vllm.v1.attention.ops.dash_attention

DASH paged KV kernels and segmented FlashAttention/FlashInfer execution.

Functions:

• cache_dash_prefill –

  Store rotated keys and integer page state in one auxiliary KV head.

• flash_dash_attention –

  Attend gathered segments and combine normalized partial outputs.

• paged_dash_decode –

  One query per request; all sequence-dependent decisions stay on GPU.

• paged_dash_prefill –

  Extend a prompt using paged history and full-prompt page coordinates.

cache_dash_prefill(k, v, ik, roles, slots, cache, sp, head_cls=None)

Store rotated keys and integer page state in one auxiliary KV head.

Source code in vllm/v1/attention/ops/dash_attention.py

def cache_dash_prefill(k, v, ik, roles, slots, cache, sp, head_cls=None):
    """Store rotated keys and integer page state in one auxiliary KV head."""
    tok, page, _, query_page = page_coordinates(roles)
    num_pages = int(page.max())
    last_image = torch.where(roles == 1)[0][-1]
    prefill = torch.arange(len(roles), device=roles.device) <= last_image
    prompt_counts = torch.zeros(num_pages + 1, device=roles.device, dtype=torch.long)
    prompt_counts.scatter_add_(0, page, prefill.long())
    state = torch.stack(
        (tok, page, torch.full_like(tok, num_pages), query_page), -1
    ).to(torch.int32)
    shifts = torch.arange(4, device=k.device) * 8
    encoded = ((state[:, :, None] >> shifts) & 255).flatten(1)
    counts = ((prompt_counts[:, None] >> shifts) & 255).to(k.dtype)
    h, dim = k.shape[1:]
    packed = torch.zeros((len(k), h + 1, 2 * dim), device=k.device, dtype=k.dtype)
    packed[:, :h, :dim] = page_rope(k, tok, page, sp, main=True)
    packed[:, :h, dim:] = v
    packed[:, h, :dim] = page_rope(ik, tok, page, sp, main=False)
    packed[:, h, dim] = roles
    packed[:, h, dim + 2 : dim + 18] = encoded
    packed[: num_pages + 1, h, dim + 18 : dim + 22] = counts
    b = cache.shape[2]
    cache[slots // b, :, slots % b] = packed
    if head_cls is not None:
        bs = int(sp["sparse_block_size"])
        cb, ch, ct, _ = cache.stride()
        _cache_summaries[(len(k) // bs,)](
            cache,
            slots,
            head_cls,
            cb,
            ch,
            ct,
            b,
            h,
            dim,
            bs,
            num_warps=4,
        )

flash_dash_attention(q, k, v, used, bias, backend, bs, local)

Attend gathered segments and combine normalized partial outputs.

k/v are contiguous [N(topkbs+local), Hkv, D]; used contains the valid prefix length of each segment, including zero for inactive rows.

Source code in vllm/v1/attention/ops/dash_attention.py

def flash_dash_attention(q, k, v, used, bias, backend, bs, local):
    """Attend gathered segments and combine normalized partial outputs.

    k/v are contiguous [N*(topk*bs+local), Hkv, D]; used contains the
    valid prefix length of each segment, including zero for inactive rows.
    """
    n, heads, dim = q.shape
    topk = bias.shape[1]
    groups = topk + 1
    device = q.device
    if q.dtype not in (torch.float16, torch.bfloat16):
        raise ValueError("DASH FlashAttention/FlashInfer require FP16 or BF16")
    if bs % 8 or local % 8:
        raise ValueError("DASH Flash backends require byte-aligned segments")
    _, _, cuq, cuk = _indptr(device, n, topk, bs, local)
    queries = q[:, None].expand(n, groups, heads, dim).reshape(-1, heads, dim)
    if backend == "FLASH_ATTN":
        from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
        from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func

        version = get_flash_attn_version(head_size=dim)
        if version is None:
            raise RuntimeError(
                "DASH FLASH_ATTN requires a supported CUDA FlashAttention"
            )
        block_table = None
        if version != 4:
            # FA2/3 use paged addressing to separate segment offsets and lengths.
            block_table = _flash_attn_block_table(device, n, topk, bs, local)
            k = k.view(-1, bs, k.shape[1], dim)
            v = v.view_as(k)
        # Empty FA2 GQA segments can leave invalid LSE; their results are masked below.
        lengths = used.clamp_min(1) if version == 2 else used
        out, lse = flash_attn_varlen_func(
            queries,
            k,
            v,
            max_seqlen_q=1,
            cu_seqlens_q=cuq,
            max_seqlen_k=max(bs, local),
            cu_seqlens_k=cuk if version == 4 else None,
            seqused_k=lengths,
            block_table=block_table,
            softmax_scale=dim**-0.5,
            causal=False,
            return_softmax_lse=True,
            fa_version=version,
        )
        lse = lse.T
    elif backend == "FLASHINFER":
        key = (device, n, topk, bs, local, heads, k.shape[1], dim, q.dtype)
        plan = _flashinfer_graph_plans.get(key)
        if plan is None:
            plan = _flashinfer_plan(*key)
            # Captured buffers must outlive eviction from the eager LRU.
            if torch.cuda.is_current_stream_capturing():
                _flashinfer_graph_plans[key] = plan
        wrapper, packed = plan
        remote = (
            torch.arange(bs, device=device)[None, None]
            < used.view(n, groups)[:, :topk, None]
        )
        nearby = (
            torch.arange(local, device=device)[None] < used.view(n, groups)[:, -1, None]
        )
        mask = torch.cat((remote.flatten(1), nearby), -1).reshape(-1, 8)
        bits = torch.arange(8, device=device)
        packed.copy_((mask.to(torch.int32) << bits).sum(-1).to(torch.uint8))
        out, lse = wrapper.run(
            queries, k, v, return_lse=True, lse_base="ln", lse_layout="NH"
        )
    else:
        raise ValueError(f"Unsupported DASH Flash backend: {backend}")
    live = used.view(n, groups, 1) > 0
    prior = torch.cat((bias, torch.zeros_like(bias[:, :1])), -1)
    logits = lse.reshape(n, groups, heads).float() + prior[:, :, None]
    weights = logits.masked_fill(~live, -torch.inf).softmax(1).nan_to_num()
    partial = out.view(n, groups, heads, dim).float().masked_fill(~live[..., None], 0)
    return (partial * weights[..., None]).sum(1).to(q.dtype)

paged_dash_decode(q, k, v, iq, ik, roles, cache, block_table, seq_lens, slots, head_cls, sp, layer_idx, max_model_len, prefix_lens=None, backend='FLASH_ATTN', cached_summaries=False, positions=None)

One query per request; all sequence-dependent decisions stay on GPU.

Source code in vllm/v1/attention/ops/dash_attention.py

def paged_dash_decode(
    q,
    k,
    v,
    iq,
    ik,
    roles,
    cache,
    block_table,
    seq_lens,
    slots,
    head_cls,
    sp,
    layer_idx,
    max_model_len,
    prefix_lens=None,
    backend="FLASH_ATTN",
    cached_summaries=False,
    positions=None,
):
    """One query per request; all sequence-dependent decisions stay on GPU."""
    n, nh, dim = q.shape
    h = k.shape[1]
    b = cache.shape[2]
    cb, ch, ct, _ = cache.stride()
    ts = block_table.stride(0)
    coords = torch.empty((n, 4), dtype=torch.int32, device=q.device)
    _coordinates[(n,)](
        cache,
        block_table,
        seq_lens,
        slots,
        roles,
        coords,
        cb,
        ch,
        ct,
        ts,
        b,
        h,
        dim,
        max_model_len,
    )
    if positions is not None:
        page = torch.minimum(positions[1], positions[2])
        prompt = torch.stack(
            (positions[0], page, positions[2], positions[1] - 1), -1
        ).to(torch.int32)
        coords = torch.where((positions[3] < 0)[:, None], prompt, coords)
    q = decode_page_rope(q, coords[:, 0], coords[:, 1], sp, main=True)
    k = decode_page_rope(k, coords[:, 0], coords[:, 1], sp, main=True)
    iq = decode_page_rope(iq, coords[:, 0], coords[:, 1], sp, main=False)
    ik = decode_page_rope(ik, coords[:, 0], coords[:, 1], sp, main=False)
    _write[(n, h + 1)](
        cache,
        k,
        v,
        ik,
        roles,
        coords,
        slots,
        cb,
        ch,
        ct,
        k.stride(0),
        k.stride(1),
        v.stride(0),
        v.stride(1),
        ik.stride(0),
        b,
        h,
        dim,
    )
    if positions is not None:
        _prefill_metadata[(n,)](
            cache,
            block_table,
            seq_lens,
            slots,
            positions,
            head_cls,
            positions.stride(0),
            cb,
            ch,
            ct,
            ts,
            b,
            h,
            dim,
            int(sp["sparse_block_size"]),
            num_warps=4,
        )
    return _paged_dash_attention(
        q,
        iq,
        coords,
        cache,
        block_table,
        seq_lens,
        slots,
        head_cls,
        sp,
        layer_idx,
        max_model_len,
        prefix_lens,
        backend,
        cached_summaries,
    )

paged_dash_prefill(q, k, v, iq, ik, roles, positions, cache, block_table, length, slots, head_cls, sp, layer_idx, backend, prefix_lens=None)

Extend a prompt using paged history and full-prompt page coordinates.

Source code in vllm/v1/attention/ops/dash_attention.py

def paged_dash_prefill(
    q,
    k,
    v,
    iq,
    ik,
    roles,
    positions,
    cache,
    block_table,
    length,
    slots,
    head_cls,
    sp,
    layer_idx,
    backend,
    prefix_lens=None,
):
    """Extend a prompt using paged history and full-prompt page coordinates."""
    n, h, dim = k.shape
    page = torch.minimum(positions[1], positions[2])
    coords = torch.stack((positions[0], page, positions[2], positions[1] - 1), -1).to(
        torch.int32
    )
    q = page_rope(q, positions[0], page, sp, main=True)
    k = page_rope(k, positions[0], page, sp, main=True)
    iq = page_rope(iq, positions[0], page, sp, main=False)
    ik = page_rope(ik, positions[0], page, sp, main=False)
    cb, ch, ct, _ = cache.stride()
    b = cache.shape[2]
    _write[(n, h + 1)](
        cache,
        k,
        v,
        ik,
        roles,
        coords,
        slots,
        cb,
        ch,
        ct,
        k.stride(0),
        k.stride(1),
        v.stride(0),
        v.stride(1),
        ik.stride(0),
        b,
        h,
        dim,
    )
    table = block_table.expand(n, -1)
    lens = torch.arange(length - n + 1, length + 1, device=q.device, dtype=torch.int32)
    _prefill_metadata[(n,)](
        cache,
        table,
        lens,
        slots,
        positions,
        head_cls,
        positions.stride(0),
        cb,
        ch,
        ct,
        table.stride(0),
        b,
        h,
        dim,
        int(sp["sparse_block_size"]),
        num_warps=4,
    )
    output = torch.empty_like(q)
    for start in range(0, n, 256):
        end = min(start + 256, n)
        prefix = torch.full_like(lens[start:end], length)
        if prefix_lens is not None:
            prefix = torch.minimum(prefix, prefix_lens)
        output[start:end] = _paged_dash_attention(
            q[start:end],
            iq[start:end],
            coords[start:end],
            cache,
            table[start:end],
            lens[start:end],
            slots[start:end],
            head_cls,
            sp,
            layer_idx,
            length,
            prefix,
            backend,
            True,
        )
    return output