vllm.models.inkling.amd.ops.gluon.rel_mha_extend_gfx950 ¶

rel_mha extend Gluon kernel for AMD GFX950.

This handles ragged, multi-token queries against a paged KV cache. The query axis is tiled into the MFMA M dimension (prefill-style): BLOCK_M query rows of a single q-head share each paged KV tile, so every KV tile is loaded once and reused across all rows in the tile. The grid is (blocks_per_req, batch, n_heads) -- all host-known sizes -- and each program self-locates its request from cu_seqlens_q / cache_seqlens in-kernel, so the launch stays CUDA-graph static with no device->host sync.

Visibility per query row depends on is_causal and the optional sliding window:

• is_causal=False: every query token attends the full visible cache, i.e. visible_kv = cache_seqlens[batch].
• is_causal=True: query tokens are a causal suffix, so the i-th query token of a request (0-indexed) attends prefix + i + 1 tokens, where prefix = cache_seqlens[batch] - query_len[batch].

Causal masking only touches the few KV tiles that reach the diagonal; the long prefix is a mask-free fast path. Sinks and sliding windows (causal or not) are applied per tile. This is the sole Gluon extend implementation.

82
83
84
85
86
87
88
89
90
91

def _select_extend_tile(max_seqlen_q: int, page_size: int) -> tuple[int, int, int]:
    """Return (BLOCK_M, BLOCK_N, NUM_WARPS) for the given max query length.

    Queries that fit in a single short tile use it (least padding, most
    occupancy); longer ones use the tall tile that covers more rows per
    shared-KV pass.
    """
    if max_seqlen_q <= _EXTEND_SHORT_Q_BLOCK_M:
        return _EXTEND_SHORT_Q_BLOCK_M, page_size, _EXTEND_SHORT_Q_NUM_WARPS
    return _EXTEND_LONG_Q_BLOCK_M, page_size, _EXTEND_LONG_Q_NUM_WARPS