[Dev] MuonClip support (non-split version) on dev branch - #2194

Merged

BoxiangW merged 31 commits into

BoxiangW:boxiangw/muon_clip_dev

Nov 18, 2025

Commits on Nov 10, 2025

Add qk-clip on dev branch
BoxiangW
committed
7391760
w/o EP working
BoxiangW
committed
28b43f7
Lint
BoxiangW
committed
13532bc
Fix MLA usage and add logging for max_attention_score
BoxiangW
committed
18b839a
Removed Attention MuonClip since it is not correct. Added MLA MuonClip
BoxiangW
committed
50db1cb
Add TE bug WAR and make sure MLA works
BoxiangW
committed
148fe7c
Fix qkclip issue
BoxiangW
committed
462e1f2
Fix bug
BoxiangW
committed
af7c58b
Fix TP usage
BoxiangW
committed
6a523cf
Fix Lint
BoxiangW
committed
8808701
Fix import and error log
BoxiangW
committed
d8425ff
lint
BoxiangW
committed
5fdf802
Address comments
BoxiangW
committed
0898ae0
Lint
BoxiangW
committed
1cfb40c
Added GQA QK Clipping
BoxiangW
committed
5bd3c05
Lint
BoxiangW
committed
87fc588
Remove comment
BoxiangW
committed
5ee6ca3
Rename max_score to max_logit
BoxiangW
committed
94fcd31
Change name
BoxiangW
committed
2e19110
Lint
BoxiangW
committed
e520c81
Fix copyright
BoxiangW
committed
4bd9acf
Move qk_clip fucntion into megatron/core
BoxiangW
committed
605aae3
Add tests for qk_clip
BoxiangW
committed
a2dc046
Lint and copyright
BoxiangW
committed
3e15d00
Add te version checks into PR
BoxiangW
committed
e13d178

Commits on Nov 12, 2025

Address comments
BoxiangW
committed
94b6589
Address comments
BoxiangW
committed
e79629d

Commits on Nov 17, 2025

DP all reduce and switch to mul_ inplace op
BoxiangW
committed
264fddb
Merge branch 'dev' into boxiangw/muon_clip_dev
BoxiangW
authored
7fd7661
Update both main_params and non
BoxiangW
committed
841d62e

Commits on Nov 18, 2025

Address comments
BoxiangW
committed
b389ebf