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 struct

Struct_msa_attention_ragged_paged_mxfp6

struct Struct_msa_attention_ragged_paged_mxfp6

Registers the mo.msa.attention.ragged.paged.mxfp6 graph op with the graph compiler.

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[kv_type: DType, //, group: Int, topk: Int, FP6_FORMAT: Int](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], output_scales: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output_scales.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], page_stride: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=page_stride.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)

mo.msa.attention.ragged.paged that emits packed MXFP6 + E8M0 scales.

AMD (gfx950) only. Same inputs and routing as the BF16 op above, but the output is the o_proj-ready MXFP8 activation: output is uint8 [num_rows, n_heads, head_dim * 3 // 4] and output_scales is float8_e8m0fnu [num_rows, n_heads * head_dim / 32], row-major -- exactly what quantize_mxfp6_amd produces and mo.matmul.dynamic.block.scaled.amd consumes. Bit-identical to running the BF16 op followed by that quantize (KERN-3384).

Only the split-K decode/spec route saves a dispatch: the combine and the quantize fuse into msa_amd_splitk_reduce_quant_mx. Prefill and the num_partitions <= 1 decode shapes still produce BF16 first (into a scratch buffer the fused route never touches) and quantize with the stock quantize_mxfp6_amd -- the same two dispatches those routes cost unfused, so the op's output contract is uniform across routes.

Deliberately a separate registration rather than a second output on the BF16 op: that op serves both vendors and BF16-o_proj configs, and this one only exists where o_proj consumes MXFP8. The routing below is the AMD half of the BF16 op's; a route added there needs a mirror here.

Parameters:

• ​group (Int): Query heads per kv-head (n_heads // n_kv_heads).
• ​topk (Int): Number of gathered KV blocks per token (d_indices stride).
• ​FP6_FORMAT (Int): FP6 element encoding, 0 for E2M3 or 1 for E3M2. Nothing downstream can recover it from the packed bytes, so it must match what the checkpoint declares.

Args:

• ​output (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec]): Quantized output [num_rows, n_heads, head_dim] FP8 e4m3.
• ​output_scales (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output_scales.static_spec]): E8M0 block scales [num_rows, n_heads * head_dim / 32].
• ​q (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=q.static_spec]): Query [num_rows, n_heads, head_dim], dtype kv_type.
• ​input_row_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input_row_offsets.static_spec]): Ragged query offsets [batch + 1] uint32.
• ​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], dtype kv_type.
• ​page_stride (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=page_stride.static_spec]): Page-to-page distance in elements; -1 when packed.
• ​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 (parity with the BF16 op).
• ​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.