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_grouped_matmul_block_scaled_mxfp6
struct Struct_grouped_matmul_block_scaled_mxfp6[FP6_FORMAT: Int = Int(0)]
MOGG wrapper for the grouped MXFP6 block-scaled matmul.
The FP6 sibling of mo.grouped.matmul.block.scaled.mxfp4, and a separate
op for the same reason the dense one is: both FP6 encodings put 24 bytes in
a lane, so lane_bytes cannot choose between them.
Preshuffled-B only. FP6's 24-byte lane fragment is plane-split (see
Shuffler.b_plane_byte_off), which the dense row-major
block_scaled_grouped_matmul_amd kernel has no path for.
Parameters
- FP6_FORMAT (
Int): 0 selects E2M3, 1 selects E3M2, matchingFP6Format.
Implemented traits
Methods
execute
static def execute[c_type: DType, a_type: DType, b_type: DType, //, target: StringSpan[ImmStaticOrigin]](c: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c.static_spec], a: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a.static_spec], b: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b.static_spec], a_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_scales.static_spec], b_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b_scales.static_spec], expert_start_indices: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_start_indices.static_spec], expert_ids: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_ids.static_spec], max_num_tokens_per_expert: UInt32, num_active_experts: UInt32, estimated_total_m: UInt32, decode_grid_m_cap: UInt32, decode_grid_m_rows: UInt32, context: DeviceContext)
Computes C = A @ B^T over expert groups with MXFP6 operands.
Parameters:
- c_type (
DType): The output tensor data type. - a_type (
DType): The packed activation data type; must be one byte wide. - b_type (
DType): The packed weight data type; must be one byte wide. - target (
StringSpan[ImmStaticOrigin]): The target GPU device.
Args:
- c (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c.static_spec]): The output tensor of shape (total_tokens, N). - a (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a.static_spec]): The packed activations, (total_tokens, K * 3 // 4). - b (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b.static_spec]): The plane-split preshuffled weights, (num_experts, N, K * 3//4). - a_scales (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_scales.static_spec]): The A scale factors, (total_tokens, K // 32). - b_scales (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b_scales.static_spec]): The B scale factors, (num_experts, N, K // 32). - expert_start_indices (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_start_indices.static_spec]): The starting token index for each expert. - expert_ids (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_ids.static_spec]): The expert ID for each group. - max_num_tokens_per_expert (
UInt32): The maximum token count for any expert. - num_active_experts (
UInt32): The number of active experts. - estimated_total_m (
UInt32): Estimated total received tokens for this GPU, used to pick the persistent vs direct kernel path. - decode_grid_m_cap (
UInt32): Decode-band gate selecting the direct kernel over the persistent one; 0 disables. - decode_grid_m_rows (
UInt32): Rows grid.y must cover per expert on the decode bands. - context (
DeviceContext): The device context pointer.