[samsja]

Performance work for DeepSeek-V4.1 Flash, stacked on #3879 (base feat/deepseek-v41).

Changes

• DeepEP works under activation checkpointing: the communication buffer is now built when the dispatcher is created. Built lazily inside the first checkpointed forward, it shifted the recompute's cached ops.

• Single-pass DeepEP local layout: one Triton gather into the padded expert-grouped layout and one fp32 score-weighted gather-sum back, with gathers in backward too (models/kernels/moe_permute.py). This replaces about 7 full copies of the routed activations per layer.

• DeepEP in the 12-node example: 9.3 s -> 5.6 s per step.

• Hopper kernels from prime-kernels, each used only when available:

  • V4.1's indexer uses dsa_indexer_topk, and its mHC uses mhc_projection.
  • model.moe.compute.backend = "prime_kernels" runs moe_experts.

• FP8 routed experts (model.moe.compute = {type = "deepgemm_fp8", backend = "prime_kernels"}): DeepGEMM blockwise FP8 grouped GEMMs (forward, dgrad, wgrad). The quantization and clamped SwiGLU are fused into the passes around them, via prime-kernels' moe_experts(fp8=True) (feat: Hopper DeepSeek-V4.1 training kernels (indexer top-k, MoE experts bf16/FP8, mHC, sparse attention bwd) prime-kernels#15). At V4.1's EP=96 shape this is ~1.4x faster than bf16. The existing DeepGEMM path is no faster than bf16, because its standalone quantization passes eat the gain.

• fp8.toml overlay for V4.1: FP8 dense linears plus FP8 experts. The 512-wide KV projections and the indexer's query projection stay in bf16; measured in FP8, they come out at 0.81x and 0.99x of bf16 speed.

• Cheaper FP8 casts for dense linears: each cast takes one reciprocal per group with an FMA correction instead of an exact divide per element, and is exhaustively bit-identical. The transposed cast pads the token dimension itself, so there are no F.pad copies. Scales are written in DeepGEMM's layout, which removes its scale-transpose kernels. About 5% less GPU time per FP8 linear.

• Sparse attention backward: uses prime-kernels' Hopper dsa_sparse_attn_bwd when it is available (~1.2x cuDNN), and falls back to cuDNN otherwise. The gain is within noise end to end, because the saved time lands in DeepEP waits.

• optim.apply_in_backward (SignSGD, SFT): applies the update to each parameter as soon as its gradient is reduced, then frees the gradient. It is exact, because sign updates ignore the loss scaling. Peak memory drops from 117.6 to 96.9 GiB at the same step time. Spending that memory on 24k tokens per GPU fits but is slower per token (242k vs 330k tokens/s).

• Copy-free grouped o_a_proj (V4 and V4.1): one GEMM per group on strided slices, instead of a batched matmul over a transposed view that copied the output in forward and the gradient in backward. Bit-identical, and 1.23x faster fwd+bwd.

