Merge GDN conv fusion (fused pre-gated-delta-rule kernels) - #13
Closed
wplf wants to merge 23 commits into
Closed
Conversation
…ining (NVIDIA#4816) Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>
Signed-off-by: conver334 <conver334@gmail.com>
…#4927) Signed-off-by: Xin Yao <xiny@nvidia.com>
…VIDIA#4894) Co-authored-by: Kunlun Li <kunlunl@nvidia.com>
…NVIDIA#3835) Signed-off-by: jinliangl <jinliangl@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
…VIDIA#5013) Co-authored-by: Xin Yao <xiny@nvidia.com>
…ned backward (NVIDIA#4908) Co-authored-by: Xin Yao <xiny@nvidia.com>
… Attention (NVIDIA#5018) Co-authored-by: Yuzhong Wang <yuzhongw@nvidia.com>
Co-authored-by: BestJuly <19769279+BestJuly@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: oliver könig <okoenig@nvidia.com>
Replace the placeholder with the all-in-one ("mega") pre-GDR path:
- Forward: a single Triton launch over a flat logical task space
(QK conv+silu+l2norm+repeat, V conv+silu, Z copy, g/beta), supporting
dense BSHD and packed THD. Replaces the streamed path's four launches.
- Backward: one fused branch kernel (QK l2norm/repeat, V layout, Z layout,
g/beta chain rule; per-head d_A_log/d_dt_bias via atomics) plus the
external causal_conv1d backward for conv input/weight gradients. Two
launches, down from five.
Outputs and saved tensors match the streamed path bit-for-bit, so the
public (query, key, value, gate, beta, g) contract and parity are preserved.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The pre_gated_delta_rule_impl selector is already exposed as a CLI arg by ArgumentGroupFactory(TransformerConfig), which auto-derives --pre-gated-delta-rule-impl from the dataclass field. Registering it manually raised an argparse "conflicting option string" error at parse time, so only a clarifying comment is kept. Recipes / passthrough can set the impl via the auto-generated arg. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Owner
Author
|
Superseded by #14, which cherry-picks/squashes only Yuzhong Wang's GDN-fusion commits onto central-dev (dropping the unrelated dev syncs), integrates the fused dispatch into the existing gdn_qkv-recompute forward, and is verified green on GB200 (fused-vs-unfused parity + recompute regression). Closing this full-branch merge in favor of #14. |
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
10 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.
Merge
yuzhongw/gdn_conv_fusion(fromyuzhongw-nvidia/Megatron-LM) intojinliangl/qwen35-vl-central-dev.Headline feature — fused pre-gated-delta-rule (GDN conv fusion)
megatron/core/fusions/fused_pre_gated_delta_rule.py(new, ~2.2k lines)megatron/core/fusions/fused_mega_pre_gated_delta_rule.py(new, ~1.1k lines) — mega-fused forward + backward kernelsmegatron/core/ssm/gated_delta_net.py— wire the fused pre-GDR path (incl. packed THD)transformer_config.py+arguments.py—--pre-gated-delta-rule-impl(auto-generated from config)tests/unit_tests/ssm/test_gated_delta_net.py— packed fused GDN path testsScope note
The source branch diverged from this base at
bd5c98f77and is 23 commits ahead / 31 commits behind. Besides the GDN-fusion commits it also carries newerdevsyncs (e.g. NVIDIA#4751 Qwen3.5-VL, NVIDIA#5018, NVIDIA#5051, NVIDIA#4894, NVIDIA#4908 …), so the base…head diff is large (~219 files). The two histories have diverged, so this merge may surface conflicts that need resolution before it can be merged.🤖 Generated with Claude Code