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 module

mega_ffn

Graph-op binding for the fused MegaFFN MoE FFN composite op (NVFP4 + MXFP8).

Target: NVIDIA SM100 (B200). Registers mo.composite.mega_ffn_nvfp4 and DTYPE-BRANCHES on the operand element type: packed uint8 activations -> mega_ffn_nvfp4_dispatch (NVFP4), float8_e4m3fn activations -> mega_ffn_mxfp8_dispatch (MXFP8). Either way it is the single-launch kernel that fuses the MoE gate/up grouped matmul + SwiGLU + re-quant (L1) and the down projection (L2) into one launch (one phase pipelined across SMs). The graph compiler emits this composite when the MegaFFN fusion fires (SM100, num_experts <= 64 per device, single-use intermediate, shared routing, bf16 out); the fusion + composite op are dtype-agnostic, so the same op carries NVFP4 (uint8) or MXFP8 (e4m3) data and this binding selects the dispatch.

Clamped-SwiGLU (swigluoai) status: the activation SELECTOR clamp_activation is plumbed end-to-end as a Bool op attribute (the emitter sets it, the MegaFFN fusion copies it from the L1 leg, and MOGG binds it to the comptime clamp_activation param here). When the selector is set, this binding supplies swigluoai's canonical clamp constants (alpha=1.702, limit=7.0) to mega_ffn_{nvfp4,mxfp8}_dispatch, so the standard swigluoai activation is numerically CORRECT through the fusion; clamp_activation=False (plain SwiGLU, the default) is also correct. The alpha/limit VALUES are supplied here rather than carried on the op because MOGG cannot bind an f32 kernel arg from an op attribute (comptime params forward only Integer/Bool/String attributes; runtime kernel args bind only from operands). A model whose clamp uses NON-standard alpha/limit would need those values carried as host f32 scalar-tensor OPERANDS on the composite ops (mirroring the masked_flash_attention scale operand) -- a .td operand addition tracked as a follow-up.

Per-expert scaling is applied PER LEG: the composite op exposes two scale arrays -- gate_up_expert_scales (L1) and down_expert_scales (L2) -- and both are forwarded to the fused kernel, which applies the L1 scale in the SwiGLU store and the L2 scale in the final-output store (grouped_1d1d_matmul_kernel.execute_epilogue). The kernel's scheduler resolves each tile's single per-slot scale to the leg-matching array by tile phase before publishing, so an MoE that emits distinct L1 and L2 expert scales is reproduced exactly (matching the two-launch chain). See KERN-3085.

Modeled on the sibling msa.mojo / linalg.mojo registrations (private //Kernels kernels registered in dedicated builtin_kernels files). MegaFFN is internal-only: unlike msa / matmul_rs it is NOT shipped in the OSS wheel, so open-source builds drop the //Kernels/src/mega_ffn dep (in api.bzl) and exclude this file from the public export (copybara).

On-chip scratch (the composite op does not expose these): c_packed and c_swiglu_scales are allocated per call via ctx.enqueue_create_buffer -- the capture-safe MOGG workspace pattern used throughout builtin_kernels (msa.mojo, attention.mojo, ep.mojo, linalg.mojo). DeviceBuffer frees are stream-ordered, so the buffers outlive the launch that uses them. They are write-then-read within the launch, so they need no initialization.

arrival_count (the cross-CTA pool-slot counters) is instead a PERSISTENT graph buffer OPERAND: the MegaFFN fusion mints it once (mo.buffer.create) and zeroes it once at setup. Under the dispatch default POST_SELF_CLEAN_UP the kernel resets every touched slot in band, so the buffer is all-zero at every launch boundary (correct under single-stream serialization). This replaces the per-launch allocate + memset this binding used to do, which sat on the launch-bound decode critical path (one memset per FFN launch). The fusion sizes the buffer to a static upper bound on total_m_blocks; the kernel only touches/resets [0, total_m_blocks), so over-allocation is safe.

Structs

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

Structs​

Structs