Dispatch the CSA Compressor gated pooling to the cudnn-frontend fused kernels (THD) - #5984
Merged
Merged
Conversation
Contributor
|
Hi @zkyue Thanks for your contribution! I think it is better to put the kernels in cudnn-frontend like other CSA/HCA fused kernels. Is it convenient for you to do so, or do you need our help to port the kernel? |
Author
|
@hxbai Sure, happy to do that — the kernels are self-contained CuTe DSL and should map cleanly onto the cudnn-frontend python API structure (we've contributed there before). Plan: I'll open a PR against cudnn-frontend porting the fwd+bwd kernels (APIBase-style wrapper, THD semantics, tests), then rework this PR into a thin dispatch that uses the cudnn-frontend implementation when available, keeping the eager fallback. Will link the kernel PR here once it's up. |
Author
|
@hxbai The cudnn-frontend kernel PR is up: NVIDIA/cudnn-frontend#427 (fwd+bwd kernels, APIBase-style wrapper, THD semantics, 40 tests; outputs bit-identical to the kernels in this PR). Once it lands I'll rework this PR into the thin dispatch that uses the cudnn-frontend implementation when available, keeping the eager fallback. |
Anerudhan
pushed a commit
to NVIDIA/cudnn-frontend
that referenced
this pull request
Jul 30, 2026
…rom Megatron-LM) (#427) * CSA: add fused Compressor forward+backward CuTe-DSL kernels Port the fused CSA/HCA Compressor gated-pooling kernels from Megatron-LM (NVIDIA/Megatron-LM#5984) into the FE-OSS Python API, per the maintainer request in NVIDIA/Megatron-LM#5984 (comment). One forward and one backward CuTe-DSL kernel fuse the THD gated-softmax pooling region (gather -> +APE -> overlap-window transform -> fp32 softmax -> gated weighted sum -> bf16 cast). Kernel math is unchanged from the Megatron original (verified bitwise on forward/dKV/dScore); the interface is reshaped to the APIBase pattern: CSACompressorForward / CSACompressorBackward classes, csa_compressor_forward_wrapper / csa_compressor_backward_wrapper, a new python/cudnn/csa package, docs, and fe_api tests (numerics vs fp32/upstream eager references and an fp64 oracle, ragged packs, static-capacity padding, run-to-run determinism, CUDA-graph capture/replay, check_support boundaries). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: vectorize forward bf16 access (32-bit, CuTe autovec_copy) Widen every bf16 load/store in the fused Compressor forward kernel from scalar (16-bit) to paired 32-bit accesses: each thread now owns 2 adjacent head dims and moves its window slices through 32-bit universal copies (cute.autovec_copy over register fragments; cute.assume makes the alignment provable). The per-lane fp32 math is unchanged and in the same order, so the forward output stays bitwise identical to the previous kernel (re-verified against the Megatron-LM original on the same 5 THD shape-mix gate as the original port). Odd head_dims keep a scalar (vec == 1) instantiation of the same kernel; the head_dim 128/512 production shapes take vec == 2. The launch schedule moves from (rows_per_cta=4, threads=128) to one row per CTA with 64-thread column groups, which widens the sub-wave grids that limit the small shapes. Measured (nsys pure kernel time, B200, ratio=4 coff=2, THD packs of 8192-token sequences, forward kernel only): shape before after speedup 1x8192 d128 6.0 us 4.5 us 1.36x 3x8192 d128 12.6 us 10.0 us 1.26x 1x8192 d512 16.0 us 12.4 us 1.29x 3x8192 d512 42.1 us 35.0 us 1.20x Alternatives measured and rejected on the same matrix (all bitwise-equal): a PTX-unpack vec2 prototype (slower than the pure-DSL version: 4.7/10.5/12.9/35.7 us), and 64/128/256-bit per-thread vectors, which cut executed instructions but lose to register pressure and occupancy (80/147/255 regs vs 48; the 256-bit variant spills). Registers stay spill-free at 48 (previous kernel: 40). Tests: numerics shape matrix extended with an odd head_dim case to cover the scalar-layout instantiation, and a check_support rejection for head_dims beyond the launch bound. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: fold dKV/dScore zero-init into the backward kernel The backward previously required zero-initialized dKV/dScore buffers: the kernel only wrote consumed positions, and two tensor-wide bf16 fill kernels (torch.zeros_like) supplied the zeros everywhere else. Those fills cost as much DRAM traffic as the kernel's own stores. The kernel now writes every position itself: consumed positions get their gradients as before (disjoint, atomic-free stores), and each never-consumed slot class gets an exact zero from a unique natural owner, keeping all stores disjoint and dKV/dScore bitwise run-to-run deterministic: - first-half columns of each segment's last block's own tokens (coff == 2; no next block consumes them) -> that last block; - per-segment tail tokens (seqlen % ratio, both halves) -> the segment's last block; - all tokens of segments with zero output blocks (seqlen < ratio) -> CTA column bidx == 0; - tokens beyond cu_seqlens[-1] (static token-capacity padding of the gradient buffers, the CUDA-graph case) -> grid-strided across CTA columns. The backward wrapper allocates grad_kv/grad_score with torch.empty_like instead of torch.zeros_like (the fp32 dAPE fill stays: it is the atomic accumulator). When total_comp == 0 the kernel cannot launch, so the wrapper falls back to zeroed allocations to preserve autograd's exact-zero semantics; the class API documents the same caveat. Measured on the backward region (kernel + remaining fills vs kernel + three fills before), B200, ratio=4 coff=2, THD packs of 8192-token sequences; nsys = sum of kernel durations per iteration, wall = CUDA events around the wrapper call: shape nsys before -> after wall before -> after 1x8192 d128 15.5 -> 12.8 us (1.21x) 56.2 -> 47.2 us (1.19x) 3x8192 d128 29.9 -> 22.2 us (1.35x) 61.8 -> 54.6 us (1.13x) 1x8192 d512 34.2 -> 22.8 us (1.50x) 67.9 -> 57.3 us (1.19x) 3x8192 d512 90.9 -> 66.0 us (1.38x) 111.6 -> 101.8 us (1.10x) The zero-writes cost the kernel 8 registers (64 -> 72, occupancy 44% -> 41% at 3x8192 d512) and a small amount of extra store traffic, repaid several times over by dropping the two whole-tensor fills. dKV/dScore remain bitwise identical to the Megatron-LM original on the 5-shape port gate and to the fp32 eager reference on every test shape. Tests: NaN-canary suite proving the kernel fully overwrites uninitialized buffers (ragged, degenerate short segments, all-tiny packs, head_dim 512, static-capacity padded rows) bitwise against zero-initialized runs and the eager reference; a zeros-fallback test for packs where no segment reaches ratio tokens. The existing CUDA-graph replay test covers the token-capacity padding class (it fails without the grid-strided sweep). Docs: csa.md updated (buffer contract, fill count, perf tables refreshed for both this commit and the forward vectorization). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * Address review feedback: docs bounds, lint, warning stacklevel - docs: list the total_comp * head_dim < 2**31 and total_comp > 0 with total_tokens < ratio support boundaries already enforced in check_support - docs: run the compressor tests from test/python per repo convention - csa/__init__.py: spell out __all__ as an explicit literal (Ruff PLE0604) - compressor/__init__.py: sort __all__ (Ruff RUF022) - compressor/api.py: add stacklevel=2 to the deterministic-mode RuntimeWarning Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: rewrite graph benchmark to fwd-only + fwd+bwd total graphs Address review feedback on the PR #427 fairness benchmark (benchmark/csa/bench_csa_compressor.py). - Drop the subtraction-derived backward-graph estimate. Graph variants are now two directly measured, symmetric columns for both eager and fused: a forward-only graph (eager under no_grad) and a forward+backward total graph. - Capture the eager fwd+bwd graph with stable, pre-allocated zero .grad buffers (the supported torch CUDA-graph pattern): the captured region zeros them in place, then runs forward + autograd backward, so every replay accumulates into a zeroed buffer -- numerically identical to a single fresh backward. A per-shape graph-vs-fresh-backward cross-check is printed (all shapes bitwise-equal). - The fused total graph captures the forward wrapper immediately followed by the backward wrapper. - Emit both the per-call and graph tables from a single run; graph speedups are eager/fused of the displayed us, truncated to one decimal (never rounded up). The backward replay is reported only as ~total-fwd (a reference), never as a column. Not collected by pytest; black -l160 clean. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * docs(csa): refresh graph table to fwd/total columns; soften wording Address review feedback on the CSA compressor performance tables (docs/fe-oss-apis/csa.md). - Replace the CUDA-graph table (which published a subtraction-derived backward column) with two directly measured columns: eager/fused forward-only graph and eager/fused forward+backward total graph. No backward-graph column; backward is noted only as ~total-fwd, never as a measured value. - Refresh both the per-call and graph tables from a single run of the benchmark harness; graph speedups are eager/fused of the displayed us, truncated to one decimal. - State that capture collapses each side's per-op launches into a single replay rather than removes launch/host overhead; say less launch-bound rather than compute-bound. - Note how the re-run reproduces the previously published per-call wall clock, and document the eager-backward capture basis (stable zero .grad buffers). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * bench(csa): guard speedup division, dedupe leaf construction, ruff nits Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: widen the validated envelope to coff in {1, 2} The kernels are already generic over coff — for coff=1 the window is the block's own ratio tokens (no overlap transform, always valid); only the check_support gate restricted the APIs to the production coff=2 form. Widen the gate to coff in {1, 2} (ratio stays gated at 4), update the error message and the envelope prose in api.py / docs / README, and parametrize the tests over both coff forms: numerics vs the fp32/upstream/fp64 eager references (the eager reference already implements coff=1 faithfully — it skips the overlap transform), static- capacity padding, replay determinism, empty outputs, NaN-canary zero-writes, deterministic-mode error paths, CUDA-graph capture, class-vs-wrapper equivalence, and check_support accept/reject boundaries (new coff=0 / coff=3 rejects; a new empty-segment shape covers both coff forms). No kernel changes. Addresses hxbai's coff=1 question on the PR. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA: add missing docstrings across the compressor Python files Bring the PR's Python docstring coverage over CodeRabbit's 80% pre-merge gate: module docstrings for the csa package/compressor __init__ files and one-line docstrings for the remaining undocumented helpers in api.py (cache/ctor/compile plumbing), compressor_sm100.py (launcher ctors), the test helpers, and the benchmark helpers. interrogate 1.7.0 over the PR's six CSA Python files: 69.4% -> 100%. Docstrings only — no logic, kernel, or launch-path changes. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> * CSA compressor: add dedicated ratio=128 kernels (forward+backward) with ratio dispatch The generic CSA compressor kernels keep the whole coff * ratio pooling window in per-thread registers; past ratio ~32 that hits the 255-register wall and spills kilobytes of local memory per thread, so ratio=128 needs dedicated kernels rather than a lifted gate. This adds them behind the same public API: check_support now admits ratio in {4, 128} x coff in {1, 2} (ratio=128 additionally gated to head_dim in {128, 512}, the hardware-validated set) and the wrappers and class APIs route to the matching kernel family by ratio transparently. Design: the forward streams the window with a chunked online softmax (one CTA per output row, per-column (max, denom, acc) triples merged in one fixed smem order); its launch schedule is bucketed by output-row count (small/default/large, all precompiled for CUDA-graph safety) and per bucket selects two-phase accumulation and an ex2.approx fast exp where measured faster. The backward stages each row's window through shared memory chunk-parallel, fuses the per-chunk den/S partial sums into the same pass, merges partials in a fixed chunk order, and stores gradients with a hoisted 1/den multiply — keeping kernel-side zero-writes to never-consumed dKV/dScore slots and fp32-atomic dAPE exactly as at ratio=4. ptxas (sm_100a): forward 32-51 registers, backward 48-128, 0 spill / 0 stack across all 16 shipped (config, schedule) kernels (reproducer: benchmark/csa/reg_probe_csa_compressor_r128.py). Numerics contract, per ratio family (docs/fe-oss-apis/csa.md): ratio=4 keeps its bitwise contract UNCHANGED — dKV/dScore bit-identical to the fp32-intermediate eager reference, at coff 1 and 2. ratio=128 is deterministic + faithful to that same fp32-intermediate eager reference: bitwise run-to-run determinism of out/dKV/dScore (NaN-prefill replay tested); the same values within final-bf16 rounding at the gate tolerances (differing elements <= max(1, 0.1%), max_abs <= 1.6e-2, calibrated on the gate's documented input distribution — bf16 deviations scale 2^k with 2^k inputs while the fp32 intermediates stay finite); the eager reference's non-finite propagation (finite inputs that overflow its fp32 intermediates poison both sides alike, with the stated caveat that the fused order saturates earlier near fp32 max — un-normalized chunk partials); and, on inputs whose fp32 intermediates stay finite in both evaluation orders, fp64-oracle parity (at least as close to an fp64 oracle as the eager reference). The reduction reorders and fast-exp buckets were each adopted on a measured win and are covered by that contract. Performance (nsys isolated GPU kernel time, 1x B200, vs the fp32-intermediate eager reference region on identical inputs): forward 13.2-44.4x, backward 7.4-21.9x across coff {1, 2} x head_dim {128, 512} x 8k-131k tokens; e.g. coff=2, d=128, 65k tokens: fwd 726.3 -> 16.4 us (44.4x), bwd 960.0 -> 61.2 us (15.7x). Full tables with methodology in docs/fe-oss-apis/csa.md. Validation (B200, CC 10.0): the committed 21-case contract gate (benchmark/csa/gate_csa_compressor_r128.py) passes 21/21 with shipped schedule coverage asserted (forward 10/10, backward 6/6 buckets); pytest fe_api/csa/test_CSA_compressor.py: 88 passed, 1 skipped at the default L0 level and 98 passed, 1 skipped with -m "L0 or L1" (both ratio families; the ratio=4 rows, including the coff=1 envelope, stay on their bitwise assertions). New tests cover the ratio=128 dispatch envelope (schedule selection at every nb_total bucket boundary; every shipped (config, schedule) kernel executed against the full contract at L1), exact-zero assertions on every never-consumed dKV/dScore slot class in the NaN-canary, CUDA-graph capture/replay for the new family, and grad_ape zeroing-ownership regressions. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com> --------- Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
…ressor The fused forward+backward kernels for the gated-pooling region of Compressor._forward_thd now live in the cudnn-frontend Python package (cudnn.csa.compressor, cudnn-frontend PR NVIDIA#427). Megatron-side this is a thin additive dispatch: maybe_compress_thd_fused() returns the pooled tensor when the frontend is importable and the configuration is supported (THD non-pre-grouped path, compress_ratio == 4, bf16 kv/score, fp32 ape, compute capability 10.0, int32 flat offsets) and None otherwise, in which case the eager region runs unchanged. SBHD and the pre-grouped CP-prep path are untouched. Frontend availability is probed by importing the concrete entry points (csa_compressor_forward_wrapper / csa_compressor_backward_wrapper), not by version comparison: nvidia-cudnn-frontend installs that predate the CSA compressor API lack cudnn.csa and keep the eager path, as does any import failure (e.g. the DSL extra missing). The dispatch also keeps the eager path under torch.use_deterministic_algorithms(True) (dAPE fp32 atomics) and under torch.compile tracing (raw-pointer launch path is not traceable). MCORE_CSA_FUSED_COMPRESSOR=0 disables the dispatch entirely. compress_ratio == 128 is functionally supported by the frontend kernels but stays on the eager path until tuned (see issue NVIDIA#5968). Addresses NVIDIA#5968. Signed-off-by: zky <kaiyue.zhou@z.ai>
The kernels themselves are validated in cudnn-frontend (PR NVIDIA#427); these tests cover the Megatron-side wiring: numerics of the dispatched fused region vs an fp32-intermediate eager reference (bitwise dKV/dScore, forward within one bf16 rounding step) and vs the verbatim upstream eager numerics (tolerance) over ragged THD packs including segments shorter than ratio; fixed_total_comp static-capacity padding rows (including that nonzero padding-row gradients are ignored); dispatch gating and eager fallback (kill switch, missing/old cudnn-frontend, deterministic mode, ratio 128, non-bf16, layout, empty output), plus Compressor._forward_thd-level integration (fused engages and matches eager, bitwise eager fallback without the frontend, gradients flow). The module is marked launch_on_gb200 for the CC 10.0 CI lane and skips cleanly without CUDA, without a cudnn-frontend that provides cudnn.csa, or off compute capability 10.0. Addresses NVIDIA#5968. Signed-off-by: zky <kaiyue.zhou@z.ai>
zkyue
force-pushed
the
feat/csa-fused-compressor
branch
from
July 30, 2026 07:36
8217542 to
e073b28
Compare
zkyue
changed the title
Author
|
cudnn-frontend #427 merged this morning ( Force-pushed |
Contributor
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Strict review passed — no significant issues found.
Scope: thin, additive dispatch shim (+732/−37) routing the THD non-pre-grouped CSA gated-pooling region to the cudnn-frontend fused kernels, with the eager region preserved verbatim as the reference/fallback.
Verified sound:
- Autograd arity —
forwardtakes 9 non-ctx args;backwardreturns exactly 9 grads (3 real + 6None). Correct. - Shape parity — fused
out.unsqueeze(1)→(total_comp, 1, head_dim), identical to the eager(kv_grouped * weights).sum(dim=1). - Gating invariants —
ratio == 4 ⟺ coff == 2 ⟺ overlapare locked together in__init__, so the fused path cannot be reached with a coff/overlap mismatch. Gates also correctly cover dtype (bf16 kv/score, fp32 ape), device (CC 10.0, cached), int32-offset overflow (total * coff * head_dim < 2**31), deterministic mode, andtorch.compiletracing. - No hot-path sync — all dispatch checks are host-side (
.shape, env read, cached capability lookup); no CPU↔GPU sync introduced. On non-CC-10.0 hardware the probe short-circuits to eager after a one-time cached import. - Import safety — the frontend import is lazy inside
_get_frontend();csa.py's import chain never fails on a missing/old/partial frontend. - Backward compat — no renamed/removed args, no checkpoint or API signature changes. H100/A100 (CC ≠ 10.0) stay bit-identical on eager. The
dsv4_hybrid_fusedfunctional case is commented out in the gb200 recipe, so no golden-value impact. - Unused-variable check — all new identifiers have real use paths;
_frontend_erroris a documented debugging aid.
One point for maintainer decision (not a defect): the fused path is on by default, so GB200 users with cudnn-frontend ≥ 1.27 get a silent numerics change (fp32-intermediate vs the eager bf16-weight multiply — documented as at-least-as-accurate against an fp64 oracle). The author has already flagged this and offered to flip to opt-in; noting only so the choice is made deliberately.
Risk: low. Additive, capability-gated, no-op against any cudnn-frontend predating the CSA API, and covered by the new dispatch/gating/fallback tests. LGTM.
hxbai
reviewed
Jul 31, 2026
hxbai
reviewed
Jul 31, 2026
hxbai
reviewed
Jul 31, 2026
…ssing-lib warning Addresses the review on NVIDIA#5984: * Gate on the frontend's validated envelope -- compress_ratio in {4, 128} x coff in {1, 2}, with ratio 128 restricted to head_dim in {128, 512} -- instead of only the two pairs Compressor happens to produce today. The dispatch follows the kernels rather than Compressor's overlap-from-ratio derivation, so a future overlap-policy change keeps the fast path instead of silently falling back to eager. Measured on B200 at ratio 128 / coff 1: 9.3-11.8x end-to-end with both sides captured into CUDA graphs, and 13.3-13.6x forward / 15.8-21.6x backward in isolated kernel time (nsys). Accuracy improves too: max absolute error against an fp64 oracle is 1.4-2.6x smaller than the eager region's across all four combinations, matching an fp32-intermediate eager reference. * Replace the MCORE_CSA_FUSED_COMPRESSOR env var with the existing config switch: Compressor.use_fused_compressor = use_fused_dsa_kernels(config), passed to the dispatch as enabled=. * Warn once per process when the device is supported but the cudnn-frontend CSA compressor API is unavailable. The device-capability check runs before the frontend probe so the warning does not fire on hardware where eager is the expected path. Unit tests cover all four (ratio, coff) combinations, the out-of-envelope head-dim and ratio fallbacks, and keep the missing-frontend fallback assertion effective by scoping the disable to the eager comparison only. Signed-off-by: zky <kaiyue.zhou@z.ai>
Author
|
@hxbai Thanks for the review — all three addressed and pushed. Everything below was measured on a B200 (CC 10.0, cudnn-frontend at #427, CUDA 13). 1. ratio 128 / coff 1Yes — and I widened the gate to the frontend's full validated envelope rather than adding just that one pair: Speed. The region captured into CUDA graphs, both sides symmetrically (forward+backward as one replay). This removes the per-op launch overhead from eager and fused alike, which is the fairest single basis for these launch-bound shapes:
Isolated kernel time (nsys, sum of kernel durations per iteration, 50 iters after 20 warmup, backward includes the dAPE zero-fill) — same methodology as the ratio-4 table in
Per-call wall clock (CUDA events, launch overhead included) is 5.0-5.3x at 1x8192/d128 and 5.8-11.1x at 1x131072/d128. The three tables are different measurement bases and are not comparable to one another. As a check on the harness, I re-ran the published ratio-4 / coff-2 configuration on the same box: it reproduces the per-call and graph tables in The Numerics. Worth stating explicitly, since the two ratio families differ: at ratio 4, dKV/dScore are bit-identical to the fp32-intermediate eager autograd; at ratio 128 it is a tolerance contract (frontend gate: differing elements <= max(1, 0.1%), max_abs <= 1.6e-2). Both are now in the module docstring. In the other direction, keeping the softmax weights in fp32 makes the fused path more accurate than the region it replaces — max absolute error against an fp64 oracle on a 4096-token pack:
Across all eight (ratio, coff, head_dim) combinations the fused error is 1.4-2.6x smaller than the eager region's and matches the fp32-intermediate reference — what remains is the single final bf16 rounding. 2. Config arg instead of the env varDone — One consequence worth confirming, because it changes the default: 3. Warning when the library is unavailableAdded, with two refinements so it stays signal rather than noise: the device-capability check now runs before the frontend probe, so the warning only fires where the kernels could actually have run (CC 10.0) instead of on every H100/A100 forward; and it is one-shot per process, with the underlying import error included. Validation
Note that the push changes the head SHA, so the CI authorization on |
Contributor
|
/ok to test bdfcca9 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30677998507 |
xuwchen
added a commit
to xuwchen/Megatron-LM
that referenced
this pull request
Aug 18, 2026
…shard mapping
GTP shards dim0 of the fused [gate|up] fc1 weight, but the swiglu checkpoint
factory used to run per-shard: each shard was chunked into [gate_i|up_i], so
the checkpoint layout depended on the save-time GTP degree and, worse, the
storage mapping it implied was interleaved -- the all-gathered weight was NOT
the logical TP-local tensor, and every consume site needed a de-interleave
permutation (plus the matching wgrad-side inverse when syncing gradients).
Those two runtime transforms cost extra kernels on every layer consume and
were the prime suspect in the GTP-vs-MFSDP perf asymmetry.
Replace the runtime-transform approach with the same storage re-layout that
Mamba/GatedDeltaNet in_proj already uses: pin the mapping "GTP shard i ==
contiguous rows [i*n, (i+1)*n) of the logical TP-local weight" via the
checkpoint wiring, and delete both runtime transforms entirely.
- New megatron/core/tensor_parallel/gtp_ckpt.py hosts the gather/slice pair
(moved verbatim from ssm/utils.py, which now re-exports them): on save,
all-gather the shards back to the unpadded TP-local tensor, THEN run the
standard apply_swiglu_sharded_factory on it -- identical on-disk format to a
non-GTP run; on load, wrap merge_fn to cat -> re-pad -> slice this rank's
contiguous rows.
- MLP.sharded_state_dict wires the pair around the swiglu factory when fc1 is
a GTP param (covers dense MLP and SharedExpertMLP). Grouped routed-expert
fc1 under EGTP keeps its NotImplementedError guard until the grouped wiring
is ported.
- The gathered weight is now logically ordered by construction, so the
consume-site de-interleave and the wgrad re-interleave (and their
transformability checks) are deleted: zero per-iteration cost, and no
transform pair to keep in sync.
Disk format is unchanged (gate/up halves of the full TP-local weight), so
checkpoints interop with 3D-parallel runs and existing checkpoints load
correctly under the new mapping.
Tests: TestGTPGatedFusedLayout now seeds shards as contiguous slices of a
broadcast logical weight and asserts elementwise fwd equality vs F.linear plus
per-rank wgrad slices (permutation-sensitive, unlike grad-norm checks); new
TestGtpFc1SwigluDcp asserts the swiglu factory shapes/offsets, merge round-trip
(pad rows re-zeroed) and DCP writer election under TP2xGTP2 with alignment
padding.
fix(gtp): give expert grads the GTP/EGTP share of the DP normalization
GTP carves its ranks out of the data-parallel axis, so DDP's 1/DP scaling
shrinks as GTP grows. Dense params recover the missing factor from the
gtp_remat AVG; expert params only ever see compensation sized to the EGTP
axis, leaving their gradients GTP/EGTP times too large. Expert layers were
therefore stepped that many times too far while dense layers were correct,
which no learning-rate change can compensate.
The failure is silent -- no error, no NaN, just expert gradients off by a
constant -- so a short run cannot see it and only convergence can. The CPU
tests therefore pin the arithmetic directly with stand-in groups: the
GTP/EGTP factor is applied to expert params (allreduce=False) and not to
dense ones, nothing happens when GTP does not exceed EGTP or the groups are
absent, calculate_per_token_loss is skipped, and params without a grad or
with requires_grad=False are left alone.
test(gtp): cover the multimodal vision->embedding->language chain ordering
Multimodal models put the vision tower before the language embedding on one
prefetch chain, and the embedding is a forward-only chain member: it links on
its on-demand forward gather but its backward (GTPEmbeddingWeight) never enters
all_gather_and_prefetch_bwd, so it produces no gather for its chain predecessor.
The per-consume gather rule from #6242 is what keeps this ordering correct —
the vision tail's backward consume finds no gather issued for it and falls back
on demand. Pin that with bitwise ground-truth comparisons on every consume, in
both directions, so a regression shows up as wrong data here rather than as a
silent stale read in a full multimodal run.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fix(gtp): import itertools.chain in the GTP sharded-param backfill
_backfill_gtp_sharded_param_map (upstream #4967) iterates
chain.from_iterable(float16_groups) but never imports chain, so the first
distributed-Muon save with an unmatched GTP param dies with
`NameError: name 'chain' is not defined` instead of backfilling.
tests/unit_tests/generalized_tensor_parallel/test_gtp_muon_dcp.py hits it
on its first optimizer sharded_state_dict call.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
feat(gtp): make the weight-remat alignment pad configurable
GTPRematConfig.pad_for_alignment defaults to 16 and is only ever overridden by the
three quantized branches in configure_gtp_remat_from_recipe, so a bf16 run pads
without any way to say otherwise. The pad rows are stripped before every GEMM, but
they are not invisible: they widen a GTP tensor's global shape in the checkpoint
(which is why a GTP checkpoint cannot be read by a non-GTP run), and on load they
are re-zeroed rather than round-tripped.
Add --gtp-remat-pad-for-alignment so a run can pin or disable it. This is what
lets an experiment separate 'the pad rows cause X' from 'something else does'
without editing the source.
CPU tests pin the precedence: each quantized recipe keeps its historical pad,
bf16 leaves the pre-existing value untouched, and an explicit value (including 0
to disable) wins over every recipe.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fix(gtp): make GTP produce a loadable fully-reshardable optimizer checkpoint
dist_ckpt_optim_fully_reshardable is the only distributed-optimizer format whose state
survives a change of parallel layout: every other format prefixes its entries with
`optimizer.distributed.dp_group_idx_N.`, pinning them to the DP grouping in force at
save time. GTP could neither write that format nor be read back from it, so a GTP run
could not hand its optimizer state to a differently-sharded run in either direction.
Two things were in the way, and fixing either alone leaves the capability unusable.
1. The save aborted. sharded_param_state_fully_reshardable matches a grad-buffer
parameter to its model ShardedTensor by object identity, and GTP puts a different
tensor in the model entry twice over: a dequantized BF16 copy for native-FP8 weights,
and a factory exposing the GATHERED tensor for GDN/Mamba in_proj. The latter raised
"Model param ...in_proj.weight ... not in model_sharded_state_dict" and stopped the
save outright. Resolve both the way _backfill_gtp_sharded_param_map already does for
the distributed-Muon path — follow the _gtp_dequant_src backlink, else rebuild the
per-shard ShardedTensor — and refuse the expert-parallel rebuild that would write
duplicate shards across EP groups. That helper documents distributed Adam as
unaffected, which holds for dp_reshardable (it never consults the model entries) but
not for this format, which does.
2. What it wrote could not be read. The model side gathers GDN in_proj back to TP-local
before splitting into [z|x|B|C|dt], so its keys match a non-GTP run; the optimizer did
not follow, and keyed the same state '...in_proj.weight' where a non-GTP run writes
'...in_proj.weight.z'. Reuse the model's own factory and gather the optimizer state
across the GTP group to the width it expects, stripping the alignment pad. The state
arrives on CPU from the DP gather while the GTP group is NCCL, so the collective runs
on the parameter's device and returns; mem_efficient is off under GTP, which means the
world tensors reach every rank and the collective is symmetric.
Name the parameter when the match fails. The error printed the tensor's values and
neither the parameter nor what the identity-keyed map actually holds, which is what a
reader needs to see that the model exposed a different tensor in its place.
CPU tests pin the resolver's contract: plain params stay the caller's problem, the
dequant backlink and the by-name factory match resolve, and the EP-unaware rebuild
refuses expert params.
Verified on the VL MTP proxy: the checkpoint saves with no dp_group_idx keys, a
same-topology resume loads and continues, and a 3D-saved checkpoint loads into a GTP run.
A no-GTP control saves cleanly in this format, which is what established the gap was
GTP's rather than the format's.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fix(gtp): checkpoint GDN input projection across GTP shards
The semantic [q|k|v|z|beta|alpha] boundaries of GatedDeltaNet's in_proj can cross
GTP shard boundaries, so a per-shard save would write a layout that depends on
the save-time GTP degree.
Save gathers the physical shards back to the logical TP-local projection before
the factory split; load pads then slices back to this rank's shard. Both reuse
the shared ssm/utils helpers that MambaMixer's in_proj already goes through, so
the checkpoint layout is identical to a non-GTP run and independent of the
topology it was written from, and the GTP rank rides in replica_id so writer
election stays correct when dp_cp_group excludes the GTP axis.
Upstream wires GTP into mamba_mixer but not gated_delta_net, which is the
attention variant Qwen3.5 uses.
Test builds a GatedDeltaNet under TP2 x GTP2 with nonzero alignment padding and
verifies the saved factory splits into the six logical TP-local sections, the
load-side merge reconstructs this rank's physical shard (pad rows re-zeroed),
and exactly one GTP peer wins DCP writer election per chunk.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
refactor(ssm): share the GTP fused-projection checkpoint helpers
Mamba's in_proj GTP checkpoint handling — gather the GTP shards back to the
logical TP-local width before the semantic split on save, re-pad and slice the
merged tensor back to this rank's shard on load — is exactly what GatedDeltaNet
needs next, so hoist both blocks into megatron/core/ssm/utils.py alongside
_split_tensor_factory and switch MambaMixer to the helpers.
Two deliberate strengthenings ride along:
- The GTP rank is folded into replica_id ((0, gtp_rank, dp_cp_rank) instead of
(0, 0, dp_cp_rank)). The gathered tensor is replicated across the GTP peers,
and when metadata['dp_cp_group'] is a GTP-excluded group (explicit
pg_collection grids pass pg_collection.dp_cp) the peers share a dp_cp rank —
without the fold, two ranks would both win DCP writer election for the same
chunk. mpu-derived grids, where dp_cp subsumes the GTP axis, are unaffected:
the extra slot only ever separates ranks that previously collided.
- The load-side merge refuses a non-2-D tensor instead of padding/slicing a
flattened buffer; only the unflattened model-weight factory is supported
(optimizer state resolves through the per-shard rebuild, never this merge).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
test(gtp): pin the zero_out_wgrad contract for GTP's dummy wgrad
The fix itself now comes from upstream #6565, which landed the same zero= pairing
in _handle_megatron_grad_accum. These two unit tests are ours and upstream has no
equivalent: one asserts the flag is forwarded to get_dummy_wgrad at all, the other
guards the consequence by running DDP's accumulation by hand and checking main_grad
stays finite. They are cheap single-process tests, unlike #6565's full MTP run.
fix(gtp): recycle the wgrad reduce-scatter inputs, not its outputs
wgrad_reduce_scatter() rebound one name to both the pool-owned full-size RS
inputs and the shard-sized RS outputs, then returned the outputs to the pool that
_wgrad_pool_get() serves full-size buffers from. Split into wgrad_inputs and
reduced_wgrads so only the pool's own buffers go back to it.
Test asserts the synchronous batched path returns the original input buffers to
the wgrad pool.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fix(gtp): restore replicated TE bias width after GTP presharding
GTP remat shards weights only, but TE sizes both weight and bias from the
pre-sharded out_features it is constructed with, leaving every bias at GTP-shard
width. Resize each bias back to the logical TP-local width after construction.
Refuses the two layouts this does not model — TE single_grouped_bias and
parameters_split — rather than silently mis-sizing them.
Upstream GTP has no bias-size handling at all, so any GTP run with
add_bias_linear would carry short biases into the first forward.
GPU-free tests cover the resize, the no-op when already correct, the
disabled-bias placeholder, the grouped case, and each refusal path. The
distributed test drives row/column/padded TE linears end to end through
forward and backward, including the bias-disabled placeholder Qwen3.5 GDN
uses.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
fix(multimodal): align Qwen3.5-VL RoPE wrappers
Match the current MCore RoPE callable contract: both fp32 wrappers now take
mla_rotary_interleaved / inverse / mla_output_remove_interleaving / max_seqlen in
the canonical parameter order and forward them through the BSHD and THD paths.
The wrappers replace the module-level apply_rotary_pos_emb, so a signature that
diverged in position 7 (max_seqlen where upstream has mla_rotary_interleaved)
could silently misbind a positional caller.
Adapted from 2a00f99b52 on feat/qwen35-gtp-clean rather than taken verbatim: this
tree has since gained a fused mRoPE THD fast path that the original body predates
and would have removed. The signature contract is kept; the stale body is not.
Guard the fused path on the new options. fused_apply_mrope_thd implements neither
the inverse rotation nor MLA output de-interleaving and takes no
mla_rotary_interleaved argument, so requesting any of them must fall through to
the unfused wrapper instead of silently computing the wrong rotation.
fix(multimodal): derive MockQwen35VLDataset samples from the sample index
The dataset drew from ambient CPU RNG and ignored idx. Ranks share the CPU seed
unless data_parallel_random_init is set, so every rank of a data-parallel group
drew the same stream and the distinct samples in a global batch collapsed to the
per-rank microbatch count — a number that changes with the parallel layout, so
two configurations being compared saw different data.
Each sample is now derived from (seed, split, idx). The split name keeps the
train/valid/test datasets apart: sharing one stream would make them emit
identical samples, trading duplication across ranks for duplication across
splits. The provider passes args.seed through, which the ambient-RNG version
got for free.
The parts are hashed with blake2b rather than added. Addition overlaps by
construction: idx=1 under seed S is idx=0 under seed S+1, and each split would
replay the previous one shifted by one index.
Since idx now determines content, out-of-range indices raise IndexError instead
of silently hashing to a plausible sample from outside the dataset.
(cherry picked from commit 716a9e7249cbcfa6a2197772a52c4cbe58522b05)
fix(gtp+mtp): zero the placeholder wgrad when zero_out_wgrad is set (#6565)
(cherry picked from commit d262dbed3d48a0b6a22b1f0ddc58c2cede2b9149)
[fix] fix GTP+DCP ckpt saving/loading for GDP module (#6503)
(cherry picked from commit 8ba5c60d2a)
Import block adapted to this base: it keeps is_using_quantization_scales, which
upstream's tree does not have, alongside the make_tp_sharded_tensor_for_checkpoint
this change adds.
[fix] GTP+DDP potential data-race: publish DDP params before GTP prefetch reads them (#6388)
Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit 87cc590ba03b7490ca2ab27dfc2372cd4d24677d)
Gated Delta Product (GDP) implementation (#6074)
(cherry picked from commit 8e70f69827, adapted for the dev base:
- hybrid_flops/gdn_layer_flops grafts re-done inside the dev-evolved
functions (DSv4/MLA/MTP-aware signatures); gdn_layer_flops keeps the
dev signature (no main-only use_gdn2)
- main-only context tests (maybe_save_dataloader_state) not carried;
tests/unit_tests/transformer/test_transformer_config.py recreated
with only the GDP tests since dev deleted the file
- _equivalent_pair hybrid fixture gains gdp_num_householder, matching
what upstream did to _make_hybrid_args)
Balance LayerWise optimizer shards by Newton-Schulz cost, not parameter size (#6379)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
(cherry picked from commit 85344907be6a9985410bf528acbedaca77b2a020)
GTP + PartialCG: Classify latent projections for partial CUDA graphs (#6446)
(cherry picked from commit 81fe7c746f6556e5d8dc24196a1358ef0898dbf7)
[fix] GTP+recompute: keep adjacent GTP recompute weights off the same gather buffer (#6407)
(cherry picked from commit e22bcb09b91375639588217c12e96242e8e48c36)
feat(gtp): opt-in GTP_remat sharding support for 'moe-latent-proj' (#6383)
(cherry picked from commit fb24e8707c94aebe28fb919e4a6253b600ad0528)
[feat] GTP+MTP (#6242)
(cherry picked from commit 79e6b6a7f525a9fb42a8adcb347534b4fc791b0a)
Perf: skip per-param copy_ dispatch in the MXFP8 param copy-back (#6094)
(cherry picked from commit 8e57bb642344f60e04bec0d7a79e5fd66c5c0023)
Enforce that the number of optimizer shards used in layout computation is the same used during the training iteration (#6048)
(cherry picked from commit 8a424b83777d4bdc8c3c857b1cb1d9ce141818af)
[GTP] Cleanup usage of pg_collection in gtp (#6234)
Co-authored-by: ykarnati <ykarnati@nvidia.com>
(cherry picked from commit 9f3fe8c925fb3588356c06300071cd077b90359a)
Populate dp process group in auto-built ProcessGroupCollection in pipeline schedules (#5901)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
(cherry picked from commit 6f7bcd48f4fccfea74fbd2385502b2fa8bb05251)
Fix GTP reduce-scatter overlap across local CUDA graphs (#6060)
(cherry picked from commit 9fa816bfd9dd994dc5acf6012fb4e841a56cfa4a)
Make Bridge Communicator aware of GTP (#6263)
(cherry picked from commit 9a3c40b25408f9c5f4ed87500989ebe299f88fea)
[feat] Add reduce-scatter-with-fp32-accumulation support for GTP (#6200)
(cherry picked from commit 23b2ff22eea5a150fec2b41b4d035afa216c8040)
Fix gradient reduction issue when EP=1, EP=TP, and EGTP != GTP (#6080)
(cherry picked from commit abf04f46e39a22f8e6f2597ba820a7712aef36c2)
Fix gradient-norm undercounting when using EP and TP (#5916)
(cherry picked from commit cd4afffa648426a959dc7cb1e24b5ce7d0c3ff54)
[fix] Give same-key GTP chain neighbours distinct gather buffers (#6207)
Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit c2fd8275d9fa13cd7d2e51c9d9370fbb6cdae02d)
Fix GTP full-iteration CUDA-graph capture regression (#6077)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
(cherry picked from commit ec53920e605b7d5cd84d182981db3c83439a4fad)
[GTP][Feat] Add one-block-ahead prefetch for GTP grouped-expert weights (#6057)
(cherry picked from commit 543b8039646f2cdb7a54eacad612849080292bd6)
[feat] Generalized Tensor Parallelism (GTP) (#4967)
Co-authored-by: Jieming Zhang <jiemingz@nvidia.com>
Co-authored-by: Deepak Narayanan <deepakn94@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Jiangfei Duan <jiangfeid@nvidia.com>
(cherry picked from commit c5ff22b7f11822e0a36526b0434708d99c8206b0)
ci(auth): treat svcnemo-autobot as internal (#6437)
[dev] [Deepseek-V4] Fix fused CSA indexer loss normalization and compact attention indices (#6349)
fix(ci): redact sensitive environment values from test output (#6408)
[Dev] Update moe recipes (#6335)
[dev] Optimize MTP contiguous packed-CP rolls (#6246)
[dev] Refactor CSA structure: Move CSA implementation helpers into csa_utils dir (#6372)
fix(optimizer): AUT-1363 avoid grad threshold check during capture (#6355)
Co-authored-by: Hristo Filaretov <hello@hgf.sh>
Co-authored-by: Michał Marcinkiewicz <43240942+mmarcinkiewicz@users.noreply.github.com>
fix(ci): AUT-1396 stabilize docker cache and external test queue (#6348)
cp: chore(ci): bump community workflow to v1.8.8 (#6345)
Co-authored-by: oliver könig <okoenig@nvidia.com>
build: AUT-1363 verify downloaded yq binaries (#6344)
fix(deps): AUT-1323 use released dev dependency sources (#6273)
Co-authored-by: NVIDIA NeMo Bot <nemo-bot@nvidia.com>
[DSv4] Use the full CSA denominator for unfused indexer loss (#5960)
Fix undefined get_data_and_context_parallel_group references (#6072)
Add standard MLite local validation (#5888)
fix(dsa): skip invalid slots in sparse indexer backward (#6166)
[Dev] Keep the mHC mapping computation in fp32 on the fused cuTile path (#6172)
[codex] Propagate SBHD padding masks across pipeline stages (#5544)
[Dev] Support full activation recompute for EP A2A overlap (VPP-stage recompute) (#5869)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Dispatch the CSA Compressor gated pooling to the cudnn-frontend fused kernels (THD) (#5984)
Fix/dsv4 cp relax assert offline packed thd (#6158)
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
fix(tensor-parallel): AUT-1177 deduplicate expert shards by expert TP (#6165)
[dev] Fix Torch checkpoint loading for updated layer-wise optimizer (#6073)
[dev] [fix] fix optimizer_cpu_offload with mark_keep_in_fp32 (#6124)
[codex] Preserve real and padded THD sequence lengths (#5541)
[dev] [fix] propagate LayerWise optimizer flag through model builders (#6082)
[Dev] Migrate main-first DSA support to dev (#6020)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Pingtian Li <pingtianl@nvidia.com>
build: serialize uv dependency installation (#6089)
[Dev] Add Triton fused mRoPE for Qwen3.5-VL (#5962)
Co-authored-by: Li Tao <lit@nvidia.com>
[dev] Align thd padding logic in eager & graph mode (#5724)
ci(actions): AUT-977 retry transient log artifact uploads (#6026)
fix(deps): AUT-960 use optional-attribution NVRX revision (#6006)
test(gdn): AUT-983 quarantine flaky chunkwise CP backward checks (#6040)
[dev] moe(perf): Support fused pre-GDR path for chunkwise CP. (#5638)
docs(skills): clarify container::lts is the older LTS PyTorch base (#6009)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[dev] moe(fix): Restore per-token MTP loss logging (#5967)
[Dev] Skip no-op BSHD padding masks in Qwen3.5-VL (#5964)
[dev] [fix] fix optimizer state offloading (#5923)
[dev]: Fix mHC boundaries in EP overlap schedule (#5471)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[Dev] Fix max_seqlen forwarding in Qwen3.5-VL vision RoPE (#5890)
Co-authored-by: Li Tao <lit@nvidia.com>
chore: nightly sync main into dev (13_07_2026) (#5784)
Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com>
Co-authored-by: Keshav Santhanam <ksanthanam@nvidia.com>
Co-authored-by: Laura Dang <lauradang.2000@gmail.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com>
Co-authored-by: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>
Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Jingyue Wu <wujingyue@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 <hollowman@opensuse.org>
Co-authored-by: muyihao <37872457+muyihao@users.noreply.github.com>
Co-authored-by: Tom Long <tolong@nvidia.com>
Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>
Co-authored-by: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com>
Co-authored-by: yeyu-nvidia <yeyu@nvidia.com>
Co-authored-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com>
Co-authored-by: Laura Dang <laurad@nvidia.com>
Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>
Co-authored-by: Lawrence McAfee <85179052+lmcafee-nvidia@users.noreply.github.com>
Co-authored-by: janEbert <janpabloe@nvidia.com>
Co-authored-by: Anil Thomas <anlthms@users.noreply.github.com>
Co-authored-by: Ritesh Patel <ripatel@nvidia.com>
Co-authored-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>
Co-authored-by: ma-jh <3364870135@qq.com>
Co-authored-by: Jenny Chen <jennifchen@nvidia.com>
Co-authored-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Jimmy Zhang <133159885+jiemingz@users.noreply.github.com>
Co-authored-by: Jingyue Wu <jingyuew@nvidia.com>
Co-authored-by: Ajay <abalasa@nvidia.com>
Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
Co-authored-by: mathemakitten <helenn@nvidia.com>
Co-authored-by: Zang Peiyu <166481866+factnn@users.noreply.github.com>
Co-authored-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-authored-by: Skand Hurkat <shurkat@nvidia.com>
Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com>
Co-authored-by: Siddharth Singh <sidsingh@nvidia.com>
Co-authored-by: Yongqiang Wang <yongqiang.seagull@gmail.com>
Co-authored-by: Minh Vu <vuhoangminh97@gmail.com>
Co-authored-by: Guihong Li <guihongl@nvidia.com>
Co-authored-by: Asadbek Xodjayev <100586658+asadbekXodjayev@users.noreply.github.com>
Co-authored-by: asadbekXodjayev <matyoqub18@gmail.com>
Co-authored-by: Björn Buschkämper <bjoern.buschkaemper@gmail.com>
Co-authored-by: Zhengmao Ye <yezhengmaolove@gmail.com>
Co-authored-by: Jon Barker <jbarker@nvidia.com>
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: Deyu Fu <deyuf@nvidia.com>
[dev] Lite: /bump lite DeepSeek v4 RL to dev (#5862)
Co-authored-by: leikaixiang.shyoshyo <leikaixiang.shyoshyo@bytedance.com>
Avoid per-segment GPU->CPU syncs in the unfused DSA THD helpers (#5884)
enable tms no cpu backup region for param buffers (#5103)
Co-authored-by: Yueming Yuan <yym022502@gmail.com>
Co-authored-by: lilei <799812479@qq.com>
Co-authored-by: Deyu Fu <deyuf@nvidia.com>
[Dev] Fix distributed optimizer state save/load for mixed-dtype param groups (#5835)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
fix: preserve static hybrid module specs
fix: resolve callable hybrid specs
fix: address post-merge unit test regressions
Fix zero-loss indexer gradients for DSv4 CP (#5809)
Merge remote-tracking branch 'origin/dev' into main2dev/13_07_2026
# Conflicts:
# megatron/training/arguments.py
# pretrain_hybrid.py
# tests/unit_tests/data/test_get_batch.py
fix: restore MLA TP fusion and DeepEP v2 SM config
[dev] Enable DeepSeek-v4 hybrid_model Part (3/N) (#5762)
[dev] Muon fp8 param gather for decoupled layout (#5470)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
fix: restore dev training-loop behavior
fix: align pretrain provider call order
fix: forward padding mask in fine-grained MoE routing
Keep split A2A schedule routing consistent with the normal forward path so dropless HybridEP excludes padded tokens in both executions.
fix: post-CI corrections (unused imports + AbsorbedMLA submodule fields)
Fixes for the first CI run's failures:
1. Remove unused imports flagged by the linting job (pylint W0611):
- dynamic_engine.py: CUDAGraphBatchDimensionBuilder, InferenceBatchDimensions
- gated_delta_net.py: ShardedTensor, ReplicaId, ShardedTensorFactory
- transformer_block.py: checkpointed_forward
2. Fix AbsorbedMLASelfAttentionSubmodules construction to match dev's split
K/V up-projection API (the merged absorbed_mla.py is dev's version, which
splits linear_kv_up_proj into linear_k_up_proj + linear_v_up_proj):
- hybrid_layer_specs.py dsa_layer + experimental_attention_variant_module_specs.py
now pass linear_k_up_proj/linear_v_up_proj to AbsorbedMLASelfAttentionSubmodules.
- The standard MLASelfAttentionSubmodules construction in hybrid_layer_specs.py
keeps linear_kv_up_proj (that class was not split).
This fixes the install-test ImportError:
"AbsorbedMLASelfAttentionSubmodules.__init__() got an unexpected keyword
argument 'linear_kv_up_proj'".
Co-Authored-By: Claude <noreply@anthropic.com>
[codex] Exclude padding tokens from MoE routing (#5542)
Co-authored-by: Xin Yao <xiny@nvidia.com>
fix: save chained optimizer state in torch checkpoints (#5684)
[Dev] Numerical fix for moe single grouped weight with fp8 fp4 primary weight and grad norm spikes (#5464)
Merge remote-tracking branch 'origin/main' into main2dev/13_07_2026
Nightly sync of main into dev (13_07_2026).
Resolution highlights:
- Kept dev's dependency triple (pyproject.toml, uv.lock, docker/Dockerfile.ci.dev),
CODEOWNERS, base-image pin (26.04, GitLab pin), and golden values. Main's
base-image bump to 26.06 (#5632) belongs in a dedicated bump PR, not the sync.
- copy-pr-bot.yaml: union of trustee lists (kept dev-only sanandaraj5597, wplf).
- Renamed args.hybrid_context_parallel -> args.dynamic_context_parallel and
get_hybrid_data_context_parallel_groups -> get_dynamic_data_context_parallel_groups
in main-originated code (training.py, pretrain_hybrid_flex.py) to match dev's
canonical naming (the deprecated hybrid arg does not exist in the merged tree).
- DSA subsystem: dsa_kernels.py resolved as a UNION of dev's kernel API
(indexer_topk, dsa_sparse_attn, ...) and main's backend-neutral hooks
(use_fused_dsa_kernels, run_fused_*); dsa.py based on main + re-added dev-only
fused_qk_topk_naive_thd/_build_causal_mask_seg for csa.py.
- MoE flex dispatcher: unioned backend Literal (deepep/deepepv2/hybridep/ncclep),
kept both _DeepepV2Manager (dev) and _NCCLEPManager (main).
- transformer_layer.py: kept dev's _forward_mlp_output_with_bias/input_ids
plumbing AND main's inter-document-masking MoE reshape
(_maybe_unflatten_for_moe/_maybe_reflatten_from_moe).
- fine_grained_activation_offload.py + checkpointing.py: took main where the
merged non-conflict common regions and callers already required main's
interface (group_offload naming; process-group threading params).
- MTP inference: kept dev's _decoder_hidden_states_cache path (consumed by
text_generation_controller); dropped main's block-scope
mtp_decoder_hidden_states approach (infra not present in merged tree).
Intentional drops (skill case (a), documented with SHA):
- _forward_mlp_postprocess's redundant _restore_token_dispatcher_attrs()-at-top
removed by main commit 5e4fe9b3c ("Optimize memory usage of partial CUDA
graphs"); replaced by main's weakref tail.
Co-Authored-By: Claude <noreply@anthropic.com>
Remove use of exec_module (#5744)
Assign BERT CODEOWNERS to GPT team (#5746)
Set is_first_microbatch when quant_recipe is configured (#5642)
Set Bert TE spec q/k_layernorm to None (#5687)
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Fix infinite recursion in abstract tokenizer special-id property aliases (#5445)
Co-authored-by: asadbekXodjayev <matyoqub18@gmail.com>
Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
Avoid X11 master port default (#5299)
Fix seq_load_balancing loss with inter-document masking and MBS > 1 (#5696)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Add forward all-gather overlap (#5513)
[codex] Restore per-microbatch MTP loss logging (#5543)
fix(fsdp): import os in safe_get_rank fallback (#4959)
remove deprecated modules from core/dist_checkpointing (#5134)
[dev] [fix] fix MTP with contiguous CP partition mode (#5706)
fix: add label shifting to mlite THD packing to align SFT loss with M… (#5682)
Co-authored-by: Xin Yao <xiny@nvidia.com>
[dev] Fix HybridEP token equalization under torch.compile without CUDA graphs (#5668)
[dev] Align Megatron Lite MoE RL execution contracts (#5694)
Keep DeepSeek V4 CSA compressor and indexer in high precision under FP8 training (#5308)
Co-authored-by: Hongxiao Bai <hongxiaob@nvidia.com>
Add NeMo Transformer audio encoder model (#5565)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Triton kernels - avoid recompilation and autotuning in prod (#5608)
[2/2] Wiring cuDNN fused DSA kernels support with THD, CP and IndexShare (GLM5.2) (#5099)
Refactor RL rollout pipeline (#5491)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Update base image to nvcr.io/nvidia/pytorch:26.06-py3 (#5632)
NCCL EP support (#5129)
Add FSDP NVTX annotations (#5704)
Remove some barriers in save_checkpoint_and_time (#5557)
Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com>
Normalize CRLF in Claude fix commands (#5712)
ci: revert unify legacy scope names (#5316) (#5709)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
deprecate common strategy (#5160)
test(determinism): add determinism tests (#5041)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
chore: rotate oncall schedule
Separate mFSDP v2 unit tests (#5640)
Fix Torch FSDP2 crash: add force_all_reduce kwarg to base finish_grad_sync (#4953)
Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>
ci: Update test configurations to unify legacy scope names (#5316)
Add microbatch context helper (#5652)
Add cspades to oncall rotation (#5695)
MoE routing analysis and metrics capture (#5220)
Fuse shared expert MLP with grouped GEMM (#5604)
add safe version of numpy.load (#5500)
[Dev] Fix FSDP backward hooks for TE-fused experts (#5636)
Fix inter-document masking crash and NaNs with TP > 1 and micro_batch_size > 1 (#5635)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Scatter embeddings for sequence parallelism in standalone LM forwards (#5628)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Pre-size the all-gather buffer for inference to max capacity (#5546)
Add int4 qat training support (#5108)
Co-authored-by: GeLee-Q <leege233@gmail.com>
Co-authored-by: fy1214 <282598660@qq.com>
Co-authored-by: Gao016 <yngao016@163.com>
Co-authored-by: yefei12 <xjtu_yefeichen@163.com>
Co-authored-by: yzlnew <yzlnew@gmail.com>
Co-authored-by: Deyu Fu <deyuf@nvidia.com>
chore(beep boop 🤖): Bump (main) (2026-07-06)
Ignore contributor DCO failures in Claude fix (#5625)
Update copy-pr-bot.yaml [skip ci]
[dev] [DeepSeek-v4] Context Parallel support (#5087)
Update golden value files for GPT-3 weekly (#5459)
Add smoke test notification functionality and update notify script (#5631)
Add NCCL symmetric-memory staging to experimental FSDP (#5440)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Fix smoke BERT/T5 test failures (#5629)
Fix Claude reaction permissions (#5613)
Document stacked dependent PR handling in split PR skill (#5496)
Optimize memory usage of partial CUDA graphs (#5451)
E2E heterogenous non colocated MiMo training (#5602)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[dev] moe(feat): support chunkwise context parallelism for GDN. (#3282)
[Main][feat] Support CUDA Graph capture offloading modules (#3697)
Co-authored-by: Xin Yao <xiny@nvidia.com>
Add Auto Quantize in ModelOpt quantize example (#4821)
Co-authored-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
[Dev] fix padding mask docstring (#5598)
fix(tensor_parallel): _reduce returns unreduced tensor for non-contig… (#5338)
[Megatron-FSDP] MaxPoolAllocator for double-buffering hybrid architectures. (#5462)
Update copy-pr-bot.yaml [skip ci]
Add /claude fix workflow for on-demand PR fixes (#4862)
chore: rotate oncall schedule
[dev] Sync Megatron Lite with the latest implementation (#5577)
Update mcore skill owners (#5586)
Update PR instructions (#5592)
Use NVIDIA inference credentials for Claude actions (#5589)
Update copy-pr-bot.yaml [skip ci]
[dev] Add experimental decoupled compact LayerWise DDP layout for Muon (#5388)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[training migration] Finish ModelBuilder integration (#5516)
Fix TEGroupedMLP pre-backward unshard in fine-grained FSDP hooks for … (#4990)
Thread dp_cp/expt_dp process groups through checkpoint load path (#5579)
Add CI duties to oncall (#5510)
Co-authored-by: oliver könig <okoenig@nvidia.com>
Deduplicate tensor-splitting utility (#5545)
Fix PR template typo (#5566)
Fix `isort` target Python version (#5567)
build: bump transformer-engine to release_v2.16.post (#5517)
Implement async scheduling for dynamic inference (#5453)
Enable Deepseek-v4 hybrid_model in dev branch Part (2/N) (#5485)
[CI] Fix `gpt_dynamic_inference_tp2_pp2_ep2_gptoss_20b_swa` tests (#5527)
Update copy-pr-bot.yaml [skip ci]
Preserve DSA output across fused inverse RoPE (#5526)
Co-authored-by: Kaixiang Lei <5780122+shyoshyo@users.noreply.github.com>
Add inter-document attention masking to GPTDataset (#5298)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
ci: pin HF_HUB_CACHE to bind-mounted cache for gpt-oss-20b inference test (#5512)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ci: Use GB300 for Github CI tests (#5520)
[dev] moe(perf): Pre-GDR kernel fusion (#5361)
Add CUDA graph training iteration test (#5417)
ci: cache-from a single coherent buildcache donor (#5509)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Update copy-pr-bot.yaml [skip ci]
test: restore G/G + lag=19 for gpt_grpo_tp4_pp1_dp2_8b throughput tests (#5514)
[Main] Generalized fix for mxfp8 param gather (#5236)
Improve default dynamic CP packing scheduler (#5154)
Add --qad-train-target {base|mtp|both} for QAD / MTP QAT (frozen-base, frozen-MTP, or co-train) (#4785)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
fix: set DATA_PATH for moe-dynamic-inference recipe (#5506)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[dev] [DeepSeek-v4] Packed Sequence (THD) support for DSv4 Hybrid Attention (#5011)
Add hybrid FSDP unit module support (#4329)
Fix NameError in is_flashinfer_min_version when check_equality=False (#4961)
chore: nightly sync main into dev (22_06_2026) (#5430)
Co-authored-by: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com>
Co-authored-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-authored-by: Pranav Thombre <pthombre@nvidia.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: janEbert <janpabloe@nvidia.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: mathemakitten <helenn@nvidia.com>
Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Shijie <505749828@qq.com>
Co-authored-by: Ajay <abalasa@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>
Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com>
Co-authored-by: Tom Long <tolong@nvidia.com>
Co-authored-by: Keshav Santhanam <ksanthanam@nvidia.com>
Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>
Co-authored-by: Jingyue Wu <wujingyue@gmail.com>
Co-authored-by: ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 <hollowman@opensuse.org>
Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Lawrence McAfee <85179052+lmcafee-nvidia@users.noreply.github.com>
Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com>
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Fix fused MLA down projection with tensor parallelism (#5383)
[split 4/4] Enable DSA CP and THD hooks (#5246)
Merge cu_seqlens across micro-batch for THD attention (#5454)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Force RL inference to CP=1 (#5423)
fix: restore fused group MLP offload in main2dev sync (#5493)
ci: auto-retry test-data download in container-build job (#5498)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[dev] Megatron Lite (4/4) shared attention (#5427)
fix sequence packing wrapper for eval (#5483)
Add Nemotron6-MoE VLM model provider for MIMO example (#5374)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Add MIMO forward step and per-token loss for hetero training (#5376)
Narrow oncall responsibilities (#5490)
fix: post-CI corrections for sync test/impl reconciliation
1) test_optimizer.py: reverted to dev. The sync kept dev's
multi_latent_attention.py (split q/kv down-proj, no
_synthesize_fused_qkv_down_weight), but auto-merged main's
test asserting the fused linear_qkv_down_proj.weight key.
2) training.py: guard the dev-only sequence_packing_scheduler config
access with getattr (lines in train_step and train()). main's new
MIMO schedule-plumbing test (#5333) passes an empty SimpleNamespace
config; the reconciled training.py keeps dev's packing path, so the
access must tolerate a config lacking the attribute. Real configs are
unaffected (getattr returns the same value).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Thread process groups through training checkpoint paths (#5486)
Clean up training.py module header (dedupe + reorganize imports/globals) (#5469)
Automated community request assignment (#5147)
Fix merges_file kwarg name in HuggingFaceTokenizer (#5406)
ci: check megatron.training imports in installation test (#5458)
[split 2/4] Scale DSA indexer loss in pipeline schedules (#5244)
build: install flash_mla from source in the CI image (#5481)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ci: launch GB200 unit tests via launch_on_gb200 marker (#5477)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add inference functions to support MCore-/MBridge- training refactor and remove legacy modelbuilder functions (#5169)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
chore: rotate oncall schedule
test: mark gated_delta_net selective-recompute test flaky_in_dev (#5476)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
test: mark TestParallelTransformerBlockCudagraphs::test_gpu_cudagraph flaky_in_dev (#5475)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add experimental Megatron-FSDP fully_shard implementation (#5387)
Support HybridModel feature specs in ModelOpt (#5354)
feat(inference): default use_coordinator to True in high-level APIs (#5326)
Add hetero grid args and MoE process groups for MIMO example (#5375)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Support SWA and sink attention in dynamic inference (#5249)
Co-authored-by: shanmugamr1992 <shanmugamr1992@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
Remove DBuffer mesh axis validation (#5441)
ci: Set test_save_verify_integrity_manifest_directly as flaky (#5468)
Add logprobs_mode (raw/processed) to inference config (#5419)
build: point flash_mla at the nv_dev branch (#5448)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rename CP batch helpers to describe balancing granularity (#5403)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Add --functional-test-name to trigger_internal_ci (#5449)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
test: mark ep_a2a_overlap activation-offloading test flaky_in_dev (#5450)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[Fix] Fix MoE router z-loss compatibility with TE CUDA Graph capture. (#5401)
Co-authored-by: yangfan.bai <yangfan.bai@shopee.com>
Add RL rollout submission and consumption granularity controls (#5306)
Add MIMO dual gradient finalization (colocated + non-colocated) (#5286)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Clean up MTP inference control flow (#5418)
Add RADIO vision encoder wrapper for MIMO example (#5397)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Disag MR3: Add heterogeneous KV/Mamba reshard planners (#5188)
Merge remote-tracking branch 'origin/main' into main2dev/22_06_2026
Nightly sync of main into dev (22_06_2026). Resolves 16 conflicts
preserving dev features (pre-push guard: 0 dropped dev lines); brings
in main's inference shard-spec API additively. Supersedes #5429.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Consistent oncall schedule (#5404)
Stabilize hybrid_2b GB200 perf test against run-to-run noise (#5364)
Support the MIMO cross-grid path in training loop (#5373)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Disag MR1: Add inference shard specs and pg-collection building (#5186)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Disag MR2: Refit into multiple destination pools and tied-embedding + UVM fixes (#5187)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Fix Mamba prefix match for chunked prefill (#4758)
chore(beep boop 🤖): Bump (main) (2026-06-22)
Add --mamba-training-ssm-states-dtype argument (#5309)
Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com>
Add MIMO runtime setup: per-role RNG seeding and DDP wrapping (#5285)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Update goldens for weekly tests after pytorch and TE bumps. (#5399)
Revert "Remove checkpoint-time GPU cache reclaim workaround (#5170)" (#5366)
Add flaky marker to fine-grained activation offloading test (#5350) (#5368)
ci: Remove sync skills workflow (#5091)
chore: rotate oncall schedule
[dev] moe(perf): Restore fused GDN THD all-to-all on dev (#5389)
[split 3/5] Refactor absorbed MLA projection handling (#5245)
Add MimoModel.zero_grad_buffer delegating to active DDP submodules (#5372)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Thread tensor-parallel group into the RADIO patch embedder (#5371)
[feat] Support fine-grained activation offloading in fused group mlp (#5082)
Remove unused distributed pytest markers (#5380)
[split 1/5] Fix packed THD RoPE under CP (#5243)
[Dev] restore DSv4 tflops calc in training and fix the packed seq case (#5358)
Add minimal DBuffer implementation (#4835)
Document agent PR commit sign-off and signing (#5381)
Update copy-pr-bot.yaml [skip ci]
Support fused MLA QKV checkpoint reload (#5310)
Profiling (#3110)
Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Make Megatron RL only materialize last token logit (#4551)
Expand the Mamba prefix caching memory safety check to include scratch space buffers (#5348)
Add full model cuda graph support for MTP inference (#4950)
Clean up pretrain_gpt.py and pretrain_hybrid.py formatting and remove module globals (#5351)
Fix memory leak with log_max_attention_logit (#4699) (#5067)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com>
[Dev] add cuda graph support for thd format training. (#4359)
Co-authored-by: Haochen Yuan <haocheny@login-eos01.eos.clusters.nvidia.com>
ci: default functional test time limit to 4h for release/weekly scopes (#5360)
[Dev] Add Megatron-FSDP weight prefetch for full recompute (#5175)
[Fix] Fix optimizer parameter override bugs. (#5213)
Co-authored-by: yangfan.bai <yangfan.bai@shopee.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
feat(ssm): whole-module 'gdn' selective recompute for GatedDeltaNet (#5296)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
chore: nightly sync main into dev (12_06_2026) (#5314)
Co-authored-by: megnvidia <mmiranda@nvidia.com>
Co-authored-by: Philip Petrakian <pgpetrak@gmail.com>
Co-authored-by: Xuanteng Huang <44627253+xuantengh@users.noreply.github.com>
Co-authored-by: Ajay <abalasa@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Pingtian Li <158665726+Wohox@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
Co-authored-by: Santosh Bhavani <santosh.bhavani@live.com>
Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Yan Xu <45385219+Connor-XY@users.noreply.github.com>
Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com>
Co-authored-by: Jingyue Wu <wujingyue@gmail.com>
Co-authored-by: lichenlu <lichenlu_8618@163.com>
Co-authored-by: peibli <lipeibao@126.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: shurkat-nvidia <shurkat@nvidia.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Lintch <44701395+returnL@users.noreply.github.com>
Co-authored-by: janEbert <janpabloe@nvidia.com>
Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>
Co-authored-by: Siddhartha Raman Sundara Raman <270218152+sraman-rgb@users.noreply.github.com>
Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
Co-authored-by: Anish Mahishi <amahishi@nvidia.com>
Co-authored-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@users.noreply.github.com>
Enable Deepseek-v4 hybrid_model in dev branch Part (1/N) (#5042)
Co-authored-by: Guihong Li <guihongl@oci-hsg-cs-001-vscode-02.cm.cluster>
Co-authored-by: Yan Xu <yxu1@nvidia.com>
Co-authored-by: hx <hongxiaob@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Thread MIMO support through the stock training loop (schedule + optimizer) (#5333)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
fix: preserve seqlen stats in train_step
Add zstandard package to Docker LTS requirements. Fix nightly failures (#5347)
Merge branch 'dev' into main2dev/12_06_2026
Fix LatentMoE theoretical memory estimate (#5145)
[dev] bump emerging optimizers to v0.3.0 (#5320)
chore(beep boop 🤖): Bump (main) (2026-06-15)
fix tflops calculation when sequence_packing_scheduler is not none (#5342)
[Dev] Add DeepEP v2 flex dispatcher backend (#4793)
Co-authored-by: Dennis(Zhenhuan) Liu <denliu@nvidia.com>
[Dev] Add MoE recipe performance summary (#5289)
Co-authored-by: Dennis Liu <denliu@denliu.nvidia.com>
Thread pg_collection through wrap_model_chunks_with_ddp (#5328)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Account for reasoning token stripping (#5313)
Inference: Cudagraph-aware admission gating in prefill scheduler (#4870)
Fix crash due to tool call at sequence length (#5302)
Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com>
Fix EP=1 inference by allocating buffers anyway (#5233)
Add code owners for optimizer-related files (#5297)
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Add moe loss normalization for RL SFT (#3956)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Allow for pre-bound socket to be passed in server (#5301)
Handle None values in sampling parameters (#5300)
Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com>
Offline Logits-Based Knowledge Distillation (#5019)
[dev]: faster implementation of mHC fused kernels (#4624)
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
fix: restore dev-only training.py features dropped by the main override
The merge takes main's training.py wholesale (per the sync skill's override
list), which silently reverted four dev-only features whose supporting core
code the merge keeps at dev's version. Each surfaced as a CI failure:
1. Dynamic context-parallel API rename. Dev renamed
get_hybrid_data_context_parallel_groups -> get_dynamic_data_context_parallel_groups
(identical signature) and args.hybrid_context_parallel -> dynamic_context_parallel.
Point training.py at dev's names. Fixes the conftest.py ImportError that
cascaded to every unit-test bucket.
2. Dynamic-CP / sequence-packing data loading. Replace main's setup-time
HybridCPDataLoaderWrapper wrap with dev's per-step wrap_data_iterator
(gated on config.sequence_packing_scheduler), which returns a
RerunDataIterator-compatible iterator. Fixes the *_cp4_dcp RerunDataIterator
assertion.
3. MTP loss logging scale. Dev's MTPLossLoggingHelper stores raw loss sums and
token counts and computes the per-token loss after reduction, so the log
scale must be 1.0; main's 1/get_num_microbatches() divided the reported
mtp_N loss by num_microbatches. Fixes moe gpt3_..._scoped_cudagraph mtp_1
loss (was 16x too small: 0.682 vs golden 10.915).
4. DSA indexer loss cross-PP reduction. dsa.py requires num_layers (and
csa_compress_ratios) so first-pipeline-stage ranks lazily initialize the
tracker and join the cross-PP all_reduce in reduce_loss_in_tracker; main's
call omitted them, so stage 0 returned early and the last stage hung on an
unmatched all_reduce. Fixes the gpt3_..._dsv4_hybrid_mhc_mtp NCCL timeout.
ci: Allow DCO check in merge queue and add DCO requirement (#5305)
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Merge remote-tracking branch 'origin/main' into main2dev/12_06_2026
# Conflicts:
# README.md
# megatron/core/optimizer/optimizer.py
# megatron/core/ssm/gated_delta_net.py
# megatron/core/transformer/multi_token_prediction.py
# megatron/core/transformer/transformer_config.py
# megatron/elastification/pretrain_hybrid_flex.py
# megatron/training/training.py
# megatron/training/utils/common_utils.py
# tests/test_utils/recipes/h100/t5.yaml
# tests/test_utils/recipes/moe2.0.yaml
# tests/unit_tests/ssm/test_gated_delta_net.py
# tests/unit_tests/transformer/test_multi_token_prediction.py
# uv.lock
Thread pg_collection into train_step reductions (#5259)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Clip mtp grads separately when mtp_detach_heads=True (#4116)
Co-authored-by: Anish Mahishi <amahishi@nvidia.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Stabilize hybrid nanov3 gb200 perf (#5295)
Enable non-deterministic results in model configuration for nemotron tests (#5239)
Thread pg_collection into get_model DDP bucket sizing (#5250)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ci: limit retries on unsuccessful test launches (#5275)
Fix fused MLA delayed weight grad hooks (#5273)
Co-authored-by: Siddhartha Raman Sundara Raman <270218152+sraman-rgb@users.noreply.github.com>
Fix Dockerfile warnings (#4856)
fix(ci): resolve t5 dataloader stall + GRPO cudagraph-memory regression (CI-validated) (#5280)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
[dev] Add experimental Megatron Lite as agentic exploration (#4885)
Co-authored-by: Deyu Fu <deyuf@nvidia.com>
Remove duplicate nccl_allocator import (#5057)
[examples] Add dynamic context parallel benchmark example (#5123)
[TE] Restore original CP group after dynamic CP forward in TEDotProductAttention (#5215)
Co-authored-by: rionawang <rionawang@tencent.com>
[Dev] Cherry-pick MTP detach heads (#5223)
[Dev] Generalized fix for mxfp8 param gather (#4994)
Co-authored-by: Xin Yao <xiny@nvidia.com>
[Dev] DeepSeek-V4-Flash recipe 20260610 (#5266)
Remove checkpoint-time GPU cache reclaim workaround (#5170)
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Add MIMO hetero topology + distributed bootstrap (examples/mimo training-loop folder) (#5260)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Update copy-pr-bot.yaml [skip ci]
ci: Allow DCO check in merge queue and add DCO requirement to Contribution guide (#5278)
Add optional group= to common_utils model/data-parallel reduction helpers (#5251)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
chore: rotate oncall schedule
Fix test_split_tensor_along_last_dim to actually assert correctness (#4710)
Co-authored-by: peibli <lipeibao@126.com>
Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
Route bridge communicator cross-grid P2P through a dedicated process group (#5234)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move uneven DTensor distributed fixture to conftest (#5237)
Fix wgrad race condition when using double buffers. (#5222)
chore: nightly sync main into dev (06_06_2026) (#5199)
Co-authored-by: Nick Schank <nick@reflection.ai>
Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>
Co-authored-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: janEbert <janpabloe@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Shivanjan Chakravorty <schakravorty846@gmail.com>
Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>
Co-authored-by: Maanu Grover <maanug@nvidia.com>
Co-authored-by: Ajay <abalasa@nvidia.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Chen Cui <chcui@nvidia.com>
Co-authored-by: Tuomas Rintamaki <trintamaki@nvidia.com>
Co-authored-by: Tyler Poon <tylerpoon@gmail.com>
Co-authored-by: Collin McCarthy <cmccarthy@nvidia.com>
Co-authored-by: Matthieu Le <matthieul@nvidia.com>
Co-authored-by: Piotr Zelasko <pzelasko@nvidia.com>
Co-authored-by: Ehsan Hosseini Asl <ehosseiniasl@nvidia.com>
Co-authored-by: Teodor-Dumitru Ene <34819528+tdene@users.noreply.github.com>
Co-authored-by: Siddharth Singh <sidsingh@nvidia.com>
Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
Co-authored-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-authored-by: Deepak Narayanan <dnarayanan@nvidia.com>
Co-authored-by: Xiaowei Ren <103958965+xrennvidia@users.noreply.github.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com>
Co-authored-by: William Dykas <wdykas@oci-hsg-cs-001-vscode-03.cm.cluster>
Co-authored-by: Fei Wu <33940270+YangFei1990@users.noreply.github.com>
Co-authored-by: Xuanteng Huang <44627253+xuantengh@users.noreply.github.com>
Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
Co-authored-by: kunlunl <kunlunl@nvidia.com>
Co-authored-by: Xuesong Ye <xuesongyey@gmail.com>
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Co-authored-by: Daisy Gao <daisyg@nvidia.com>
Co-authored-by: Pingtian Li <158665726+Wohox@users.noreply.github.com>
Co-authored-by: gautham-kollu <gkollu@nvidia.com>
Co-authored-by: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com>
Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: lichenlu <lichenlu_8618@163.com>
Co-authored-by: peibli <lipeibao@126.com>
Co-authored-by: Yan Xu <45385219+Connor-XY@users.noreply.github.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: Xin Yao <xiny@nvidia.com>
Co-authored-by: Gao Deng <160076886+gdengk@users.noreply.github.com>
Co-authored-by: Gao Deng <gdeng@login-lyris02.lyris.clusters.nvidia.com>
Co-authored-by: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com>
Co-authored-by: Pavel Gein <pavelgejn@yandex.ru>
Co-authored-by: Gao Deng <gdeng@login-lyris01.lyris.clusters.nvidia.com>
Co-authored-by: Eric Harper <eharper@nvidia.com>
Co-authored-by: Robin Zhang <robinz@nvidia.com>
Co-authored-by: Shanmugam Ramasamy <111910568+shanmugamr1992@users.noreply.github.com>
Co-authored-by: Charlie Truong <chtruong@nvidia.com>
Co-authored-by: Abhishree Thittenamane <47577437+athitten@users.noreply.github.com>
Co-authored-by: Abhishree Thittenamane <athittenaman@cw-dfw-cs-001-login-01.cm.cluster>
Co-authored-by: root <root@pool0-01849.cm.cluster>
Co-authored-by: Jingyue Wu <wujingyue@gmail.com>
Co-authored-by: Jianbin Chang <shjwudp@gmail.com>
Co-authored-by: Kamran Jafari <kjafarisadeg@nvidia.com>
Co-authored-by: Nan Zheng <80790206+nanz-nv@users.noreply.github.com>
Co-authored-by: Li Tao <lit@nvidia.com>
Co-authored-by: Zhongbo Zhu <zhongboz@nvidia.com>
Co-authored-by: Li Jinliang <jinliangl@nvidia.com>
Co-authored-by: Robert Kirby <ArEsKay3@users.noreply.github.com>
Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>
Co-authored-by: Cole Hawkins <chawkins@nvidia.com>
Co-authored-by: Qi Zhang <qizhang@nvidia.com>
Co-authored-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Co-authored-by: tongliu <tongliu@nvidia.com>
Co-authored-by: sraman-rgb <sraman@nvidia.com>
Co-authored-by: Siddhartha Raman S <sraman@login-lyris02.lyris.clusters.nvidia.com>
Co-authored-by: Siddhartha Raman S <sraman@login-lyris01.lyris.clusters.nvidia.com>
Co-authored-by: Qiyu Wan <39144338+WanZzzzzz@users.noreply.github.com>
Co-authored-by: Devil1716 <149754374+Devil1716@users.noreply.github.com>
Co-authored-by: Zijie Yan <zijiey@nvidia.com>
Co-authored-by: Dennis Liu <denliu@nvidia.com>
Co-authored-by: Shifang Xu <shifangx@nvidia.com>
Co-authored-by: Li Ding <liding@nvidia.com>
Co-authored-by: sraman <sraman@users.noreply.github.com>
Co-authored-by: sraman-rgb <270218152+sraman-rgb@users.noreply.github.com>
Co-authored-by: Yi-Fu Wu <yifu.wu@gmail.com>
Co-authored-by: Gerald Shen <geshen@nvidia.com>
Co-authored-by: nvskills-svc-account <svc-nvskills-signing@nvidia.com>
Co-authored-by: Hongbin Liu <lhb8125@users.noreply.github.com>
Co-authored-by: Yuzhong Wang <yuzhongw@computelab-frontend-3.nvidia.com>
Co-authored-by: Simiao Zhang <56124251+conver334@users.noreply.github.com>
Co-authored-by: Shaurya Singh <sshaurya914@gmail.com>
Co-authored-by: Kevin <23258141+kaimo455@users.noreply.github.com>
Co-authored-by: mokai <mokai@baidu.com>
Co-authored-by: Moozy <108285604+Moozy23232@users.noreply.github.com>
Co-authored-by: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com>
Co-authored-by: Keshav Santhanam <ksanthanam@nvidia.com>
Co-authored-by: Mike Chrzanowski <mchrzanowski@nvidia.com>
Co-authored-by: Mike Chrzanowski <mchrzanowski@gcp-nrt-cs-001-login-001.cm.cluster>
Co-authored-by: mathemakitten <helenn@nvidia.com>
Co-authored-by: Jenny Chen <jennifchen@nvidia.com>
Co-authored-by: Yu Yao <54727607+yaoyu-33@users.noreply.github.com>
Co-authored-by: kajalj22 <kajalj@nvidia.com>
Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Add named layouts to HyperCommGrid for heterogeneous parallelism (#5148)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Fix CUDA IMA in fsdp_double_buffer when an FSDP unit's bucket doesn't fit the pool (#4810)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Apply MIMO SP/CP sharding with explicit groups and enable THD in non-colocated path (#5150)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Merge remote-tracking branch 'origin/dev' into main2dev/06_06_2026
# Conflicts:
# megatron/training/arguments.py
Fuse per-sequence AlltoAll into a unified one in GDN forward (#4913)
[Dev] Add separate toggle for varlen input padding for HybridEP in THD training (#5048)
Co-authored-by: Xin Yao <xiny@nvidia.com>
docs: fix install guide NGC container anchor (#5224)
fix: wrap mtp logging comment
Add mtp_detach_heads config to detach MTP head inputs (#3456)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ci: add smoke tests (#5143)
varlendataset for thd e2e and benchmark (#4832)
Fix bug with Megatron-FSDP zero counter not working with decoupled gradients. (#4802)
Merge remote-tracking branch 'origin/dev' into main2dev/06_06_2026
# Conflicts:
# megatron/core/models/gpt/gpt_model.py
# megatron/core/ssm/mamba_mixer.py
# megatron/core/transformer/multi_token_prediction.py
# tests/unit_tests/transformer/test_multi_token_prediction.py
docs: Update Latest News in README.md (#3790)
fix: post-CI corrections
Minor improvements for Dynamic-cp (#4226)
fix(combined-1f1b): release loss-node input storage after combined backward (#4909)
fix(elastification): align with get_batch + utils refactors (#5194)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Enable selective recompute for `norm_out` in GDN layers (#4715)
Avoid stat syscall in rerun result validation (#5107)
chore(beep boop 🤖): Bump (main) (2026-06-08)
AI aided audit for Nvidia Style guidance (#5141)
Co-authored-by: Philip Petrakian <pgpetrak@gmail.com>
nvidia style guide audit for getting started folder (#5168)
Merge current dev into nightly sync
Resolve MoE op-fuser conflicts and carry over #5029 Deyu/Philip fixes.
[dev] [DeepSeek-v4] Add ClampedSwiGLU to MoE mlp_op_fuser and add force balance to hash routing (#5130)
fix: post-CI corrections (docs dup fields, nvrx version gate, utils re-export)
- Remove merge-introduced duplicate definitions of use_transformer_engine_op_fuser…
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
4 participants
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.
What (reworked per review)
Addresses #5968. Reworked as discussed with @hxbai: the fused kernels themselves have
moved to cudnn-frontend — NVIDIA/cudnn-frontend#427 (merged into
developatb950af1, 2026-07-30) adds them behind the FE-OSSAPIBasepattern (cudnn.csa.compressor,CSACompressorForward/Backward+high-level wrappers) — and this PR is now the thin Megatron-side dispatch only:
Compressor._forward_thd(THD packed, non-pre-grouped path) tries the fused fastpath for the gated-softmax pooling region — gather-index build, gather,
+ APE,overlap-window transform (
coff == 2), fp32 softmax, gated weighted sum — and keepsthe existing eager implementation as the semantic reference and fallback for every
unsupported configuration. One fused compute kernel per direction (plus a small
dAPEzero-init in backward) instead of ~40 forward / ~50 backward eager launches.the pre-grouped CP-prep input path are untouched.
Compared to the previous revision of this PR, the ~1000-line kernel module and the
benchmark harness are gone (they live in cudnn-frontend #427, harness included in its
test/bench assets); what remains is a ~240-line dispatch shim, the
csa.pyedit, and atrimmed test file: +732 / −37 vs
dev(previously +1930 / −37).Dependency
CSA: add fused Compressor forward+backward CuTe-DSL kernels (ported from Megatron-LM) cudnn-frontend#427, merged into
developatb950af1(2026-07-30). Thefused path activates when the installed cudnn-frontend provides
cudnn.csa— i.e. asource install of
develop>=b950af1, or the upcomingnvidia-cudnn-frontend1.27pip release (expected early August) — otherwise the dispatch silently keeps the eager
path. This includes the frontend's
nvidia-cutlass-dslruntime dependency (thefrontend's
cutedslextra). No hard dependency is added to Megatron: availabilityis probed at dispatch time by importing the concrete entry points
(
cudnn.csa.compressor.csa_compressor_forward_wrapper/..._backward_wrapper) —capability detection, not version comparison.
nvidia-cudnn-frontendinstalls thatpredate the API simply lack
cudnn.csa, and any import failure (missing package,missing DSL extra, partial install) silently keeps the eager path;
csa.py's importchain never fails because of the frontend.
torch.autograd.Functionaround thefrontend's forward/backward wrappers); the frontend APIs are pure
kernels-plus-validation.
Dispatch / gating (unchanged semantics vs the previous revision)
The fast path engages only for: THD packed non-pre-grouped path,
compress_ratio == 4(
coff == 2, the production CSA/HCA configuration), bf16kv/score, fp32ape,compute capability 10.0 (the frontend's validated envelope), int32 flat offsets
(
total_tokens * coff * head_dim < 2**31). Everything else — SBHD, the pre-groupedCP-prep path,
compress_ratio == 128(functionally supported by the frontend kernels,stays on eager until tuned), missing/old frontend, other devices — keeps the eager
implementation.
MCORE_CSA_FUSED_COMPRESSOR=0disables the dispatch entirely. Thedispatch also keeps the eager path under
torch.use_deterministic_algorithms(True)(
dAPEis accumulated with fp32 atomics in the fused backward) and undertorch.compiletracing (the frontend launch path takes raw pointers; eager lets thecompiler fuse the region itself). CUDA graphs: capture-compatible after a one-step
eager warmup per
(ratio, head_dim, coff)configuration (a first call that would JITunder capture raises loudly, per the frontend); the dispatch passes the caller's static
fixed_total_compcapacity through without device synchronization.Numerics
The numerics contract is the frontend kernels' (measured and tested in
NVIDIA/cudnn-frontend#427; original analysis in #5968): all arithmetic fp32 with a
single final bf16 rounding,
mul.rn/fma.rnpinned in PTX. vs an fp32-intermediateeager reference:
dKV/dScorebit-identical, forward within one bf16 roundingstep. Not bit-identical to the current eager region (which rounds softmax weights to
bf16 and multiplies in bf16) but at least as accurate against an fp64 oracle on every
tested output. Forward,
dKV,dScoreare bitwise run-to-run deterministic;dAPEisnot (fp32 atomics) — hence the deterministic-mode fallback above. Incoming gradients on
static-capacity padding rows are ignored (they are tail padding, not consumed
downstream); never-consumed input elements get exact zeros, matching autograd. These
gates are re-verified here at the dispatch level (see Tests).
Performance
The kernels this PR dispatches to are the ones measured in NVIDIA/cudnn-frontend#427
(same B200 methodology as before; the port also picked up two additional optimization
commits there — forward 32-bit vectorized access and backward kernel-side zero-writes —
so the numbers are better than this PR's previous revision). From #427, isolated GPU
kernel time (nsys, sum of kernel durations, 50 iterations after 20 warmup; THD packs of
8192-token sequences,
ratio = 4,coff = 2, bf16, vs the verbatim eager region ofCompressor._forward_thd):Wall-clock numbers, the hardware-ceiling audit (DRAM traffic within 1% of algorithmic
bytes), and per-commit bitwise gates are in #427's description.
Tests
tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py— trimmed to the Megatron-side wiring (kernel-level coverage — zero-write ownership,
CUDA-graph capture/replay, determinism,
coff == 1, fp64 oracle — lives in #427's testsuite):
before (
dKV/dScorebitwise vs an fp32-intermediate reference; forward within onebf16 rounding step; tolerance vs verbatim upstream numerics) over ragged THD packs
including segments shorter than
ratio;fixed_total_compstatic-capacity padding through the dispatch (forward semantics +ignored padding-row gradients);
cudnn.csa→ silent eager, verified bitwise-identical to the kill-switch path at theCompressor._forward_thdlevel), deterministic mode (dispatch falls back; enablingdeterministic mode between forward and backward raises loudly in the frontend),
ratio == 128, non-bf16, unexpected layout, empty output;Compressor._forward_thdintegration: fused engages (spy), matches eager, gradientsflow to inputs and
ape.The module keeps
pytestmark = pytest.mark.launch_on_gb200; without CUDA, without acudnn-frontend that provides
cudnn.csa, or off CC 10.0 every test skips cleanly (theGB200 CI lane needs
nvidia-cudnn-frontend>= 1.27 — the first release containing#427 — plus its
cutedslextra to actually exercise the fused path; until then thetests skip rather than fail).
Full runs on 1×B200 with cudnn-frontend
develop@b950af1(the #427 merge)installed from source: the new test file
8 passed; the related suites (
test_attention_variant_csa.py,test_csa_cp_layout_kernels.py,test_csa_cp_utils.py) with the fused dispatch live:123 passed, 1 skipped (needs ≥2 ranks), 0 failed. With no cudnn-frontend installed:
8 skipped, related suites unaffected (eager path is untouched).
Notes
MCORE_CSA_FUSED_COMPRESSOR=0), as in the previous revision — happy to flip toopt-in if you prefer.
develop@b950af1, 2026-07-30), sothe ordering constraint from the previous revision is satisfied; this PR remains a
no-op (eager everywhere) against any cudnn-frontend that predates the CSA API.
dAPEreduction for theratio=4 backward, wider arch enablement after validation. (The dedicated
compress_ratio == 128kernels, previously listed here, landed in [QUESTION] Does the zero-out operation conduct during the model initialization and affect reproducibility? #427 itself.)