From 1d01e7a7c154959c8d9907c942fc0e7e1b63be27 Mon Sep 17 00:00:00 2001 From: zq Date: Mon, 3 Aug 2026 15:23:03 +0800 Subject: [PATCH 1/8] fix: guard PRE_MLP NVFP4 fusion for dense layers Signed-off-by: zq --- .../_torch/models/modeling_exaone_moe.py | 33 ++++++++++----- tensorrt_llm/_torch/models/modeling_glm.py | 33 ++++++++++----- .../modeling/test_pre_mlp_nvfp4_fusion.py | 41 +++++++++++++++++++ 3 files changed, 85 insertions(+), 22 deletions(-) create mode 100644 tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py diff --git a/tensorrt_llm/_torch/models/modeling_exaone_moe.py b/tensorrt_llm/_torch/models/modeling_exaone_moe.py index 40ae3653d6e0..67596c64b457 100644 --- a/tensorrt_llm/_torch/models/modeling_exaone_moe.py +++ b/tensorrt_llm/_torch/models/modeling_exaone_moe.py @@ -433,17 +433,28 @@ def forward_mlp( residual: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: if self.fusion_config.PRE_MLP_FUSION: - act_fp4, act_sf, residual = self.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, - residual=residual, - norm_weight=self.post_attention_layernorm.weight, - scale=self.mlp.gate_up_proj.input_scale, - eps=self.post_attention_layernorm.variance_epsilon, - ), - ) - hidden_states = Fp4QuantizedTensor(act_fp4, act_sf) + if self.mlp.gate_up_proj.has_nvfp4: + act_fp4, act_sf, residual = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + scale=self.mlp.gate_up_proj.input_scale, + eps=self.post_attention_layernorm.variance_epsilon, + ), + ) + hidden_states = Fp4QuantizedTensor(act_fp4, act_sf) + else: + hidden_states, residual = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + ), + ) else: hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 2572ea548e48..7845c6e14acd 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/tensorrt_llm/_torch/models/modeling_glm.py @@ -774,17 +774,28 @@ def forward_mlp( spec_metadata: Optional[SpecMetadata] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: if self.fusion_config.PRE_MLP_FUSION: - act_fp4, act_sf, residual = self.allreduce( - hidden_states, - all_reduce_params=AllReduceParams( - fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, - residual=residual, - norm_weight=self.post_attention_layernorm.weight, - scale=self.mlp.gate_up_proj.input_scale, - eps=self.post_attention_layernorm.variance_epsilon, - ), - ) - hidden_states = Fp4QuantizedTensor(act_fp4, act_sf) + if self.mlp.gate_up_proj.has_nvfp4: + act_fp4, act_sf, residual = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + scale=self.mlp.gate_up_proj.input_scale, + eps=self.post_attention_layernorm.variance_epsilon, + ), + ) + hidden_states = Fp4QuantizedTensor(act_fp4, act_sf) + else: + hidden_states, residual = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + ), + ) else: # No fusion # We need to add twoshot allreduce here to avoid modifying MLA logic diff --git a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py new file mode 100644 index 000000000000..83660dae47a2 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +@pytest.mark.parametrize( + "model_file", + [ + "modeling_deepseekv3.py", + "modeling_glm.py", + "modeling_exaone_moe.py", + ], +) +def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(model_file: str) -> None: + source = (_REPO_ROOT / "tensorrt_llm" / "_torch" / "models" / model_file).read_text() + start = source.index(" def forward_mlp") + end = source.index(" hidden_states = self.mlp", start) + pre_mlp_branch = source[start:end] + + guard = "if self.mlp.gate_up_proj.has_nvfp4:" + scale_access = "scale=self.mlp.gate_up_proj.input_scale" + + assert guard in pre_mlp_branch + assert pre_mlp_branch.index(guard) < pre_mlp_branch.index(scale_access) From 27d8384a64aa65e8e6de611d8ae0506c3238c9c4 Mon Sep 17 00:00:00 2001 From: zq Date: Thu, 6 Aug 2026 01:08:57 +0800 Subject: [PATCH 2/8] test: assert PRE_MLP NVFP4 fallback branch Signed-off-by: zq --- .../modeling/test_pre_mlp_nvfp4_fusion.py | 44 +++++++++++++++---- 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py index 83660dae47a2..bd7f01f1c16d 100644 --- a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import ast from pathlib import Path import pytest @@ -20,6 +21,32 @@ _REPO_ROOT = Path(__file__).resolve().parents[4] +def _attribute_path(node: ast.AST) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + prefix = _attribute_path(node.value) + return f"{prefix}.{node.attr}" if prefix else node.attr + return "" + + +def _contains_attribute(node: ast.AST, attribute_path: str) -> bool: + return any(_attribute_path(child) == attribute_path for child in ast.walk(node)) + + +def _forward_mlp_has_nvfp4_branch(model_file: str) -> ast.If: + source = (_REPO_ROOT / "tensorrt_llm" / "_torch" / "models" / model_file).read_text() + module = ast.parse(source) + for node in ast.walk(module): + if isinstance(node, ast.FunctionDef) and node.name == "forward_mlp": + for child in ast.walk(node): + if (_attribute_path(getattr(child, "test", ast.Constant(None))) + == "self.mlp.gate_up_proj.has_nvfp4"): + assert isinstance(child, ast.If) + return child + raise AssertionError(f"{model_file} does not guard PRE_MLP NVFP4 fusion") + + @pytest.mark.parametrize( "model_file", [ @@ -29,13 +56,14 @@ ], ) def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(model_file: str) -> None: - source = (_REPO_ROOT / "tensorrt_llm" / "_torch" / "models" / model_file).read_text() - start = source.index(" def forward_mlp") - end = source.index(" hidden_states = self.mlp", start) - pre_mlp_branch = source[start:end] + nvfp4_branch = _forward_mlp_has_nvfp4_branch(model_file) - guard = "if self.mlp.gate_up_proj.has_nvfp4:" - scale_access = "scale=self.mlp.gate_up_proj.input_scale" + assert _contains_attribute(nvfp4_branch, "self.mlp.gate_up_proj.input_scale") + assert _contains_attribute(nvfp4_branch, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4") - assert guard in pre_mlp_branch - assert pre_mlp_branch.index(guard) < pre_mlp_branch.index(scale_access) + 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) From 6dbe66172a64c5fdd9647dd38e8afcc9b5aa5028 Mon Sep 17 00:00:00 2001 From: zq Date: Fri, 7 Aug 2026 01:50:22 +0800 Subject: [PATCH 3/8] test: add PRE_MLP NVFP4 guard to B200 test list Signed-off-by: zq --- tests/integration/test_lists/test-db/l0_b200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 51f7f6d821e3..bc15f9250851 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -86,6 +86,7 @@ l0_b200: - 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 - unittest/llmapi/test_deepseek_v4_tokenizer.py - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py From 2bdc269eb075dd33f1e447f96c81691ef52480ac Mon Sep 17 00:00:00 2001 From: zq Date: Sat, 8 Aug 2026 02:36:45 +0800 Subject: [PATCH 4/8] test: address PRE_MLP NVFP4 review feedback Signed-off-by: zq --- .../test_lists/test-db/l0_b200.yml | 1 - .../test_lists/test-db/l0_cpu_arm.yml | 1 + .../test_lists/test-db/l0_cpu_x86.yml | 1 + .../modeling/test_pre_mlp_nvfp4_fusion.py | 111 ++++++++++++++++-- 4 files changed, 103 insertions(+), 11 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index bc15f9250851..51f7f6d821e3 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -86,7 +86,6 @@ l0_b200: - 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 - unittest/llmapi/test_deepseek_v4_tokenizer.py - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py diff --git a/tests/integration/test_lists/test-db/l0_cpu_arm.yml b/tests/integration/test_lists/test-db/l0_cpu_arm.yml index 9e21ccf8b5fc..14edb20a25a6 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_arm.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_arm.yml @@ -17,3 +17,4 @@ l0_cpu_arm: - unittest/executor/test_event_loop_error_broadcast.py - unittest/others/test_http_utils_fail_fast.py - unittest/llmapi/test_bench_async.py + - unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py diff --git a/tests/integration/test_lists/test-db/l0_cpu_x86.yml b/tests/integration/test_lists/test-db/l0_cpu_x86.yml index 76d344b77adb..0afa8c9118e0 100644 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ b/tests/integration/test_lists/test-db/l0_cpu_x86.yml @@ -18,3 +18,4 @@ l0_cpu_x86: - unittest/executor/test_multi_frontend_routing.py - unittest/executor/test_event_loop_error_broadcast.py - unittest/llmapi/test_bench_async.py + - unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py diff --git a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py index bd7f01f1c16d..fe6e5e4113df 100644 --- a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -14,12 +14,11 @@ # limitations under the License. import ast +import textwrap from pathlib import Path import pytest -_REPO_ROOT = Path(__file__).resolve().parents[4] - def _attribute_path(node: ast.AST) -> str: if isinstance(node, ast.Name): @@ -34,16 +33,59 @@ def _contains_attribute(node: ast.AST, attribute_path: str) -> bool: return any(_attribute_path(child) == attribute_path for child in ast.walk(node)) -def _forward_mlp_has_nvfp4_branch(model_file: str) -> ast.If: - source = (_REPO_ROOT / "tensorrt_llm" / "_torch" / "models" / model_file).read_text() +def _mentions_nvfp4_flag(node: ast.AST) -> bool: + for child in ast.walk(node): + if isinstance(child, ast.Attribute) and child.attr == "has_nvfp4": + return True + if isinstance(child, ast.Name) and child.id == "has_nvfp4": + return True + if isinstance(child, ast.Constant) and child.value == "has_nvfp4": + return True + return False + + +def _forward_mlp_nvfp4_aliases(function_node: ast.FunctionDef) -> set[str]: + aliases: set[str] = set() + for stmt in function_node.body: + if isinstance(stmt, ast.Assign) and _mentions_nvfp4_flag(stmt.value): + for target in stmt.targets: + if isinstance(target, ast.Name): + aliases.add(target.id) + elif isinstance(stmt, ast.AnnAssign) and stmt.value is not None and _mentions_nvfp4_flag(stmt.value): + target = stmt.target + if isinstance(target, ast.Name): + aliases.add(target.id) + return aliases + + +def _find_repo_root(start: Path) -> Path: + for candidate in (start, *start.parents): + if (candidate / "tensorrt_llm" / "_torch" / "models").exists(): + return candidate + raise AssertionError(f"Could not locate repo root from {start}") + + +def _forward_mlp_has_nvfp4_branch(model_file: str, repo_root: Path) -> ast.If: + source = (repo_root / "tensorrt_llm" / "_torch" / "models" / model_file).read_text() module = ast.parse(source) + for node in ast.walk(module): if isinstance(node, ast.FunctionDef) and node.name == "forward_mlp": + aliases = _forward_mlp_nvfp4_aliases(node) for child in ast.walk(node): - if (_attribute_path(getattr(child, "test", ast.Constant(None))) - == "self.mlp.gate_up_proj.has_nvfp4"): - assert isinstance(child, ast.If) - return child + if not isinstance(child, ast.If): + continue + if not (_mentions_nvfp4_flag(child.test) or ( + isinstance(child.test, ast.Name) and child.test.id in aliases)): + continue + if not _contains_attribute(child, "self.mlp.gate_up_proj.input_scale"): + continue + if not _contains_attribute(child, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4"): + continue + if not any(_contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM") + for false_branch_node in child.orelse): + continue + return child raise AssertionError(f"{model_file} does not guard PRE_MLP NVFP4 fusion") @@ -55,8 +97,9 @@ def _forward_mlp_has_nvfp4_branch(model_file: str) -> ast.If: "modeling_exaone_moe.py", ], ) -def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(model_file: str) -> None: - nvfp4_branch = _forward_mlp_has_nvfp4_branch(model_file) +def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(pytestconfig, model_file: str) -> None: + repo_root = _find_repo_root(pytestconfig.rootpath) + nvfp4_branch = _forward_mlp_has_nvfp4_branch(model_file, repo_root) assert _contains_attribute(nvfp4_branch, "self.mlp.gate_up_proj.input_scale") assert _contains_attribute(nvfp4_branch, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4") @@ -67,3 +110,51 @@ def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(model_file: str) -> N assert any( _contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM") for false_branch_node in nvfp4_branch.orelse) + + +def test_forward_mlp_has_nvfp4_branch_supports_alias_and_getattr(tmp_path) -> None: + repo_root = tmp_path / "repo" + model_dir = repo_root / "tensorrt_llm" / "_torch" / "models" + model_dir.mkdir(parents=True) + model_file = model_dir / "modeling_alias.py" + model_file.write_text( + textwrap.dedent( + ''' + class Dummy: + def forward_mlp(self): + if self.fusion_config.PRE_MLP_FUSION: + gate_up_proj = self.mlp.gate_up_proj + has_nvfp4 = getattr(gate_up_proj, "has_nvfp4", False) + if has_nvfp4: + act_fp4, act_sf, residual = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + scale=self.mlp.gate_up_proj.input_scale, + eps=self.post_attention_layernorm.variance_epsilon, + ), + ) + hidden_states = Fp4QuantizedTensor(act_fp4, act_sf) + else: + hidden_states, residual = self.allreduce( + hidden_states, + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=self.post_attention_layernorm.weight, + eps=self.post_attention_layernorm.variance_epsilon, + ), + ) + ''' + ) + ) + + nvfp4_branch = _forward_mlp_has_nvfp4_branch("modeling_alias.py", repo_root) + + assert _contains_attribute(nvfp4_branch, "self.mlp.gate_up_proj.input_scale") + assert _contains_attribute(nvfp4_branch, "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) From 15dce66fd004f5a907d40c7f06c9a78435f87e90 Mon Sep 17 00:00:00 2001 From: zq Date: Sat, 8 Aug 2026 02:47:33 +0800 Subject: [PATCH 5/8] test: collect nested PRE_MLP NVFP4 aliases Signed-off-by: zq --- .../modeling/test_pre_mlp_nvfp4_fusion.py | 29 ++++++++++++------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py index fe6e5e4113df..87fac7a9655c 100644 --- a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -44,17 +44,26 @@ def _mentions_nvfp4_flag(node: ast.AST) -> bool: return False +def _collect_nvfp4_aliases(node: ast.AST, aliases: set[str]) -> None: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): + return + if isinstance(node, ast.Assign) and _mentions_nvfp4_flag(node.value): + for target in node.targets: + if isinstance(target, ast.Name): + aliases.add(target.id) + elif isinstance(node, ast.AnnAssign) and node.value is not None and _mentions_nvfp4_flag(node.value): + target = node.target + if isinstance(target, ast.Name): + aliases.add(target.id) + + for child in ast.iter_child_nodes(node): + _collect_nvfp4_aliases(child, aliases) + + def _forward_mlp_nvfp4_aliases(function_node: ast.FunctionDef) -> set[str]: aliases: set[str] = set() for stmt in function_node.body: - if isinstance(stmt, ast.Assign) and _mentions_nvfp4_flag(stmt.value): - for target in stmt.targets: - if isinstance(target, ast.Name): - aliases.add(target.id) - elif isinstance(stmt, ast.AnnAssign) and stmt.value is not None and _mentions_nvfp4_flag(stmt.value): - target = stmt.target - if isinstance(target, ast.Name): - aliases.add(target.id) + _collect_nvfp4_aliases(stmt, aliases) return aliases @@ -124,8 +133,8 @@ class Dummy: def forward_mlp(self): if self.fusion_config.PRE_MLP_FUSION: gate_up_proj = self.mlp.gate_up_proj - has_nvfp4 = getattr(gate_up_proj, "has_nvfp4", False) - if has_nvfp4: + use_nvfp4 = getattr(gate_up_proj, "has_nvfp4", False) + if use_nvfp4: act_fp4, act_sf, residual = self.allreduce( hidden_states, all_reduce_params=AllReduceParams( From 80c96cd28b2cc9dddb4f3703f62114580d779d1a Mon Sep 17 00:00:00 2001 From: zq Date: Sat, 8 Aug 2026 15:05:47 +0800 Subject: [PATCH 6/8] test: tighten PRE_MLP NVFP4 alias detection Signed-off-by: zq --- .../modeling/test_pre_mlp_nvfp4_fusion.py | 141 ++++++++++++++---- 1 file changed, 116 insertions(+), 25 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py index 87fac7a9655c..4dcb7ab0494a 100644 --- a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -44,9 +44,8 @@ def _mentions_nvfp4_flag(node: ast.AST) -> bool: return False -def _collect_nvfp4_aliases(node: ast.AST, aliases: set[str]) -> None: - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): - return +def _nvfp4_aliases_assigned_by_statement(node: ast.AST) -> set[str]: + aliases: set[str] = set() if isinstance(node, ast.Assign) and _mentions_nvfp4_flag(node.value): for target in node.targets: if isinstance(target, ast.Name): @@ -55,16 +54,54 @@ def _collect_nvfp4_aliases(node: ast.AST, aliases: set[str]) -> None: target = node.target if isinstance(target, ast.Name): aliases.add(target.id) + return aliases + - for child in ast.iter_child_nodes(node): - _collect_nvfp4_aliases(child, aliases) +def _is_pre_mlp_nvfp4_branch(node: ast.If, aliases: set[str]) -> bool: + test_uses_nvfp4 = _mentions_nvfp4_flag(node.test) or ( + isinstance(node.test, ast.Name) and node.test.id in aliases) + if not test_uses_nvfp4: + return False + if not _contains_attribute(node, "self.mlp.gate_up_proj.input_scale"): + return False + if not _contains_attribute(node, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4"): + return False + return any( + _contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM") + for false_branch_node in node.orelse) -def _forward_mlp_nvfp4_aliases(function_node: ast.FunctionDef) -> set[str]: - aliases: set[str] = set() - for stmt in function_node.body: - _collect_nvfp4_aliases(stmt, aliases) - return aliases +def _scan_statement_list_for_nvfp4_branch(statements: list[ast.stmt], aliases: set[str]) -> ast.If | None: + visible_aliases = set(aliases) + for stmt in statements: + if isinstance(stmt, ast.If): + if _is_pre_mlp_nvfp4_branch(stmt, visible_aliases): + return stmt + for branch in (stmt.body, stmt.orelse): + branch_match = _scan_statement_list_for_nvfp4_branch(branch, visible_aliases) + if branch_match is not None: + return branch_match + elif isinstance(stmt, (ast.For, ast.AsyncFor, ast.While, ast.With, ast.AsyncWith)): + branch_match = _scan_statement_list_for_nvfp4_branch(stmt.body, visible_aliases) + if branch_match is not None: + return branch_match + if isinstance(stmt, (ast.For, ast.AsyncFor, ast.While)): + branch_match = _scan_statement_list_for_nvfp4_branch(stmt.orelse, visible_aliases) + if branch_match is not None: + return branch_match + elif isinstance(stmt, ast.Try): + for branch in (stmt.body, stmt.orelse, stmt.finalbody, *(handler.body for handler in stmt.handlers)): + branch_match = _scan_statement_list_for_nvfp4_branch(branch, visible_aliases) + if branch_match is not None: + return branch_match + elif isinstance(stmt, ast.Match): + for case in stmt.cases: + branch_match = _scan_statement_list_for_nvfp4_branch(case.body, visible_aliases) + if branch_match is not None: + return branch_match + + visible_aliases.update(_nvfp4_aliases_assigned_by_statement(stmt)) + return None def _find_repo_root(start: Path) -> Path: @@ -80,24 +117,13 @@ def _forward_mlp_has_nvfp4_branch(model_file: str, repo_root: Path) -> ast.If: for node in ast.walk(module): if isinstance(node, ast.FunctionDef) and node.name == "forward_mlp": - aliases = _forward_mlp_nvfp4_aliases(node) - for child in ast.walk(node): - if not isinstance(child, ast.If): - continue - if not (_mentions_nvfp4_flag(child.test) or ( - isinstance(child.test, ast.Name) and child.test.id in aliases)): - continue - if not _contains_attribute(child, "self.mlp.gate_up_proj.input_scale"): - continue - if not _contains_attribute(child, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4"): - continue - if not any(_contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM") - for false_branch_node in child.orelse): - continue - return child + nvfp4_branch = _scan_statement_list_for_nvfp4_branch(node.body, set()) + if nvfp4_branch is not None: + return nvfp4_branch raise AssertionError(f"{model_file} does not guard PRE_MLP NVFP4 fusion") + @pytest.mark.parametrize( "model_file", [ @@ -167,3 +193,68 @@ def forward_mlp(self): assert any( _contains_attribute(false_branch_node, "AllReduceFusionOp.RESIDUAL_RMS_NORM") for false_branch_node in nvfp4_branch.orelse) + + +def _write_model_fixture(repo_root: Path, model_name: str, forward_mlp_body: str) -> None: + model_dir = repo_root / "tensorrt_llm" / "_torch" / "models" + model_dir.mkdir(parents=True) + (model_dir / model_name).write_text( + "class Dummy:\n" + textwrap.indent(textwrap.dedent(forward_mlp_body), " ") + ) + + +def test_forward_mlp_alias_must_be_defined_before_candidate_branch(tmp_path) -> None: + repo_root = tmp_path / "repo" + _write_model_fixture( + repo_root, + "modeling_alias_use_before_assignment.py", + ''' + def forward_mlp(self): + if use_nvfp4: + self.allreduce( + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + scale=self.mlp.gate_up_proj.input_scale, + ), + ) + else: + self.allreduce( + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + ), + ) + use_nvfp4 = self.mlp.gate_up_proj.has_nvfp4 + ''', + ) + + with pytest.raises(AssertionError, match="does not guard PRE_MLP NVFP4 fusion"): + _forward_mlp_has_nvfp4_branch("modeling_alias_use_before_assignment.py", repo_root) + + +def test_forward_mlp_alias_must_not_leak_from_sibling_branch(tmp_path) -> None: + repo_root = tmp_path / "repo" + _write_model_fixture( + repo_root, + "modeling_alias_sibling_branch.py", + ''' + def forward_mlp(self): + if self.fusion_config.PRE_MLP_FUSION: + use_nvfp4 = self.mlp.gate_up_proj.has_nvfp4 + if use_nvfp4: + self.allreduce( + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4, + scale=self.mlp.gate_up_proj.input_scale, + ), + ) + else: + self.allreduce( + all_reduce_params=AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + ), + ) + ''', + ) + + with pytest.raises(AssertionError, match="does not guard PRE_MLP NVFP4 fusion"): + _forward_mlp_has_nvfp4_branch("modeling_alias_sibling_branch.py", repo_root) From cd6d0ff2166488a103b8f86f75944533e36f7a24 Mon Sep 17 00:00:00 2001 From: zq Date: Sun, 9 Aug 2026 16:48:59 +0800 Subject: [PATCH 7/8] test: strengthen PRE_MLP NVFP4 guard test Signed-off-by: zq --- .../modeling/test_pre_mlp_nvfp4_fusion.py | 37 +++++++++++-------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py index 4dcb7ab0494a..455742c882ba 100644 --- a/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -19,6 +19,11 @@ import pytest +import tensorrt_llm._torch.models as _models + + +_MODELS_DIR = Path(_models.__file__).parent + def _attribute_path(node: ast.AST) -> str: if isinstance(node, ast.Name): @@ -104,15 +109,8 @@ def _scan_statement_list_for_nvfp4_branch(statements: list[ast.stmt], aliases: s return None -def _find_repo_root(start: Path) -> Path: - for candidate in (start, *start.parents): - if (candidate / "tensorrt_llm" / "_torch" / "models").exists(): - return candidate - raise AssertionError(f"Could not locate repo root from {start}") - - -def _forward_mlp_has_nvfp4_branch(model_file: str, repo_root: Path) -> ast.If: - source = (repo_root / "tensorrt_llm" / "_torch" / "models" / model_file).read_text() +def _forward_mlp_has_nvfp4_branch(model_file: str, models_dir: Path = _MODELS_DIR) -> ast.If: + source = (models_dir / model_file).read_text() module = ast.parse(source) for node in ast.walk(module): @@ -120,7 +118,13 @@ def _forward_mlp_has_nvfp4_branch(model_file: str, repo_root: Path) -> ast.If: nvfp4_branch = _scan_statement_list_for_nvfp4_branch(node.body, set()) if nvfp4_branch is not None: return nvfp4_branch - raise AssertionError(f"{model_file} does not guard PRE_MLP NVFP4 fusion") + raise AssertionError( + f"{model_file} does not guard PRE_MLP NVFP4 fusion. " + "Expected a branch that gates RESIDUAL_RMS_NORM_QUANT_NVFP4 on has_nvfp4, " + "uses self.mlp.gate_up_proj.input_scale in that branch, and falls back to " + "RESIDUAL_RMS_NORM otherwise; the has_nvfp4 check supports direct attributes, " + "getattr(..., \"has_nvfp4\", ...), or an alias assigned before the branch." + ) @@ -132,9 +136,8 @@ def _forward_mlp_has_nvfp4_branch(model_file: str, repo_root: Path) -> ast.If: "modeling_exaone_moe.py", ], ) -def test_pre_mlp_nvfp4_fusion_guards_unquantized_dense_mlp(pytestconfig, model_file: str) -> None: - repo_root = _find_repo_root(pytestconfig.rootpath) - nvfp4_branch = _forward_mlp_has_nvfp4_branch(model_file, repo_root) +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") @@ -186,7 +189,7 @@ def forward_mlp(self): ) ) - nvfp4_branch = _forward_mlp_has_nvfp4_branch("modeling_alias.py", repo_root) + nvfp4_branch = _forward_mlp_has_nvfp4_branch("modeling_alias.py", model_dir) assert _contains_attribute(nvfp4_branch, "self.mlp.gate_up_proj.input_scale") assert _contains_attribute(nvfp4_branch, "AllReduceFusionOp.RESIDUAL_RMS_NORM_QUANT_NVFP4") @@ -227,8 +230,9 @@ def forward_mlp(self): ''', ) + model_dir = repo_root / "tensorrt_llm" / "_torch" / "models" with pytest.raises(AssertionError, match="does not guard PRE_MLP NVFP4 fusion"): - _forward_mlp_has_nvfp4_branch("modeling_alias_use_before_assignment.py", repo_root) + _forward_mlp_has_nvfp4_branch("modeling_alias_use_before_assignment.py", model_dir) def test_forward_mlp_alias_must_not_leak_from_sibling_branch(tmp_path) -> None: @@ -256,5 +260,6 @@ def forward_mlp(self): ''', ) + model_dir = repo_root / "tensorrt_llm" / "_torch" / "models" with pytest.raises(AssertionError, match="does not guard PRE_MLP NVFP4 fusion"): - _forward_mlp_has_nvfp4_branch("modeling_alias_sibling_branch.py", repo_root) + _forward_mlp_has_nvfp4_branch("modeling_alias_sibling_branch.py", model_dir) From 816787021c49bb99bbcd8fa9d0b7417d2936a499 Mon Sep 17 00:00:00 2001 From: zq Date: Mon, 10 Aug 2026 05:37:10 +0800 Subject: [PATCH 8/8] test: drop obsolete split CPU stage entries Signed-off-by: zq --- .../test_lists/test-db/l0_cpu_arm.yml | 20 ------------------ .../test_lists/test-db/l0_cpu_x86.yml | 21 ------------------- 2 files changed, 41 deletions(-) delete mode 100644 tests/integration/test_lists/test-db/l0_cpu_arm.yml delete mode 100644 tests/integration/test_lists/test-db/l0_cpu_x86.yml diff --git a/tests/integration/test_lists/test-db/l0_cpu_arm.yml b/tests/integration/test_lists/test-db/l0_cpu_arm.yml deleted file mode 100644 index 14edb20a25a6..000000000000 --- a/tests/integration/test_lists/test-db/l0_cpu_arm.yml +++ /dev/null @@ -1,20 +0,0 @@ -version: 0.0.1 -l0_cpu_arm: -- condition: - ranges: - system_gpu_count: - gte: 0 - lte: 0 - wildcards: - linux_distribution_name: ubuntu* - cpu: aarch64 - terms: - stage: pre_merge - backend: generic - orchestrator: mpi - tests: - - unittest/executor/test_rpc.py - - unittest/executor/test_event_loop_error_broadcast.py - - unittest/others/test_http_utils_fail_fast.py - - unittest/llmapi/test_bench_async.py - - unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py diff --git a/tests/integration/test_lists/test-db/l0_cpu_x86.yml b/tests/integration/test_lists/test-db/l0_cpu_x86.yml deleted file mode 100644 index 0afa8c9118e0..000000000000 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ /dev/null @@ -1,21 +0,0 @@ -version: 0.0.1 -l0_cpu_x86: -- condition: - ranges: - system_gpu_count: - gte: 0 - lte: 0 - wildcards: - linux_distribution_name: ubuntu* - cpu: x86_64 - terms: - stage: pre_merge - backend: generic - orchestrator: mpi - tests: - - unittest/executor/test_rpc.py - - unittest/others/test_http_utils_fail_fast.py - - unittest/executor/test_multi_frontend_routing.py - - unittest/executor/test_event_loop_error_broadcast.py - - unittest/llmapi/test_bench_async.py - - unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py