[Fusion] Manual ActivationQuantFusionPass initial application

tests/compile/fusions_e2e/common.py

@@ -43,6 +43,45 @@ class AttentionBackendCase(NamedTuple):
 """Are we running on Blackwell, a lot of tests depend on it"""
 
 
+def nvfp4_kernel_exposes_input_quant_key() -> bool:
+    """Check if the NVFP4 kernel selected on this platform exposes input_quant_key.
+
+    FlashInferCuteDslNvFp4LinearKernel does not expose input_quant_key() due to
+    layout incompatibility between the manual fusion kernel (silu_and_mul_nvfp4_quant)
+    output format and the cutedsl backend's expected input format.
+
+    FlashInferCutlassNvFp4LinearKernel does expose input_quant_key() and supports
+    manual fusion.
+    """
+    if not current_platform.is_cuda():
+        return False
+
+    try:
+        from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
+            FlashInferCuteDslNvFp4LinearKernel,
+        )
+
+        is_supported, _ = FlashInferCuteDslNvFp4LinearKernel.is_supported()
+        if is_supported:
+            return False
+    except ImportError:
+        pass
+
+    try:
+        from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
+            FlashInferCutlassNvFp4LinearKernel,
+        )
+
+        is_supported, _ = FlashInferCutlassNvFp4LinearKernel.is_supported()
+        if is_supported:
+            return True
+    except ImportError:
+        pass
+
+    # Fallback: assume no manual fusion support
+    return False
+
+
 def custom_ops_combos(*custom_ops: str) -> Iterable[str]:
     """Generate all combinations of custom ops for parametrization."""
     custom_ops_lists = [[f"-{op}", f"+{op}"] for op in custom_ops]

tests/compile/fusions_e2e/test_tp1_quant.py

@@ -14,6 +14,7 @@
     Matches,
     custom_ops_combos,
     is_blackwell,
