[Bugfix] Stop fusing scale_ue8m0=True group quant into ue8m0-unaware fused kernels

tests/compile/passes/test_fusion.py

@@ -676,3 +676,74 @@ def test_aiter_fusion_rmsnorm_gated_quant_no_gdn_layers(
         model_fused(x, z)
 
         assert fusion_pass.matched_count == 0
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="V1 attention only on CUDA")
+def test_fusion_rmsnorm_group_quant_skips_ue8m0():
+    """scale_ue8m0=True group quant must stay unfused.
+
+    The fused op ``rms_norm_per_block_quant`` has no ue8m0 parameter and
+    always emits raw float32 group scales. Replacing a ``scale_ue8m0=True``
+    ``per_token_group_fp8_quant`` with it silently drops the power-of-two
+    scale contract of the DeepGEMM path, so the pass must not match those
+    graphs. Regression test for the pattern-registration fix; before it,
+    the pass matched and the compiled model produced non-power-of-two
+    scales (contract break measured as 128/128 groups on SM120).
+    """
+    from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(torch.bfloat16)
+
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=torch.bfloat16),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["none", "+rms_norm", "+quant_fp8"],
+            backend="eager",
+            pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True),
+        ),
+    )
+
+    with vllm.config.set_current_vllm_config(vllm_config):
+
+        class Model(torch.nn.Module):
+            def __init__(self, hidden: int):
+                super().__init__()
+                self.norm = RMSNorm(hidden, eps=1e-6)
+                self.quant = QuantFP8(
+                    static=False,
+                    group_shape=GroupShape(1, 128),
+                    column_major_scales=False,
+                    use_ue8m0=True,
+                    compile_native=False,
+                )
+
+            def forward(self, x, residual):
+                y, res = self.norm(x, residual)
+                q, s = self.quant(y)
+                deq = q.to(torch.float32) * s.repeat_interleave(128, dim=1).to(
+                    torch.float32
+                )
+                return deq, s, res
+
+        model = Model(512)
+        fusion_pass = RMSNormQuantFusionPass(vllm_config)
+        backend = TestBackend(
+            NoOpEliminationPass(vllm_config), fusion_pass, PostCleanupPass(vllm_config)
+        )
+
+        x = torch.randn(32, 512)
+        residual = torch.randn(32, 512)
+        result = model(x.clone(), residual.clone())
+        result2 = torch.compile(model, backend=backend)(x.clone(), residual.clone())
+
+        # The ue8m0 graph must not be rewritten...
+        assert fusion_pass.matched_count == 0
+        # ...and the compiled model must keep emitting power-of-two scales.
+        for scales in (result[1], result2[1]):
+            log2_scales = torch.log2(scales.float())
+            assert torch.equal(log2_scales, log2_scales.round()), (
+                "expected power-of-two group scales on the scale_ue8m0=True path"
+            )
+        torch.testing.assert_close(result[0], result2[0], atol=1e-6, rtol=0)

tests/compile/passes/test_silu_mul_quant_fusion.py

@@ -368,3 +368,71 @@ def test_fusion_silu_and_mul_quant(
 
         # In post-nodes, fused kernels should be present and quant op should not
         backend.check_after_ops(model.ops_in_model_after())
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="ue8m0 group quant is CUDA-only"
+)
+def test_fusion_silu_and_mul_group_quant_skips_ue8m0():
+    """scale_ue8m0=True group quant must stay unfused.
+
+    ``silu_and_mul_per_block_quant`` has no ue8m0 parameter and always emits
+    raw float32 group scales, so fusing a ``scale_ue8m0=True``
+    ``per_token_group_fp8_quant`` would silently drop the power-of-two scale
+    contract of the DeepGEMM path. Regression test for the
+    pattern-registration fix; before it, the pass matched and the compiled
+    model produced non-power-of-two scales (128/128 groups on SM120).
+    """
+    torch.set_default_device("cuda")
+    torch.set_default_dtype(torch.bfloat16)
+
+    config = VllmConfig(
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["none", "+silu_and_mul", "+quant_fp8"],
+            backend="eager",
+            pass_config=PassConfig(fuse_act_quant=True, eliminate_noops=True),
+        ),
+    )
+
+    with set_current_vllm_config(config):
+
+        class Model(torch.nn.Module):
+            def __init__(self, hidden: int):
+                super().__init__()
+                self.silu_and_mul = SiluAndMul()
+                self.quant = QuantFP8(
+                    static=False,
+                    group_shape=GroupShape(1, 128),
+                    column_major_scales=False,
+                    use_ue8m0=True,
+                    compile_native=False,
+                )
+
+            def forward(self, x):
+                y = self.silu_and_mul(x)
+                q, s = self.quant(y)
+                deq = q.to(torch.float32) * s.repeat_interleave(128, dim=1).to(
+                    torch.float32
+                )
+                return deq, s
+
+        model = Model(512)
+        fusion_pass = ActivationQuantFusionPass(config)
+        backend = TestBackend(
+            NoOpEliminationPass(config), fusion_pass, PostCleanupPass(config)
+        )
+
+        x = torch.randn(32, 1024)
+        result = model(x.clone())
+        result2 = torch.compile(model, backend=backend)(x.clone())
+
+        # The ue8m0 graph must not be rewritten...
+        assert fusion_pass.matched_count == 0
+        # ...and the compiled model must keep emitting power-of-two scales.
+        for scales in (result[1], result2[1]):
+            log2_scales = torch.log2(scales.float())
+            assert torch.equal(log2_scales, log2_scales.round()), (
+                "expected power-of-two group scales on the scale_ue8m0=True path"
+            )
+        torch.testing.assert_close(result[0], result2[0], atol=1e-6, rtol=0)

