Add QK layernorm support for dot-product attention in MambaModel - #4067

Merged

Phlip79 merged 21 commits into

NVIDIA:mainfrom

Phlip79:philip/add-qk-norm

Apr 17, 2026

Commits on Mar 31, 2026

Add QK layernorm support for dot-product attention in MambaModel

Phlip79
and
claude
committed
301d12e
Add tests for MambaModel QK layernorm spec functions

Phlip79
and
claude
committed
367b8a8
Fix linting
Phlip79
committed
1c46827
Fix linting again
Phlip79
committed
b28a11f
Ran isort
Phlip79
committed
97607d6
Update qk layernorm settings and tests
Phlip79
committed
bdf2d12
Fix linting
Phlip79
committed
488e448
Fix linting again
Phlip79
committed
e9c20a4
Fix TENorm isinstance check in QK layernorm test

Phlip79
and
claude
committed
d16170f

Commits on Apr 3, 2026

Use blueprint system instead
Phlip79
committed
2ee490b
Update comments and failure
Phlip79
committed
c778f9c
Merge branch 'main' of https://github.com/NVIDIA/Megatron-LM into philip/add-qk-norm
Phlip79
committed
cb2c31e

Commits on Apr 7, 2026

Update ordering
Phlip79
committed
046f443

Commits on Apr 10, 2026

Add TODO

Phlip79
and
yuzhongw-nvidia
authored
cd26bb4

Commits on Apr 16, 2026

Refactor QK norm resolution: explicit branches + ValueError

Phlip79
and
claude
committed
d09c4ca
Merge remote-tracking branch 'origin/philip/add-qk-norm' into philip/add-qk-norm
Phlip79
committed
861fe3b
Update test_l2_qk_norm to set qk_l2_norm on config

Phlip79
and
claude
committed
a49723f

Commits on Apr 17, 2026

Set qk_l2_norm on shared test config instead of duplicating

Phlip79
and
claude
committed
e16a762
Add tests for config-driven QK norm fallback and spec/config mismatch

Phlip79
and
claude
committed
236c50f
Fix black formatting in test_qk_layernorm_spec_config_mismatch_raises

Phlip79
and
claude
committed
c8a313d
Fix TENorm isinstance check and add L2Norm config-fallback test

Phlip79
and
claude
committed
1ca9061