• Fix: FP8 / MXFP8 linear conversion now swaps only plain nn.Linear. Subclasses with their own forward (V4's block-diagonal o_a_proj) were silently replaced and broke DeepGEMM downstream.

• Engram: the tables are plain nn.Embeddings, row-sharded by a new AllToAllEmbeddingParallel style next to EmbeddingParallel. Deduplicated ids go to their owners and only those rows come back. RowwiseParallel would instead reduce-scatter ~19 GB of partial rows per table per rank at 96 GPUs. The custom table class is gone from the model.

• DeepEP as one custom op (prime_rl::deepep_moe):

  • Forward and backward are pipelined over token chunks, with communication outside autograd.
  • Activation checkpointing saves the op's outputs, so recompute neither resends tokens nor reruns experts.
  • Parity with the torch dispatcher is ~3e-3 relative on outputs, input and weight gradients, with and without chunking.
  • It does not speed up a 12-node step (see below), but chunking no longer costs anything.

• Optional FP8 token dispatch ([model.moe.dispatch] fp8 = true, DeepEP with deepgemm_fp8 experts): tokens are quantized per 1x128 block with power-of-two scales before the forward dispatch, then dequantized exactly on arrival to the bf16 the experts quantize again. Forward dispatch traffic and the received tokens kept for backward halve; gradients stay bf16. In a microbenchmark at the trainer's shapes (64 ranks) a dispatch takes 7.5 ms instead of 14.2 ms. End to end: 4.8 s -> 4.7 s per step and 4 GiB less peak memory, with the loss curve unchanged within noise.

• In-place RoPE for V4.1's query and for the gradient of the attention output: only the 64 rotary channels of each 512-wide head are read and written, instead of copying the whole head (6.3x faster per call at 16k tokens, bit-identical). Within noise end to end: like the faster attention backward, the saved compute turns into DeepEP waits.

• Fix: DeepEP's local experts run in row slices when routing sends a rank more token-expert pairs than prime-kernels' FP8 path addresses (32-bit offsets, ~419k rows at V4.1's width). Before, that rank raised and every other rank hung in DeepEP until its GPU timeout.

• Pipeline parallelism (model.pp, pp_schedule, pp_stages_per_rank, pp_layers_per_stage): stage build in model.py, trainer/pipeline.py (pooled receive buffers, an async 1F1B schedule, a no-grad shape probe for static shapes), prune_to_pipeline_stage / pipeline_stage_forward for V4.1, PP-aware data ranks and grad clipping across the PP group. Losses match pp = 1 exactly.

• Kept expert activations (model.moe.dispatch.keep_expert_activations): the routed experts' activations are kept for backward instead of recomputed (−5.8% step on a 5-layer EP64 proxy, +3.2 GB per layer).

• FP8 expert weights quantized once per optimizer step, and FP8 weight gradients accumulated straight into FSDP's fp32 buffer (prime-kernels moe_experts_fp8_quantize_weights / moe_experts_fp8_backward_accumulate): +8.7% and +9.1% on a PP2×4 proxy.

• dispatch.weight_grads_after_combine: each chunk's gradient is sent back before its weight-gradient GEMMs run, so those GEMMs hide the send. Bitwise-identical gradients; −2.5 ms per layer at EP48.

• moe.reduce_local_expert_grads_once: when each expert FSDP group is a single rank, its per-micro-batch FSDP work (cast, reduce copies) runs once per step. −4.2 ms per layer at EP48; same up to fp32 summation order; ~0.65 GiB more peak memory per layer.

• dispatch.fp8_grad (off by default, changes numerics): FP8 backward dispatch, halving its bytes. −2.5 ms per layer at EP48.

• Spread balanced routing: debug.force_balanced_routing now strides experts by 97, so a token's 6 experts land on ~4–5 nodes like the real router. The old round-robin sent every token to one rank, which hid nearly all inter-node traffic.

• prime_rl::dsv4_linear_rope: V4.1's query projection and its RoPE as one op that rotates the projection output in place and keeps no unrotated copy for backward. Bitwise identical; −0.46 ms and −1 GiB per layer. It is in the default selective-AC save list; runs passing an explicit --model.ac.targets must add it. RoPE, mHC and q-norm ops are also saveable by selective AC.

• mHC stream-gradient sum folded into prime-kernels' projection backward when the installed build has hyper_connection_streams (bitwise identical, −0.47 ms per layer); otherwise the old path.

• Fused Triton cross-entropy gradient in the chunked LM head: 131.7 -> 96.1 ms per 16k tokens, exact.

• dispatch.type = "mega_moe" (off by default): prime-mega-moe's fused BF16 Hopper Mega MoE as a backend (distributed/mega_moe.py), one custom op with gradients for x, top-k weights, routed and shared expert weights. dispatch.num_sms caps its SMs: at the full count it deadlocks against FSDP's NCCL kernels. On one node it is slower than DeepEP + FP8 experts (97–105 vs 80.3 ms per layer), because the backward recomputes the forward and it is BF16.

• debug.random_init draws per-parameter seeded N(0, 0.02). Before, it materialized zeros and kernels ran 1.5–3.5x too fast. Setup order is now AC -> FSDP -> compile.

Pipeline-parallel push (PP16 x EP8 plan of record)

• MFU is 6ND only: perf.py drops the attention term; FLOPs per token = 6 x active matmul params (LoRA keeps its 4/2/6 split). PerfCounter no longer takes seq_len.
• Async pipeline executor (AsyncPipelineSchedule): runs a precomputed op order and never waits on a stage transfer before it is needed. Schedules Async1F1B (with pp_warmup_step) and DualPipeV. pp_layers_per_stage takes multiples of 0.5, so a stage boundary can cut a V4.1 layer between its attention and MoE blocks. Sent outputs that no backward reads are freed early. Grad clipping and the loss broadcast handle several stages per rank and V schedules. Schedule-order unit tests are added.
• Stage transports (pp_transport): nccl (default), copy_engine, and put. put uses one-sided ncclPutSignal / ncclWaitSignal on zero-CTA 2-rank communicators, so a transfer in flight holds no SMs and no longer slows the persistent kernels. It needs NCCL >= 2.30, and nvidia-nccl-cu13 is now overridden to 2.32.3. pp_transport_ctas and pp_gemm_sms pin SM budgets for nccl.
• PP activation offload (pp_activation_offload, default off): keeps in-flight micro-batches' saved activations in pinned host memory between forward and backward. target = "fp8" instead keeps them on the GPU as FP8 with 1x128 scales. That changes numerics, is unverified, and measured slower than fp8t.
• 1F1B-W (model.pp_defer_expert_weight_grads, default off): FP8 Mega MoE weight-gradient GEMMs run after the stage posts its input-gradient send. Same GEMMs, order and accumulators. On 6 nodes, steps 1–3 are bitwise and later steps are within FP8 noise. +1–3 GB.
• Master offload (model.pp_offload_masters, default off; functionally validated; memory/step cost unvalidated, needs a run with expert DP > 1; keep default off): the sharded fp32 masters stay in pinned host memory between a step's first and last forward. The reduced-grad pool is released after the optimizer step. Frees 1.13 GB per MoE block at D=12; same values.
• PP activation offload v2 (functionally validated; memory/step cost unvalidated, needs a run with expert DP > 1; keep default off):
  • Reloads allocate from the compute stream's pool.
  • every (default 1) offloads one micro-batch in k.
  • paced (default off) throttles warmup forwards to the swap-outs.
• Validation status: 1F1B-W was validated on 6 nodes (bitwise for steps 1–3, then within FP8 noise). The wgrad-operand pool change is bitwise. The per-half AC modes were unit-checked. The PutEdge cleanup was parity-checked on 6 nodes (identical losses, per-stage times within 1 ms).
• fp8_transposed_on_demand also accepts a list of pipeline stages. model.ac.layer_modes accepts [attention, MoE] pairs for layers that a stage boundary cuts. PutEdge checks for NCCL >= 2.30 at runtime. The FP8 Mega MoE wgrad operand is allocated from the compute stream's pool (bitwise).
• FP8 Mega MoE (moe.dispatch.fp8, default off): prime-mega-moe's SM90 FP8 kernels for the routed experts. The forward keeps FP8 inputs and the bf16 gate/up output, so the backward does not recompute the forward. Weight gradients go through DeepGEMM K-grouped FP8 GEMMs into FSDP's fp32 accumulators. Sub-flags, each default off:
  • capacity_factor
  • wgrad_tile_scales (changes numerics)
  • free_bf16_expert_weights (same numerics; −3.4 GB per layer at EP8)
  • fp8_transposed_on_demand (bitwise; −1.7 GB per layer, ~0.95 ms per MoE backward)
• Dense FP8 wgrad into fp32 (model.quantization.accumulate_wgrad_fp32, default off; 250-step study: within run-to-run noise): each dense FP8 linear's weight-gradient GEMM writes straight into FSDP's fp32 accumulator. This removes the bf16 rounding per micro-batch. −0.52 ms/layer on a 6-layer EP8 stage bench. Before a reducing post-backward, the group's other gradients are upcast to fp32, so FSDP's reduce-scatter sees one dtype.
• Fused FP8 shared expert (always on, bitwise): V4/V4.1 FP8 shared experts run their gate and up projections as one FP8 GEMM, with a Triton clamped-SwiGLU kernel that also emits the down projection's FP8 input and a backward kernel that writes the dgrad/wgrad FP8 operands directly. The routed and shared outputs are added inside the mHC write-back kernel.
  • Evidence: on 1 GPU, outputs and all gradients are torch.equal under every AC mode. On 8 GPUs, the stage forward and the shared-expert fp32 grads are equal.
  • −0.66 ms/layer. Not yet run inside a PP job.
• Fused expert wgrad across micro-batches (moe.dispatch.fused_wgrad_micro_batches, FP8 Mega MoE; default 1 = today's behaviour; recommended 2 on H200): n > 1 holds a layer's wgrad operands, then adds n micro-batches with one K-grouped GEMM launch per weight.
  • The int can instead be a per-stage list.
  • Bitwise equal to one launch per micro-batch.
  • +1.17 GB per MoE layer per held micro-batch, which is a net loss at 80 GB.
  • Track B validated it in PP on 6 nodes: losses equal within FP8 noise, −2.0% cycle at n = 2.
  • Required prime-mega-moe commit: 5827a5b (branch feat/sm90-fp8-mega-moe-v2, also on kern2/moe), which includes b6460ed and 125a11a. _C must be rebuilt for b6460ed.
    • b6460ed prefetches the fp32 accumulator tile into L2 and uses 128x144 dW2 tiles: bitwise, −0.55 ms/layer at n ≥ 2.
    • 5827a5b orders the FP8 Mega MoE forward and backward tasks expert-major whenever no expert exceeds 16 M blocks (checked per launch), so each expert's weights are read about once. It is bitwise and saves −0.53 ms/layer on top of b6460ed. It applies to the default path too, not just n ≥ 2. prime-mega-moe is not a declared dependency (it is supplied as prime_mega_moe out of band), so that commit is not pinned in pyproject.toml.
• Sparse-attention slot list built once per index source (always on, bitwise): V4.1's index-source layer stores the window + top-k int32 slot list in the shared attention state, and consumer layers reuse it instead of rebuilding it in every forward and recompute. ~−0.11 ms/layer of GPU time, within noise in an untraced stage A/B. Pipeline boundaries now carry the int32 slot list instead of the int64 picks.
• Experimental: single-stream residual for V4.1 (residual_type = "mhc" | "gated" | "layerscale" on the V4.1 config; default "mhc", bitwise unchanged; quality ablation pending): gated is mHC's post gate at hc_mult = 1, run by a fused Triton op (prime_rl::dsv41_gated_residual, a default selective-AC save target). Not checkpoint-compatible with DeepSeek's mHC gates.
  • −5.4% per layer.
  • Stage transfers drop from 335 to 84 MB.
  • 250-step smoke test: within noise of mHC. layerscale was worse.
• One-pass FP8 dual cast for FP8-linear gradients (always on, bitwise; falls back to Triton): with prime-kernels' fp8_dual_cast, grad_output's per-token and transposed casts take one pass, and the q_b gradient's inverse RoPE is fused into it. −0.17 ms/layer (cast kernels 0.93 → 0.77 ms) once prime-kernels has fp8_dual_cast (prime-kernels 2ef8b4c, PR update qwen #15).
• Activation-checkpointing levels: model.ac.mode adds full_moe (also recomputes the FP8 Mega MoE), projections, attention and matmul. model.ac.layer_modes sets a mode per layer, indexed by full-model layer so it works across pipeline stages. The default (full) is unchanged.
• Engram:
  • Fix: the gate multiplies q_weight * k_weight in fp32, as DeepSeek's reference does. This changes numerics toward the reference.
  • The gate's wkv projection is no longer checkpointed by default. model.ac.recompute_engram (default off) restores the recompute on full / full_moe layers. This keeps 51 KB/token more per engram layer than before.
  • Lookups are routed ahead of the forward.
  • Later micro-batches add table rows straight into the gradient (fp32 summation order differs from autograd's).
  • model.engram_offload (default off, prototype) keeps the table shards and their optimizer state in host memory. Checkpointing offloaded tables is not implemented.
• Optimizer:
  • One-pass Triton SignSGD update, bitwise with the torch ops.
  • optim.embedding_update for Muon (default off): V4.1's parameter split, with the Sinkhorn-balanced momentum update for embeddings, LM head and engram tables (MuonWithSinkhorn).
• DeepSeek Muon recipe (optim.recipe = "deepseek", default "dion" unchanged; requires embedding_update): Muon as in the V4/V4.1 reports. It uses Nesterov momentum, 10 hybrid Newton-Schulz iterations in bf16, update RMS 0.18 and AdamW eps 1e-20. Each expert, query head, o_a group, engram wkv K/V and mHC mix is orthogonalized as its own matrix. The engram gate q/k weights are treated as norm weights with weight decay.
  • Distribution: each expert is orthogonalized on its expert-DP shard owner. Dense matrices go to owner ranks via a bf16 all-to-all, overlapped with the local work.
  • Evidence:
    • 11/11 GPU tests, including Sinkhorn.
    • The 8-GPU sharded step is bitwise equal to a single GPU.
    • A 250-step proxy is within run-to-run noise.
  • Cost:
    • Optimizer step ~150 ms per 43 s step.
    • +4 B/param of momentum.
    • Triton Newton-Schulz is off by default, because it kept gradients alive.
• Kernels:
  • The mHC post backward is one pass.
  • dsv41_sparse_attn_rope fuses V4.1's sparse attention with the in-place inverse RoPE of its output, keeping only the rotary channels the backward needs.

Required kernel commits

• prime-kernels 2ef8b4c (feat/hopper-dsv41-kernels, feat: Hopper DeepSeek-V4.1 training kernels (indexer top-k, MoE experts bf16/FP8, mHC, sparse attention bwd) prime-kernels#15): out_tail, the register-sourced dKV MMAs and fp8_dual_cast.
• prime-mega-moe 5827a5b (feat/sm90-fp8-mega-moe-v2, also on kern2/moe) for the FP8 Mega MoE.

Known numerics

• Default block compile on AC-none layers fuses the shared-expert SwiGLU without bf16 rounding. That is more accurate against fp64, but it differs from eager.

Pipeline results (8-node emulation of PP16 stages 0–7, 8k tokens/GPU/micro-batch, M = 320, 1F1B)

MFU is 6·N·D (16.13B active, 989 TFLOP/s). Days are for 20T tokens. All rows are measured on the emulation and extrapolated with the measured-cycle method.

setup	MFU	days 1,536 GPUs	days 1,024 GPUs	peak mem/GPU
H200, pp_transport = nccl	28.9%	51.0	76.4	115.8 GB
H200, pp_transport = put	34.0%	43.4	65.1	117.5 GB
80 GB cap, put	31.6%	46.7	70.0	79.5 GB
80 GB cap, put + fp8_transposed_on_demand	31.6% (full PP16: +1.2 pt over the 30.4% projected baseline)	46.7	70.0	78.7 GB

With put, the pipeline cycle drops by 24.3 ms (14.9%) on H200. The step-10 loss and grad norm match nccl.

Results with spread routing (16k tokens/GPU, random init, selective AC, FP8, FP8 dispatch)

Per-layer cost is the marginal between two layer counts of reduced proxies; the 40-layer projection adds it to the measured step. MFU is 6·N·D with 16.13B active params at 989 TFLOP/s.

layout	per layer	40-layer step	MFU	days for 20T on 1k H200
EP64, 8 nodes	~104 ms	~4.3 s	~37%	~62
EP48, 6 nodes, + weight_grads_after_combine, reduce_local_expert_grads_once, fp8_grad	96.7 ms	—	39.4%	~57.6
EP8, 1 node, 64-expert proxy (no InfiniBand: the ceiling with ideal inter-node communication)	80.3 ms	3.39 s	47.3%	~48

• EP64 is InfiniBand-bound: ~150 GB per GPU per step against ~143 GB at the target step time, and the forward DeepEP kernels already run back to back at ~46 GB/s.
• PP4×EP8 on 8 nodes: 23–28% MFU, from the bubble, torch 1F1B running adjacent stages in lockstep, and stage imbalance.

The 12-node table below predates spread routing and the random-init fix, so it understates communication.

Earlier results (12x8 H200, 16k tokens/GPU, real weights, fake data, round-robin routing)

config	step	tok/s
#3879 (torch EP, overlap_chunks = 2)	9.3 s	170k
+ DeepEP, single-pass layout, prime-kernels kernels (bf16)	5.65 s	278k
+ FP8 experts	5.1 s	308k
+ FP8 experts + FP8 dense (fp8.toml)	4.9 s	321k
+ cheaper FP8 casts, copy-free grouped linear, bf16 KV/indexer projections	4.8 s	327k
+ FP8 token dispatch (dispatch.fp8 = true)	4.7 s	334k

These results use fake.toml, which also forces round-robin expert routing (force_balanced_routing). With the real router the 12-node config does not fit yet: on real math data (INTELLECT-3-SFT-10K) a rank runs out of memory in the first step. On 2 nodes with the first 4 layers, real data gives a max expert load ~20x the mean (Max Vio 18-20) and 1.1-2.0 s steps against 0.72 s with forced balancing.

• FP8 numerics: FP8 shifts the loss curve on this fake-data run. For example, at step 5 the FP8 run is at 0.256 against 0.183 for bf16. That's why it is an overlay, not the default.
• Where the time goes (trace of the bf16 config, rank 0): compute stream busy 4.47 s of 5.72 s.
  • Dense GEMMs: 1.71 s. Microbenchmarks put every large V4.1 projection at 600–720 TFLOP/s.
  • Expert GEMMs: 1.45 s.
  • Sparse attention: 0.65 s.
  • Glue: ~1 s.
  • Waits on DeepEP: ~1.2 s idle.
• Why chunking does not hide DeepEP:
  • With 2 chunks, idle drops to 0.85 s only if the persistent expert GEMMs leave DeepEP its SMs, and that costs as much compute as it saves.
  • Each chunk adds a cross-rank cached_notify barrier.
  • Expert load is balanced across ranks (expert GEMM time 1,433–1,465 ms), so it is not skew.
• Ruled out:
  • Disabling activation checkpointing: ~8 GiB per layer at 16k.
  • Fusing the small projections: ~20 ms per step.
  • Tuning DeepEP's send chunk sizes: the best of 144 configs is within 2% of the built-in ones.
  • EP 8 with torch all-to-all instead of EP across all GPUs: 2.5x slower on a 10-layer proxy, because FSDP then moves each GPU's 48 experts per layer every step (~13.6 GB per layer per GPU).
  • EP 8 with DeepEP's NVLink-only kernels: same result (2.8 s vs 1.1 s for EP 32 on the 10-layer proxy). Token traffic inside a node is cheap (~2.4 ms per dispatch at V4.1's shapes); the expert FSDP traffic is not.
  • Selective activation checkpointing at 40 layers: OOM with the default targets; saving only the sparse-attention outputs fits at 132.5 GiB but is slower (4.9-5.0 s).

Tried and parked: two-batch overlap

Branch exp/deepseek-v41-tbo (not part of this PR).

How it works:

• Each packed row splits at its middle document boundary into two micro-batches that move through every layer together.
• DeepEP is split into launch / experts / combine-wait ops whose backward roles mirror the forward, so the backward transfers overlap too.

Results:

• It is correct: losses match.
• Two 8k documents per row, bf16: 5.66 s -> 5.48 s, and only once the expert GEMMs are capped to leave DeepEP its 20 SMs. A cuBLAS SM carveout on top makes it slower.
• FP8: 4.95 s -> 5.03 s, even with DeepGEMM capped.

Why so little: in the trace, DeepEP's internode kernels start a median of 12–16 ms after launch, because they need all their SMs co-resident and wait behind persistent or cooperative GEMMs.

Testing

• 12-node runs above.
• tests/unit/test_configs.py (203 passed).
• prime-kernels tests (52 passed).
• DeepEP vs torch dispatch parity on 2 GPUs under activation checkpointing.
• Op-level parity for dsv4_linear_rope (bf16 and FP8; eager, compiled, SAC), the mHC stream-gradient path, weight_grads_after_combine and the Mega MoE backend; prime-kernels mhc_projection tests (16 passed).
• Run on this push (CPU node, no GPU), on the pushed tree:
  • ruff check and ruff format --check on every changed file: clean.
  • tests/unit/test_configs.py, tests/unit/train/test_perf.py, tests/unit/train/test_pipeline_orders.py: all passed (test_configs, test_pipeline_orders, test_checkpointing and test_deepseek_v4 rerun after the latest commits; test_deepseek_v4's 8 GPU-only tests fail here with no NVIDIA driver, the same 8 as on the base).
  • tests/unit/train/models/test_checkpointing.py: 2 failures, with the same 2 on the base commit. prime_rl::fp8_indexer is not registered without a GPU.
  • tests/unit/train/test_sinkhorn.py and test_muon.py: GPU-only, so they fail here (no NVIDIA driver). Track optim ran them on H200: 11/11 pass.

🤖 Generated with Claude Code