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

Mojo function

mla_indexer_ragged_float8_paged

def mla_indexer_ragged_float8_paged[dtype: DType, KCollectionT: KVCollectionT, num_heads: Int, depth: Int, top_k: Int, mask_str: StringSpan[ImmStaticOrigin], scores_budget_bytes: Int = Int((mul get_defined_int[StringSpan("MLA_INDEX_SCORES_BUDGET_MB"), Int(512)](), 1048576)), kpool: Int = Int(1)](output_indices: TileTensor[.int32, Engine=output_indices.Engine, address_space=output_indices.address_space, linear_idx_type=output_indices.linear_idx_type], q: TileTensor[dtype, Engine=q.Engine, address_space=q.address_space, linear_idx_type=q.linear_idx_type], q_s: TileTensor[.float32, Engine=q_s.Engine, address_space=q_s.address_space, linear_idx_type=q_s.linear_idx_type], input_row_offsets: TileTensor[.uint32, Engine=input_row_offsets.Engine, address_space=input_row_offsets.address_space, linear_idx_type=input_row_offsets.linear_idx_type], k_collection: KCollectionT, layer_idx: UInt32, ctx: DeviceContext)

Compute FP8 indexed attention scores using paged KV cache and return top-k indices.

This function:

Computes FP8 matmul between q and cached k (with scales), aggregated across heads
Applies the specified mask (causal, etc.)
Computes top-k indices per token (scores are summed across all heads)

Parameters:

dtype (DType): Element type of the q query tensor, an FP8 dtype.
KCollectionT (KVCollectionT): Type of the KV collection holding cached K values and K scales.
num_heads (Int): Number of attention heads per token.
depth (Int): Per-head key dimension (head size) in elements.
top_k (Int): Requested number of top-scoring key indices to select per token.
mask_str (StringSpan[ImmStaticOrigin]): Name of the mask to apply, either MaskName.NULL or MaskName.CAUSAL.
scores_budget_bytes (Int): Peak bytes the transient score matrix may occupy. Longer batches are scored a row-window at a time to stay under it (see the chunking below). Exposed so tests can force a window small enough to exercise the multi-chunk path on toy shapes.
kpool (Int): Tokens per pooled cache row. 1 scores one row per token; k > 1 scores one pooled key per k consecutive tokens, so every candidate count and the caller's top_k are pool-granular.

Args:

output_indices (TileTensor[.int32, Engine=output_indices.Engine, address_space=output_indices.address_space, linear_idx_type=output_indices.linear_idx_type]): Dense output tensor for top-k indices [total_seq_len, top_k]. Invalid positions (where there are fewer than top_k valid keys due to causal masking or shorter sequences) are filled with -1.
q (TileTensor[dtype, Engine=q.Engine, address_space=q.address_space, linear_idx_type=q.linear_idx_type]): Query tensor [total_seq_len, num_heads, head_dim] in FP8.
q_s (TileTensor[.float32, Engine=q_s.Engine, address_space=q_s.address_space, linear_idx_type=q_s.linear_idx_type]): Query scales [total_seq_len, num_heads] in float32.
input_row_offsets (TileTensor[.uint32, Engine=input_row_offsets.Engine, address_space=input_row_offsets.address_space, linear_idx_type=input_row_offsets.linear_idx_type]): Ragged row offsets for queries [batch_size + 1].
k_collection (KCollectionT): KV collection containing cached K values and K scales. K scales are accessed via k_cache.scales (quantization_granularity=head_size).
layer_idx (UInt32): Layer index for retrieving cache.
ctx (DeviceContext): Device context.