+    nvfp4_kernel_exposes_input_quant_key,
 )
 from .models import (
     FLASHINFER_ATTN,
@@ -77,7 +78,6 @@ def test_tp1_fp8_fusions(
     inductor_graph_partition: bool,
     use_deepgemm: bool,
     run_e2e_fusion_test,
-    monkeypatch,
 ):
     if use_deepgemm and not current_platform.is_cuda():
         pytest.skip("DeepGemm only supported on CUDA")
@@ -109,7 +109,7 @@ def test_tp1_fp8_fusions(
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
             fuse_norm_quant=True,
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             enable_qk_norm_rope_fusion=True,
         ),
@@ -119,7 +119,6 @@ def test_tp1_fp8_fusions(
 
     matches_check = [
         "rms_quant_fusion",
-        "act_quant_fusion",
         "norm_rope_fusion",
         "attn_quant_fusion",
     ]
@@ -168,6 +167,12 @@ def test_tp1_fp4_fusions(
     inductor_graph_partition: bool,
     run_e2e_fusion_test,
 ):
+    if nvfp4_kernel_exposes_input_quant_key():
+        pytest.skip(
+            "NVFP4 kernel exposes input_quant_key; manual fusion fires "
+            "instead of compiler pass-based fusion"
+        )
+
     matches = matches_fn(n_layers)
 
     # Reduce size of model and skip weight loading time
@@ -181,13 +186,13 @@ def test_tp1_fp4_fusions(
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
             fuse_norm_quant=True,
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             enable_qk_norm_rope_fusion=True,
         ),
     )
 
-    matches_check = ["act_quant_fusion", "attn_quant_fusion", "norm_rope_fusion"]
+    matches_check = ["attn_quant_fusion", "norm_rope_fusion"]
 
     run_e2e_fusion_test(
         model_name,

tests/compile/fusions_e2e/test_tp2_ar_rms.py

@@ -14,6 +14,7 @@
     Matches,
     custom_ops_combos,
     is_blackwell,
+    nvfp4_kernel_exposes_input_quant_key,
 )
 from .models import (
     FLASHINFER_ATTN,
@@ -68,7 +69,6 @@ def test_tp2_ar_rms_fp8_fusions(
     custom_ops: str,
     inductor_graph_partition: bool,
     run_e2e_fusion_test,
-    monkeypatch,
 ):
     matches = matches_fn(n_layers)
 
@@ -89,7 +89,7 @@ def test_tp2_ar_rms_fp8_fusions(
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
             fuse_norm_quant=True,
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             enable_qk_norm_rope_fusion=True,
             fuse_allreduce_rms=True,
@@ -98,7 +98,6 @@ def test_tp2_ar_rms_fp8_fusions(
 
     matches_check = [
         "rms_quant_fusion",
-        "act_quant_fusion",
         "norm_rope_fusion",
         "attn_quant_fusion",
         "ar_rms_fusion",
@@ -139,8 +138,13 @@ def test_tp2_ar_rms_fp4_fusions(
     custom_ops: str,
     inductor_graph_partition: bool,
     run_e2e_fusion_test,
-    monkeypatch,
 ):
+    if nvfp4_kernel_exposes_input_quant_key():
+        pytest.skip(
+            "NVFP4 kernel exposes input_quant_key; manual fusion fires "
+            "instead of compiler pass-based fusion"
+        )
+
     matches = matches_fn(n_layers)
 
     # Reduce size of model and skip weight loading time
@@ -154,14 +158,13 @@ def test_tp2_ar_rms_fp4_fusions(
         use_inductor_graph_partition=inductor_graph_partition,
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             fuse_allreduce_rms=True,
         ),
     )
 
     matches_check = [
-        "act_quant_fusion",
         "attn_quant_fusion",
         "ar_rms_fusion",
     ]

tests/compile/fusions_e2e/test_tp2_async_tp.py

@@ -14,6 +14,7 @@
     Matches,
     custom_ops_combos,
     is_blackwell,
+    nvfp4_kernel_exposes_input_quant_key,
 )
 from .models import (
     FLASHINFER_ATTN,
@@ -61,7 +62,7 @@ def test_tp2_async_tp_fp8_fusions(
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
             fuse_norm_quant=True,
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             enable_qk_norm_rope_fusion=True,
             enable_sp=True,
@@ -74,7 +75,6 @@ def test_tp2_async_tp_fp8_fusions(
 
     matches_check = [
         "rms_quant_fusion",
-        "act_quant_fusion",
         "norm_rope_fusion",
         "attn_quant_fusion",
         "sequence_parallel",
@@ -114,6 +114,12 @@ def test_tp2_async_tp_nvfp4_fusions(
     inductor_graph_partition: bool,
     run_e2e_fusion_test,
 ):
+    if nvfp4_kernel_exposes_input_quant_key():
+        pytest.skip(
+            "NVFP4 kernel exposes input_quant_key; manual fusion fires "
+            "instead of compiler pass-based fusion"
+        )
+
     # NVFP4 currently wires the all-gather + GEMM path only.
     matches = matches_fn(n_layers)._replace(async_tp=n_layers * 2)
 
@@ -127,7 +133,7 @@ def test_tp2_async_tp_nvfp4_fusions(
         use_inductor_graph_partition=inductor_graph_partition,
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             enable_sp=True,
             fuse_gemm_comms=True,
@@ -138,7 +144,6 @@ def test_tp2_async_tp_nvfp4_fusions(
     )
 
     matches_check = [
-        "act_quant_fusion",
         "attn_quant_fusion",
         "sequence_parallel",
         "async_tp",
@@ -245,7 +250,7 @@ def test_tp2_sp_ar_rms_fp8_fusions(
         custom_ops=custom_ops.split(","),
         pass_config=PassConfig(
             fuse_norm_quant=True,
-            fuse_act_quant=True,
+            fuse_act_quant=False,
             fuse_attn_quant=True,
             enable_qk_norm_rope_fusion=True,
             enable_sp=True,
@@ -258,7 +263,6 @@ def test_tp2_sp_ar_rms_fp8_fusions(
 
     matches_check = [
         "rms_quant_fusion",
-        "act_quant_fusion",
         "norm_rope_fusion",
         "attn_quant_fusion",
         "ar_rms_fusion",