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

kn_generate_ps_metadata

def kn_generate_ps_metadata(seqlens_qo_indptr: List[Int32], pages_kv_indptr: List[Int32], context_lens: List[Int32], num_heads_k: Int32, qhead_granularity: Int32, qlen_granularity: Int32, kvlen_granularity: Int32, block_size: Int32, is_causal: Bool, available_tgs: Int32, cluster_id: Int32) -> PsMetadata

Faithful port of v1_2_host.cuh:61-241 (single-cluster: cluster_id=0, current_work_idx=0 -- the nkv=1 prefill case).

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): Number of KV heads assigned to this cluster; each head receives an identical work split.
qhead_granularity (Int32): Number of query heads packed per work-item; sets the stride used to build q_head_range.
qlen_granularity (Int32): Query tile width in tokens; one work-item covers this many tokens of a single head.
kvlen_granularity (Int32): KV split unit width in tokens; must be a multiple of block_size. Drives the even split across thread-groups.
block_size (Int32): Paged-attention block size in tokens; the granularity of pages_kv_indptr and the KV block count.
is_causal (Bool): When true, each query tile's effective KV length is masked to the causal boundary; when false, the full KV length is used.
available_tgs (Int32): Number of persistent thread-groups to split the total KV units across evenly.
cluster_id (Int32): Index of this cluster; offsets the KV head index used to build q_head_range.

Returns:

PsMetadata