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

flare_mla_prefill

def flare_mla_prefill[rank: Int, cache_t: KVCacheT, mask_t: MHAMask, dtype: DType, output_type: DType, //](output: TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment], v: LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment], k_rope: cache_t, mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], cache_row_offsets: TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, linear_idx_type=cache_row_offsets.linear_idx_type], scale: Float32, ctx: DeviceContext, q_max_seq_len: OptionalReg[Int] = None, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]] = None)

MLA prefill kernel that would only be called in the optimized compute graph. Only supports ragged Q/K/V inputs.

The Q input has a shape of [seq_len, num_heads, q_depth]. The K and V input has a shape of [cache_len, num_heads, depth]. The K_rope input is retrieved from the KV cache, with a shape of [cache_len, 1, q_depth - depth].

Specifically, for DeepSeek V2/3, depth = 128 and q_depth = 192.

When computing attention scores (Q @ K), each head of K is smaller than Q head. The missing 64 elements of K are retrieved from the K cache, and broadcasted to all the heads. This kernel also handles that output has reduced dimension compared to input Q.

This kernel handles batches with different valid lengths (i.e., before the padding). Such lengths are passed in valid_length argument.

Parameters:

rank (Int): Tensor rank of q and output (inferred). Must be 3 for ragged inputs.
cache_t (KVCacheT): Paged KV cache type backing k_rope (inferred). Carries the KV layout, dtype, and head geometry.
mask_t (MHAMask): Mask functor type applied to attention scores (inferred).
dtype (DType): Element type of q (inferred). Must be bfloat16 or float8_e4m3fn.
output_type (DType): Element type of output (inferred). bfloat16 when q is float8_e4m3fn, else matches q.

Args:

output (TileTensor[output_type, Storage=output.Storage, linear_idx_type=output.linear_idx_type]): Output tensor with ragged rank-3 shape [total_q_tokens, num_heads, depth]. Dtype is bfloat16 when q is float8_e4m3fn, else matches q.
q (TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type]): Query tensor with ragged rank-3 shape [total_q_tokens, num_heads, q_depth]. For DeepSeek V2/3, q_depth is 192.
k (LayoutTensor[element_layout=k.element_layout, layout_int_type=k.layout_int_type, linear_idx_type=k.linear_idx_type, masked=k.masked, alignment=k.alignment]): Key tensor (nope part) with shape [cache_len, num_heads, depth]. For DeepSeek V2/3, depth is 128.
v (LayoutTensor[element_layout=v.element_layout, layout_int_type=v.layout_int_type, linear_idx_type=v.linear_idx_type, masked=v.masked, alignment=v.alignment]): Value tensor with shape [cache_len, num_heads, depth]. Same depth as k.
k_rope (cache_t): Paged KV cache operand providing the rope part of K, with shape [cache_len, 1, q_depth - depth]. For DeepSeek V2/3, the last dim is 64.
mask_functor (mask_t): Mask functor instance applied to attention scores.
valid_length (TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type]): Per-batch uint32 tensor of cumulative row offsets into the ragged Q layout; offsets[b] to offsets[b+1] gives batch b's token range.
cache_row_offsets (TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, linear_idx_type=cache_row_offsets.linear_idx_type]): Per-batch uint32 tensor of row offsets into the ragged KV layout, used to build the ragged K/V operands.
scale (Float32): Softmax scale factor applied to QK^T.
ctx (DeviceContext): Device context used to enqueue the kernel.
q_max_seq_len (OptionalReg[Int]): Optional maximum query sequence length (tokens per batch); defaults to the cache's max_prompt_length when None.
cache_offsets (OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]]): Optional per-batch uint32 tensor of starting offsets into the paged KV cache; None when the cache is contiguous.

