diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 1aac4b2bec4..e9c8d0deaa7 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -272,12 +272,10 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): token_num=16, eps=1e-6, dtype: torch.dtype = torch.bfloat16, - use_triton_quant: bool = False, ): super().__init__() self.hidden_size = hidden_size self.eps = eps - self.use_triton_quant = use_triton_quant assert hidden_size % self.quant_group_size == 0, ( f"hidden_size ({hidden_size}) must be a multiple of " f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant" @@ -289,10 +287,6 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): ] def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - if self.use_triton_quant: - return torch.ops.vllm.triton_per_token_group_quant_fp8( - rms, self.quant_group_size - ) return torch.ops.vllm.rocm_aiter_group_fp8_quant.default( rms, self.quant_group_size ) @@ -339,11 +333,7 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module): def ops_in_model_before(self): return [ torch.ops.vllm.all_reduce.default, - ( - torch.ops.vllm.triton_per_token_group_quant_fp8.default - if self.use_triton_quant - else torch.ops.vllm.rocm_aiter_group_fp8_quant.default - ), + torch.ops.vllm.rocm_aiter_group_fp8_quant.default, ] def ops_in_model_after(self): @@ -646,7 +636,6 @@ def all_reduce_fusion_pass_on_test_model( @multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize("use_triton_quant", [True, False]) @pytest.mark.parametrize("batch_size", [8]) @pytest.mark.parametrize("seq_len", [8]) @pytest.mark.parametrize("hidden_size", [128]) @@ -663,7 +652,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): """Sibling of ``test_all_reduce_fusion_pass_replace`` for the new @@ -676,9 +664,9 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( * ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual, single ``rms`` consumer) * ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with- - residual, DSv3.2 indexer fan-out; parametrized over both - ``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant`` - producers). + residual, DSv3.2 indexer fan-out; parametrized over + ``rocm_aiter_group_fp8_quant`` + producer). """ with monkeypatch.context() as m: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -703,7 +691,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace( hidden_size, dtype, enable_rms_norm_custom_op, - use_triton_quant, monkeypatch, ), nprocs=nprocs, @@ -721,7 +708,6 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( hidden_size: int, dtype: torch.dtype, enable_rms_norm_custom_op: bool, - use_triton_quant: bool, monkeypatch: pytest.MonkeyPatch, ): set_random_seed(0) @@ -749,10 +735,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( custom_ops = [] if enable_rms_norm_custom_op: custom_ops.append("+rms_norm") - # ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip`` - # only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at - # the call site). The patterns in this PR are robust to both Triton and - # rocm_aiter forms; we always enable +quant_fp8 so the matcher's example + # We always enable +quant_fp8 so the matcher's example # trace finds the same form the test model uses. custom_ops.append("+quant_fp8") @@ -783,9 +766,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model( ) token_num = batch_size * seq_len - model = test_model_cls( - hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant - ) + model = test_model_cls(hidden_size, token_num, dtype=dtype) hidden_states = torch.randn((token_num, hidden_size), requires_grad=False) diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 92d1902b2c2..591b014d9e2 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -195,8 +195,6 @@ class TestModel(torch.nn.Module): # Blockwise path if self.use_aiter_fusion and self.use_aiter_quant_op: return [rocm_aiter_ops.get_group_quant_op()] - if self.use_aiter_fusion: - return [torch.ops.vllm.triton_per_token_group_quant_fp8.default] else: if self.use_aiter_quant_op: return [rocm_aiter_ops.get_per_token_quant_op()] diff --git a/tests/compile/passes/test_silu_mul_quant_fusion.py b/tests/compile/passes/test_silu_mul_quant_fusion.py index bc134ed427a..7d291cc5044 100644 --- a/tests/compile/passes/test_silu_mul_quant_fusion.py +++ b/tests/compile/passes/test_silu_mul_quant_fusion.py @@ -158,13 +158,6 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): input_dtype=dtype, ) - if not current_platform.is_fp8_fnuz(): - kernel = self.w8a8_block_fp8_linear.kernel - orig_quant = kernel.quant_fp8 - kernel.quant_fp8 = lambda *a, use_triton=False, **kw: orig_quant( - *a, use_triton=True, **kw - ) - self.enable_silu_mul_custom_op = self.silu_and_mul.enabled() def forward(self, x): @@ -175,9 +168,7 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module): def ops_in_model_before(self): return [ SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul, - rocm_aiter_ops.get_group_quant_op() - if current_platform.is_fp8_fnuz() - else torch.ops.vllm.triton_per_token_group_quant_fp8.default, + rocm_aiter_ops.get_group_quant_op(), ] def ops_in_model_after(self): diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index f7ff7df66cd..1722b524eeb 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -1416,8 +1416,7 @@ class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern( The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent with the sibling patterns above), which traces both ``QuantFP8.forward_hip`` and ``forward_native`` paths and so matches whichever op the call site - lowers to (``vllm.triton_per_token_group_quant_fp8`` or - ``vllm.rocm_aiter_group_fp8_quant``). + lowers to (``vllm.rocm_aiter_group_fp8_quant``). """ def __init__( diff --git a/vllm/model_executor/layers/quantization/input_quant_fp8.py b/vllm/model_executor/layers/quantization/input_quant_fp8.py index e8810919c20..2eb34630aa6 100644 --- a/vllm/model_executor/layers/quantization/input_quant_fp8.py +++ b/vllm/model_executor/layers/quantization/input_quant_fp8.py @@ -139,11 +139,6 @@ class QuantFP8(CustomOp): scale_ub: torch.Tensor | None = None, use_triton: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - if self.is_group_quant and use_triton: - assert scale is None, "Dynamic group quantization does not use scale" - - return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size) - use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous() use_aiter_per_tensor_quant = ( use_aiter_quant and self.group_shape.is_per_tensor() diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 83e56a4567b..2e4fbdf4c64 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -34,7 +34,6 @@ from vllm.utils.deep_gemm import ( transform_sf_into_required_layout, ) from vllm.utils.platform_utils import get_device_name_as_file_name -from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -45,39 +44,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool: return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz -def _triton_per_token_group_quant_fp8_impl( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - return per_token_group_quant_fp8( - x, group_size, column_major_scales=False, use_ue8m0=False - ) - - -def _triton_per_token_group_quant_fp8_fake( - x: torch.Tensor, - group_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - M, N = x.shape - x_fp8 = torch.empty((M, N), dtype=current_platform.fp8_dtype(), device=x.device) - out_bs = torch.empty( - ( - M, - (N + group_size - 1) // group_size, - ), - dtype=torch.float32, - device=x.device, - ) - return x_fp8, out_bs - - -direct_register_custom_op( - "triton_per_token_group_quant_fp8", - _triton_per_token_group_quant_fp8_impl, - fake_impl=_triton_per_token_group_quant_fp8_fake, -) - - def input_to_float8( x: torch.Tensor, dtype: torch.dtype | None = None ) -> tuple[torch.Tensor, torch.Tensor]: