For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

build_ps_metadata

def build_ps_metadata(seqlens_qo_indptr: List[Int32], pages_kv_indptr: List[Int32], context_lens: List[Int32], num_heads_k: Int32, gqa_ratio: Int32, tile_q: Int32, tile_kv: Int32, block_size: Int32, is_causal: Bool, available_tgs: Int32) -> PsMetadata

Port of get_ps_metadata_v1_2_host (v1_2_host.cuh:265-314): the host wrapper that GCD-clusters heads across TGs, then calls the per-cluster kn_generate_ps_metadata and concatenates.

For MLA-prefill this is MHA (gqa_ratio==1, one head per work-item): the work-item Q tile is qlen_granularity = tile_q // gqa_ratio TOKENS of ONE head (token-major; the 256 MMA rows are 256 tokens, NOT 16 tok x 16 head), and q_head_range's low 16 bits = the head index (= cluster_id).

Args:

seqlens_qo_indptr (List[Int32]): Prefix-sum array of length batch+1; entry i is the global query-token offset where batch i begins.
pages_kv_indptr (List[Int32]): Prefix-sum of KV block counts per batch, length batch+1; entry i is the block offset where batch i's KV pages begin.
context_lens (List[Int32]): Per-batch KV token lengths, length batch; the number of cached KV tokens for each sequence.
num_heads_k (Int32): Total number of KV heads; GCD-clustered across available_tgs thread-groups.
gqa_ratio (Int32): Ratio of query heads to KV heads; for MLA-prefill MHA this is 1. Sets qhead_granularity and divides tile_q for qlen_granularity.
tile_q (Int32): Query tile size in tokens before head division; one work-item covers tile_q // gqa_ratio tokens of one head.
tile_kv (Int32): KV tile size in blocks; sets kvlen_granularity via max(tile_kv, block_size).
block_size (Int32): Paged-attention block size in tokens; the granularity of pages_kv_indptr.
is_causal (Bool): When true, each query tile's effective KV length is masked to the causal boundary.
available_tgs (Int32): Number of persistent thread-groups; GCD-clustered with num_heads_k to form clusters.

Returns:

PsMetadata