Skip to content

MuonClip support (non-split version) #1929

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

Jump to bottom
Open
wants to merge 23 commits into
base: main
Choose a base branch
from
Open

MuonClip support (non-split version) #1929

wants to merge 23 commits into from

Conversation

@BoxiangW
Copy link
Contributor

@BoxiangW BoxiangW commented Oct 24, 2025

Added MLA and MHA(GQA) clipping support

@BoxiangW
Add qk-clip
0ef342c 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
w/o EP working
a5d2ea8 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Lint
978b99b 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Fix MLA usage and add logging for max_attention_score
b77a6d6 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Removed Attention MuonClip since it is not correct. Added MLA MuonClip
a83b68c 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Add TE bug WAR and make sure MLA works
eae6a5a 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Fix qkclip issue
c370cdd 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Fix bug
6f8f77d 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Fix TP usage
527075f 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Fix Lint
3cebc1f 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Fix import and error log
e3c468f 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
lint
ef995ea 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Address comments
7f5f1e2 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Lint
65fc829 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Added GQA QK Clipping
75ee38c 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Lint
a1c0ff6 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW BoxiangW requested review from a team as code owners October 24, 2025 22:00
@copy-pr-bot
Copy link

copy-pr-bot bot commented Oct 24, 2025

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.

@BoxiangW BoxiangW self-assigned this Oct 24, 2025
@BoxiangW
Copy link
Contributor Author

BoxiangW commented Oct 24, 2025

TE's NVIDIA/TransformerEngine#2195 (2.9.0) is needed for this PR

@BoxiangW
Remove comment
762bdb5 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Rename max_score to max_logit
ba5c073 
Signed-off-by: Boxiang Wang <[email protected]>
@BoxiangW
Change name
7917e68 
Signed-off-by: Boxiang Wang <[email protected]>
@skyw
Copy link
Contributor

skyw commented Oct 28, 2025

TE's NVIDIA/TransformerEngine#2195 (2.9.0) is needed for this PR

It has been merged.

@BoxiangW BoxiangW added Expert Review Apply this label to indicate that your PR is ready for expert review. Run tests labels Oct 31, 2025
@BoxiangW BoxiangW added this to the Core 0.15 milestone Oct 31, 2025
@BoxiangW
Copy link
Contributor Author

BoxiangW commented Nov 3, 2025

/ok to test 7917e68

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@skyw skyw Awaiting requested review from skyw

@mkhona-nvidia mkhona-nvidia Awaiting requested review from mkhona-nvidia

At least 1 approving review is required to merge this pull request.

Assignees

@BoxiangW BoxiangW

Labels

Expert Review Apply this label to indicate that your PR is ready for expert review. Run tests

Projects

None yet

Milestone

Core 0.15

Development

Successfully merging this pull request may close these issues.

2 participants

@BoxiangW @skyw