vllm.models.step5.nvidia.ops.sparse_decode_attention

DSA decode attention read directly from paged KV by selected regions.

Replaces gather-to-scratch plus a ragged FlashInfer call: each program takes all query heads of one KV group (GQA packing) for one split of a row's selected regions, loads the regions' pages in place (dequantizing FP8), and a combine kernel merges splits and clears inactive rows.

Functions:

• sparse_decode_attention –

  Attention for (token, KV group) rows over their selected regions.

• sparse_decode_splits –

  Splits per row for ~256 programs (about 1.7 waves on B300).

_gate(Gate, token, head, mask, GATED, GATE_ROW, GATE_HEAD)

Headwise sigmoid gate, rounded like the unfused gate kernel (1 if off).

Source code in vllm/models/step5/nvidia/ops/sparse_decode_attention.py

@triton.jit
def _gate(Gate, token, head, mask, GATED: tl.constexpr, GATE_ROW, GATE_HEAD):
    """Headwise sigmoid gate, rounded like the unfused gate kernel (1 if off)."""
    if GATED:
        gate = tl.load(
            Gate + token.to(tl.int64) * GATE_ROW + head * GATE_HEAD, mask, other=0.0
        ).to(tl.float32)
        sigmoid = tl.div_rn(1.0, 1.0 + tl.extra.cuda.libdevice.exp(-gate))
        return sigmoid.to(Gate.dtype.element_ty).to(tl.float32)
    return tl.full(head.shape, 1.0, tl.float32)

sparse_decode_attention(query, cache, regions, block_table, positions, active, output, *, scale, query_len=1, splits, clear_inactive=True, k_scale=None, v_scale=None, gate=None)

Attention for (token, KV group) rows over their selected regions.

With gate ([tokens, G * heads]) the headwise sigmoid gate is applied to the output, as the separate gate kernel would.

query/output: [rows, heads, 192] with rows = tokens * G; cache: [pages, G, block, 384] (K then V); regions: [rows, R] region ids (-1 unused); positions/active: per row. Rows of one request are G * query_len consecutive rows sharing a table row.

Source code in vllm/models/step5/nvidia/ops/sparse_decode_attention.py

def sparse_decode_attention(
    query: torch.Tensor,
    cache: torch.Tensor,
    regions: torch.Tensor,
    block_table: torch.Tensor,
    positions: torch.Tensor,
    active: torch.Tensor,
    output: torch.Tensor,
    *,
    scale: float,
    query_len: int = 1,
    splits: int,
    clear_inactive: bool = True,
    k_scale: torch.Tensor | None = None,
    v_scale: torch.Tensor | None = None,
    gate: torch.Tensor | None = None,
) -> None:
    """Attention for (token, KV group) rows over their selected regions.

    With ``gate`` ([tokens, G * heads]) the headwise sigmoid gate is applied
    to the output, as the separate gate kernel would.

    ``query``/``output``: [rows, heads, 192] with rows = tokens * G;
    ``cache``: [pages, G, block, 384] (K then V); ``regions``: [rows, R]
    region ids (-1 unused); ``positions``/``active``: per row. Rows of one
    request are ``G * query_len`` consecutive rows sharing a table row.
    """
    rows, heads, dim = query.shape
    groups = cache.shape[1]
    assert dim == 192 and cache.shape[-1] == 384 and heads <= 16
    assert query.stride(2) == output.stride(2) == cache.stride(3) == 1
    assert regions.stride(1) == 1 and rows == regions.shape[0]
    fp8 = cache.dtype in (torch.uint8, torch.float8_e4m3fn)
    if fp8:
        assert k_scale is not None and v_scale is not None
        cache = cache.view(torch.float8_e4m3fn)
    if not rows:
        return
    partial = torch.empty(
        (rows * splits * heads if splits > 1 else 1, 192),
        dtype=torch.float32,
        device=query.device,
    )
    lse = torch.empty(partial.shape[0], dtype=torch.float32, device=query.device)
    _sparse_decode_split[(rows, splits)](
        query,
        cache,
        regions,
        block_table,
        positions,
        active,
        k_scale,
        v_scale,
        partial,
        lse,
        output,
        gate,
        scale,
        output.stride(0),
        output.stride(1),
        clear_inactive,
        gate is not None,
        gate.stride(0) if gate is not None else 0,
        gate.stride(1) if gate is not None else 0,
        query.stride(0),
        query.stride(1),
        cache.stride(0),
        cache.stride(1),
        cache.stride(2),
        regions.stride(0),
        block_table.stride(0),
        groups,
        groups * query_len,
        heads,
        16,
        regions.shape[1],
        cache.shape[2],
        splits,
        fp8,
        _PDL,
        num_warps=4,
        launch_pdl=_PDL,
    )
    if splits > 1:
        _sparse_decode_combine[(rows, heads)](
            partial,
            lse,
            active,
            output,
            gate,
            output.stride(0),
            output.stride(1),
            heads,
            splits,
            triton.next_power_of_2(splits),
            clear_inactive,
            groups,
            gate is not None,
            gate.stride(0) if gate is not None else 0,
            gate.stride(1) if gate is not None else 0,
            _PDL,
            num_warps=4,
            launch_pdl=_PDL,
        )

sparse_decode_splits(rows)

Splits per row for ~256 programs (about 1.7 waves on B300).

Graphs are captured per batch size, so the split cannot follow context length; each split still covers only the row's live regions.

Source code in vllm/models/step5/nvidia/ops/sparse_decode_attention.py

def sparse_decode_splits(rows: int) -> int:
    """Splits per row for ~256 programs (about 1.7 waves on B300).

    Graphs are captured per batch size, so the split cannot follow context
    length; each split still covers only the row's live regions.
    """
    splits = 1
    while splits < 32 and rows * splits * 2 <= 256:
        splits *= 2
    return splits