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

struct Struct_grouped_matmul_block_scaled

MOGG wrapper for grouped block-scaled matrix multiplication.

Provides graph compiler integration for block-scaled grouped matmul operations used in Mixture of Experts (MoE) layers on SM100 GPUs.

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[c_type: DType, a_type: DType, b_type: DType, scales_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], a_scale_offsets: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_scale_offsets.static_spec], expert_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_scales.static_spec], estimated_total_m: UInt32, num_active_experts: UInt32, context: DeviceContext)

Executes grouped block-scaled matrix multiplication.

Computes C = A @ B^T for multiple expert groups where A and B are block-scaled. uint8 operands are nibble-packed 4-bit E2M1, so their rows are K // 2 wide; float8_e4m3fn operands are unpacked. The operand and scale dtypes together select the UMMA kind (see block_scaled_umma_kind): NVFP4, MXFP4, MXFP8, or the mixed W4A8 pair.

Parameters:

• ​c_type (DType): The output tensor data type.
• ​a_type (DType): The input A data type. Constraints: Must be uint8 (NVFP4/MXFP4) or float8_e4m3fn (MXFP8/W4A8).
• ​b_type (DType): The input B data type. Constraints: Must equal a_type, except for W4A8, which pairs a float8_e4m3fn A with a uint8 B.
• ​scales_type (DType): The scale factor data type. Constraints: Must be float8_e4m3fn (NVFP4) or float8_e8m0fnu (MXFP4/MXFP8/W4A8).
• ​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 input tensor of shape (total_tokens, K) for float8_e4m3fn or (total_tokens, K // 2) for packed uint8.
• ​b (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b.static_spec]): The weight tensor of shape (num_experts, N, K) for float8_e4m3fn or (num_experts, N, K // 2) for packed uint8.
• ​a_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_scales.static_spec]): The A scale factors in tcgen05 5D layout.
• ​b_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b_scales.static_spec]): The B scale factors in tcgen05 6D layout.
• ​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.
• ​a_scale_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_scale_offsets.static_spec]): The starting scale index for each expert.
• ​expert_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_scales.static_spec]): The per-expert scaling factors for the epilogue.
• ​estimated_total_m (UInt32): The estimated total number of tokens.
• ​num_active_experts (UInt32): The number of active experts.
• ​context (DeviceContext): The device context pointer.