fix unit tests

Fridah-nv · Fridah-nv · commit 40ef0686a27d · 2025-08-26T23:45:44.000Z
Signed-off-by: Frida Hou &lt;201670829+Fridah-nv@users.noreply.github.com&gt;

add unit test for quant fusion and quant sharding

Signed-off-by: Frida Hou &lt;201670829+Fridah-nv@users.noreply.github.com&gt;
diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
@@ -82,6 +82,8 @@ transforms:
     stage: post_load_fusion
   fuse_fp8_gemms:
     stage: post_load_fusion
+  fuse_quant:
+    stage: post_load_fusion
   fuse_allreduce_residual_rmsnorm:
     stage: post_load_fusion
   fuse_collectives:
@@ -91,8 +93,6 @@ transforms:
   fuse_rmsnorm:
     stage: post_load_fusion
     backend: flashinfer
-  fuse_quant:
-    stage: post_load_fusion
   ############################################################################################
   # SWITCH TO CACHED+FLATTENED ATTENTION + INITIALIZE CACHES
   ############################################################################################
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_quant.py
@@ -6,6 +6,8 @@
     cutlass_fp4_scale_to_modelopt_fp4_scale,
 )
 
+from .quant import QUANT_LINEAR_OPS, QUANT_OPS
+
 # ===== Enums =====
 FORMAT_FP8 = 0
 FORMAT_NVFP4 = 1
@@ -285,3 +287,11 @@ def torch_fake_quant_fp4_linear(
     weight_zp: List[torch.Tensor],
 ) -> torch.Tensor:
     return torch.ops.aten.linear(input, weight_quantized.repeat(1, 2).to(input.dtype), bias)
+
+
+CUSTOM_QUANT_LINEAR_OPS = [
+    torch.ops.auto_deploy.torch_fake_quant_fp8_linear,
+    torch.ops.auto_deploy.torch_fake_quant_fp4_linear,
+]
+QUANT_LINEAR_OPS.extend(CUSTOM_QUANT_LINEAR_OPS)
+QUANT_OPS.extend(CUSTOM_QUANT_LINEAR_OPS)
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -249,8 +249,6 @@ def is_linear_op(node: Node, include_quantization: bool = False) -> bool:
 
     if include_quantization:
         lin_ops.update(QUANT_LINEAR_OPS)
-        lin_ops.update([torch.ops.auto_deploy.torch_fake_quant_fp8_linear])
-        lin_ops.update([torch.ops.auto_deploy.torch_fake_quant_fp4_linear])
     return is_op(node, lin_ops)
 
 
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py
@@ -98,10 +98,14 @@ def create(quant_type_or_node: Union[str, Node], is_bmm: bool = False):
         for q in [
             FP4QuantizationImpl,
             FP8QuantizationImpl,
-            FP8BMMQuantizationImpl,
         ]:
             if is_op(quant_type_or_node, q.custom_op()):
                 return q
+        for q in [
+            FP8BMMQuantizationImpl,
+        ]:
+            if is_op(quant_type_or_node, q.target_op()):
+                return q
         return None
 
     @staticmethod
diff --git a/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py b/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py
@@ -253,6 +253,28 @@ def forward(self, x):
         return torch.bmm(x, dynamic_weights)
 
 
