WIP to map custom quant op to real implementation using pattern matcher

Fridah-nv · Fridah-nv · commit d14d5f455ba5 · 2025-08-25T20:05:04.000Z
Signed-off-by: Frida Hou &lt;201670829+Fridah-nv@users.noreply.github.com&gt;

Update fusion to pass in args instead of kwrargs

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
@@ -91,6 +91,8 @@ 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/transform/library/fuse_quant.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_quant.py
@@ -0,0 +1,282 @@
+from typing import Tuple
+
+import torch
+from torch.fx import GraphModule
+
+from ...models.factory import ModelFactory
+from ...shim.interface import CachedSequenceInterface
+from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern
+from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry
+
+
+def _fp8_ref_pattern_1(
+    x: torch.Tensor,
+    w_fp8: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+):
+    # Matches: torch_fake_quant_fp8_linear(input, weight_fp8, bias, [in_s], [w_s], [], [])
+    return torch.ops.auto_deploy.torch_fake_quant_fp8_linear.default(
+        x,
+        w_fp8,
+        None,
+        input_scale=[input_scale],
+        weight_scale=[weight_scale],
+        input_zp=[],
+        weight_zp=[],
+    )
+
+
+def _fp8_ref_repl_1(
+    x: torch.Tensor,
+    w_fp8: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+):
+    # Map lists -> scalars for fused op
+    # in_s = input_scale[0]
+    # w_s = weight_scale[0]
+    return torch.ops.auto_deploy.torch_quant_fp8_linear(
+        x,
+        w_fp8,
+        None,
+        input_scale=input_scale,
+        weight_scale=weight_scale,
+    )
+
+
+def _fp8_ref_pattern_2(
+    x: torch.Tensor,
+    w_fp8: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+):
+    # Matches: torch_fake_quant_fp8_linear(input, weight_fp8, bias, [in_s], [w_s], [], [])
+    return torch.ops.auto_deploy.torch_fake_quant_fp8_linear.default(
+        x,
+        w_fp8,
+        bias,
+        input_scale=[input_scale],
+        weight_scale=[weight_scale],
+        input_zp=[],
+        weight_zp=[],
+    )
+
+
+def _fp8_ref_repl_2(
+    x: torch.Tensor,
+    w_fp8: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+):
+    # Map lists -> scalars for fused op
+    # in_s = input_scale[0]
+    # w_s = weight_scale[0]
+    return torch.ops.auto_deploy.torch_quant_fp8_linear(
+        x,
+        w_fp8,
+        bias,
+        input_scale=input_scale,
+        weight_scale=weight_scale,
+    )
+
+
+# NVFP4: reference (search) and fused (replacement)
+def _fp4_ref_pattern_1(
+    x: torch.Tensor,
+    w_fp4: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+):
+    # Matches: torch_fake_quant_fp4_linear(x, w_fp4, bias, [s_in2], [cutlass_scale, alpha], [], [])
+    return torch.ops.auto_deploy.torch_fake_quant_fp4_linear(
+        x,
+        w_fp4,
+        None,
+        input_scale=[input_scale],
+        weight_scale=[weight_scale, alpha],
+        input_zp=[],
+        weight_zp=[],
+    )
+
+
+def _fp4_ref_repl_1(
+    x: torch.Tensor,
+    w_fp4: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+):
+    return torch.ops.auto_deploy.torch_quant_fp4_linear(
+        x,
+        w_fp4,
+        bias=None,
+        input_scale=input_scale,
+        weight_scale=weight_scale,
+        alpha=alpha,
+    )
+
+
+def _fp4_ref_pattern_2(
+    x: torch.Tensor,
+    w_fp4: torch.Tensor,
+    bias: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+):
+    # Matches: torch_fake_quant_fp4_linear(x, w_fp4, bias, [s_in2], [cutlass_scale, alpha], [], [])
+    return torch.ops.auto_deploy.torch_fake_quant_fp4_linear(
+        x,
+        w_fp4,
+        bias,
+        input_scale=[input_scale],
+        weight_scale=[weight_scale, alpha],
+        input_zp=[],
+        weight_zp=[],
+    )
+
+
+def _fp4_ref_repl_2(
+    x: torch.Tensor,
+    w_fp4: torch.Tensor,
+    bias: torch.Tensor | None,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    alpha: torch.Tensor,
+):
+    return torch.ops.auto_deploy.torch_quant_fp4_linear(
+        x,
+        w_fp4,
+        bias=bias,
+        input_scale=input_scale,
+        weight_scale=weight_scale,
+        alpha=alpha,
+    )
+
+
+def _register_quant_linear_patterns(patterns: ADPatternMatcherPass) -> None:
+    """
+    Register the FP8 and FP4 patterns with robust dummy args and minimal ignores.
+    """
+    # Use harmless meta tensors; no dtype/device constraints during tracing.
+    # Shapes mirror your unit tests but can be arbitrary as long as tracing succeeds.
+    x_fp8 = torch.randn(3, 16, device="meta", dtype=torch.float16)
+    w_fp8 = torch.randn(32, 16, device="meta", dtype=torch.float16)  # dtype not enforced in trace
+    bias32 = torch.randn(32, device="meta", dtype=torch.float32)
+    one = torch.tensor(1.0, device="meta", dtype=torch.float32)
+
+    dummy_args_fp8 = [
+        x_fp8,
+        w_fp8,
+        one,
+        torch.tensor(0.5, device="meta", dtype=torch.float32),
+    ]
+
+    dummy_args_fp8_2 = [
+        x_fp8,
+        w_fp8,
+        bias32,
+        one,
+        torch.tensor(0.5, device="meta", dtype=torch.float32),
+    ]
+
+    register_ad_pattern(
+        search_fn=_fp8_ref_pattern_1,
+        replace_fn=_fp8_ref_repl_1,
+        patterns=patterns,
+        dummy_args=dummy_args_fp8,
+        # No special scalar_workaround or op_ignore_types needed here.
+    )
+    register_ad_pattern(
+        search_fn=_fp8_ref_pattern_2,
+        replace_fn=_fp8_ref_repl_2,
+        patterns=patterns,
+        dummy_args=dummy_args_fp8_2,
+        # No special scalar_workaround or op_ignore_types needed here.
+    )
+
+    # FP4 dummy args
+    N = 32
+    K_packed = 32  # weight is packed by 2 FP4 per byte
+    K_eff = 2 * K_packed  # <- effective K after repeat(1, 2) in the fake impl
+
+    x_fp4 = torch.randn(3, K_eff, device="meta", dtype=torch.float16)  # was 3 x 32, must be 3 x 64
+    w_fp4 = torch.randint(0, 255, (N, K_packed), device="meta", dtype=torch.uint8)
+
+    s_in2 = torch.tensor(0.01, device="meta", dtype=torch.float32)
+    alpha = torch.tensor(1.2345, device="meta", dtype=torch.float32)
+
+    # Optional: give a realistic-length CUTLASS scale vector (one uint8 per 16-wide block)
+    # num_blocks = N * (K_eff // 16)
+    cutlass_len = N * (K_eff // 16)  # 32 * (64/16) = 128
+    cutlass_vec = torch.randint(0, 255, (cutlass_len,), device="meta", dtype=torch.uint8)
+
+    dummy_args_fp4_1 = [
+        x_fp4,
+        w_fp4,
+        s_in2,  # input_scale list
+        cutlass_vec,
+        alpha,  # weight_scale list: [per-block vec, alpha]
+    ]
+
+    dummy_args_fp4_2 = [
+        x_fp4,
+        w_fp4,
+        torch.randn(N, device="meta", dtype=torch.float16),  # bias
+        s_in2,  # input_scale list
+        cutlass_vec,
+        alpha,  # weight_scale list: [per-block vec, alpha]
+    ]
+
+    register_ad_pattern(
+        search_fn=_fp4_ref_pattern_1,
+        replace_fn=_fp4_ref_repl_1,
+        patterns=patterns,
+        dummy_args=dummy_args_fp4_1,
+    )
+
+    register_ad_pattern(
+        search_fn=_fp4_ref_pattern_2,
+        replace_fn=_fp4_ref_repl_2,
+        patterns=patterns,
+        dummy_args=dummy_args_fp4_2,
+    )
+
+
+@TransformRegistry.register("fuse_quant")
+class FuseQuant(BaseTransform):
+    """
+    Use ADPatternMatcherPass to rewrite reference quantized linear ops into fused ones:
+
+      FP8:
+        torch_fake_quant_fp8_linear(x, w_fp8, bias, [in_s], [w_s], [], [])
+        -> torch_quant_fp8_linear(x, w_fp8, bias=bias, input_scale=in_s, weight_scale=w_s)
+
+      FP4 (NVFP4):
+        torch_fake_quant_fp4_linear(x, w_fp4, bias, [s_in2], [cutlass_vec, alpha], [], [])
+        -> torch_quant_fp4_linear(x, w_fp4, bias=bias, input_scale=s_in2,
+                                  weight_scale=cutlass_vec, alpha=alpha)
+    """
+
+    def _apply(
+        self,
+        gm: GraphModule,
+        cm: CachedSequenceInterface,
+        factory: ModelFactory,
+        shared_config: SharedConfig,
+    ) -> Tuple[GraphModule, TransformInfo]:
+        patterns = ADPatternMatcherPass()
+        _register_quant_linear_patterns(patterns)
+        num_matches = patterns.apply(gm.graph)
+
+        info = TransformInfo(
+            skipped=(num_matches == 0),
+            num_matches=num_matches,
+            is_clean=False,
+            has_valid_shapes=False,
+        )
+        return gm, info
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fusion.py
@@ -125,6 +125,10 @@ def fuse_rule(
     ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
         raise NotImplementedError
 
+    def build_custom_args_for_linear(self, scale_getattrs: Dict[str, Node]) -> Tuple[object, ...]:
+        """Return the *positional* tail after bias for the fused call."""
+        raise NotImplementedError
+
     def _insert_fused_quant_gemm(
         self, gm: GraphModule, idx: int, parent_node: Node, linear_nodes: List[Node]
     ):
@@ -163,16 +167,17 @@ def split_output(tensor: torch.Tensor) -> Tuple[torch.Tensor, ...]:
 
             # For each kwarg group (e.g., input_scale, weight_scale[, alpha]),
             # create a list of get_attr nodes in the same structure the op expects.
-            for group, kwarg_name in zip(self.scale_groups, ["input_scale", "weight_scale"]):
-                fused_kwargs[kwarg_name] = [
-                    gm.graph.create_node("get_attr", f"{key_fused}_{name}") for name in group
-                ]
+            scale_getattrs: Dict[str, Node] = {
+                name: gm.graph.create_node("get_attr", f"{key_fused}_{name}")
+                for name in flat_scale_names
+            }
+            custom_tail_args = self.build_custom_args_for_linear(scale_getattrs)
 
         # add new linear node + split node
         with gm.graph.inserting_before(linear_nodes[0]):
             fused_linear_node = gm.graph.call_function(
                 get_op_overload_packet(linear_nodes[0].target),
-                args=(parent_node, get_param_node, None),
+                args=(parent_node, get_param_node, None, *custom_tail_args),
                 kwargs=fused_kwargs,
             )
             split_node = gm.graph.call_function(split_output, args=(fused_linear_node,))
@@ -286,6 +291,15 @@ def fuse_rule(
             "input_scale": input_scale[0].clone(),
         }
 
+    def build_custom_args_for_linear(self, scale_getattrs: Dict[str, Node]) -> Tuple[object, ...]:
+        # (..., bias, input_scale(list), weight_scale(list), input_zp(list), weight_zp(list))
+        return (
+            [scale_getattrs["input_scale"]],
+            [scale_getattrs["weight_scale"]],
+            [],
+            [],
+        )
+
     def _apply(
         self,
         gm: GraphModule,
@@ -323,6 +337,15 @@ def fuse_rule(
             "input_scale": input_scale[0].clone(),
         }
 
+    def build_custom_args_for_linear(self, scale_getattrs: Dict[str, Node]) -> Tuple[object, ...]:
+        # (..., bias, input_scale(list), weight_scale(list-with-alpha), input_zp(list), weight_zp(list))
+        return (
+            [scale_getattrs["input_scale"]],
+            [scale_getattrs["weight_scale"], scale_getattrs["alpha"]],
+            [],
+            [],
+        )
+
     def _apply(
         self,
         gm: GraphModule,
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py b/tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py
@@ -77,12 +77,10 @@ def _insert_quantized_linear(
         for scale_name in quantization_impl.scale_names():
             scales[scale_name] = gm.graph.create_node("get_attr", modname + "." + scale_name)
 
-    custom_kwargs = quantization_impl.build_custom_kwargs_for_linear(
-        scales,
-    )
+    custom_args = quantization_impl.build_custom_args_for_linear(scales)
 
     node.target = quantization_impl.custom_op()
-    node.kwargs = {**node.kwargs, **custom_kwargs}
+    node.args = (*node.args, *custom_args)
 
 
 def _insert_quantized_bmm(
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py b/tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py
