fix: guard PRE_MLP NVFP4 fusion for dense layers - #17191
Conversation
Signed-off-by: zq <zhouquan1511@163.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe Exaone-MoE and GLM MLP paths now select quantized or standard residual RMSNorm fusion from ChangesNVFP4 pre-MLP fusion
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py`:
- Around line 40-41: Strengthen the test around the guarded pre-MLP behavior by
verifying that `input_scale` is accessed only within the true `has_nvfp4`
branch, rather than relying on source-text ordering. Also exercise the false
`has_nvfp4` path with a raising `input_scale` sentinel and assert the
`RESIDUAL_RMS_NORM` fallback behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5280d550-7e2a-4ed6-90b4-2161dcd865e1
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_exaone_moe.pytensorrt_llm/_torch/models/modeling_glm.pytests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py
Signed-off-by: zq <zhouquan1511@163.com>
|
Thanks for the review — I pushed Changes made:
Validation:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py (1)
37-47: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winScope the AST lookup to the PRE_MLP fusion branch.
_forward_mlp_has_nvfp4_branch()returns the first matchinghas_nvfp4branch inforward_mlp. It does not verify that the branch is nested underself.fusion_config.PRE_MLP_FUSION. A later unrelated branch could satisfy the test while PRE_MLP fusion is unguarded. Restrict the search to the expected outer branch before scanning its children.This keeps the regression check tied to the PRE_MLP fusion requirement.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py` around lines 37 - 47, Update _forward_mlp_has_nvfp4_branch to first locate the self.fusion_config.PRE_MLP_FUSION outer branch within forward_mlp, then search only that branch’s descendants for self.mlp.gate_up_proj.has_nvfp4. Return the matching ast.If and fail the assertion if the expected PRE_MLP branch or nested NVFP4 guard is absent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py`:
- Around line 58-69: Add test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp
to a suitable CI and QA test-list entry under
tests/integration/test_lists/test-db, ensuring all three model parameters are
covered. Then run pytest tests/unittest/ to validate the updated coverage.
---
Nitpick comments:
In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py`:
- Around line 37-47: Update _forward_mlp_has_nvfp4_branch to first locate the
self.fusion_config.PRE_MLP_FUSION outer branch within forward_mlp, then search
only that branch’s descendants for self.mlp.gate_up_proj.has_nvfp4. Return the
matching ast.If and fail the assertion if the expected PRE_MLP branch or nested
NVFP4 guard is absent.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0669592d-7a22-4622-9391-27e7cc5b4cd3
📒 Files selected for processing (1)
tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py
| def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(model_file: str) -> None: | ||
| nvfp4_branch = _forward_mlp_has_nvfp4_branch(model_file) | ||
|
|
||
| assert _contains_attribute(nvfp4_branch, "self.mlp.gate_up_proj.input_scale") | ||
| assert _contains_attribute(nvfp4_branch, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4") | ||
|
|
||
| for false_branch_node in nvfp4_branch.orelse: | ||
| assert not _contains_attribute(false_branch_node, "self.mlp.gate_up_proj.input_scale") | ||
| assert not _contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4") | ||
| assert any( | ||
| _contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM") | ||
| for false_branch_node in nvfp4_branch.orelse) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
test_name='test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp'
test_file='test_pre_mlp_nvfp4_fusion.py'
for list_dir in \
tests/integration/test_lists/test-db \
tests/integration/test_lists/qa; do
echo "== $list_dir =="
rg -n --fixed-strings "$test_name" "$list_dir" || true
rg -n --fixed-strings "$test_file" "$list_dir" || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 238
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== changed files =="
git status --short
git diff --stat
echo "== test file outline =="
ast-grep outline tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py
echo "== test file context =="
sed -n '1,140p' tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py
echo "== test-list files =="
find tests/integration/test_lists -maxdepth 2 -type f \( -path '*/test-db/*' -o -path '*/qa/*' \) -print | sort | head -80
echo "== relevant list entries =="
rg -n -i 'pre[_-]mlp|nvfp4|fusion' tests/integration/test_lists/test-db tests/integration/test_lists/qa || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== test-list documentation =="
sed -n '1,220p' tests/integration/test_lists/test-db/README.md
sed -n '1,220p' tests/integration/test_lists/qa/README.md
echo "== list and CI references to unit-test paths =="
rg -n --glob '*.yml' --glob '*.yaml' --glob '*.md' --glob '*.py' \
'tests/integration/test_lists|unittest/' .github tests scripts 2>/dev/null \
| head -160 || true
echo "== exact file and function references in all tracked files =="
rg -n --fixed-strings \
'test_pre_mlp_nvfp4_fusion.py|test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp' \
. --glob '!tests/integration/test_lists/**' || true
echo "== deterministic AST check for the three model files =="
python3 - <<'PY'
import ast
from pathlib import Path
models = [
"modeling_deepseekv3.py",
"modeling_glm.py",
"modeling_exaone_moe.py",
]
root = Path("tensorrt_llm/_torch/models")
def path(node):
if isinstance(node, ast.Name):
return node.id
if isinstance(node, ast.Attribute):
prefix = path(node.value)
return f"{prefix}.{node.attr}" if prefix else node.attr
return ""
def contains(node, expected):
return any(path(child) == expected for child in ast.walk(node))
for model in models:
module = ast.parse((root / model).read_text())
branches = []
for function in ast.walk(module):
if isinstance(function, ast.FunctionDef) and function.name == "forward_mlp":
for child in ast.walk(function):
if path(getattr(child, "test", ast.Constant(None))) == \
"self.mlp.gate_up_proj.has_nvfp4":
branches.append(child)
print(model, "matching_branches=", len(branches))
for branch in branches:
true_checks = [
contains(branch, "self.mlp.gate_up_proj.input_scale"),
contains(branch, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4"),
]
false_checks = [
not contains(node, "self.mlp.gate_up_proj.input_scale")
and not contains(
node, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4")
for node in branch.orelse
]
fallback = any(
contains(node, "AllReduceFusionOp.RESIDUAL_RMS_NORM")
for node in branch.orelse
)
print(" true_checks=", true_checks,
"false_checks=", false_checks,
"fallback=", fallback)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 24521
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== modeling-directory coverage in test lists =="
rg -n --fixed-strings \
'unittest/_torch/modeling' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa \
|| true
echo "== exact target path coverage in test lists =="
rg -n --fixed-strings \
'unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa \
|| true
echo "== relevant tests README guidance =="
sed -n '135,175p' tests/README.mdRepository: NVIDIA/TensorRT-LLM
Length of output: 10503
Add the unit test to the CI test list.
test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp is absent from both CI and QA lists. Add it to a suitable tests/integration/test_lists/test-db/*.yml entry and run pytest tests/unittest/.
Coverage summary: one modified test with three model parameters. It checks NVFP4 and fallback branches. Coverage verdict: insufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py` around lines 58
- 69, Add test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp to a suitable
CI and QA test-list entry under tests/integration/test_lists/test-db, ensuring
all three model parameters are covered. Then run pytest tests/unittest/ to
validate the updated coverage.
Source: Path instructions
Signed-off-by: zq <zhouquan1511@163.com>
|
Thanks for the review — I pushed Changes made:
Validation:
Note: plain root-level pytest in this local Windows environment is blocked by unrelated test-harness imports ( |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Fix looks right. I confirmed the predicate holds up: apply_quant_config_exclude_modules (tensorrt_llm/_torch/models/modeling_utils.py:619) replaces the excluded module's quant_config, so Linear.has_nvfp4 really does go False for model.layers.0*-style exclusions, and the fallback still issues an allreduce, keeping disable_attn_allreduce consistent. The GLM/ExaOne change is a verbatim copy of the guard already on main in modeling_deepseekv3.py:1607, which is the right call for consistency.
Worth stating in the description that the crash reported in #15516 is in the DeepSeekV3 path and is already fixed on main — this PR is parity for the two other models with the same pattern, plus a test that also pins the existing DeepSeekV3 guard. As written it reads like the fix for the reported crash.
Follow-up, out of scope here: modeling_llama.py:838 gates the same gate_up_proj.input_scale read on the layer-level self.is_nvfp4, so it has the identical latent failure for checkpoints with submodule-level exclusions.
Comments below are on the test, not the fix.
| - unittest/_torch/attention/sparse/deepseek_v4/test_compressor_tf32.py TIMEOUT (15) | ||
| - unittest/_torch/attention/sparse/test_sparse_mla_forward.py TIMEOUT (60) | ||
| - unittest/_torch/modeling/test_modeling_deepseekv4.py | ||
| - unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py |
There was a problem hiding this comment.
This test is pure ast.parse over source files — no GPU, no torch, 0.02s locally. Registering it on the B200 pre-merge list spends scarce Blackwell capacity on something that can't fail differently there. tests/integration/test_lists/test-db/l0_cpu_x86.yml (and l0_cpu_arm.yml) is the right home.
| from pathlib import Path | ||
|
|
||
| import pytest | ||
|
|
There was a problem hiding this comment.
parents[4] hardcodes the test's depth in the tree and assumes a source checkout rooted there. Moving the file one level silently points at the wrong path (or FileNotFoundError), and if a CI stage runs against an installed tensorrt_llm wheel this asserts on repo source that isn't necessarily what's imported. Derive it from the package instead:
import tensorrt_llm._torch.models as _models
_MODELS_DIR = Path(_models.__file__).parentand drop the "tensorrt_llm" / "_torch" / "models" join in _forward_mlp_has_nvfp4_branch.
| def _contains_attribute(node: ast.AST, attribute_path: str) -> bool: | ||
| return any(_attribute_path(child) == attribute_path for child in ast.walk(node)) | ||
|
|
||
|
|
There was a problem hiding this comment.
_forward_mlp_has_nvfp4_branch matches on the exact attribute path self.mlp.gate_up_proj.has_nvfp4, so any equivalent refactor — hoisting the check to a local, getattr(gu, "has_nvfp4", False) as modeling_gemma4.py:831 does, or inverting to if not ... — fails the test with a message claiming the guard is missing. That's the accepted cost of AST-pinning something that needs TP>1 + an NVFP4 checkpoint to exercise, but the raise AssertionError(f"{model_file} does not guard PRE_MLP NVFP4 fusion") message should say what shape it was looking for so the next person doesn't go hunting for a regression that isn't there.
Signed-off-by: zq <zhouquan1511@163.com>
|
Thanks for the review — I pushed Changes made:
Validation:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py`:
- Around line 47-58: Update _forward_mlp_nvfp4_aliases to recursively inspect
nested control-flow statements while skipping nested functions and classes,
collecting matching assignment and annotated-assignment targets throughout
forward_mlp. Rename the fixture alias has_nvfp4 to a distinct name such as
use_nvfp4 so alias tracking is exercised.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d19b771b-af32-44f1-995b-05bff941973f
📒 Files selected for processing (3)
tests/integration/test_lists/test-db/l0_cpu_arm.ymltests/integration/test_lists/test-db/l0_cpu_x86.ymltests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py
Signed-off-by: zq <zhouquan1511@163.com>
|
Thanks for the follow-up — I pushed Changes made:
Validation:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py`:
- Around line 63-66: Update _forward_mlp_nvfp4_aliases and its callers to
preserve source order and branch scope, accepting an alias only when its
assignment dominates the candidate if use_nvfp4 branch; reject
use-before-assignment and assignments from sibling branches. Add fixtures
covering both invalid cases and register the test file in the l0_cpu_arm.yml and
l0_cpu_x86.yml test lists.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4ba9b23b-40cc-44c7-b41d-49bcea27bc3c
📒 Files selected for processing (1)
tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py
Signed-off-by: zq <zhouquan1511@163.com>
|
Thanks for the review — I pushed Changes made:
Validation:
|
3 participants
Summary
gate_up_proj.has_nvfp4module state.Fixes #15516
Test Plan
python -m pytest --confcutdir=tests/unittest/_torch/modeling tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py -qpython -m ruff check tensorrt_llm/_torch/models/modeling_glm.py tensorrt_llm/_torch/models/modeling_exaone_moe.py tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.pypython -m py_compile tensorrt_llm/_torch/models/modeling_glm.py tensorrt_llm/_torch/models/modeling_exaone_moe.py tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.pyDev Engineer Review
has_nvfp4.RESIDUAL_RMS_NORMpath.input_scaleaccess.QA Engineer Review
test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp, parameterized for DeepSeekV3, GLM, and ExaOne.test_forward_mlp_has_nvfp4_branch_supports_alias_and_getattr.tests/integration/test_lists/test-db/l0_cpu_arm.ymltests/integration/test_lists/test-db/l0_cpu_x86.yml