vllm.v1.attention.ops.triton_prefill_attention

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

Functions:

• context_attention_fwd –

  q, k, v: [b * s, head, head_dim]

_fit_block_to_shared_memory(block, head_dim, element_size) cached

Halve the tile until its Q, K and P tiles fit in shared memory.

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

@cache
def _fit_block_to_shared_memory(block: int, head_dim: int, element_size: int) -> int:
    """Halve the tile until its Q, K and P tiles fit in shared memory."""
    max_shared_mem = get_max_shared_memory_bytes()
    while block > 16 and element_size * block * (2 * head_dim + block) > max_shared_mem:
        block //= 2
    return block

_prefer_narrow_kv_tile()

RDNA3/RDNA4 prefer a narrower KV tile than the shared default.

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

def _prefer_narrow_kv_tile() -> bool:
    """RDNA3/RDNA4 prefer a narrower KV tile than the shared default."""
    if not current_platform.is_rocm():
        return False
    from vllm.platforms.rocm import on_gfx1x

    # RDNA3/RDNA4 only: on_gfx1x() excludes gfx10xx and gfx1250.
    return on_gfx1x()

context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len, is_causal=True, softmax_scale=None, sliding_window_q=None, sliding_window_k=None, sinks=None, sinks_bias_key0=False)

q, k, v: [b * s, head, head_dim] b_start_loc: [b] b_seq_len: [b] out: [b * s, head, head_dim]

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

def context_attention_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    b_start_loc: torch.Tensor,
    b_seq_len: torch.Tensor,
    max_input_len: int,
    is_causal: bool = True,
    softmax_scale: float | None = None,
    sliding_window_q: int | None = None,
    sliding_window_k: int | None = None,
    sinks: torch.Tensor | None = None,
    sinks_bias_key0: bool = False,
):
    """q, k, v: [b * s, head, head_dim]
    b_start_loc: [b]
    b_seq_len: [b]
    out: [b * s, head, head_dim]
    """
    Lq, Lk, _ = q.shape[-1], k.shape[-1], v.shape[-1]
    BLOCK = get_block_size(q.dtype)
    if Lk >= 512:
        BLOCK = min(BLOCK, 32)
    if current_platform.is_cuda():
        BLOCK = _fit_block_to_shared_memory(
            BLOCK, triton.next_power_of_2(Lk), q.element_size()
        )

    sm_scale = 1.0 / (Lq**0.5) if softmax_scale is None else softmax_scale
    # rescale with 1/ln(2) for triton exp2
    sm_scale *= RCP_LN2

    batch, head = b_seq_len.shape[0], q.shape[1]
    kv_group_num = q.shape[1] // k.shape[1]
    if sinks is not None:
        assert sinks.shape[0] == head, "Sinks must be num_query_heads size"

    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
    num_warps = 4 if Lk <= 64 else 8

    # BLOCK_M, num_warps and num_stages stay at the shared defaults; min()
    # leaves dtypes whose default tile is already 32, such as float32, alone.
    BLOCK_N = min(BLOCK, 32) if _prefer_narrow_kv_tile() else BLOCK

    sliding_window_q = sliding_window_q if sliding_window_q is not None else 0
    sliding_window_k = sliding_window_k if sliding_window_k is not None else 0

    _fwd_kernel[grid](
        q,
        k,
        v,
        sinks if sinks is not None else q,
        sm_scale,
        b_start_loc,
        b_seq_len,
        o,
        q.stride(0),
        q.stride(1),
        k.stride(0),
        k.stride(1),
        v.stride(0),
        v.stride(1),
        o.stride(0),
        o.stride(1),
        kv_group_num=kv_group_num,
        BLOCK_M=BLOCK,
        BLOCK_DMODEL=triton.next_power_of_2(Lk),
        BLOCK_N=BLOCK_N,
        IS_CAUSAL=is_causal,
        SLIDING_WINDOW_Q=sliding_window_q,
        SLIDING_WINDOW_K=sliding_window_k,
        USE_SINKS=sinks is not None,
        SINKS_BIAS_KEY0=sinks_bias_key0,
        num_warps=num_warps,
        num_stages=1,
        Lk=Lk,
    )