{"slug": "enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus", "title": "Enabling DeepSeek-V4-Flash Training on AMD Instinct MI355X GPUs with Primus", "summary": "DeepSeek-AI released the DeepSeek-V4 series on April 24, 2026, including the MIT-licensed DeepSeek-V4-Flash model with 284B total parameters (13B activated) and a one-million-token context window, and a blog post details how to enable its pretraining on AMD Instinct MI355X GPUs using the Primus framework. The model interleaves three attention types across 43 layers, routes tokens through 256 experts, and uses hyper-connection blocks, requiring custom kernel work to run efficiently in BF16.", "body_md": "# Enabling DeepSeek-V4-Flash Training on AMD Instinct MI355X GPUs with Primus[#](#enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus)\n\nDeepSeek-AI released the [DeepSeek-V4 series](https://arxiv.org/abs/2606.19348) on\nApril 24, 2026: a preview pair of MIT-licensed Mixture-of-Experts models, with\n[DeepSeek-V4-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash) at 284B\ntotal parameters (13B activated) and a one-million-token context window. Flash\npushes sparse attention further than any open-weight model before it: its 43\ntransformer layers interleave three different attention types, each layer sits\ninside a hyper-connection block rather than a plain residual, and every token is\nrouted through 256 experts. Each of those choices breaks an assumption baked into\nstock Megatron-LM training code.\n\nThis blog walks you through enabling end-to-end DeepSeek-V4-Flash pretraining in\nPrimus on AMD Instinct™ MI355X GPUs. You will learn what the architecture looks\nlike layer by layer, which knobs Primus exposes to configure it, and — where most\nof the engineering went — the kernel work that took the model from *it runs* to\n*it runs fast*. All of it is BF16 training — see [the endnotes](#endnotes) for\nwhat that leaves out. Every config, launch script, and benchmark referenced here\nships in the open-source Primus repository, so you can reproduce the run\nyourself.\n\n## The DeepSeek-V4-Flash architecture[#](#the-deepseek-v4-flash-architecture)\n\nDeepSeek-V4-Flash keeps the skeleton you already know from DeepSeek-V3 — a Transformer stack with DeepSeekMoE feed-forward layers and a Multi-Token Prediction head. What changed sits underneath: how attention reads the KV cache, and how residual connections carry signal between blocks. Pretraining also moves to the Muon optimizer for most parameters, keeping AdamW for the embedding, the prediction head, and the RMSNorm weights.\n\nFigure 1 shows a single block. The overall shape is familiar: 43 layers, hidden size 4,096, a 129,280-token vocabulary, and a MoE layer in every block with one shared expert alongside 256 routed experts, six of which activate per token. Two details already depart from V3 — every sub-layer is wrapped in manifold-constrained hyper-connections (mHC) rather than a plain residual, and the first three MoE layers route tokens by hash instead of through the learned router.\n\nThe attention module is where most of the change lives. All 64 query heads read a single 512-dimensional latent that serves as both key and value, making the layer multi-query rather than multi-head. Queries arrive through a low-rank path of rank 1,024 that the sparse-selection indexer shares. And because 64 heads of 512 dimensions is a wide tensor to project back down to 4,096, the output projection splits into 8 groups that each pass through a 1,024-dimensional bottleneck.\n\n### Three attention types, interleaved[#](#three-attention-types-interleaved)\n\nV4’s headline change is that not every layer attends the same way. A per-layer compression ratio picks one of three paths, fixed when the model is defined.\n\nThe first two layers run dense attention over a 128-token sliding window, a local warm-up before any compression kicks in. The remaining 41 layers alternate between Compressed Sparse Attention (CSA) and Heavily Compressed Attention (HCA), which works out to 21 CSA layers and 20 HCA layers. The MTP layer reuses the dense type.\n\nThe payoff shows up in the lower half of Figure 2. At a one-million-token context, a query in an HCA layer reads roughly 7,900 KV entries and a query in a CSA layer reads 640 — against a million for dense attention. Aggregated over the model, that is what lets DeepSeek report V4-Flash at about 10% of DeepSeek-V3.2’s single-token inference FLOPs with 7% of the KV cache.\n\nFigure 3 shows how each path gets there. CSA pools every 4 tokens into one KV entry, then a lightweight “lightning indexer” scores every pooled entry and keeps the best 512 for the attention itself. HCA pools far more aggressively — 128 tokens per entry — but skips selection and attends densely over everything it produced. Both add the same 128-token sliding-window branch so a query can still see recent tokens at full resolution, and both share one learned pooling operator: a softmax over the group, biased by a learnable per-position term, used to weight the sum.\n\nThe asymmetry worth remembering is that CSA’s groups overlap. Each compressed entry pools its own four tokens plus the previous four, which is why CSA needs four KV-side projections where HCA needs two. That extra projection work reappears later when we break down kernel time.\n\n### Hyper-connections in place of the residual[#](#hyper-connections-in-place-of-the-residual)\n\nInstead of `x + F(x)`\n\n, each sub-layer sits between a pair of mHC mixers\noperating on four parallel residual streams. The first mixer collapses those\nfour streams into the single tensor the sub-layer consumes; the second expands\nthe result back out and combines it with the streams coming in. That\ncombination matrix is projected onto the doubly-stochastic manifold by 20\nSinkhorn-Knopp iterations, which bounds its spectral norm at 1 and keeps signal\npropagation non-expansive across all 43 layers.\n\nOne ordering detail matters if you are porting this: the RMSNorm sits *after*\nthe collapse, not before it. Several published diagrams of V4 get this\nbackwards.\n\nTaken one at a time, none of these changes is exotic. Taken together they mean you cannot train V4 by pointing stock Megatron-LM at a new config file, which is where Primus comes in.\n\n## Enabling DeepSeek-V4 in Primus[#](#enabling-deepseek-v4-in-primus)\n\nPrimus describes a model as a chain of YAML files, each overriding the one below it. For V4-Flash that chain has three links:\n\n```\nprimus/configs/models/megatron/llama_base.yaml     generic decoder defaults\n  └─ deepseek_v4_base.yaml                         everything the V4 family shares\n       └─ deepseek_v4_flash.yaml                   Flash-specific shapes\n```\n\n`deepseek_v4_base.yaml`\n\nis where the V4 vocabulary enters Primus. These are the\nknobs that have no equivalent in a V3 config:\n\nField |\nFlash value |\nWhat it controls |\n|---|---|---|\n|\n|\nPer-layer attention type; 43 decoder entries plus one for MTP |\n|\n|\nHow many compressed entries the lightning indexer keeps |\n|\n|\nIndexer scoring shape |\n|\n|\nThe local branch every layer type carries |\n|\n|\nmHC residual streams and Sinkhorn-Knopp iterations |\n|\n|\nGrouped low-rank output projection |\n|\n|\nHow many leading MoE layers use hash routing |\n|\n|\nV4’s router scoring |\n|\n|\nClamped SwiGLU, for FP8 and FP4 stability |\n\nOne field does more than configure: `model_type: deepseek_v4`\n\nis what routes\nthe build away from the standard GPT path and into\n`primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_builders.py`\n\n,\nwhich assembles the per-layer specs from `compress_ratios`\n\n. Change nothing else\nand Primus would happily build a V3-shaped model with V4 numbers in it; this\nline is what makes the hybrid attention stack real.\n\nOn top of the model config sits an experiment config carrying the training\nhyperparameters, the parallelism, and the kernel selection:\n`examples/megatron/configs/MI355X/deepseek_v4_flash-BF16-pretrain.yaml`\n\n, with\nan FP8 variant beside it. The parts specific to V4 are short:\n\n```\nmodules:\n  pre_trainer:\n    framework: megatron\n    model: ${PRIMUS_MODEL:deepseek_v4_flash}.yaml\n    overrides:\n      tensor_model_parallel_size: ${PRIMUS_TP:1}\n      pipeline_model_parallel_size: ${PRIMUS_PP:4}\n      expert_model_parallel_size: ${PRIMUS_EP:8}\n\n      # attention kernels, chosen per path\n      use_v4_attention_backend: ${PRIMUS_USE_V4_ATTENTION_BACKEND:turbo}\n      use_v4_csa_attention_backend: ${PRIMUS_USE_V4_CSA_ATTENTION_BACKEND:turbo}\n      use_v4_fp8_indexer: ${PRIMUS_USE_V4_FP8_INDEXER:false}\n      use_v4_compiled_sinkhorn: ${PRIMUS_USE_V4_COMPILED_SINKHORN:false}\n\n      # MoE acceleration\n      use_turbo_grouped_gemm: true\n      use_turbo_deepep: ${PRIMUS_USE_TURBO_DEEPEP:true}\n```\n\nThree things there are worth pointing out. The attention backend is selected\nseparately for the dense and HCA layers (`use_v4_attention_backend`\n\n) and for\nthe CSA layers (`use_v4_csa_attention_backend`\n\n), because CSA’s indexer and\ntop-k selection make it a different kernel problem — we come back to that when\nwe look at performance. The MoE lines shown here are one of two alternatives:\nsetting `USE_TURBO_MEGA_MOE=True`\n\nswaps in MegaMoE, which fuses the\ncommunication into the grouped GEMM and turns DeepEP off automatically, since\nthe two patch the same layer. And every V4 knob reads through an environment\nvariable with a default, which is what makes it practical to sweep one\noptimization at a time.\n\n## Memory and performance projection[#](#memory-and-performance-projection)\n\nBefore booking cluster time it helps to know where the parameter budget and the memory budget actually go. Primus Projection derives both from the model config without running a training step, and for a model shaped like V4-Flash the answers are lopsided in ways that are worth seeing before you start tuning.\n\n### Where the parameters go[#](#where-the-parameters-go)\n\nTwo things stand out in Figure 4:\n\n**MoE holds 95.7% of the model.** All 43 attention modules together come to 1.70% — 4.94B parameters against the MoE stack’s 278.15B. Inside a single MoE layer the concentration is sharper still: the 256 routed experts are 99.59% of it, the shared expert 0.39%, and the router gate 0.02%.**The total is 290.80B, not the 284B on the model card.** The difference is the 6.61B MTP module, which the published figure leaves out.\n\n### Where the memory goes[#](#where-the-memory-goes)\n\nFigure 5 shows rank 0 of the PP=4, EP=8 configuration at sequence length 4,096, assuming 11 transformer layers on the rank plus the embedding. Three takeaways:\n\n**It fits, with about 4% to spare.** 257.62 GiB per rank against the 268.2 GiB an MI355X exposes from its 288 GB of HBM3E, leaving roughly 10.6 GiB of headroom.**FP32 state, not activations, is the largest bucket.** The gradient buffer, the optimizer’s main parameter copy, and the two moments come to 142.67 GiB, 55.4% of the total — seven times the 20.38 GiB of BF16 weights they shadow, and more than the 94.56 GiB of activations.**Activations are the only bucket you can trade against compute.** Recompute buys memory back by paying for a second forward pass, and how much of it a run needs depends on everything else that run is doing. The shipped four-node configuration ends up needing none — but only because the kernel work below frees the memory first. Where the layers land and how much to recompute are tuned together, and every layout measured below puts ten layers on stage 0, not the eleven assumed here;[that section](#pipeline-layout-and-recompute)measures both.\n\nThese numbers are projections derived from the config rather than measurements, and they are an upper bound: the measured peak on the shipped four-node configuration is 242.98 GiB — 260.9 GB in the units the sections below use — roughly 15 GiB under the projection.\n\n## Performance optimizations[#](#performance-optimizations)\n\nThe sections below follow the order in which we switched these on, and\n[the ladder at the end](#stacking-the-optimizations) measures what each one is\nworth on a four-node run.\n\n### Kernel fusions[#](#kernel-fusions)\n\nDeepSeek-V4 brings in a lot of new machinery — mHC on every sub-layer, a compressor and an indexer on every compressed layer, two new routers. Written the obvious way, each of those is a chain of small elementwise operations, and PyTorch dispatches every one as its own kernel with a full HBM round trip. None of them is expensive on paper. Together they dominated our first working build, and each intermediate they materialize is memory you do not get back.\n\nSo we fused them. The table below is what ships today; each row replaces an eager chain with a single forward kernel and, where a backward is needed, a single backward kernel.\n\nFusion |\nWhat the eager path does |\nFused into |\nWritten in |\n|---|---|---|---|\nSWA / CSA / HCA attention |\nSeparate K and V paths, a split CSA pool kernel, and a sliding-window branch joined afterwards — see |\nOne single-latent sparse-MLA kernel per layer type, gathering the selected KV entries in-kernel and folding in the window branch and the softmax sink |\nTriton, Gluon, FlyDSL |\nRMSNorm |\nCast to fp32, square, mean, add eps, rsqrt, scale, cast back, optional weight multiply — an 8-op chain |\nOne kernel pair, covering every non-TE RMSNorm site in the model body |\nTriton |\nInterleaved partial RoPE |\nA 9-op chain ending in a |\nOne kernel pair |\nTriton |\nSinkhorn-Knopp |\n39 fp32 reductions over a 4×4 matrix — one priming column normalization plus 19 row/column pairs — each its own launch |\nOne kernel pair that keeps the entire trajectory in registers |\nTriton, after a |\nHyper-connection glue |\nThree slices, three fused multiply-adds, two sigmoids, a softmax and two eps adds — about 8 launches |\nOne kernel |\nTriton |\nHyper-connection collapse |\nA broadcast multiply that materializes a full |\nOne kernel that contracts |\nTriton |\nHyper-connection expand |\nAn outer product, a contraction over |\nOne kernel |\nTriton |\nCompressor pooling |\nAdd the positional bias, cast, softmax over the window, cast back, multiply, reduce — about 5 launches |\nOne forward kernel that reduces in fp32 and handles both the CSA window of 8 and the HCA window of 128 |\nTriton |\nIndexer scoring tail |\nReLU, per-head multiply, sum over heads, mask allocation, mask add, cast — about 5 ATen launches |\nOne kernel that materializes the causal mask inline, with no |\nTriton |\nMoE router tail |\nScore function, gather, sum, clamp, divide, scale, then two scatters |\nOne kernel |\nTriton |\nGrouped expert weight stack |\n|\nOne kernel, single pass |\nTriton |\n\nNone of these is a headline optimization on its own, which is exactly why they are easy to leave on the table. Switching all of them on at once is the single largest step in the whole ladder: it nearly doubles end-to-end throughput and frees 22 GB of memory at the same time, because every intermediate that no longer gets written is also memory that no longer gets allocated.\n\n### Attention kernels for the three layer types[#](#attention-kernels-for-the-three-layer-types)\n\nV4-Flash runs [three different kinds of attention](#three-attention-types-interleaved),\npicked per layer by `compress_ratio`\n\n, and each one hands the kernel a different\nproblem. Primus implements all three as fused kernels in several backends and\nhas tuned each of them:\n\n**eager**— a plain PyTorch path. Slow, but it is the reference the parity tests compare against.** Triton**— the first production backend, and the portable one.** Gluon**— Triton’s experimental Gluon dialect, gfx950 only. It exposes the warp-level pipeline, so the kernel can be scheduled explicitly instead of leaving the decision to the compiler.**FlyDSL**— the fastest of the four. FlyDSL gives fine-grained control over instruction scheduling and software pipelining on MI355X, which is exactly what the compressed layer types need: their inner loop is a gather over a sparse set of KV entries, and hiding that latency behind MFMA issue is a scheduling problem more than a math problem. The FlyDSL DeepSeek-V4 attention kernels live in Primus-Turbo.\n\nThe tables below are single-GPU MI355X measurements at sequence length 4,096,\nmicro-batch 1, BF16, attention sink on, 128-token sliding window. Each cell is\nmedian latency in milliseconds and the achieved TFLOP/s.[1]\n\n#### Forward[#](#forward)\n\nModel |\nLayer type |\nTriton |\nGluon |\nFlyDSL |\n|---|---|---|---|---|\nV4-Flash |\nSWA (cr = 0) |\n0.30 | 230.0 |\n0.28 | 248.3 |\n|\nV4-Flash |\nCSA (cr = 4) |\n0.87 | 397.1 |\n0.66 | 523.6 |\n|\nV4-Flash |\nHCA (cr = 128) |\n0.38 | 223.9 |\n0.33 | 263.2 |\n|\nV4-Pro |\nSWA (cr = 0) |\n0.58 | 236.2 |\n0.51 | 269.0 |\n|\nV4-Pro |\nCSA (cr = 4) |\n2.78 | 444.3 |\n1.92 | 645.1 |\n|\nV4-Pro |\nHCA (cr = 128) |\n0.72 | 238.6 |\n0.61 | 280.9 |\n|\n\n#### Backward[#](#backward)\n\nModel |\nLayer type |\nTriton |\nGluon |\nFlyDSL |\n|---|---|---|---|---|\nV4-Flash |\nSWA (cr = 0) |\n1.16 | 148.4 |\n1.13 | 152.0 |\n|\nV4-Flash |\nCSA (cr = 4) |\n5.93 | 144.9 |\n3.99 | 215.1 |\n|\nV4-Flash |\nHCA (cr = 128) |\n1.67 | 128.8 |\n1.54 | 139.2 |\n|\nV4-Pro |\nSWA (cr = 0) |\n1.81 | 190.1 |\n1.70 | 202.2 |\n|\nV4-Pro |\nCSA (cr = 4) |\n10.74 | 287.8 |\n8.52 | 362.9 |\n|\nV4-Pro |\nHCA (cr = 128) |\n2.47 | 174.2 |\n2.27 | 189.4 |\n|\n\nTo reproduce any column, set the two selectors in the experiment config. The dense and HCA layers read the first field, the CSA layers the second:\n\nConfig field |\nTriton |\nGluon |\nFlyDSL |\n|---|---|---|---|\n|\n|\n|\n|\n|\n|\n|\n|\n\nGetting the FlyDSL column to those numbers took work in both directions. Sparse-MLA is one of the harder attention kernels to make fast: the KV cache is a single 512-dimensional latent serving as both key and value, and each query reads only a sparse top-k subset of it, so on top of the usual attention math the kernel pays a gather/scatter tax — two passes over an intermediate tensor that dense flash attention never touches.\n\n**The forward pass is latency-bound rather than throughput-bound.** With `exp2`\n\nand MFMA issue roughly balanced there is no occupancy left to buy, so every gain\ncame from shortening or overlapping the serial `QK → softmax → PV`\n\nchain:\n\nBatching two adjacent tiles into\n\n`K=32`\n\ndoubles MFMA depth and halves the read-after-write chain.Moving to one work-group per token at\n\n`BLOCK_H=128`\n\n. Under a shared latent, two work-groups per token each store their own copy of it — pure redundancy. Storing it once cuts roughly a quarter of the work.Exploiting softmax’s shift invariance: take the first key pair’s maximum as a fixed bound, and the rescale factor becomes a constant 1 that the compiler folds away. That buys no-max speed at pure-accumulation precision, worth about 13% on its own. It is now the default path.\n\n**The backward pass splits into three kernels**, each with a different bound and\na different fix:\n\n**dQ** takes the largest share and is pinned to single-wave occupancy by register pressure. Its bottleneck is HBM latency on the KV gather, and the only way to hide it is to keep the per-tile`QK → softmax → PV`\n\ninterleaving. It is the healthiest of the three, running 1.4–2.4× faster than the Triton version.**interm** is a head-dimension contraction GEMM. Replacing its LDS staging with a hand-rolled 16×16 in-register transpose through`ds_bpermute`\n\n, then moving to`K=32`\n\nMFMA, halves the instruction count.**delta**— the`rowsum(O·dO)`\n\nreduction — was a standalone, fully serial micro-kernel. Inlining it into dQ removes an entire launch. Batching kv blocks then let dQ and interm fuse as well, which also drops the HBM round trip for the intermediate tensor.\n\nAll six shapes are numerically correct in both directions at BF16.\n\n### Expert parallelism: DeepEP and the grouped GEMM[#](#expert-parallelism-deepep-and-the-grouped-gemm)\n\nWith 256 experts and six of them active per token, the MoE layer is where both the FLOPs and the communication live. Two independent optimizations sit on that path, and they are worth separating because they are almost always enabled together and then reported as one.\n\n**DeepEP** replaces the token dispatch and combine. Expert parallelism has to\nsend every token to whichever rank owns its experts and bring the results back;\nthe stock path does that as a pair of all-to-all collectives with the\npermutation and its inverse done in PyTorch around them. DeepEP does the\npermutation, the transfer and the reverse in dedicated kernels, producing the\ntoken layout the GEMM wants rather than assembling it afterwards.\n\n**The Turbo grouped GEMM** replaces the per-expert loop. At EP=8 every rank\nholds 32 of the 256 experts, and each of them multiplies a different number of\ntokens by its own weights. Issued as 32 separate GEMMs most are far too small to\nfill the GPU, and the launch overhead alone is comparable to the math. The\ngrouped GEMM issues all of them as one kernel over a ragged batch.\n\nBoth live in the experiment config:\n\n```\nenable_primus_turbo: true\nuse_turbo_deepep: true\nuse_turbo_grouped_gemm: true\n```\n\nSwitched on one after the other rather than in one step, the grouped GEMM is\nworth about six times what DeepEP is worth on this model: DeepEP adds 1.3%,\nand the grouped GEMM on top of it another 7.9%\n([rungs 4 and 5 of the ladder below](#stacking-the-optimizations)). Enabling\nthem together, as most configurations do, would put that gain on the wrong\nfeature.\n\n### MegaMoE: fusing communication into the grouped GEMM[#](#megamoe-fusing-communication-into-the-grouped-gemm)\n\nBoth optimizations above make one half of the expert path faster, but they leave it in two halves: the communication still runs next to the GEMM rather than inside it. The usual way to hide it is to put the transfer on its own stream and overlap it with the math. That is awkward to orchestrate, and the two streams then compete for the same compute units and memory bandwidth — the overlap gives back part of what it saves.\n\nPrimus overlaps inside the kernel instead. Data movement and math interleave at the instruction level rather than racing as separate streams, which is what FlyDSL’s fine-grained control over the pipeline makes possible. This ships today as single-node fusion on MI355X; the same approach extends to much larger EP degrees on the next generation of rack-scale systems.\n\nAs Figure 6 shows, MegaMoE — the FlyDSL layer that replaces the native\n`MoELayer`\n\n— collapses that chain into two kernels: `dispatch_grouped_gemm`\n\nfuses the token dispatch all-to-all into the first grouped GEMM, and\n`grouped_gemm_combine`\n\nfuses the second grouped GEMM into the combine and the\nweighted reduce. With a fused router in front and SwiGLU in between, the whole\nexpert path becomes `dispatch_grouped_gemm → SwiGLU → grouped_gemm_combine`\n\n.\n\n**Two stages, so the DDP collectives still overlap.** Primus-Turbo does expose\nall of this as a single fused op. Primus deliberately drives it as two stages\ninstead, each owning one weight and each wrapped in a tiny weight module that\ncomputes nothing:\n\n```\nMegaMoEExperts\n├── fc1_weight : MegaMoEWeightModule   # w1 [g, 2I, H]  gate + up\n└── fc2_weight : MegaMoEWeightModule   # w2 [g, H, I]   down\n\nFORWARD (in order)                   BACKWARD (in order)\n─────────────────────────────        ─────────────────────────────\nw1 = fc1_weight()                    stage2.backward -> dW2\n  hook: all-gather(w1), wait           hook: reduce-scatter(dW2) ─┐\nstage1: dispatch + GEMM1     ─┐                                   │ overlap\nw2 = fc2_weight()             │ ovl  stage1.backward -> dW1 ───────┘\n  hook: all-gather(w2) ───────┘        hook: reduce-scatter(dW1)\nstage2: SwiGLU + GEMM2 + combine         (overlaps the next layer)\n```\n\nThose modules exist to be hook sites. The distributed optimizer overlaps two collectives at module and parameter granularity, and neither can overlap anything if the expert path is one opaque call:\n\n`overlap_param_gather`\n\nrides the forward pre-hook, which fires per module. A single call site taking both weights means both all-gathers have to land before any compute starts. Split,`w2`\n\n’s gather is issued at`fc2_weight`\n\nand hides under stage 1.`overlap_grad_reduce`\n\nrides the grad hook, which fires when a parameter’s`.grad`\n\nappears. One fused autograd node emits`dW1`\n\nand`dW2`\n\ntogether at the end of the layer backward. Split,`dW2`\n\nlands early and its reduce-scatter hides under stage 1’s backward.\n\nThe split is purely at the Python and autograd level — the kernels themselves are unchanged.\n\n**Configuration.** Two flags turn it on, and MegaMoE is EP-only:\n\n```\nenable_primus_turbo: true\nuse_turbo_mega_moe: true      # EP-only, TP=1, BF16\ntensor_model_parallel_size: 1\nadd_bias_linear: false\n```\n\nThe replacement is applied only when `enable_primus_turbo`\n\nand\n`use_turbo_mega_moe`\n\nare both set, `tensor_model_parallel_size == 1`\n\n,\n`params_dtype == bf16`\n\n, and an EP process group exists. Anything else asserts.\nSequence-level and global aux loss, z-loss, sinkhorn and input jitter are\nunsupported — only the standard `aux_loss`\n\n— and aux-loss-free expert bias raises\n`NotImplementedError`\n\n.\n\n**What it buys.** The expert-parallel intra-node all-to-all is fused into the\nFlyDSL grouped-GEMM kernel, so the ideal cost becomes `max(comm, gemm)`\n\nrather\nthan their sum. In practice the fused kernel holds at least 85% of that\nperfect-overlap roofline, 90% or better in most cases, with only 0.3–0.5 ms of\noverhead left over.\n\nFigure 7 measures what that is worth. These are times for the MoE module on its\nown.[1]\n\n### Pipeline layout and recompute[#](#pipeline-layout-and-recompute)\n\nThe last two knobs are not kernels. They decide how the 43 transformer layers, the embedding, the MTP module and the loss are spread across the four pipeline stages, and how much activation memory is traded back for recompute.\n\nThe default split — 10 layers on stage 0, which also carries the embedding, and\n11 on each of the others — looks fair and is not. The last stage also carries\nthe MTP module and the loss, while 1F1B leaves stage 0 holding four microbatches\nin flight where the last stage holds one. Moving two layers off the last stage\nonto the middle two, `Et*10|t*12|t*12|t*9mL`\n\n, evens out the time per stage and\nshortens the pipeline bubble; stage 0 keeps its 10 layers either way, because it\nis the one under activation pressure. Recompute is the other half — the\nconservative starting point checkpoints the first three layers of every stage —\nand the optimizations above free enough memory to stop paying that tax.\n\nMeasured one at a time against the same reference:[2]\n\nChange |\nLayout |\nRecompute |\nTFLOP/s |\nGain |\nPeak memory |\n|---|---|---|---|---|---|\nreference |\n|\n3 |\n1167.2 |\n— |\n217.2 GB |\nlayout only |\n|\n3 |\n1273.4 |\n+9.1% |\n225.9 GB |\nrecompute only |\n|\n0 |\n1255.7 |\n+7.6% |\n260.9 GB |\nboth |\n|\n0 |\n1378.8 |\n+18.1% |\n260.9 GB |\n\nThey contribute almost equally, and doing both is worth 1.5 points more than the sum of doing each alone: dropping recompute frees time that an unbalanced pipeline would partly give back as bubble, and rebalancing the pipeline has little to fill unless recompute stops taking the time. The memory bill is dominated by recompute: dropping it costs 35 to 44 GB depending on the layout, against the 8.7 GB the rebalance adds on its own. Once recompute is off, both layouts peak at the same 260.9 GB of the 288 GB an MI355X provides, since stage 0 holds 10 layers and the most microbatches in flight either way. That is also why recompute cannot be dropped first: it only fits once the fusions and the MoE work have given the memory back.\n\n**Neither knob is a one-time decision.** Those four rows isolate what each one is\nworth at a single point in the project; they are not the method that produced the\nshipped values. Layout and recompute were retuned continuously throughout V4\ndevelopment, because every kernel that landed moved the target: a fusion that\nfrees 8 GB changes which layout balances best, and a faster attention kernel\nchanges which stage sits on the critical path. Recompute is not a switch either.\nPrimus exposes the granularity, the per-stage layer count, an explicit list of\nglobal layer ids, and a per-module selection, so how much to recompute is a\nsearch over a space rather than a boolean. Zero is where that search happens to\nland for four nodes with everything else on; eight nodes, or a different set of\nkernels, land elsewhere.\n\nRunning that search by hand does not scale past a handful of configurations, and it has to be redone every time the kernels change. We are building an auto-tuner that chooses the pipeline layout and a fine-grained recompute plan together, and will open-source it in Primus as it matures.\n\n### Stacking the optimizations[#](#stacking-the-optimizations)\n\nEvery section above reports what one optimization is worth in isolation. The number that decides whether a run is practical is what they are worth together, and that is not the same thing — each one changes the balance the next one sees.\n\nSo we measured the whole ladder end to end on four nodes: start from a build\nwith every optimization switched off, turn on exactly one thing per rung, keep\neverything already on, and hold the shapes fixed at global batch 256,\nmicro-batch 1, sequence length 4,096, all in BF16, with router load balancing\nforced to uniform so the expert GEMM shapes do not drift between rungs. Ten\niterations per rung, averaged over iterations 4 to 10.[2]\n\nFigure 8 plots that climb, and the table below gives the exact number each rung lands on:\n\nStage |\nOptimization |\nWhat it changes |\nTFLOP/s/GPU |\nThis step |\nCumulative |\n|---|---|---|---|---|---|\n0 |\nBaseline |\nEvery optimization off: unfused elementwise chains, first-generation Triton attention, the native MoE layer, an even pipeline split with three recomputed layers per stage |\n439.5 |\n— |\n— |\n1 |\nKernel fusions |\nThe fusions listed above, plus Megatron’s permutation, cross-entropy and gradient-accumulation fusions |\n875.4 |\n+99.2% |\n+99.2% |\n2 |\nGluon attention |\nSparse-MLA moves to the Gluon dialect ( |\n917.0 |\n+4.8% |\n+108.6% |\n3 |\nFlyDSL attention |\nSparse-MLA moves again, to the FlyDSL kernels in Primus-Turbo |\n954.3 |\n+4.1% |\n+117.1% |\n4 |\nDeepEP |\nToken dispatch and combine become dedicated kernels instead of PyTorch permutation around two all-to-all collectives |\n966.5 |\n+1.3% |\n+119.9% |\n5 |\nTurbo grouped GEMM |\nThe 32 local expert GEMMs issue as one ragged-batch kernel |\n1042.6 |\n+7.9% |\n+137.2% |\n6 |\nMegaMoE |\nReplaces both of the above: the all-to-all is fused into the grouped GEMM rather than sitting next to it |\n1167.2 |\n+12.0% |\n+165.6% |\n7 |\nPipeline layout and recompute |\nLayers rebalanced to 10/12/12/9, recompute dropped to zero |\n1378.8 |\n+18.1% |\n+213.7% |\n\n## Reproduce: training DeepSeek-V4-Flash[#](#reproduce-training-deepseek-v4-flash)\n\nEverything above ships in the open-source\n[Primus](https://github.com/AMD-AGI/Primus/tree/8e24522d3ccf9be38411385a38bb881261378eb9)\nrepository, and the four-node configuration in this blog is the default — you do\nnot have to reassemble the optimizations by hand.\n\nThe runs here pin Primus at commit\n[ 8e24522](https://github.com/AMD-AGI/Primus/commit/8e24522d3ccf9be38411385a38bb881261378eb9),\nwhich is where this launcher landed on\n\n`main`\n\n. Build the container from the\nDockerfile at that same commit —\n[, which puts PyTorch, Megatron-LM, Primus-Turbo and the FlyDSL kernels on a ROCm base — and point](https://github.com/AMD-AGI/Primus/blob/8e24522d3ccf9be38411385a38bb881261378eb9/.github/workflows/docker/Dockerfile)\n\n`.github/workflows/docker/Dockerfile`\n\n`DOCKER_IMAGE`\n\nat it; the launcher requires that variable.One thing the image does not settle is the code. Any Primus image — including\none built from that Dockerfile — ships its own snapshot of the repository under\n`/workspace/Primus`\n\n, and that snapshot is not necessarily this commit. Check\n`8e24522`\n\nout on the host and mount it over that path, so the image supplies\nthe environment and your checkout supplies the code.\n\nEverything else is one launcher:\n[ examples/deepseek-v4/run_deepseek_v4_flash.sh](https://github.com/AMD-AGI/Primus/blob/8e24522d3ccf9be38411385a38bb881261378eb9/examples/deepseek-v4/run_deepseek_v4_flash.sh).\nRun it with no flags for rung 7. Its header documents the rest: one switch per\noptimization family, so any rung of the ladder is a single variable away, and a\ndry-run mode that resolves a combination and prints what it means before you\nspend an allocation on it.\n\nA healthy four-node run settles at roughly 8.5 s per iteration and 1,370–1,385\nTFLOP/s per GPU, with peak memory around 261 GB of the 288 GB on each MI355X.[2]\nThe launcher is tuned for four nodes; another node count needs its own\n`PRIMUS_PP`\n\nand `PRIMUS_PP_LAYOUT`\n\n.\n\n## Summary[#](#summary)\n\nIn this blog you explored what it takes to train DeepSeek-V4-Flash end to end in\nPrimus on AMD Instinct MI355X GPUs. You read the architecture layer by layer —\nthree interleaved attention types, manifold-constrained hyper-connections in\nplace of the plain residual, and a 256-expert MoE in every block — and saw why\nnone of it drops into stock Megatron-LM unchanged. You saw which knobs Primus\nexposes to describe that shape in YAML, what Primus Projection says about the\nparameter and memory budget before you book a single node, and then — where most\nof the engineering went — the kernel work that took the model from *it runs* to\n*it runs fast*.\n\nNo single change got it there. Fusing the small operations V4 introduces — mHC on every sub-layer, a compressor and an indexer on every compressed layer, two new routers — was the largest single step, nearly doubling throughput and freeing 22 GB at once. Moving the three attention types from Triton to Gluon and then to the FlyDSL sparse-MLA kernels added another 9% to end-to-end throughput and delivered up to a 2.3× speedup on the CSA backward pass alone, through scheduling and pipelining rather than new math. On the expert path, DeepEP and the Turbo grouped GEMM each accelerate one half of it, until MegaMoE replaces both by fusing the expert-parallel all-to-all into the GEMM itself. And the last 18% was not a kernel at all: rebalancing the pipeline to 10/12/12/9 layers and switching recompute off, which only fits because the kernel work freed the memory first.\n\nTogether they take a four-node run from 439.5 to 1,378.8 TFLOP/s per GPU — 3.1× —\nwith the model holding at 261 GB of the 288 GB each GPU provides. All of it is in\nthe [Primus repository](https://github.com/AMD-AGI/Primus) and on by default:\nattach a four-node allocation and run the launcher.\n\nThree threads continue from here, and we will cover them in future posts as they land. FP8 is the nearest — the experiment config already sits beside the BF16 one and the numerical pieces it leans on are in place, so what is left is coverage and stability rather than enablement. The pipeline-layout and recompute search we ran by hand for this blog is becoming an auto-tuner that plans both together, and we will open-source it in Primus as it matures. And MegaMoE’s in-kernel overlap, which today supports single-node fusion, is the piece that extends to much larger expert-parallel degrees on the next generation of rack-scale systems. Each of these will land in Primus before it appears in a blog, so the repository is the place to watch.\n\n## Acknowledgments[#](#acknowledgments)\n\nWe would like to express our sincere gratitude to the following teams and individuals for their invaluable contributions and collaboration, their expertise and support have been instrumental in advancing the progress of this project: Felix Li from the FlyDSL Team, and Wen Chen and Ye Wang from the TE Team.\n\n## Additional Resources[#](#additional-resources)\n\n[AMD Instinct™ MI355X GPUs](https://www.amd.com/en/products/accelerators/instinct/mi350/mi355x.html): Product page for the accelerators every measurement in this blog runs on.[DeepSeek-V4 technical report](https://arxiv.org/abs/2606.19348): The architecture this enablement follows, including the compressed-attention and hyper-connection definitions.[DeepSeek-V4-Flash model card](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash): Published weights and reference configuration for the model.[Primus](https://github.com/AMD-AGI/Primus): The training framework holding the V4 model definition, the experiment configs, and the launcher used here.[Primus-Turbo](https://github.com/AMD-AGI/Primus-Turbo): AMD’s operator library, where the FlyDSL sparse-MLA attention and MegaMoE kernels live.[Megatron-LM](https://github.com/NVIDIA/Megatron-LM): The backend Primus builds on, and the stock code path the V4 builders replace.[DeepEP](https://github.com/deepseek-ai/DeepEP): Expert-parallel dispatch and combine library, measured at rung 4 of the ladder.[Triton](https://github.com/triton-lang/triton): Compiler for the portable kernel backend, and home of the experimental Gluon dialect used for the gfx950 attention kernels.[Primus Projection: Estimate Memory and Performance Before You Train](https://rocm.blogs.amd.com/software-tools-optimization/primus-projection/README.html): The tool behind the parameter and memory projections in this blog.[MoE Training Best Practices on AMD GPUs](https://rocm.blogs.amd.com/software-tools-optimization/primus-moe-package/README.html): Broader MoE training guidance that complements the V4-specific work here.[Porting High-Performance HIP Kernels to FlyDSL](https://rocm.blogs.amd.com/software-tools-optimization/porting-hip-flydsl/README.html): Background on the FlyDSL programming model behind the fastest attention backend.\n\n## Endnotes[#](#endnotes)\n\n[1] Test Environment\n\nSingle-GPU kernel latency — the attention backend tables and the MegaMoE MoE-module times — was measured on one AMD Instinct MI355X GPU of an 8-GPU node with BF16 precision, sequence length 4,096 and micro-batch 1. Server manufacturers may vary configurations, which can yield different results. Performance may also vary based on the use of the latest drivers and optimizations.\n\nAMD system configuration:\n\nDual AMD EPYC 9575F 64-core processor\n\n8× AMD Instinct MI355X GPUs, 288 GB HBM3E per GPU\n\n1 NUMA node per socket\n\nSystem model: Supermicro AS-4126GS-NMR-LCC, system BIOS 1.4a\n\nHost OS: Ubuntu 22.04.5 LTS with Linux kernel 6.8.0-107-generic\n\nHost GPU driver: ROCm 7.0.1 + amdgpu 6.14.14\n\nVBIOS version: 113-M355-01-1K1-010C\n\nPyTorch 2.12.0\n\nAMD ROCm 7.14 software in the container\n\nPrimus-Turbo 0.3.2, FlyDSL 0.2.4, Triton 3.7.0, Transformer Engine 2.14.0\n\n[2] Test Environment\n\nEnd-to-end pretraining throughput (TFLOP/s per GPU) — the optimization ladder, the pipeline layout and recompute comparison, and the four-node figures in the reproduce section — was measured on 4 MI355X nodes (32 GPUs total) with BF16 precision, TP=1, PP=4, EP=8, global batch 256, micro-batch 1 and sequence length 4,096, averaged over iterations 4 to 10 of a 10-iteration run. Server manufacturers may vary configurations, which can yield different results. Performance may also vary based on the use of the latest drivers and optimizations.\n\nAMD system configuration:\n\nDual AMD EPYC 9575F 64-core processor per node\n\n32× AMD Instinct MI355X GPUs across 4 nodes, 288 GB HBM3E per GPU\n\n1 NUMA node per socket\n\nSystem model: Supermicro AS-4126GS-NMR-LCC, system BIOS 1.4a\n\nHost OS: Ubuntu 22.04.5 LTS with Linux kernel 6.8.0-107-generic\n\nHost GPU driver: ROCm 7.0.1 + amdgpu 6.14.14\n\nVBIOS version: 113-M355-01-1K1-010C\n\nPyTorch 2.12.0\n\nAMD ROCm 7.14 software in the container\n\nPrimus-Turbo 0.3.2, FlyDSL 0.2.4, Triton 3.7.0, Transformer Engine 2.14.0\n\n[3] Scope of these measurements\n\nEverything measured here is one configuration, and it is worth being explicit\nabout where its edges are. The optimizer is AdamW in BF16, not the Muon that\nDeepSeek used for V4 pretraining — Primus wires Muon in behind `OPTIMIZER=muon`\n\n,\nbut that is not what these numbers measure. The indexer distillation loss that\ntrains CSA’s selector is off, which also leaves the indexer parameters frozen:\nthe right setting for loading an already-trained indexer, or for measuring what\nthe kernels cost, and the wrong one for pretraining from scratch, where it has\nto be on. FP8 is not measured here either, though the pieces that path leans on\nalready ship — an `E4M3`\n\npath for the indexer QK, and the clamped SwiGLU that\ngives FP8 and FP4 their numerical headroom.\n\n## Disclaimers[#](#disclaimers)\n\nThe information presented in this document is for informational purposes only and may contain technical inaccuracies, omissions, and typographical errors. The information contained herein is subject to change and may be rendered inaccurate for many reasons, including but not limited to product and roadmap changes, component and motherboard version changes, new model and/or product releases, product differences between differing manufacturers, software changes, BIOS flashes, firmware upgrades, or the like. Any computer system has risks of security vulnerabilities that cannot be completely prevented or mitigated. AMD assumes no obligation to update or otherwise correct or revise this information. However, AMD reserves the right to revise this information and to make changes from time to time to the content hereof without obligation of AMD to notify any person of such revisions or changes. THIS INFORMATION IS PROVIDED ‘AS IS.” AMD MAKES NO REPRESENTATIONS OR WARRANTIES WITH RESPECT TO THE CONTENTS HEREOF AND ASSUMES NO RESPONSIBILITY FOR ANY INACCURACIES, ERRORS, OR OMISSIONS THAT MAY APPEAR IN THIS INFORMATION. AMD SPECIFICALLY DISCLAIMS ANY IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR ANY PARTICULAR PURPOSE. IN NO EVENT WILL AMD BE LIABLE TO ANY PERSON FOR ANY RELIANCE, DIRECT, INDIRECT, SPECIAL, OR OTHER CONSEQUENTIAL DAMAGES ARISING FROM THE USE OF ANY INFORMATION CONTAINED HEREIN, EVEN IF AMD IS EXPRESSLY ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. Third-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF THIRD-PARTY CONTENT. Illustrations may have been created using generative AI and reviewed by AMD. AMD, the AMD Arrow logo, AMD Instinct, AMD ROCm, and combinations thereof are trademarks of Advanced Micro Devices, Inc. PyTorch is a registered trademark of Meta Platforms, Inc. Other product names used in this publication are for identification purposes only and may be trademarks of their respective companies. © 2026 Advanced Micro Devices, Inc. All rights reserved", "url": "https://wpnews.pro/news/enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus", "canonical_source": "https://rocm.blogs.amd.com/software-tools-optimization/primus-deepseek-v4/README.html", "published_at": "2026-09-03 00:00:00+00:00", "updated_at": "2026-09-03 16:54:25.397173+00:00", "lang": "en", "topics": ["large-language-models", "ai-research", "ai-infrastructure", "ai-products"], "entities": ["DeepSeek-AI", "DeepSeek-V4-Flash", "AMD Instinct MI355X", "Primus", "Megatron-LM", "DeepSeek-V3", "DeepSeek-V3.2"], "alternates": {"html": "https://wpnews.pro/news/enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus", "markdown": "https://wpnews.pro/news/enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus.md", "text": "https://wpnews.pro/news/enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus.txt", "jsonld": "https://wpnews.pro/news/enabling-deepseek-v4-flash-training-on-amd-instinct-mi355x-gpus-with-primus.jsonld"}}