perf(gated_delta_net): fold q/k L2-norm into the gated_delta_rule kernel

[yuchenwang3]

What

Move GatedDeltaNet q/k L2 normalization into the FLA gated_delta_rule kernel instead of materializing normalized q/k before the call.

The current GDN and GDN2 implementations share _prepare_input_for_gated_delta_rule, but GDN2's kernel path does not enable in-kernel normalization. The rebased change therefore makes deferral an explicit helper option:

• GDN defers normalization and passes use_qk_l2norm_in_kernel=self.use_qk_l2norm.
• GDN2 keeps its existing caller-side normalization and continues to pass False to its kernel.

This preserves GDN2 behavior added after the original PR rather than silently removing its normalization during the rebase.

Why

Caller-side normalization leaves an additional pre-split normalized query_key activation live for backward alongside the q/k tensors the gated-delta kernel already needs. Letting FLA own normalization removes that extra caller-side activation and keeps the normalization state with the kernel-owned q/k, which saves activation memory in long-context GDN-heavy hybrids.

The numerical contract is unchanged:

• both caller-side FLA l2norm and the in-kernel path use eps=1e-6;
• normalization is per head over the last dimension, so GQA head replication does not change the result;
• use_qk_l2norm=False remains a no-op.

Validation

The original change ran in Qwen3.5-35B-A3B 128K SFT on 16×B200. The current-main replay passes Black, isort, Ruff, Pylint, py_compile, and git diff --check.

The final 398a42b4 head was also exercised with an 8-rank torchrun on 8×B200:

• the GDN/GDN2 kernel flag-vs-explicit-normalization parity test passed on every rank;
• the BF16 end-to-end GatedDeltaNet.forward comparison passed on every rank, covering outputs, input gradients, all parameter gradients, and both routing flags;
• all 12 parameterized deterministic-mode GDN/GDN2 cases passed on every rank.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuchenwang3]

@xuantengh pinging you as the most recent committer on gated_delta_net.py. This folds the explicit q/k l2norm into the FLA gated_delta_rule kernel via use_qk_l2norm_in_kernel=True, dropping a materialized intermediate at long context; it ran in real Qwen3.5-35B-A3B hybrid training at 128K. One-file change, open since Jun 17 — could you review or /ok to test?

[xuantengh]

Thanks for the contribution! We had a refactor to the GDN module recently in #5843, would you please rebase that and apply this to the new file?

[yuchenwang3]

Done, @xuantengh — merged latest main and re-applied the change onto the refactored layout from #5843:

• common.py: the explicit l2norm() in _GDNBase._prepare_input_for_gated_delta_rule is removed (replaced by a NOTE explaining the fold and the eps=1e-6 equivalence).
• gdn.py: the gated_delta_rule call in GatedDeltaNet.forward now passes use_qk_l2norm_in_kernel=self.use_qk_l2norm.

Net diff vs main is the same two-line semantic change as before (+7/−4 with the comment block). Ready for another look.

[yuchenwang3]

@xuantengh one more follow-up before I refresh this again: current main added GDN2 sharing _prepare_input_for_gated_delta_rule. I've kept GDN2's caller-side normalization and only defer GDN to the kernel. Does that split match what you'd expect? If so I'll update the branch and request CI.

[yuchenwang3]

Rebased onto current main and adjusted the shared input helper added since the original PR: GDN now defers q/k normalization to its kernel, while GDN2 keeps the existing caller-side normalization. Local formatting, syntax, and routing checks pass. @xuantengh, could you re-review this head when you get a chance?

[xuantengh]

Sorry for the delayed reply. Would you please add a test to tests/unit_tests/ssm/test_gated_delta_net.py to validate toggling use_qk_l2norm can produce on-par results?

[yuchenwang3]

Done, @xuantengh. I added a small kernel-level parity check plus a BF16 end-to-end test that runs the real GatedDeltaNet.forward path. The latter compares the old caller-normalized route with kernel-owned normalization across outputs, input gradients, and every parameter gradient, and also asserts that both the helper and kernel receive the in-kernel flag and that raw q/k reach the kernel.

The new tests pass on all 8 ranks on 8×B200. While running the existing deterministic matrix I also caught and fixed a current FLA API mismatch in the newly activated fallback (l2norm has no dim argument); all 12 deterministic GDN/GDN2 cases now pass on every rank. The final head is 398a42b4 and DCO is green. Could you take another look and run /ok to test 398a42b4af753c76ffeeb14198d9b41a1bb2de1a when you get a chance?

[xuantengh]

Hi @yuchenwang3 thanks for your contribution. Generally LGTM, but currently there is no way to specify use_qk_norm_in_kernel at the module spec level. So would you please also add this as an optional knob (default to False, aligned with upstream FLA interface signature) to TransformerConfig and ModuleSpec. And we can toggle this via config in test, rather than the install_route_probe hook.

[yuchenwang3]

Done, @xuantengh — I rebased onto current main and added TransformerConfig.use_qk_l2norm_in_kernel with a False default, using the exact upstream FLA argument name. The value is forwarded through the GDN ModuleSpec.params and remains available from the config for other construction paths.

The BF16 parity test now builds the caller-normalized and kernel-normalized modules by toggling that config value through the real ModuleSpec path; the route-probe hook is gone. It still compares outputs, input gradients, and every parameter gradient. The module-spec test file is 47/47 locally, and formatting/static checks pass.

The updated head is edf2ff50cbfdba177c5f27eb8d01a82d792a6bcd. Could you take another look and run /ok to test edf2ff50cbfdba177c5f27eb8d01a82d792a6bcd?

[yuchenwang3]

The current head edf2ff50cbfdba177c5f27eb8d01a82d792a6bcd has now completed the reviewer-requested focused validation on 8×B200 with flash-linear-attention==0.5.1 (HAVE_FLA=True, HAVE_FLA_GDN2=True). Every rank finished with 15 passed, 0 failed, 0 skipped: the GDN/GDN2 normalization parity cases, the BF16 end-to-end config/ModuleSpec route, and all 12 deterministic GDN/GDN2 cases. The logs show the Blackwell FLA/TileLang kernel path executing; there were no tracebacks, NaNs, or OOMs, and the source tree remained clean.

Could you take another look and run /ok to test edf2ff50cbfdba177c5f27eb8d01a82d792a6bcd?