+FP8_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+
+class FakeFP8Linear(nn.Linear):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        device = self.weight.device
+        weight_scale = torch.max(torch.abs(self.weight)).to(torch.float).to(device) / FP8_MAX
+        self.weight = nn.Parameter((self.weight / weight_scale).to(torch.float8_e4m3fn))
+        self.register_buffer(
+            "input_scale", torch.tensor(1.0, device=self.weight.device, dtype=torch.float)
+        )
+        self.register_buffer("weight_scale", weight_scale)
+        if self.bias is not None:
+            self.bias = nn.Parameter(self.bias.to(torch.half))
+
+    def forward(self, x):
+        return torch.ops.auto_deploy.torch_fake_quant_fp8_linear(
+            x, self.weight, self.bias, [self.input_scale], [self.weight_scale], [], []
+        )
+
+
 def generate_dynamic_shapes(max_batch_size, max_seq_len):
     dynamic_shapes = (
         {
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py
@@ -9,6 +9,7 @@
 import torch.nn.functional as F
 from _dist_test_utils import get_device_counts
 from _graph_test_helpers import run_sharding_pattern_detection_test, run_test_transformed_gm
+from _model_test_utils import FakeFP8Linear
 
 import tensorrt_llm._torch.auto_deploy.distributed.common as dist_common
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
@@ -77,6 +78,19 @@ def forward(self, x):
         return self.linear2(y)
 
 
+class FP8MLP(nn.Module):
+    def __init__(self, in_features, out_features, bias=False):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.linear1 = FakeFP8Linear(in_features, 4 * in_features, bias=bias)
+        self.linear2 = FakeFP8Linear(4 * in_features, out_features, bias=bias)
+
+    def forward(self, x):
+        y = F.relu(self.linear1(x))
+        return self.linear2(y)
+
+
 def _run_job(
     model_cls: nn.Module,
     dist_op_expected: str,
@@ -99,6 +113,8 @@ def _run_job(
             hidden_size=num_features,
             num_key_value_heads=num_key_value_heads,
         ).to(device="cuda", dtype=torch.float16)
+    elif model_cls == FP8MLP:
+        model = model_cls(num_features, num_features, bias=bias).to("cuda")
     else:
         model = model_cls(num_features, num_features, bias=bias).to(
             device="cuda", dtype=torch.float16
@@ -231,7 +247,7 @@ def _run_pattern_detection_job(
                             min_local_shape=min_local_shape,
                         )
                     )
-        elif model_cls == MLP:
+        elif model_cls == MLP or model_cls == FP8MLP:
             for node in gm.graph.nodes:
                 if is_linear_op(node, include_quantization=True):
                     # linear1 should be sharded on dim=0, add_dist=False, min_local_shape=1
@@ -291,6 +307,7 @@ def _run_pattern_detection_job(
     "model_cls, dist_op_expected",
     (
         (MLP, "torch_dist_all_reduce"),
+        (FP8MLP, "torch_dist_all_reduce"),
         (nn.Linear, "torch_dist_all_gather"),
         (GQA_Block, "torch_dist_all_reduce"),
     ),
@@ -302,12 +319,13 @@ def test_sharding(model_cls: Type[nn.Module], dist_op_expected: str, bias: bool,
     )
 
 
-@pytest.mark.parametrize("world_size", [1, 8])
+@pytest.mark.parametrize("world_size", [8])
 @pytest.mark.parametrize("bias", [False, True])
 @pytest.mark.parametrize(
     "model_cls, dist_op_expected",
     (
         (MLP, "torch_dist_all_reduce"),
+        (FP8MLP, "torch_dist_all_reduce"),
         (nn.Linear, "torch_dist_all_gather"),
         (GQA_Block, "torch_dist_all_reduce"),
     ),
diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gemm_fusion.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_gemm_fusion.py
@@ -8,9 +8,9 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from _graph_test_helpers import count_buffers, run_test_transformed_gm
+from _model_test_utils import FakeFP8Linear
 from _torch_test_utils import all_close, fp8_compatible, reset_parameters
 
-from tensorrt_llm._torch.auto_deploy.custom_ops.quant import FP8Linear
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
 from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_linear_op
@@ -74,7 +74,7 @@ def __init__(self, **kwargs):
 
 class FusableModel1_M_FP8(FusableModel1_M):
     def __init__(self, **kwargs):
-        super().__init__(**{"cls": FP8Linear, **kwargs})
+        super().__init__(**{"cls": FakeFP8Linear, **kwargs})
 
 
 class FusableModel1_L(FusableModel1):
@@ -89,7 +89,7 @@ def __init__(self, **kwargs):
 
 class FusableModel1_XL_FP8(FusableModel1_XL):
     def __init__(self, **kwargs):
-        super().__init__(**{"cls": FP8Linear, **kwargs})
+        super().__init__(**{"cls": FakeFP8Linear, **kwargs})
 
 
 class FusableModel2(FusableModel):
@@ -112,7 +112,7 @@ def num_gemms_after_fusion(self) -> int:
 
 class FusableModel2_FP8(FusableModel2):
     def __init__(self, **kwargs):
-        super().__init__(**{"cls": FP8Linear, **kwargs})
+        super().__init__(**{"cls": FakeFP8Linear, **kwargs})
 
 
 class FusableModel3(FusableModel):
@@ -137,7 +137,7 @@ def num_gemms_after_fusion(self) -> int:
 
 class FusableModel3_FP8(FusableModel3):
     def __init__(self, **kwargs):
-        super().__init__(**{"cls": FP8Linear, **kwargs})
+        super().__init__(**{"cls": FakeFP8Linear, **kwargs})
 
 
 class FusableModel4(FusableModel):
@@ -168,7 +168,7 @@ def num_gemms_after_fusion(self) -> int:
 
 class FusableModel4_FP8(FusableModel4):
     def __init__(self, **kwargs):
-        super().__init__(**{"cls": FP8Linear, **kwargs})
+        super().__init__(**{"cls": FakeFP8Linear, **kwargs})
 
 
 # TODO: consider adding test cases for classic GQA and MLP layers
@@ -262,6 +262,9 @@ def test_fusion(get_model: Callable[[], TestModel], dtype: str):
             "fuse_gemms": {
                 "stage": "post_load_fusion",
             },
+            "fuse_fp8_gemms": {
+                "stage": "post_load_fusion",
+            },
         },
     )(None, gm)
 
diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_fusion.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_fusion.py
diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_quantization_utils.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_quantization_utils.py
