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 struct
Struct_msa_attention_ragged_paged
struct Struct_msa_attention_ragged_paged
Registers the mo.msa.attention.ragged.paged graph op with the graph compiler.
Implemented traits
Methods
execute
static def execute[kv_type: DType, //, group: Int, topk: Int](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], q: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec], input_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec], cache_row_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_row_offsets.static_spec], total_context_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=total_context_length.static_spec], kv_blocks: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=kv_blocks.static_spec], cache_lengths: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_lengths.static_spec], kv_lookup_table: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=kv_lookup_table.static_spec], max_prompt_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_prompt_length.static_spec], max_cache_length: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_cache_length.static_spec], msa_scalar_args: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=msa_scalar_args.static_spec], layer_idx: UInt32, d_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=d_indices.static_spec], scale: Float32, ctx: DeviceContext)
Block-sparse MHA for SM100 (BF16 or FP8 e4m3, head_dim 128).
The KV cache dtype (kv_type, inferred from q / kv_blocks) selects
BF16 or native FP8 e4m3; q and kv_blocks must share it and the
kernel accumulates in FP32. FP8 is scale-free (no per-block dequant
scales), matching the msa_sm100_* FP8 path. The output is always BF16.
Gathers topk KV blocks per (kv head, query token) using the block ids
in d_indices. Dispatches to the decode kernel when
kv_collection.max_seq_length == 1 (one query token per sequence) and to
the prefill kernel otherwise.
Decode uses NullMask + an SM-fill split-K heuristic
(get_mha_decoding_max_num_partitions clamped by topk): `num_partitions
1
runs the block-major fwd over partitioned KV bands then combines via the sharedmha_splitk_reduce;num_partitions == 1takes the no-combineNoPartitionpath. Prefill uses the device-CSR plan/run path (msa_sm100_prefill_plan+msa_sm100_prefill_run`): the run is pure-device, but the plan sizes its buffers from the per-batch cu-seqlens on host, so one D2H readback + sync per call is unavoidable while this stays a single stateless op.
Routing is purely by the runtime query length
max_q_len = kv_collection.max_seq_length (the max new query tokens):
== 1 decode, 2 / 3 / 4 sparse speculative decode (one CTA per draft
token, real per-token causal, capture-stable over-launch -- see the
module docstring; spec_max_seq_len is bound to the matched length per
branch), and > 4 prefill. A short 2-4 prefill is correctly handled by
the spec path, so no prefill/spec disambiguation is needed. Spec decode
derives each draft token's logical query position in-kernel from
cache_lengths + tok_in_seq (mirrors the prefill use_causal path), so
no q_positions array is built or passed.
Parameters:
- group (
Int): Query heads per kv-head (n_heads // n_kv_heads); assertsgroup <= MMA_Min the kernel. - topk (
Int): Number of gathered KV blocks per token (d_indicesstride).
Args:
- output (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec]): Output[num_rows, n_heads, head_dim]BF16. - q (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec]): Query[num_rows, n_heads, head_dim], dtypekv_type(BF16 or FP8 e4m3;num_rows== total_q on prefill, batch on decode). - input_row_offsets (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec]): Ragged query offsets[batch + 1]uint32 (1 token/seq on decode). - cache_row_offsets (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_row_offsets.static_spec]): Ragged valid cache offsets[batch + 1]uint32. - total_context_length (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=total_context_length.static_spec]): Total context length of the current batch. - kv_blocks (
ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=kv_blocks.static_spec]): Main-KV paged blocks[num_blocks, 2, num_layers, page_size, n_kv_heads, head_dim], dtypekv_type(BF16 or FP8 e4m3, scale-free). - cache_lengths (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=cache_lengths.static_spec]): Main-KV cache lengths[batch]uint32. - kv_lookup_table (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=kv_lookup_table.static_spec]): Main-KV page table[batch, max_pages]uint32. - max_prompt_length (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_prompt_length.static_spec]): Main-KV max prompt (query) length[1]uint32. - max_cache_length (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=max_cache_length.static_spec]): Main-KV max cache length[1]uint32. - msa_scalar_args (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=msa_scalar_args.static_spec]): On-device scalar arguments for the MSA decode msa_scalar_args[0] = batch_size msa_scalar_args[1] = max_cache_valid_length. - layer_idx (
UInt32): Layer index for the main-KV cache. - d_indices (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=d_indices.static_spec]): Selected block ids[n_kv_heads, num_rows, topk]int32. - scale (
Float32): QK scale. - ctx (
DeviceContext): Device context.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!