AMD Primus adds FP8/MXFP8 kernels, megakernel fusion and DeepSeek-V3 training at 1024 GPUs

AMD has published a technical deep-dive on MoE training optimizations added to its Primus framework for AMD Instinct MI355X GPUs. The work builds on an earlier post about MoE training best practices and addresses four simultaneous bottlenecks: expert GEMM throughput, all-to-all communication, memory capacity, and host launch overhead. New features include FP8 and MXFP8 grouped GEMM kernels, a fused "MegaMoE" megakernel that embeds all-to-all operations inside grouped GEMMs, and DeepEP-based dispatch with 1F1B pipeline overlap. The team demonstrated end-to-end training of DeepSeek-V3 on 256 to 1024 GPUs including a 1700-step uninterrupted run. Both Megatron-LM and JAX (MaxText) backends are supported. Primus Projection, a planning tool for Megatron-LM memory and pipeline layout, is also documented.

What Primus adds for MoE training on MI355X

• Low-precision expert GEMMs: FP8 and MXFP8 grouped GEMM kernels are implemented at both kernel and end-to-end levels, with quantized-weight caching, targeting the FFN-dominant FLOP profile of MoE models.
• Fused MoE megakernel (MegaMoE): A single kernel fuses each all-to-all communication step with the grouped GEMM that produces or consumes its data, operating at tile granularity rather than as separate operations.
• Communication optimizations: DeepEP dispatch handles all-to-all dispatch and combine phases; 1F1B (one-forward-one-backward) pipeline scheduling overlaps all-to-all communication with computation.
• Memory and planning tools: A precision-aware optimizer, pipeline layout tuning, fine-grained recomputation, and Primus Projection (an ahead-of-time planning tool for Megatron-LM) address memory ceilings that limit micro-batch size.
• Host overhead reduction: Sync-free MoE execution, NUMA binding, launch tuning, and pipeline warm-up target the launch overhead that dominates when layers are small or expert granularity is fine.
• Framework coverage: Both Megatron-LM and JAX (MaxText) backends are supported, with the JAX path bringing the same grouped-GEMM and DeepEP primitives to that ecosystem.
• Production validation: DeepSeek-V3 training was demonstrated at scale from 256 to 1024 MI355X GPUs, including a 1700-step uninterrupted run, with support also confirmed for Qwen3-235B-A22B, Qwen3-30B-A3B, GLM, MiniMax, GPT-OSS, and Mixtral.

Why it matters

These optimizations strengthen AMD ROCm as a viable open-platform alternative to CUDA for training frontier-scale MoE models, directly supporting open-weight architectures like DeepSeek-V3 and Qwen3 that are widely used in the open-source AI community. By publishing kernel-level details, releasing code through GitHub, and supporting both PyTorch/Megatron-LM and JAX ecosystems, AMD reduces vendor lock-in for large-model training infrastructure and gives Linux-based AI clusters more hardware choice.