vllm.v1.attention.ops.triton_decode_attention

Memory-efficient attention for decoding. It supports page size >= 1.

Functions:

• decode_attention_fwd_verify –

  Causal MLA verify over a per-request block table.

• select_verify_config –

  Launch config from MI300X (gfx942) sweeps, bf16, page 16.

• verify_num_kv_splits –

  Static upper bound on KV splits (safe to bake into a cudagraph).

• warmup_decode_attention_fwd_verify –

  Compile every launch variant serving can reach for these tensors.

_storage_view(t, shape)

Contiguous shape over t's storage, or a new tensor if too small.

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

def _storage_view(t, shape):
    """Contiguous ``shape`` over ``t``'s storage, or a new tensor if too small."""
    if math.prod(shape) > t.untyped_storage().nbytes() // t.element_size():
        return t.new_empty(shape)
    strides = [1]
    for size in reversed(shape[1:]):
        strides.insert(0, strides[0] * size)
    return torch.as_strided(t, shape, strides, 0)

decode_attention_fwd_verify(q, k_buffer, v_buffer, o, lse, req_to_token, seq_lens, real_seq_lens, qlen, attn_logits, num_kv_splits, sm_scale, page_size, k_scale=None)

Causal MLA verify over a per-request block table.

Parameters:

• q

  (Tensor) –

  [B * qlen, H, Lk] queries, verify tokens of a request contiguous.

• k_buffer

  (Tensor) –

  paged latent cache, [num_pages, page_size, 1, Lk] (or 3D [num_pages * page_size, 1, Lk]).

• v_buffer

  (Tensor) –

  k_buffer[..., :Lv] (V is read from K; only its width is used).

• o

  (Tensor) –

  [B * qlen, H, Lv] output.

• lse

  (Tensor) –

  [B * qlen, H] fp32 natural-log LSE (-inf for rows with no key).

• req_to_token

  (Tensor) –

  [B, max_pages] per-request block table (not repeated).

• seq_lens

  (Tensor) –

  [B] int32 kernel KV lengths L.

• real_seq_lens

  (Tensor | None) –

  [B] int32 or None; row limit is max(min(L - qlen + 1 + t, L_real), 0).

• qlen

  (int) –

  verify tokens per request.

• attn_logits

  (Tensor | None) –

  fp32 workspace with >= B * qlen * H * num_kv_splits * (Lv + 1) elements, laid out [B * qlen, H, num_kv_splits, Lv + 1]. Unused, and may be None, when num_kv_splits is 1.

• num_kv_splits

  (int) –

  static cap on KV splits per request (see verify_num_kv_splits); shorter requests use fewer, longer splits.

• sm_scale

  (float) –

  softmax scale.

• page_size

  (int) –

  tokens per KV cache page.

• k_scale

  (Tensor | None, default: None ) –

  scalar tensor dequantizing an FP8 cache; ignored for bf16.

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

