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

mla_prefill_decode_graph_bf16

def mla_prefill_decode_graph_bf16[collection_t: KVCollectionT, //, mask_str: StringSlice[ImmStaticOrigin], kv_input_fn: def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[DType.bfloat16, width], target: StringSlice[ImmStaticOrigin] = StringSlice("cpu"), sparse_mla: Bool = False, sparse_indices_stride: Int = Int(0)](output: TileTensor[DType.bfloat16, Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[DType.bfloat16, Storage=q.Storage, linear_idx_type=q.linear_idx_type], input_row_offsets: TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type], freqs_cis: TileTensor[Storage=freqs_cis.Storage, linear_idx_type=freqs_cis.linear_idx_type], kv_norm_gamma: TileTensor[Storage=kv_norm_gamma.Storage, linear_idx_type=kv_norm_gamma.linear_idx_type], kv_collection: collection_t, layer_idx: UInt32, scale: Float32, epsilon: Float32, buffer_row_offsets: TileTensor[DType.uint32, Storage=buffer_row_offsets.Storage, linear_idx_type=buffer_row_offsets.linear_idx_type], cache_offsets: TileTensor[DType.uint32, Storage=cache_offsets.Storage, linear_idx_type=cache_offsets.linear_idx_type], buffer_length: Int, max_seq_len: Int, w_k: TileTensor[DType.bfloat16, Storage=w_k.Storage, linear_idx_type=w_k.linear_idx_type], w_uk: TileTensor[DType.bfloat16, Storage=w_uk.Storage, linear_idx_type=w_uk.linear_idx_type], w_uv: TileTensor[DType.bfloat16, Storage=w_uv.Storage, linear_idx_type=w_uv.linear_idx_type], scalar_args_buf: TileTensor[DType.int64, Storage=scalar_args_buf.Storage, linear_idx_type=scalar_args_buf.linear_idx_type], ctx: DeviceContext, d_indices: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, indices_stride: Int = Int(0), topk_lengths: OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]] = None, attn_sink_ptr: OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]] = None, num_partitions_in: Optional[Int] = None)

BF16 MLA prefill/decode graph.

Dispatches to prefill or decode based on max sequence length in the batch.

Parameters:

collection_t (KVCollectionT): Type of the KV collection (inferred).
mask_str (StringSlice[ImmStaticOrigin]): Mask variant.
kv_input_fn (def[width: Int](IndexList[Int(2)]) capturing thin -> SIMD[DType.bfloat16, width]): Input lambda function to load the KV latent values. Shape: [tot_seq_len, cache_head_dim]. Where cache_head_dim = kv_lora_rank
- qk_rope_head_dim.
target (StringSlice[ImmStaticOrigin]): Target device (defaults to "cpu").
sparse_mla (Bool): Whether to use sparse MLA (defaults to False).
sparse_indices_stride (Int): Row stride of the sparse decode index buffer (defaults to 0).

Args:

output (TileTensor[DType.bfloat16, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Output tensor of shape [tot_seq_len, num_heads, v_head_dim].
q (TileTensor[DType.bfloat16, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Combined query tensor containing both nope and rope parts. Shape: [tot_seq_len, num_heads, qk_nope_head_dim + qk_rope_head_dim].
input_row_offsets (TileTensor[DType.uint32, Storage=input_row_offsets.Storage, linear_idx_type=input_row_offsets.linear_idx_type]): Indicates where each request starts and ends in q. Shape: [num_batches + 1].
freqs_cis (TileTensor[Storage=freqs_cis.Storage, linear_idx_type=freqs_cis.linear_idx_type]): Precomputed RoPE frequency values for rotary position embeddings. Shape: [max_seq_len, qk_rope_head_dim].
kv_norm_gamma (TileTensor[Storage=kv_norm_gamma.Storage, linear_idx_type=kv_norm_gamma.linear_idx_type]): RMSNorm gamma weights for normalizing the KV cache. Shape: [kv_lora_rank].
kv_collection (collection_t): Paged KV Cache object.
layer_idx (UInt32): Layer index.
scale (Float32): Scale for the attention calculation.
epsilon (Float32): Small constant for numerical stability in RMSNorm.
buffer_row_offsets (TileTensor[DType.uint32, Storage=buffer_row_offsets.Storage, linear_idx_type=buffer_row_offsets.linear_idx_type]): Indicates where each request's KV latent values should be stored in the contiguous K buffer. This is a 1D tensor of shape [num_batches + 1].
cache_offsets (TileTensor[DType.uint32, Storage=cache_offsets.Storage, linear_idx_type=cache_offsets.linear_idx_type]): Indicates the starting token position in the KV cache from which to copy KV latent values for each request. This is a 1D tensor of shape [num_batches + 1].
buffer_length (Int): The total number of tokens in the KV cache. Scalar.
max_seq_len (Int): Maximum sequence length in the batch, used to select prefill versus decode.
w_k (TileTensor[DType.bfloat16, Storage=w_k.Storage, linear_idx_type=w_k.linear_idx_type]): Weight matrix for up-projecting the latent cache to full K. Shape: [num_heads * qk_nope_head_dim, kv_latent_dim].
w_uk (TileTensor[DType.bfloat16, Storage=w_uk.Storage, linear_idx_type=w_uk.linear_idx_type]): Weight matrix for projecting the non-rope part of each query head to KV latent space. Shape: [num_heads, kv_latent_dim, qk_nope_head_dim].
w_uv (TileTensor[DType.bfloat16, Storage=w_uv.Storage, linear_idx_type=w_uv.linear_idx_type]): Weight matrix for projecting the output of the attention back to each head's original space. Shape: [num_heads, v_head_dim, kv_latent_dim].
scalar_args_buf (TileTensor[DType.int64, Storage=scalar_args_buf.Storage, linear_idx_type=scalar_args_buf.linear_idx_type]): Packed MLA dispatch metadata buffer.
ctx (DeviceContext): Device context.
d_indices (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional device pointer to packed int32 physical KV row indices for sparse decode (defaults to None).
indices_stride (Int): Stride between batch rows in d_indices (defaults to 0).
topk_lengths (OptionalReg[Pointer[Int32, MutAnyOrigin, _safe=False]]): Optional per-batch valid top-k counts (defaults to None).
attn_sink_ptr (OptionalReg[Pointer[Float32, MutAnyOrigin, _safe=False]]): Optional per-batch attention sink weights (defaults to None).
num_partitions_in (Optional[Int]): Capturable-graph num_partitions override (defaults to None).