Skip to content

[OMNIML-2932] Fusing pre_quant_scale for NVFP4 AWQ #421

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 11 commits into
base: main
Choose a base branch
from
Open

[OMNIML-2932] Fusing pre_quant_scale for NVFP4 AWQ #421

wants to merge 11 commits into from
+332 −3

Conversation

@meenchen
Copy link
Contributor

@meenchen meenchen commented Oct 9, 2025 •
edited
Loading

What does this PR do?

Type of change: ?

Overview:

This PR and NVIDIA/TensorRT-LLM#8698 enable NVFP4 AWQ deployment for TRT-LLM. Specifically, this PR fuses pre_quant_scale in following two cases:

  • For MLP, pre_quant_scale of gate_proj layer is fused into up_proj's weight, so we don't need an extra handle in downstream fused moe kernels.
  • For attention, we will try to fuse the pre_quant_scale of o_proj to v_proj if their dimensions match, which means we will skip fusion for MQA/GQA models.

Usage

# Add a code snippet demonstrating how to use this

Testing

unit test, e2e test for Qwen3 dense and moe models.

Before your PR is "Ready for review"

  • Make sure you read and follow Contributor guidelines and your commits are signed.
  • Is this change backward compatible?: Yes/No
  • Did you write any new necessary tests?: Yes/No
  • Did you add or update any necessary documentation?: Yes/No
  • Did you update Changelog?: Yes/No

Additional Information

@copy-pr-bot
Copy link

copy-pr-bot bot commented Oct 9, 2025

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@coderabbitai
Copy link
Contributor

coderabbitai bot commented Oct 9, 2025 •
edited
Loading

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch weimingc/fuse_pqs

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@codecov
Copy link

codecov bot commented Oct 9, 2025 •
edited
Loading

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.44%. Comparing base (32d168c) to head (a591330).
⚠️ Report is 39 commits behind head on main.

Additional details and impacted files 
@@            Coverage Diff             @@
##             main     #421      +/-   ##
==========================================
+ Coverage   73.38%   73.44%   +0.06%     
==========================================
  Files         180      180              
  Lines       17934    18147     +213     
==========================================
+ Hits        13160    13328     +168     
- Misses       4774     4819      +45

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@meenchen meenchen self-assigned this Oct 14, 2025
@meenchen
pattern-based fusion
9bd8e41 
Signed-off-by: weimingc <[email protected]>
@meenchen
fix GQA
13042fa 
Signed-off-by: weimingc <[email protected]>
@meenchen
minor
d1c5d19 
Signed-off-by: weimingc <[email protected]>
@meenchen
unit test
e599d43 
Signed-off-by: weimingc <[email protected]>
@meenchen
fix doc
26f2eb7 
Signed-off-by: weimingc <[email protected]>
@meenchen
revert unintended change
c5d9682 
Signed-off-by: weimingc <[email protected]>
@meenchen
minor
6dd1b87 
Signed-off-by: weimingc <[email protected]>
@meenchen
resmooth
a5a6e39 
Signed-off-by: weimingc <[email protected]>
@meenchen meenchen requested a review from cjluo-nv October 27, 2025 19:29
@meenchen meenchen mentioned this pull request Oct 27, 2025
Open
1 task
@meenchen
fix moe fusion
6020e94 
Signed-off-by: weimingc <[email protected]>
@meenchen meenchen changed the title Pattern-based fusion for pre_quant_scale Fusing pre_quant_scale for NVFP4 AWQ Nov 3, 2025
@meenchen meenchen changed the title Fusing pre_quant_scale for NVFP4 AWQ [OMNIML-2932] Fusing pre_quant_scale for NVFP4 AWQ Nov 3, 2025
@meenchen meenchen marked this pull request as ready for review November 3, 2025 23:39
@meenchen meenchen requested a review from a team as a code owner November 3, 2025 23:39
@meenchen
fix test
9339223 
Signed-off-by: weimingc <[email protected]>
@meenchen
only fuse for nvfp4 awq
a591330 
Signed-off-by: weimingc <[email protected]>
cjluo-nv
cjluo-nv reviewed Nov 4, 2025
View reviewed changes
modelopt/torch/export/quant_utils.py
kv_head_dim = linear_fuse_into.weight.shape[0] // num_kv_heads
n_rep = pre_quant_scale.numel() // num_kv_heads // kv_head_dim

# Reshape:(num_kv_heads, n_rep, kv_head_dim)
Copy link
Collaborator

@cjluo-nv cjluo-nv Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what's n_rep here?

cjluo-nv
cjluo-nv reviewed Nov 4, 2025
View reviewed changes
modelopt/torch/export/quant_utils.py
old_pre_quant_scale = module.input_quantizer._pre_quant_scale
module.weight = nn.Parameter(
module.weight
* old_pre_quant_scale.to(
Copy link
Collaborator

@cjluo-nv cjluo-nv Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we want to cast to fp32 for this manipulation?

cjluo-nv
cjluo-nv reviewed Nov 4, 2025
View reviewed changes
tests/gpu/torch/export/test_quant_utils.py
@@ -0,0 +1,193 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Copy link
Collaborator

@cjluo-nv cjluo-nv Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need update https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py#L43 as well?

cjluo-nv
cjluo-nv reviewed Nov 4, 2025
View reviewed changes
modelopt/torch/export/quant_utils.py
.reshape(-1)
)

def _update_pre_quant_scale(module, new_pre_quant_scale):
Copy link
Collaborator

@cjluo-nv cjluo-nv Nov 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we merge duplicated code with line 1090?

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

Reviewers

@cjluo-nv cjluo-nv cjluo-nv left review comments

@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1

@realAsma realAsma Awaiting requested review from realAsma

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

Assignees

@meenchen meenchen

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@meenchen @cjluo-nv