def decode_attention_fwd_verify(
    q,
    k_buffer,
    v_buffer,
    o,
    lse,
    req_to_token,
    seq_lens,
    real_seq_lens,
    qlen,
    attn_logits,
    num_kv_splits,
    sm_scale,
    page_size,
    k_scale=None,
):
    """Causal MLA verify over a per-request block table.

    Args:
        q (torch.Tensor): [B * qlen, H, Lk] queries, verify tokens of a request
            contiguous.
        k_buffer (torch.Tensor): paged latent cache, [num_pages, page_size, 1, Lk]
            (or 3D [num_pages * page_size, 1, Lk]).
        v_buffer (torch.Tensor): k_buffer[..., :Lv] (V is read from K; only its
            width is used).
        o (torch.Tensor): [B * qlen, H, Lv] output.
        lse (torch.Tensor): [B * qlen, H] fp32 natural-log LSE (-inf for rows
            with no key).
        req_to_token (torch.Tensor): [B, max_pages] per-request block table
            (not repeated).
        seq_lens (torch.Tensor): [B] int32 kernel KV lengths L.
        real_seq_lens (torch.Tensor | None): [B] int32 or None; row limit is
            max(min(L - qlen + 1 + t, L_real), 0).
        qlen (int): verify tokens per request.
        attn_logits (torch.Tensor | None): fp32 workspace with >= B * qlen * H *
            num_kv_splits * (Lv + 1) elements, laid out
            [B * qlen, H, num_kv_splits, Lv + 1].
            Unused, and may be None, when num_kv_splits is 1.
        num_kv_splits (int): static cap on KV splits per request (see
            verify_num_kv_splits); shorter requests use fewer, longer splits.
        sm_scale (float): softmax scale.
        page_size (int): tokens per KV cache page.
        k_scale (torch.Tensor | None): scalar tensor dequantizing an FP8 cache;
            ignored for bf16.

    """
    cfg = select_verify_config(q.shape[0] // qlen, qlen, q.shape[1])
    Lk = k_buffer.shape[-1]
    Lv = v_buffer.shape[-1]
    assert k_buffer.shape[-2] == 1, "verify path expects a single latent KV head"
    num_rows, num_heads = q.shape[0], q.shape[1]
    assert num_rows % qlen == 0
    batch = num_rows // qlen
    if num_kv_splits == 1:
        # Stage 1 writes o / lse directly; o only stands in for the pointer.
        attn_logits, mid_strides = o, (0, 0, 0)
    else:
        assert attn_logits.shape[-2] == num_kv_splits
        mid_strides = attn_logits.stride()[:3]
    if k_scale is None:
        k_scale = torch.ones((), dtype=torch.float32, device=q.device)
    BLOCK_DMODEL = triton.next_power_of_2(Lv)
    BLOCK_DPE = triton.next_power_of_2(Lk - Lv) if Lk > Lv else 0
    assert Lv == BLOCK_DMODEL and Lk == BLOCK_DMODEL + BLOCK_DPE
    has_real = real_seq_lens is not None
    real = real_seq_lens if has_real else seq_lens

    hip_kw = {}
    if is_hip_:
        hip_kw = dict(
            waves_per_eu=cfg["waves_per_eu"],
            matrix_instr_nonkdim=cfg["matrix_instr_nonkdim"],
            kpack=cfg["kpack"],
        )
    grid = (verify_row_blocks(qlen, num_heads, cfg["BLOCK_H"]) * batch * num_kv_splits,)
    _fwd_verify_kernel_stage1[grid](
        q,
        k_buffer,
        sm_scale * _LOG2E,
        req_to_token,
        seq_lens,
        real,
        attn_logits,
        o,
        lse,
        req_to_token.stride(0),
        q.stride(0),
        q.stride(1),
        _page_stride(k_buffer, page_size),
        k_buffer.stride(-3),
        *mid_strides,
        o.stride(0),
        o.stride(1),
        lse.stride(0),
        lse.stride(1),
        k_scale,
        batch,
        QLEN=qlen,
        NUM_HEADS=num_heads,
        HAS_REAL=has_real,
        BLOCK_DMODEL=BLOCK_DMODEL,
        BLOCK_DPE=BLOCK_DPE,
        BLOCK_N=cfg["BLOCK_N"],
        BLOCK_H=cfg["BLOCK_H"],
        NUM_KV_SPLITS=num_kv_splits,
        SPLIT_ALIGN=cfg["SPLIT_ALIGN"],
        MIN_SPLIT_LEN=cfg["MIN_SPLIT_LEN"],
        PAGE_SIZE=page_size,
        Lv=Lv,
        QSPLIT=bool(cfg["QSPLIT"]),
        SINGLE=num_kv_splits == 1,
        num_warps=cfg["num_warps"],
        num_stages=cfg["num_stages"],
        **hip_kw,
    )
    if num_kv_splits == 1:
        return  # stage 1 wrote o / lse
    _fwd_verify_kernel_stage2[(num_rows, num_heads)](
        attn_logits,
        o,
        lse,
        seq_lens,
        real,
        attn_logits.stride(0),
        attn_logits.stride(1),
        attn_logits.stride(2),
        o.stride(0),
        o.stride(1),
        lse.stride(0),
        lse.stride(1),
        QLEN=qlen,
        HAS_REAL=has_real,
        NUM_KV_SPLITS=num_kv_splits,
        SPLIT_ALIGN=cfg["SPLIT_ALIGN"],
        MIN_SPLIT_LEN=cfg["MIN_SPLIT_LEN"],
        BLOCK_S=cfg["S2_BLOCK_S"],
        BLOCK_DV=triton.next_power_of_2(Lv),
        Lv=Lv,
        num_warps=cfg["s2_num_warps"],
        num_stages=1,
    )

select_verify_config(batch, qlen, num_heads)

Launch config from MI300X (gfx942) sweeps, bf16, page 16.

Few rows (B * qlen * H <= 640): 32-row tiles with more KV splits fill the GPU. Otherwise 64-row tiles (Q split in two 256-wide halves, next K/V tile prefetched) halve the KV re-reads per request.

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

def select_verify_config(batch: int, qlen: int, num_heads: int) -> dict[str, Any]:
    """Launch config from MI300X (gfx942) sweeps, bf16, page 16.

    Few rows (B * qlen * H <= 640): 32-row tiles with more KV splits fill the
    GPU. Otherwise 64-row tiles (Q split in two 256-wide halves, next K/V
    tile prefetched) halve the KV re-reads per request.
    """
    cfg = dict(VERIFY_DEFAULT_CONFIG)
    if batch * qlen * num_heads <= _VERIFY_SMALL_MAX_ROWS:
        cfg.update(_VERIFY_SMALL_CONFIG)
    else:
        cfg.update(_VERIFY_LARGE_CONFIG)
    return cfg

verify_num_kv_splits(max_seq_len, batch, qlen, num_heads, sm_count, block_h=VERIFY_DEFAULT_CONFIG['BLOCK_H'], occupancy=1.0, min_work_per_split=64, max_splits=8)

Static upper bound on KV splits (safe to bake into a cudagraph).

The kernel splits each request's own length into max(cdiv(L, splits), MIN_SPLIT_LEN) chunks, so this only caps the split count. occupancy is the number of resident stage-1 programs per CU. If one split of every request fits in a single dispatch round, use as many splits as fit in that round; otherwise pick n <= max_splits minimising rounds(n) / n (the idle tail of the last round), preferring fewer splits within 5% of the best.

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

def verify_num_kv_splits(
    max_seq_len: int,
    batch: int,
    qlen: int,
    num_heads: int,
    sm_count: int,
    block_h: int = VERIFY_DEFAULT_CONFIG["BLOCK_H"],
    occupancy: float = 1.0,
    min_work_per_split: int = 64,
    max_splits: int = 8,
) -> int:
    """Static upper bound on KV splits (safe to bake into a cudagraph).

    The kernel splits each request's own length into
    max(cdiv(L, splits), MIN_SPLIT_LEN) chunks, so this only caps the split
    count. ``occupancy`` is the number of resident stage-1 programs per CU.
    If one split of every request fits in a single dispatch round, use as many
    splits as fit in that round; otherwise pick n <= max_splits minimising
    rounds(n) / n (the idle tail of the last round), preferring fewer splits
    within 5% of the best.
    """
    programs_per_split = max(1, batch) * verify_row_blocks(qlen, num_heads, block_h)
    ideal = triton.cdiv(max(1, max_seq_len), max(1, min_work_per_split))
    slots = max(1, int(occupancy * sm_count))
    if programs_per_split <= slots:
        n = slots // programs_per_split
    else:
        cost = [
            (triton.cdiv(n * programs_per_split, slots) / n, n)
            for n in range(1, max(1, min(ideal, max_splits)) + 1)
        ]
        best = min(c for c, _ in cost)
        n = min(n for c, n in cost if c <= best * 1.05)
    return max(1, min(ideal, n, sm_count * 2))

warmup_decode_attention_fwd_verify(q, k_buffer, o, lse, req_to_token, has_real, qlens, attn_logits, sm_scale, page_size, k_scale=None)

Compile every launch variant serving can reach for these tensors.

A launch specializes on qlen, the row-count config and a single split, but graph capture and profiling only run the full verify length. Each variant runs once on scratch tensors with the strides of the real ones; the split scratch reuses attn_logits' storage, whose size the launch specializes on.

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

def warmup_decode_attention_fwd_verify(
    q,
    k_buffer,
    o,
    lse,
    req_to_token,
    has_real,
    qlens,
    attn_logits,
    sm_scale,
    page_size,
    k_scale=None,
):
    """Compile every launch variant serving can reach for these tensors.

    A launch specializes on qlen, the row-count config and a single split, but
    graph capture and profiling only run the full verify length. Each variant
    runs once on scratch tensors with the strides of the real ones; the split
    scratch reuses attn_logits' storage, whose size the launch specializes on.
    """
    num_heads, Lv = q.shape[1], o.shape[-1]
    if attn_logits is None:
        attn_logits = q.new_empty(0, dtype=torch.float32)

    def scratch(t, rows):
        return torch.empty_strided(
            (rows, *t.shape[1:]), t.stride(), dtype=t.dtype, device=t.device
        )

    for qlen in qlens:
        for batch in {1, _VERIFY_SMALL_MAX_ROWS // (qlen * num_heads) + 1}:
            rows = batch * qlen
            seq_lens = torch.full((batch,), qlen, dtype=torch.int32, device=q.device)
            block_table = torch.empty_strided(
                (batch, req_to_token.shape[1]),
                req_to_token.stride(),
                dtype=req_to_token.dtype,
                device=req_to_token.device,
            ).zero_()
            for num_kv_splits in (1, 2):
                decode_attention_fwd_verify(
                    scratch(q, rows).zero_(),
                    k_buffer,
                    k_buffer[..., :Lv],
                    scratch(o, rows),
                    scratch(lse, rows),
                    block_table,
                    seq_lens,
                    seq_lens if has_real else None,
                    qlen,
                    None
                    if num_kv_splits == 1
                    else _storage_view(
                        attn_logits, (rows, num_heads, num_kv_splits, Lv + 1)
                    ),
                    num_kv_splits,
                    sm_scale,
                    page_size,
                    k_scale=k_scale,
                )