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_swiglu_nvfp4

struct Struct_grouped_matmul_swiglu_nvfp4

MOGG wrapper for fused grouped NVFP4 matmul + SwiGLU + NVFP4 quant.

Fuses the MoE gate/up grouped matmul, SwiGLU activation, and per-block NVFP4 quantization into a single SM100 kernel. The caller must pre-permute the weight b and its scale tile b_scales on the N axis with sigma(2i)=i, sigma(2i+1)=D+i (where D = moe_dim, N = 2D).

Implemented traits

AnyType, Deinitable, Movable

Methods

execute

static def execute[c_type: DType, a_type: DType, b_type: DType, scales_type: DType, row_scales_type: DType, //, clamp_activation: Bool, has_a_row_scales: Bool, target: StringSpan[ImmStaticOrigin]](c_packed: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c_packed.static_spec], c_swiglu_scales: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c_swiglu_scales.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], a_row_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_row_scales.static_spec], c_input_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=c_input_scales.static_spec], estimated_total_m: UInt32, num_active_experts: UInt32, swiglu_alpha: Float32, swiglu_limit: Float32, context: DeviceContext)

Executes fused grouped NVFP4 matmul + SwiGLU + NVFP4 quant.

Computes (c_packed, c_swiglu_scales) = quantize_nvfp4(silu(C[..., even]) * C[..., odd], c_input_scales) where C = A @ B^T for multiple expert groups. Because B is sigma-permuted on N, adjacent matmul-output columns carry (gate, up) pairs that the epilogue consumes in-place.

Parameters:

• ​c_type (DType): The output tensor data type.
• ​a_type (DType): The input A data type. Constraints: Must be uint8.
• ​b_type (DType): The input B data type. Constraints: Must be uint8.
• ​scales_type (DType): The scale factor data type. Constraints: Must be float8_e4m3fn.
• ​row_scales_type (DType): The per-row input scale data type. Constraints: Must be bfloat16 when has_a_row_scales.
• ​clamp_activation (Bool): Whether to clamp the activation (swigluoai).
• ​has_a_row_scales (Bool): Whether a_row_scales holds per-row input scales.
• ​target (StringSpan[ImmStaticOrigin]): The target GPU device.

Args:

• ​c_packed (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c_packed.static_spec]): Packed NVFP4 output of shape (total_tokens, D // 2).
• ​c_swiglu_scales (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=c_swiglu_scales.static_spec]): 5D FP8 SF tile in tcgen05 layout for the output.
• ​a (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a.static_spec]): The input tensor of shape (total_tokens, K // 2).
• ​b (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=b.static_spec]): The sigma-permuted weight of shape (num_experts, 2D, K // 2).
• ​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 sigma-permuted 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.
• ​a_row_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_row_scales.static_spec]): Per-row input scales of shape (total_tokens,), multiplied into each row with its expert scale before the SwiGLU. Unread unless has_a_row_scales.
• ​c_input_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=c_input_scales.static_spec]): Per-expert SiLU input scale (= 1/output_inv_scale).
• ​estimated_total_m (UInt32): The estimated total number of tokens.
• ​num_active_experts (UInt32): The number of active experts.
• ​swiglu_alpha (Float32): The alpha value for the clamped activation.
• ​swiglu_limit (Float32): The limit value for the clamped activation.
• ​context (DeviceContext): The device context pointer.