perf(deepseek-v4): fuse and enable the indexer distillation loss - #992
Open
lhzhang333 wants to merge 3 commits into
Open
perf(deepseek-v4): fuse and enable the indexer distillation loss#992lhzhang333 wants to merge 3 commits into
lhzhang333 wants to merge 3 commits into
Conversation
…nels
Building the KL target eagerly gathers the selected pool rows into a
[B, S, K, head_dim] tensor -- 2.1 GB per CSA layer per microbatch at V4-Flash
widths -- writes it to HBM and reads it straight back into the GEMM, then runs
the per-head softmax and head sum as a chain of elementwise kernels around it.
It cost more than the attention it is imitating.
indexer_distill_target indexes the pool inside the kernel, so the gather never
exists and the softmax and head sum happen in registers. It is also more
accurate than the path it replaces: tl.dot accumulates in fp32 where the eager
einsum lands in bf16 first (1.8e-07 vs 2.5e-03 against an fp64 reference).
indexer_distill_kl fuses the KL tail with an analytic backward, so nothing
beyond the inputs is saved for it.
PRIMUS_V4_INDEXER_COMPILE additionally hands the indexer to torch.compile:
training it spends most of its time in the autograd engine walking the ~20
nodes its forward builds, not in arithmetic.
Measured on 4 nodes / 32 MI355X, 10 iterations, coeff=1e-2, against a coeff=0
baseline of 8590.4 ms/iter and 1371.8 TFLOP/s:
eager 10555.3 ms 1116.5 TFLOP/s +22.88%
+ fused target and KL 9056.4 ms 1301.3 TFLOP/s +5.42%
+ compiled indexer 8792.2 ms 1340.4 TFLOP/s +2.35%
Both kernels are opt-out and fall back to the eager body on shapes they do not
cover, so a broken kernel build degrades rather than fails.
(cherry picked from commit f7a10a3)
…t softmax A CSA layer takes one softmax over the concatenation of the sliding-window keys, the sparse compressed entries and the per-head sink. The distribution it places on the compressed entries is therefore a conditional of that joint softmax, and the share of a head's attention that reaches them -- exp(compressed_lse - full_lse) -- is below one and different for every head. The target was instead renormalising each head over the compressed entries alone, which throws that share away and head-sums as if every head weighted the compressed branch equally. On a model that deliberately mixes a sliding window with sparse compressed attention that is the wrong objective: a head that spends almost all of its attention on the window got the same vote on which compressed entries to select as one that actually depends on them. The reference implementation carries the log mass of the non-compressed part into the denominator, and this does the same. That term is free inside the existing target kernel; the cost is computing the mass, a [S, window] score matrix per head. Eagerly that is 2.000 ms -- more than five times the target kernel for a quarter of the arithmetic, because four fifths of the fp32 logit tensor it materialises is masked away. indexer_distill_window_lse keeps the scores in registers and loads only the key band a query block spans: 0.193 ms. End to end on 4 nodes the fix costs 39.4 ms/iter (+0.44%), inside the +-0.5% run-to-run spread: 9017.0 ms with it off against 9056.4 ms with it on. The L1 renormalisation floor moves from 1e-10 to float32 tiny at the same time. With the joint denominator a row's mass is the share of attention reaching the compressed entries, which for an almost entirely local query is legitimately far below the H a per-branch softmax would produce, and 1e-10 would clip it. PRIMUS_V4_DISTILL_NONCOMP_LSE=0 restores the previous behaviour for A/B purposes. Tests pin the target against a joint softmax written out the long way, and cover the empty-row and local-branch-dominates cases that decide whether the arithmetic can produce a NaN. (cherry picked from commit 0af35aa)
… by default The loss was computed, scaled and written into the MoE aux-loss tracker every step, and never shown. Two things were in the way. DeepseekV4HybridLayer did not pass layer_number to build_module for the attention, so every attention module came up as layer 0 -- the "unnumbered" sentinel that log_indexer_distill_loss rejects, since 0 would otherwise index the tracker's last slot. The value was dropped before it was ever recorded. Nothing in the forward depends on layer_number, so only a test keeps it wired. Even recorded, it would not have been printed: training_log passes track_moe_metrics an explicit track_names list built from the MoE router options, and reduce_aux_losses_tracker_across_ranks iterates exactly that list, so a key outside it is written and then zeroed. A patch appends the key, gated on both the model type and the coefficient so it cannot affect another model. Both predicates read args, which is identical on every rank, so the ranks agree on whether the key joins the reduction -- disagreeing would hang that collective. The launcher now defaults the coefficient on. It is part of the recipe rather than an optimization: topk is not differentiable, so without the loss the CSA lightning indexer never receives a gradient and a from-scratch run selects compressed entries at random. Its fused kernels follow PRIMUS_OPT_FUSION, so the ladder's stage 0 still means every optimization off. (cherry picked from commit ce54244)
lhzhang333
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
August 17, 2026 12:43
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| @pytest.fixture(autouse=True) | ||
| def _args_from_ctx(monkeypatch): | ||
| """Point the patch module's ``get_args`` at the fake context.""" | ||
| import primus.backends.megatron.patches.deepseek_v4_indexer_loss_patches as mod |
| """A second install leaves the first wrapper in place.""" | ||
| import megatron.training.training as training_module | ||
|
|
||
| import primus.backends.megatron.patches.deepseek_v4_indexer_loss_patches as mod |
| full CSA forward. | ||
| """ | ||
| monkeypatch.delenv("PRIMUS_V4_INDEXER_TRAINABLE", raising=False) | ||
| import primus.backends.megatron.core.transformer.indexer_distill_loss as idl |
Contributor
There was a problem hiding this comment.
Pull request overview
Improves DeepSeek-V4 CSA indexer distillation by (1) correcting the teacher target to match the layer’s joint softmax (window + sparse + sink), (2) substantially reducing overhead via Triton fusions for the target, KL tail, and window LSE, and (3) ensuring the loss is actually surfaced in training logs and wired correctly by layer number.
Changes:
- Fix target definition to use the joint-softmax conditional (via
noncompressed_lse) and add extensive unit coverage for tricky masking/NaN edge cases. - Add Triton kernels for fused target construction, fused KL(attention||indexer) (with backward), and fused sliding-window log-sum-exp.
- Ensure the loss is reported (training-log patch +
layer_numberpropagation) and enable the loss by default in the V4 flash launcher recipe.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit_tests/megatron/transformer/deepseek_v4/test_v4_indexer_distill_loss.py | Refactors helpers and adds thorough tests pinning the corrected joint-softmax target and Triton/eager parity. |
| tests/unit_tests/megatron/patches/test_v4_indexer_loss_logging_patch.py | Adds tests ensuring the logging patch gates correctly and wraps track_moe_metrics safely. |
| primus/backends/megatron/patches/deepseek_v4_indexer_loss_patches.py | New patch to append indexer_distill_loss to Megatron aux-loss tracking keys when enabled. |
| primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_distill_window_lse.py | New Triton kernel to compute sliding-window + sink log-mass efficiently for the joint denominator. |
| primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_distill_target.py | New Triton kernel to build the KL target without materializing the huge gathered tensor. |
| primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/indexer_distill_kl.py | New Triton fused KL tail with analytic backward to reduce autograd/kernels overhead. |
| primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/init.py | Documents the new distillation-loss Triton modules. |
| primus/backends/megatron/core/transformer/indexer_distill_loss.py | Implements corrected joint-softmax target path, fused dispatch, and noncompressed_lse. |
| primus/backends/megatron/core/transformer/deepseek_v4_attention.py | Wires joint-denominator inputs into the loss call; adds optional indexer compile and diagnostic freeze. |
| primus/backends/megatron/core/models/deepseek_v4/deepseek_v4_block.py | Threads layer_number through attention construction so per-layer logging indexes correctly. |
| examples/deepseek-v4/run_deepseek_v4_flash.sh | Enables the distillation loss by default for flash recipe and documents reproduction/perf toggles. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+222
to
+228
| def can_use_triton_target(*, query: torch.Tensor, topk_idxs: torch.Tensor) -> bool: | ||
| """Whether the fused kernel covers this shape / configuration.""" | ||
| if os.environ.get(_ENABLE_ENV, "1") != "1": | ||
| return False | ||
| if not query.is_cuda: | ||
| return False | ||
|
|
Comment on lines
+188
to
+196
| has_sink = sink is not None | ||
| if has_sink: | ||
| sink_t = sink.detach().to(torch.float32) | ||
| if sink_t.numel() != H: | ||
| raise ValueError(f"sink must hold {H} head values, got {sink_t.numel()}") | ||
| sink_t = sink_t if sink_t.stride(-1) == 1 else sink_t.contiguous() | ||
| else: | ||
| # Triton needs a real pointer for the unused argument. | ||
| sink_t = q |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The CSA lightning indexer is trained by a KL against the attention distribution
over the compressed entries it selected. Three things were wrong with it: the
eager implementation cost more than the attention it imitates, the target it
matched against was the wrong distribution, and the loss it produced was never
shown. This branch fixes all three.
ffd8d8dbperf: fuse the KL target and the KL tail into Triton kernelsda3d4b2ffix: build the target from the layer's joint softmax, add the sliding-window LSE kernelbb464e1efeat: report the loss in the training log, default the coefficient onWhy the target was wrong
A CSA layer takes one softmax over the concatenation of the sliding-window keys,
the sparse compressed entries and the per-head sink. The distribution it places
on the compressed entries is therefore a conditional of that joint softmax, and
the share of a head's attention that reaches them --
exp(compressed_lse - full_lse)-- is below one and different for every head.The target was instead renormalising each head over the compressed entries alone,
which throws that share away and head-sums as if every head weighted the
compressed branch equally. On a model that deliberately mixes a sliding window
with sparse compressed attention that is the wrong objective: a head that spends
almost all of its attention on the window got the same vote on which compressed
entries to select as one that actually depends on them.
Behaviour change reviewers should notice
examples/deepseek-v4/run_deepseek_v4_flash.shnow defaultsPRIMUS_V4_INDEXER_DISTILL_LOSS_COEFFto1e-2instead of0.0. This is partof the recipe rather than an optimization:
topkis not differentiable, sowithout the loss the indexer never receives a gradient and a from-scratch run
selects compressed entries at random. It costs +2.52% per iteration, so the
throughput numbers in that script's speedup ladder now need
PRIMUS_V4_INDEXER_DISTILL_LOSS_COEFF=0to reproduce; the comment block says so.run_deepseek_v4.shstill defaults it to0.0, so no other V4 launcher changesbehaviour.
Performance
4 nodes / 32x MI355X, 10 iterations, TP=1 PP=4 EP=8, GBS=256, seq=4096,
MOE_FORCE_LB_TYPE=uniform, imagepr-927. Each rung adds one switch to the oneabove it, so every kernel's contribution is separable.
Eager to default is −17.00% per iteration, and the loss now costs +2.52%
over having it off, down from +23.53%.
Peak memory (
max reserved, the three PP ranks that report it): turning the losson costs about 10 GB per rank (242.9/241.6/196.0 GB to 252.6/252.3/203.1 GB) --
that is the unfrozen indexer's gradients and optimizer state, and it is the same
whether the kernels are on or off. The fused target removes a 2.1 GB gather per
CSA layer per microbatch, but as a short-lived temporary the caching allocator
just reuses that space, so the win shows up in bandwidth rather than in the peak.
Test plan
Unit tests, in
pr-927on one MI355X node:test_v4_indexer_distill_loss.py-- 51 passed. Extended by this branch topin the target against a joint softmax written out the long way, and to
cover the empty-row and local-branch-dominates cases that decide whether
the arithmetic can produce a NaN.
test_v4_indexer_loss_logging_patch.py-- 15 passed (new file).tests/unit_tests/backends/megatron/patches/-- 21 passed.test_training_log_patches.py-- 7 passed.tests/unit_tests/megatron/transformer/deepseek_v4/-- 11 failed, 605passed, 81 skipped. Same directory on
maingives 11 failed, 578 passed,81 skipped, and the sorted
FAILEDidsdiffempty: the same 11 cases.The 27 extra passes are this branch's new cases. Of those 11, nine are
FlyDSL hitting
lld invocation failedwhen lowering to gfx950 (thetraining path works around it via
runner/helpers/patches/11_fix_lld_stub.sh, which a baredocker runofpytest does not apply), and two are
'DeepseekV4MoE' object has no attribute 'mega_moe_experts'.No impact on anything else:
coeff=0,indexer_distill_enabledis false, socompute_indexer_distill_lossis never called and the kernels are not evenimported. This branch and
mainprint an identicallm lossat iteration 1and identical peak memory on all three reporting ranks.
configuration is not bitwise reproducible (MegaMoE's all-to-all and RCCL's
reductions are not order-deterministic). Re-running one commit twice gives
a max relative
lm lossdifference of 3.69e-06 against 7.97e-06 across thetwo branches -- same order, and at iteration 9 the repeat sits closer to
mainthan the original does, so the difference has no direction.model_type == "deepseek_v4"andcoeff > 0. Both predicates readargs, which is identical on every rank,so the ranks agree on whether the key joins
reduce_aux_losses_tracker_across_ranks-- disagreeing would hang thatcollective.
_triton_common/, which has nopackage-level re-export; each is imported by full submodule path, so
adding files there cannot affect an existing kernel at import time.
Reproduce