def flare_mla_prefill[rank: Int, mask_t: MHAMask, dtype: DType, //](output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: TileTensor[Storage=k.Storage, linear_idx_type=k.linear_idx_type], v: TileTensor[Storage=v.Storage, linear_idx_type=v.linear_idx_type], k_rope: TileTensor[Storage=k_rope.Storage, linear_idx_type=k_rope.linear_idx_type], mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], cache_row_offsets: TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, linear_idx_type=cache_row_offsets.linear_idx_type], scale: Float32, ctx: DeviceContext, q_max_seq_len: OptionalReg[Int] = None, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]] = None)

def flare_mla_prefill[rank: Int, mask_t: MHAMask, dtype: DType, //](output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], q: TileTensor[dtype, Storage=q.Storage, linear_idx_type=q.linear_idx_type], k: TileTensor[Storage=k.Storage, linear_idx_type=k.linear_idx_type], v: TileTensor[Storage=v.Storage, linear_idx_type=v.linear_idx_type], k_rope: TileTensor[Storage=k_rope.Storage, linear_idx_type=k_rope.linear_idx_type], k_rope_scales: TileTensor[Storage=k_rope_scales.Storage, linear_idx_type=k_rope_scales.linear_idx_type], mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], cache_row_offsets: TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, linear_idx_type=cache_row_offsets.linear_idx_type], scale: Float32, ctx: DeviceContext, q_max_seq_len: OptionalReg[Int] = None, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]] = None)

def flare_mla_prefill[rank: Int, mask_t: MHAMask, dtype: DType, scale_dtype: DType, //](output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], q_nope: TileTensor[dtype, Storage=q_nope.Storage, linear_idx_type=q_nope.linear_idx_type], q_rope: TileTensor[Storage=q_rope.Storage, linear_idx_type=q_rope.linear_idx_type], q_scale: TileTensor[scale_dtype, Storage=q_scale.Storage, linear_idx_type=q_scale.linear_idx_type], k: TileTensor[Storage=k.Storage, linear_idx_type=k.linear_idx_type], k_scales: TileTensor[scale_dtype, Storage=k_scales.Storage, linear_idx_type=k_scales.linear_idx_type], v: TileTensor[Storage=v.Storage, linear_idx_type=v.linear_idx_type], k_rope: TileTensor[Storage=k_rope.Storage, linear_idx_type=k_rope.linear_idx_type], mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], cache_row_offsets: TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, linear_idx_type=cache_row_offsets.linear_idx_type], scale: Float32, ctx: DeviceContext, q_max_seq_len: OptionalReg[Int] = None, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]] = None)

def flare_mla_prefill[rank: Int, cache_t: KVCacheT, mask_t: MHAMask, dtype: DType, scale_dtype: DType, //](output: TileTensor[Storage=output.Storage, linear_idx_type=output.linear_idx_type], q_nope: TileTensor[dtype, Storage=q_nope.Storage, linear_idx_type=q_nope.linear_idx_type], q_rope: TileTensor[Storage=q_rope.Storage, linear_idx_type=q_rope.linear_idx_type], q_scale: TileTensor[scale_dtype, Storage=q_scale.Storage, linear_idx_type=q_scale.linear_idx_type], k: TileTensor[Storage=k.Storage, linear_idx_type=k.linear_idx_type], k_scales: TileTensor[scale_dtype, Storage=k_scales.Storage, linear_idx_type=k_scales.linear_idx_type], v: TileTensor[Storage=v.Storage, linear_idx_type=v.linear_idx_type], k_rope: cache_t, mask_functor: mask_t, valid_length: TileTensor[DType.uint32, Storage=valid_length.Storage, linear_idx_type=valid_length.linear_idx_type], cache_row_offsets: TileTensor[DType.uint32, Storage=cache_row_offsets.Storage, linear_idx_type=cache_row_offsets.linear_idx_type], scale: Float32, ctx: DeviceContext, q_max_seq_len: OptionalReg[Int] = None, cache_offsets: OptionalReg[LayoutTensor[DType.uint32, Layout.row_major(Int(-1)), MutAnyOrigin]] = None)