vllm/compilation/passes/fusion/act_quant_fusion.py

@@ -299,22 +299,25 @@ def __init__(self, config: VllmConfig) -> None:
             self.register(SiluMulNvfp4QuantPattern())
 
         if current_platform.is_cuda():
+            # NOTE: is_e8m0 (scale_ue8m0=True) variants are intentionally not
+            # registered: silu_and_mul_per_block_quant has no ue8m0 parameter
+            # and always emits raw float32 group scales, so the replacement
+            # would silently drop the power-of-two scale contract of the
+            # DeepGEMM path. See the matching note in rms_quant_fusion.py.
             for (
                 quant_key,
                 is_scale_transposed,
-                is_e8m0,
                 is_tma_aligned,
             ) in itertools.product(
                 [kFp8Dynamic128Sym, kFp8Dynamic64Sym],
                 [False, True],
-                [True, False],
                 [False, True],
             ):
                 self.register(
                     SiluMulBlockQuantPattern(
                         quant_key,
                         is_scale_transposed=is_scale_transposed,
-                        is_e8m0=is_e8m0,
+                        is_e8m0=False,
                         is_tma_aligned=is_tma_aligned,
                     )
                 )

vllm/compilation/passes/fusion/rms_quant_fusion.py

@@ -644,29 +644,36 @@ def __init__(self, config: VllmConfig) -> None:
             RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)
 
             # Only register group quant patterns on CUDA/ROCm where the C++ op exists
+            # NOTE: is_e8m0 (scale_ue8m0=True) variants are intentionally not
+            # registered. The fused op rms_norm_per_block_quant has no ue8m0
+            # parameter and always emits raw float32 group scales, so replacing
+            # a scale_ue8m0=True per_token_group_fp8_quant with it silently
+            # drops the power-of-two scale contract that the DeepGEMM path
+            # relies on (see the TODO(quant-rms-fusion) note in
+            # tests/compile/passes/test_fusion.py and #40552). Keep those graphs
+            # unfused until the kernel supports UE8M0 scales.
             for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:
                 for has_col_major_scales in [True, False]:
-                    for is_e8m0 in [True, False]:
-                        for is_tma_aligned in [False, True]:
-                            # Fuse fused_add_rms_norm + fp8 group quant
-                            FusedAddRMSNormGroupQuantPattern(
-                                epsilon,
-                                FP8_DTYPE,
-                                group_shape=group_shape,
-                                is_e8m0=is_e8m0,
-                                has_col_major_scales=has_col_major_scales,
-                                is_tma_aligned=is_tma_aligned,
-                            ).register(self.patterns)
-
-                            # Fuse rms_norm + fp8 group quant
-                            RMSNormGroupQuantPattern(
-                                epsilon,
-                                FP8_DTYPE,
-                                group_shape=group_shape,
-                                is_e8m0=is_e8m0,
-                                has_col_major_scales=has_col_major_scales,
-                                is_tma_aligned=is_tma_aligned,
-                            ).register(self.patterns)
+                    for is_tma_aligned in [False, True]:
+                        # Fuse fused_add_rms_norm + fp8 group quant
+                        FusedAddRMSNormGroupQuantPattern(
+                            epsilon,
+                            FP8_DTYPE,
+                            group_shape=group_shape,
+                            is_e8m0=False,
+                            has_col_major_scales=has_col_major_scales,
+                            is_tma_aligned=is_tma_aligned,
+                        ).register(self.patterns)
+
+                        # Fuse rms_norm + fp8 group quant
+                        RMSNormGroupQuantPattern(
+                            epsilon,
+                            FP8_DTYPE,
+                            group_shape=group_shape,
+                            is_e8m0=False,
+                            has_col_major_scales=has_col_major_scales,
+                            is_tma_aligned=is_tma_aligned,
+                        ).register(self.patterns)
 
         self.dump_patterns(config, self.patterns)