Is your feature request related to a problem? Please describe.
MCore's MXFP8 MoE inference path currently fuses squared-ReLU with MXFP8 quantization, but SwiGLU uses two kernel launches and materializes a BF16 activation tensor between them. This adds memory traffic and temporary storage on the expert hot path. The limitation is also documented in PR #7300, which retained separate SwiGLU activation and MXFP8 quantization.
@NVIDIA/mcore-oncall
Describe the solution you'd like
Add a padding-aware Triton kernel that performs SwiGLU, BF16-compatible rounding, MXFP8 block quantization, and 32x4x4 scale swizzling in one launch. Select it for non-batch-invariant MXFP8 SwiGLU execution while preserving the existing separate path behind disable_fused_quant_kernels.
The fused result should be bit-exact with the existing padded_swiglu followed by MXFP8Tensor.from_bf16, including E4M3 data, E8M0 scales, padding behavior, and deterministic replay.
Describe alternatives you've considered
Keep the existing two-kernel route. It is correct, but retains the extra launch and BF16 intermediate. Reusing the squared-ReLU kernel is not possible because SwiGLU reads the gate and up halves and applies SiLU multiplication before quantization.
Additional context
On an RTX 5090 with PyTorch 2.10.0 and Triton 3.6.0, an activation-plus-quantization microbenchmark shows 1.13x to 1.28x speedup across representative MoE shapes. For M=8192, K=2688, fusion also avoids an approximately 42 MiB BF16 intermediate.
The RTX 5090 can validate the Triton activation/quantization kernel, scale layout, padding, and determinism. PyTorch's MXFP8 scaled_grouped_mm does not currently support SM120, so full MCore MXFP8 MoE execution still requires supported datacenter Blackwell hardware or project CI.
Is your feature request related to a problem? Please describe.
MCore's MXFP8 MoE inference path currently fuses squared-ReLU with MXFP8 quantization, but SwiGLU uses two kernel launches and materializes a BF16 activation tensor between them. This adds memory traffic and temporary storage on the expert hot path. The limitation is also documented in PR #7300, which retained separate SwiGLU activation and MXFP8 quantization.
@NVIDIA/mcore-oncall
Describe the solution you'd like
Add a padding-aware Triton kernel that performs SwiGLU, BF16-compatible rounding, MXFP8 block quantization, and 32x4x4 scale swizzling in one launch. Select it for non-batch-invariant MXFP8 SwiGLU execution while preserving the existing separate path behind
disable_fused_quant_kernels.The fused result should be bit-exact with the existing
padded_swiglufollowed byMXFP8Tensor.from_bf16, including E4M3 data, E8M0 scales, padding behavior, and deterministic replay.Describe alternatives you've considered
Keep the existing two-kernel route. It is correct, but retains the extra launch and BF16 intermediate. Reusing the squared-ReLU kernel is not possible because SwiGLU reads the gate and up halves and applies SiLU multiplication before quantization.
Additional context
On an RTX 5090 with PyTorch 2.10.0 and Triton 3.6.0, an activation-plus-quantization microbenchmark shows 1.13x to 1.28x speedup across representative MoE shapes. For
M=8192, K=2688, fusion also avoids an approximately 42 MiB BF16 intermediate.The RTX 5090 can validate the Triton activation/quantization kernel, scale layout, padding, and determinism. PyTorch's MXFP8
scaled_grouped_mmdoes not currently support SM120, so full MCore MXFP8 MoE execution still requires supported datacenter Blackwell hardware or project CI.