[port] Add GDN2 (Gated DeltaNet-2) attention variant (upstream #5765)

[HJSang]

Adds experimental_attention_variant="gdn2": channel-wise decay/erase/write gates on top of the existing Gated Delta Net (GDN) layer family, backed by fla.ops.gdn2.chunk.chunk_gdn2 (flash-linear-attention >= 0.5.1) with a torch-native deterministic-mode fallback.

Not a literal cherry-pick

This is a hand-adapted port of upstream 0c1fc4c5c9 (PR NVIDIA#5765), not a direct cherry-pick. Our common.py/gdn.py base already carries a later, independently-cherry-picked refactor (3549dc62a, upstream PR NVIDIA#5843) that generalizes headwise/chunkwise context parallelism, selective recompute of the output norm, and the fused streamed pre-gated-delta-rule path — none of which existed yet when 0c1fc4c5c9 was authored upstream. GatedDeltaNet2 is written against that refactored base (subclassing GatedDeltaNet to reuse its forward()/CP/recompute plumbing) rather than against 0c1fc4c5c9's pre-refactor gdn2.py.

Deliberately excluded

• Upstream PR Fused GDN attention support NVIDIA/Megatron-LM#6645 ("Fused GDN attention support") and its revert ci: Revert recently merged changes as they cause CI unit test failures.  NVIDIA/Megatron-LM#7326: Fused GDN attention support NVIDIA/Megatron-LM#6645 was merged to upstream/main and then reverted for CI failures, so it is not part of main's current working state.
• GDN2 chunkwise context parallelism and gdn_pre_gated_delta_rule_fusion: neither existed in upstream's original GDN2 commit; both are rejected with a clear assertion in _setup_variant_attrs rather than silently mishandled.
• Dedicated GDN2 unit test coverage: upstream's test changes targeted the pre-refactor flat test_gated_delta_net.py, which our tree had already split into tests/unit_tests/ssm/gated_delta_net/{test_gdn.py,...}. Porting that coverage is left as a follow-up.

Validation

Not yet run. py_compile is clean on every touched file and the static cross-references (base-class method/attribute signatures, config field wiring) have been checked by hand, but this has not been executed — no unit tests, no training smoke test. Needs a GPU box with flash-linear-attention >= 0.5.1 (fla.ops.gdn2) before merging.

Series

Independent of #port/kda (opened separately), but port/kda is stacked on top of this branch and will show this commit in its diff until this PR merges.

🤖 Generated with Claude Code

https://claude.ai/code/session_01UawmJvYe7WkXApRw26ZFMd

[HJSang]

Update 2026-10-05. Rebased onto current miles-main (fd15ee20a) to pick up #102's new is_gated_delta_net_variant/GATED_DELTA_NET_VARIANTS (now the canonical home for this check, replacing the duplicate copy this PR originally carried in experimental_attention_variant_module_specs.py) and #108's DSA top-k work; both files now reconcile cleanly. Extended GATED_DELTA_NET_VARIANTS to include "gdn2".

Also fixed a bug from the original push: gdn2.py's adapted content (subclassing GatedDeltaNet, headwise/chunkwise-CP-aware) had been written to disk but never re-staged after the initial cherry-pick auto-added the original pre-adaptation upstream blob, so the first version of this PR actually contained the wrong, incompatible gdn2.py. Caught and fixed before any review/merge.

Mergeable against miles-main again. Still unvalidated on GPU (see PR description).

[HJSang]

Closing — GDN2 isn't adopted by any open-source model yet (no model attribution found upstream or in this repo; see discussion in this session). Nothing concrete to validate against right now. Can reopen and rebase if/when a model actually ships with it.

Note: #110 (KDA) is unaffected — Kimi does use KDA, so that one stays open. Its branch (port/kda) carries its own copy of the GDN2 commit, so closing this PR doesn't affect it.