vllm.models.step5.amd.ops.csa_paged

Paged Step-5 CSA kernels; each cache has its own block table and page size.

Functions:

• sparse_attention_paged –

  Attention over each query's selected regions of the paged K/V, into

sparse_attention_paged(query, key_cache, value_cache, regions, block_table, token_req, positions, output, *, region_size, scale, k_scale=None, v_scale=None)

Attention over each query's selected regions of the paged K/V, into output. key_cache/value_cache are [pages, page_tokens, kv_heads, head_dim] views; token_req and positions locate every query row.

Source code in vllm/models/step5/amd/ops/csa_paged.py

def sparse_attention_paged(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    regions: torch.Tensor,
    block_table: torch.Tensor,
    token_req: torch.Tensor,
    positions: torch.Tensor,
    output: torch.Tensor,
    *,
    region_size: int,
    scale: float,
    k_scale: torch.Tensor | None = None,
    v_scale: torch.Tensor | None = None,
) -> None:
    """Attention over each query's selected regions of the paged K/V, into
    ``output``. ``key_cache``/``value_cache`` are ``[pages, page_tokens, kv_heads,
    head_dim]`` views; ``token_req`` and ``positions`` locate every query row."""
    tokens, heads, head_dim = query.shape
    if tokens == 0:
        return
    groups = key_cache.shape[2]
    heads_per_group = heads // groups
    fp8 = key_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
    assert value_cache.dtype == key_cache.dtype
    assert not fp8 or (k_scale is not None and v_scale is not None)
    _attention_paged_kernel[(tokens, groups)](
        query,
        key_cache,
        value_cache,
        regions,
        block_table,
        token_req,
        positions,
        output,
        k_scale,
        v_scale,
        scale,
        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),
        regions.stride(0),
        regions.stride(1),
        block_table.stride(0),
        output.stride(0),
        output.stride(1),
        HEADS_PER_GROUP=heads_per_group,
        HEAD_DIM=head_dim,
        PAGE_TOKENS=key_cache.shape[1],
        REGION_SIZE=region_size,
        SLOTS=regions.shape[-1],
        BLOCK_H=max(16, triton.next_power_of_2(heads_per_group)),
        BLOCK_D=triton.next_power_of_2(head_dim),
        # 8 regions x 8 tokens x 256 padded dims fits gfx942's 64 KiB LDS.
        BLOCK_S=8,
        num_stages=1,
        FP8=fp8,
    )