PR comments: add _custom_op suffix

Signed-off-by: Luka Govedič <lgovedic@redhat.com>

Author: ProExpertProg

tests/compile/test_fusion.py

@@ -71,8 +71,8 @@ def __init__(
                 act_quant_group_shape=group_shape,
             )
 
-        self.enable_rms_norm = self.norm[0].enabled()
-        self.enable_quant_fp8 = self.fp8_linear.quant_fp8.enabled()
+        self.enable_rms_norm_custom_op = self.norm[0].enabled()
+        self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
 
     def forward(self, x):
         # avoid having graph input be an arg to a pattern directly
@@ -107,21 +107,25 @@ def ops_in_model_after(self):
     def ops_in_model_before(self):
         return (
             [QUANT_OPS[self.quant_key]]
-            if self.enable_quant_fp8
+            if self.enable_quant_fp8_custom_op
             else [torch.ops.aten.reciprocal]
         )
 
     def ops_in_model_before_partial(self):
-        return [RMS_OP, RMS_ADD_OP] if self.enable_rms_norm else [torch.ops.aten.rsqrt]
+        return (
+            [RMS_OP, RMS_ADD_OP]
+            if self.enable_rms_norm_custom_op
+            else [torch.ops.aten.rsqrt]
+        )
 
 
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("hidden_size", [64])
 @pytest.mark.parametrize("num_tokens", [257])
 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
 @pytest.mark.parametrize("static", [True, False])
-@pytest.mark.parametrize("enable_rms_norm", [True, False])
-@pytest.mark.parametrize("enable_quant_fp8", [True, False])
+@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
+@pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False])
 # cuda_force_torch used to test torch code path on platforms that
 # cutlass_fp8_supported() == True.
 @pytest.mark.parametrize(
@@ -136,8 +140,8 @@ def test_fusion_rmsnorm_quant(
     num_tokens,
     eps,
     static,
-    enable_rms_norm,
-    enable_quant_fp8,
+    enable_rms_norm_custom_op,
+    enable_quant_fp8_custom_op,
     cuda_force_torch,
 ):
     torch.set_default_device("cuda")
@@ -146,9 +150,9 @@ def test_fusion_rmsnorm_quant(
     maybe_create_device_identity()  # needed for certain non-cutlass fp8 paths
 
     custom_ops = []
-    if enable_rms_norm:
+    if enable_rms_norm_custom_op:
         custom_ops.append("+rms_norm")
-    if enable_quant_fp8:
+    if enable_quant_fp8_custom_op:
         custom_ops.append("+quant_fp8")
     vllm_config = VllmConfig(
         model_config=ModelConfig(dtype=dtype),
@@ -195,7 +199,7 @@ def test_fusion_rmsnorm_quant(
         # there's a risk that the fused add doesn't get included in the
         # replacement and only the rms part gets fused with quant.
         # Hence, we check only 2 add nodes are left (final fused rmsnorm add).
-        if not enable_rms_norm:
+        if not enable_rms_norm_custom_op:
             n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))
             # 7 = 1 (RMS) + 3x2 (3xRMS_ADD, 2 each)
             assert n_add_nodes(backend.graph_pre_pass) == 7

tests/compile/test_fusion_all_reduce.py

@@ -194,7 +194,7 @@ def ops_in_model_before(self):
 
 @multi_gpu_test(num_gpus=2)
 @pytest.mark.parametrize(
-    "test_model, enable_quant_fp8",
+    "test_model, enable_quant_fp8_custom_op",
     [
         (TestAllReduceRMSNormModel, False),
         (TestAllReduceRMSNormStaticQuantFP8Model, True),
@@ -206,7 +206,7 @@ def ops_in_model_before(self):
 @pytest.mark.parametrize("seq_len", [8])
 @pytest.mark.parametrize("hidden_size", [64])
 @pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("enable_rms_norm", [True, False])
+@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
 @pytest.mark.skipif(
     not find_spec("flashinfer")
@@ -220,8 +220,8 @@ def test_all_reduce_fusion_pass_replace(
     seq_len: int,
     hidden_size: int,
     dtype: torch.dtype,
-    enable_rms_norm,
-    enable_quant_fp8,
+    enable_rms_norm_custom_op,
+    enable_quant_fp8_custom_op,
 ):
     num_processes = 2
     if (
@@ -243,8 +243,8 @@ def run_torch_spawn(fn, nprocs):
                 seq_len,
                 hidden_size,
                 dtype,
-                enable_rms_norm,
-                enable_quant_fp8,
+                enable_rms_norm_custom_op,
+                enable_quant_fp8_custom_op,
             ),
             nprocs=nprocs,
         )
@@ -260,8 +260,8 @@ def all_reduce_fusion_pass_on_test_model(
     seq_len: int,
     hidden_size: int,
     dtype: torch.dtype,
-    enable_rms_norm,
-    enable_quant_fp8,
+    enable_rms_norm_custom_op,
+    enable_quant_fp8_custom_op,
 ):
     current_platform.seed_everything(0)
 
@@ -284,9 +284,9 @@ def all_reduce_fusion_pass_on_test_model(
     initialize_model_parallel(tensor_model_parallel_size=world_size)
 
     custom_ops = []
-    if enable_rms_norm:
+    if enable_rms_norm_custom_op:
         custom_ops.append("+rms_norm")
-    if enable_quant_fp8:
+    if enable_quant_fp8_custom_op:
         custom_ops.append("+quant_fp8")
 
     vllm_config = VllmConfig(