Skip to content

[dev] Refactor: extract and split common logic between GDN & GDN2 (cherry-pick #5843) #6088

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuzhongw-nvidia merged 22 commits into from
Aug 17, 2026
Merged

[dev] Refactor: extract and split common logic between GDN & GDN2 (cherry-pick #5843) #6088

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
a46b579
refactor(ssm): replace GDN files with PR 5765
yuzhongw-nvidia Aug 3, 2026
c615423
chore(ssm): reapply dev GDN delta after PR 5765 split
yuzhongw-nvidia Aug 3, 2026
05307f4
Update common.py
yuzhongw-nvidia Aug 4, 2026
b530893
fix(ssm): address GDN review comments
yuzhongw-nvidia Aug 10, 2026
2ea70d3
chore(ssm): sync latest PR 5765 updates
yuzhongw-nvidia Aug 11, 2026
d2ac2bc
chore(deps): refresh uv lock after dev rebase
yuzhongw-nvidia Aug 11, 2026
13b1952
Update megatron/core/ssm/gated_delta_net/common.py
yuzhongw-nvidia Aug 11, 2026
674a3dc
fix(ssm): address GDN follow-up review comments
yuzhongw-nvidia Aug 11, 2026
e43376c
fix(ssm): use runtime CP size in GDN paths
yuzhongw-nvidia Aug 11, 2026
3244f7e
fix(ssm): clean up GDN refactor review nits
yuzhongw-nvidia Aug 11, 2026
9ef0226
fix(ssm): narrow GDN norm recompute region
yuzhongw-nvidia Aug 11, 2026
782e792
chore(deps): drop unrelated GDN dependency bumps
yuzhongw-nvidia Aug 11, 2026
40fb1a9
fix(ssm): revert over-scoped GDN review changes
yuzhongw-nvidia Aug 11, 2026
b3cc28b
Apply suggestions from code review
yuzhongw-nvidia Aug 11, 2026
57af510
chore(ssm): apply autoformat fixes
yuzhongw-nvidia Aug 11, 2026
07ff7a2
Update __init__.py
yuzhongw-nvidia Aug 11, 2026
de038b1
fix(ssm): skip deterministic GDN test for CP
yuzhongw-nvidia Aug 12, 2026
080ed23
fix(ssm): keep torch GDN rule qkv names
yuzhongw-nvidia Aug 12, 2026
5484feb
fix(ssm): restore pre-GDR tuple contract
yuzhongw-nvidia Aug 13, 2026
e398a42
test(ssm): split gated delta net unit tests
yuzhongw-nvidia Aug 13, 2026
fe54f6d
test(ssm): group gated delta net tests
yuzhongw-nvidia Aug 14, 2026
1aa8afe
fix(deps): install TileLang for FLA on Hopper
yuzhongw-nvidia Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1,453 changes: 0 additions & 1,453 deletions megatron/core/ssm/gated_delta_net.py
View file
Open in desktop

This file was deleted.

36 changes: 36 additions & 0 deletions megatron/core/ssm/gated_delta_net/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

"""Gated Delta Net (GDN) layer exports.

This package replaces the former ``megatron/core/ssm/gated_delta_net.py`` module
at the same import path; the names below preserve that module's public surface.
"""

from megatron.core.ssm.gated_delta_net.common import (
HAVE_FLA,
GatedDeltaNetSubmodules,
_build_head_perm_for_split_sections,
_build_thd_cp_a2a_perm,
_split_tensor_factory,
causal_conv1d,
chunk_gated_delta_rule,
get_parameter_local_cp,
l2norm,
tensor_a2a_cp2hp,
tensor_a2a_hp2cp,
)
Comment on lines +9 to +21

@claude claude Bot Aug 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The package docstring says these names "preserve that module's public surface", but the old gated_delta_net.py also exported _build_head_perm_for_split_sections and _build_thd_cp_a2a_perm (both were imported from it by tests/unit_tests/ssm/test_gated_delta_net.py before this PR, and the test had to be updated to reach into .common). Re-exporting them here would make the path compatible for any out-of-tree importer:

from megatron.core.ssm.gated_delta_net.common import (
    HAVE_FLA,
    GatedDeltaNetSubmodules,
    _build_head_perm_for_split_sections,
    _build_thd_cp_a2a_perm,
    _split_tensor_factory,
    causal_conv1d,
    chunk_gated_delta_rule,
    get_parameter_local_cp,
    l2norm,
    tensor_a2a_cp2hp,
    tensor_a2a_hp2cp,
)

(plus the two names in __all__). Not blocking if you consider the underscore-prefixed helpers private.

from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet, torch_chunk_gated_delta_rule

__all__ = [
"HAVE_FLA",
"GatedDeltaNet",
"GatedDeltaNetSubmodules",
"_split_tensor_factory",
"causal_conv1d",
"chunk_gated_delta_rule",
"get_parameter_local_cp",
"l2norm",
"tensor_a2a_cp2hp",
"tensor_a2a_hp2cp",
"torch_chunk_gated_delta_rule",
]
Comment thread
yuzhongw-nvidia marked this conversation as resolved.
Loading
Loading