For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/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, ImplicitlyDeletable

Methods

execute

static def execute[c_type: DType, a_type: DType, b_type: DType, scales_type: DType, //, clamp_activation: Bool, target: StringSlice[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], 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; kept zero in band by POST_SELF_CLEAN_UP).

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).
• ​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.
• ​target (StringSlice[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.
• ​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. Kept zero at every launch boundary by the POST_SELF_CLEAN_UP in-band reset; sized to a static upper bound on total_m_blocks (the kernel only touches [0, total_m_blocks)).
• ​context (DeviceContext): The device context.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit