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/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 9e21ccf8b5fc..000000000000 --- a/tests/integration/test_lists/test-db/l0_cpu_arm.yml +++ /dev/null @@ -1,19 +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 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 76d344b77adb..000000000000 --- a/tests/integration/test_lists/test-db/l0_cpu_x86.yml +++ /dev/null @@ -1,20 +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 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..455742c882ba --- /dev/null +++ b/tests/unittest/_torch/modeling/test_pre_mlp_nvfp4_fusion.py @@ -0,0 +1,265 @@ +# 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. + +import ast +import textwrap +from pathlib import Path + +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): + 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 _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 _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): + 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) + return 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 _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 _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): + if isinstance(node, ast.FunctionDef) and node.name == "forward_mlp": + 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. " + "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." + ) + + + +@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: + 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) + + +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 + 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( + 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", 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") + 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 + ''', + ) + + 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", model_dir) + + +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, + ), + ) + ''', + ) + + 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", model_dir)