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_mega_ffn_nvfp4

struct Struct_mega_ffn_nvfp4

MOGG wrapper for the fused single-launch MegaFFN NVFP4 MoE FFN.

Lowers the mo.composite.mega_ffn_nvfp4 composite op (gate/up GMM + SwiGLU + NVFP4 re-quant fused with the down GMM, all in one launch) to mega_ffn_nvfp4_dispatch on SM100 GPUs. The clamped-SwiGLU activation selector clamp_activation is bound from the same-named Bool op attribute; the alpha/limit values are a follow-up (see the module docstring).

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_gate_up_a_row_scales: Bool, target: StringSpan[ImmStaticOrigin]](output: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec], hidden_states: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=hidden_states.static_spec], gate_up_weight: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_weight.static_spec], gate_up_a_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_a_scales.static_spec], gate_up_b_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_b_scales.static_spec], down_weight: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=down_weight.static_spec], down_b_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=down_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], gate_up_expert_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_expert_scales.static_spec], gate_up_a_row_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_a_row_scales.static_spec], down_expert_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=down_expert_scales.static_spec], c_input_scales: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=c_input_scales.static_spec], estimated_total_m: UInt32, gate_up_num_active_experts: UInt32, down_num_active_experts: UInt32, arrival_count: ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=arrival_count.static_spec], context: DeviceContext)

Executes the fused single-launch MegaFFN NVFP4 MoE FFN.

Computes out = down((swiglu(gate_up(hidden_states))) re-quantized to NVFP4) for the active MoE experts in one kernel launch. a, both weights, and the down intermediate are NVFP4 (4-bit packed as uint8); scales are float8_e4m3fn in tcgen05 layout.

The clamped-SwiGLU selector clamp_activation is forwarded to the dispatch (its alpha/limit values are a follow-up; see the module docstring). Both per-leg scale arrays are forwarded: gate_up_expert_scales drives the L1 SwiGLU store and down_expert_scales drives the L2 final-output store. The on-chip c_packed / c_swiglu_scales scratch is allocated per call (capture-safe enqueue_create_buffer); the arrival_count pool-slot counters are a persistent graph buffer operand (zeroed once at setup; thereafter POST_SELF_CLEAN_UP rotates them per launch rather than returning them to zero).

Parameters:

• ​c_type (DType): The output tensor data type (bfloat16).
• ​a_type (DType): The input A / activation element type. uint8 (packed NVFP4) selects the NVFP4 dispatch; float8_e4m3fn selects MXFP8.
• ​b_type (DType): The weight element type (matches a_type: uint8 for NVFP4 or float8_e4m3fn for MXFP8).
• ​scales_type (DType): The block scale-factor dtype (float8_e4m3fn for NVFP4, float8_e8m0fnu for MXFP8).
• ​row_scales_type (DType): The L1 per-row input scale dtype (bfloat16 when has_gate_up_a_row_scales).
• ​clamp_activation (Bool): Activation flavor for the fused L1 SwiGLU epilogue. False = plain SwiGLU; True = clamped (swigluoai). Bound from the op's clamp_activation attribute.
• ​has_gate_up_a_row_scales (Bool): Whether gate_up_a_row_scales holds per-row input scales for the L1 leg. Constraints: NVFP4 only.
• ​target (StringSpan[ImmStaticOrigin]): The target GPU device.

Args:

• ​output (ManagedTensorSlice[IOSpec[_, _].Output, static_spec=output.static_spec]): Final output (M_total, N2) bf16.
• ​hidden_states (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=hidden_states.static_spec]): Token activations (M_total, K1 // 2) uint8 (packed NVFP4).
• ​gate_up_weight (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_weight.static_spec]): Pre-permuted gate/up weights (E, N1, K1 // 2) uint8; N1 == 2 * moe_dim.
• ​gate_up_a_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_a_scales.static_spec]): Token (A) 5D E4M3 scale tile for L1; its leading dim is the scale-block count reused by the SwiGLU intermediate scratch.
• ​gate_up_b_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_b_scales.static_spec]): Pre-permuted W13 6D E4M3 scale tile.
• ​down_weight (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=down_weight.static_spec]): Down-projection weights (E, N2, moe_dim // 2) uint8 (packed NVFP4).
• ​down_b_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=down_b_scales.static_spec]): W2 6D E4M3 scale tile.
• ​expert_start_indices (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_start_indices.static_spec]): Per-expert prefix-sum token offsets (E + 1,) uint32.
• ​expert_ids (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=expert_ids.static_spec]): Active expert IDs (E,) int32 (-1 = masked).
• ​a_scale_offsets (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=a_scale_offsets.static_spec]): Per-expert 128-row scale-block offsets (E,) uint32.
• ​gate_up_expert_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_expert_scales.static_spec]): L1 (gate+up) per-expert scaling (E,) f32. Applied in the fused SwiGLU store.
• ​gate_up_a_row_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=gate_up_a_row_scales.static_spec]): L1 per-row input scales (M_total,), applied with gate_up_expert_scales before the SwiGLU. Unread unless has_gate_up_a_row_scales.
• ​down_expert_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=down_expert_scales.static_spec]): L2 / final-output per-expert scaling (E,) f32. Applied in the down store. May differ from gate_up_expert_scales.
• ​c_input_scales (ManagedTensorSlice[IOSpec[_, _].Input, static_spec=c_input_scales.static_spec]): L1 SwiGLU per-expert input scale (tensor_sf) (E,) f32, used for the NVFP4 re-quant of the intermediate.
• ​estimated_total_m (UInt32): Estimated total non-padded token count (the avg_m gate numerator).
• ​gate_up_num_active_experts (UInt32): Active expert slots for L1.
• ​down_num_active_experts (UInt32): Active expert slots for L2 (must equal gate_up_num_active_experts; the kernel walks one shared expert list).
• ​arrival_count (ManagedTensorSlice[IOSpec[_, _].MutableInput, static_spec=arrival_count.static_spec]): Persistent cross-CTA pool-slot counters (uint32, rank 1), minted + zeroed once by the MegaFFN fusion. Under POST_SELF_CLEAN_UP a launch accumulates into the slot half its generation parity selects and clears the half the launch before it used, so this buffer holds the last launch's counts rather than zeros between launches -- do not assume an all-zero buffer or reset it per launch, which would pin the generation. Sized to a static upper bound on total_m_blocks, which also bounds the high-water span a launch clears.
• ​context (DeviceContext): The device context.