vllm.models.step5.nvidia.ops.swa_decode

Sliding-window GQA decode with split KV and a fused headwise gate.

FA2 runs one CTA per query head and scans the whole window serially; its Python wrapper rejects split KV, and FA4 has no head_dim 192 on SM100. Here one program handles all query heads of a KV head (the GQA group) for one window split, and a combine kernel merges splits and applies the gate.

Functions:

• swa_decode_gated –

  Decode rows of query_len tokens per request over a causal window.

• swa_decode_plan –

  Return (pack, splits) for a decode batch.

• swa_decode_splits –

  Split the window until (token, KV head) programs reach ~2 waves.

swa_decode_gated(query, key_cache, value_cache, block_table, seq_lens, gate, output, *, scale, window, query_len, splits, pack=1)

Decode rows of query_len tokens per request over a causal window.

query/output: [tokens, heads, 192]; key_cache/value_cache: [pages, block, kv_heads, 192] views; gate: [tokens, heads]. Rows of requests with no visible keys (padding) produce zeros.

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

def swa_decode_gated(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
    gate: torch.Tensor,
    output: torch.Tensor,
    *,
    scale: float,
    window: int,
    query_len: int,
    splits: int,
    pack: int = 1,
) -> None:
    """Decode rows of ``query_len`` tokens per request over a causal window.

    ``query``/``output``: [tokens, heads, 192]; ``key_cache``/``value_cache``:
    [pages, block, kv_heads, 192] views; ``gate``: [tokens, heads]. Rows of
    requests with no visible keys (padding) produce zeros.
    """
    tokens, heads, dim = query.shape
    kv_heads = key_cache.shape[2]
    group = heads // kv_heads
    assert dim == 192 and group <= 16 and heads == group * kv_heads
    assert query.stride(2) == key_cache.stride(3) == value_cache.stride(3) == 1
    assert output.stride(2) == 1 and tokens == seq_lens.shape[0] * query_len
    assert query_len % pack == 0 and pack * group <= 64
    if not tokens:
        return
    partial = torch.empty(
        (tokens * kv_heads * splits * group if splits > 1 else 1, 192),
        dtype=torch.float32,
        device=query.device,
    )
    lse = torch.empty(partial.shape[0], dtype=torch.float32, device=query.device)
    _swa_decode_split[(tokens // pack * kv_heads, splits)](
        query,
        key_cache,
        value_cache,
        block_table,
        seq_lens,
        partial,
        lse,
        gate,
        output,
        scale,
        output.stride(0),
        output.stride(1),
        gate.stride(0),
        gate.stride(1),
        query.stride(0),
        query.stride(1),
        key_cache.stride(0),
        key_cache.stride(1),
        key_cache.stride(2),
        value_cache.stride(0),
        value_cache.stride(1),
        value_cache.stride(2),
        block_table.stride(0),
        kv_heads,
        group,
        pack,
        max(16, triton.next_power_of_2(pack * group)),
        query_len,
        key_cache.shape[1],
        window,
        splits,
        64,
        _PDL,
        num_warps=4,
        launch_pdl=_PDL,
    )
    if splits > 1:
        _swa_decode_combine[(tokens, heads)](
            partial,
            lse,
            gate,
            output,
            output.stride(0),
            output.stride(1),
            gate.stride(0),
            gate.stride(1),
            kv_heads,
            group,
            splits,
            triton.next_power_of_2(splits),
            _PDL,
            num_warps=4,
            launch_pdl=_PDL,
        )

swa_decode_plan(tokens, kv_heads, group, query_len, num_sms, window)

Return (pack, splits) for a decode batch.

MTP verify tokens of a request are packed into one program when (token, KV head) programs already fill about a wave, since their windows overlap; the window is then split to fill the GPU.

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

def swa_decode_plan(
    tokens: int, kv_heads: int, group: int, query_len: int, num_sms: int, window: int
) -> tuple[int, int]:
    """Return (pack, splits) for a decode batch.

    MTP verify tokens of a request are packed into one program when
    (token, KV head) programs already fill about a wave, since their windows
    overlap; the window is then split to fill the GPU.
    """
    pack = (
        query_len
        if query_len > 1 and tokens * kv_heads >= 128 and query_len * group <= 64
        else 1
    )
    return pack, swa_decode_splits(tokens // pack * kv_heads, num_sms, window)

swa_decode_splits(programs, num_sms, window)

Split the window until (token, KV head) programs reach ~2 waves.

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

def swa_decode_splits(programs: int, num_sms: int, window: int) -> int:
    """Split the window until (token, KV head) programs reach ~2 waves."""
    splits = 1
    while (
        splits < 8
        and programs * splits * 2 <= 2 * num_sms
        and window // (splits * 2) >= 64
    ):
        splits *= 2
    return splits