Remove triton per group quant [ROCm] [Bugfix] (#49621)
Signed-off-by: Andy Friedrich <afriedri@amd.com>
afriedri committed
Jul 28, 2026 at 00:53 UTC
f472ab0a4c0987d2badc2f2aeecd38fee4180a4c
6 files changed
+8
-78
tests/compile/passes/distributed/test_fusion_all_reduce.py
+6
-25
@@ -272,12 +272,10 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
272
token_num=16,
273
eps=1e-6,
274
dtype: torch.dtype = torch.bfloat16,
275
- use_triton_quant: bool = False,
275
):
276
super().__init__()
277
self.hidden_size = hidden_size
278
self.eps = eps
280
- self.use_triton_quant = use_triton_quant
279
assert hidden_size % self.quant_group_size == 0, (
280
f"hidden_size ({hidden_size}) must be a multiple of "
281
f"quant_group_size ({self.quant_group_size}) for per-group FP8 quant"
@@ -289,10 +287,6 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
287
]
288
289
def _group_quant(self, rms: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
292
- if self.use_triton_quant:
293
- return torch.ops.vllm.triton_per_token_group_quant_fp8(
294
- rms, self.quant_group_size
295
- )
290
return torch.ops.vllm.rocm_aiter_group_fp8_quant.default(
291
rms, self.quant_group_size
292
)
@@ -339,11 +333,7 @@ class TestAiterAllReduceRMSNormGroupQuantFP8Model(torch.nn.Module):
333
def ops_in_model_before(self):
334
return [
335
torch.ops.vllm.all_reduce.default,
342
- (
343
- torch.ops.vllm.triton_per_token_group_quant_fp8.default
344
- if self.use_triton_quant
345
- else torch.ops.vllm.rocm_aiter_group_fp8_quant.default
346
- ),
336
+ torch.ops.vllm.rocm_aiter_group_fp8_quant.default,
337
]
338
339
def ops_in_model_after(self):
@@ -646,7 +636,6 @@ def all_reduce_fusion_pass_on_test_model(
636
637
638
@multi_gpu_test(num_gpus=2)
649
-@pytest.mark.parametrize("use_triton_quant", [True, False])
639
@pytest.mark.parametrize("batch_size", [8])
640
@pytest.mark.parametrize("seq_len", [8])
641
@pytest.mark.parametrize("hidden_size", [128])
@@ -663,7 +652,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
652
hidden_size: int,
653
dtype: torch.dtype,
654
enable_rms_norm_custom_op: bool,
666
- use_triton_quant: bool,
655
monkeypatch: pytest.MonkeyPatch,
656
):
657
"""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(
664
* ``AiterAllreduceFusedAddRMSNormGroupQuantFP8Pattern`` (with-residual,
665
single ``rms`` consumer)
666
* ``AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern`` (with-
679
- residual, DSv3.2 indexer fan-out; parametrized over both
680
- ``triton_per_token_group_quant_fp8`` and ``rocm_aiter_group_fp8_quant``
681
- producers).
667
+ residual, DSv3.2 indexer fan-out; parametrized over
668
+ ``rocm_aiter_group_fp8_quant``
669
+ producer).
670
"""
671
with monkeypatch.context() as m:
672
m.setenv("VLLM_ROCM_USE_AITER", "1")
@@ -703,7 +691,6 @@ def test_rocm_aiter_all_reduce_rmsnorm_group_quant_fp8_fusion_pass_replace(
691
hidden_size,
692
dtype,
693
enable_rms_norm_custom_op,
706
- use_triton_quant,
694
monkeypatch,
695
),
696
nprocs=nprocs,
@@ -721,7 +708,6 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
708
hidden_size: int,
709
dtype: torch.dtype,
710
enable_rms_norm_custom_op: bool,
724
- use_triton_quant: bool,
711
monkeypatch: pytest.MonkeyPatch,
712
):
713
set_random_seed(0)
@@ -749,10 +735,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
735
custom_ops = []
736
if enable_rms_norm_custom_op:
737
custom_ops.append("+rms_norm")
752
- # ``triton_per_token_group_quant_fp8`` is emitted by ``QuantFP8.forward_hip``
753
- # only when QuantFP8 is enabled as a custom op (and ``use_triton=True`` at
754
- # the call site). The patterns in this PR are robust to both Triton and
755
- # rocm_aiter forms; we always enable +quant_fp8 so the matcher's example
738
+ # We always enable +quant_fp8 so the matcher's example
739
# trace finds the same form the test model uses.
740
custom_ops.append("+quant_fp8")
741
@@ -783,9 +766,7 @@ def rocm_aiter_group_quant_fusion_pass_on_test_model(
766
)
767
768
token_num = batch_size * seq_len
786
- model = test_model_cls(
787
- hidden_size, token_num, dtype=dtype, use_triton_quant=use_triton_quant
788
- )
769
+ model = test_model_cls(hidden_size, token_num, dtype=dtype)
770
771
hidden_states = torch.randn((token_num, hidden_size), requires_grad=False)
772
tests/compile/passes/test_fusion.py
-2
@@ -195,8 +195,6 @@ class TestModel(torch.nn.Module):
195
# Blockwise path
196
if self.use_aiter_fusion and self.use_aiter_quant_op:
197
return [rocm_aiter_ops.get_group_quant_op()]
198
- if self.use_aiter_fusion:
199
- return [torch.ops.vllm.triton_per_token_group_quant_fp8.default]
198
else:
199
if self.use_aiter_quant_op:
200
return [rocm_aiter_ops.get_per_token_quant_op()]
tests/compile/passes/test_silu_mul_quant_fusion.py
+1
-10
@@ -158,13 +158,6 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
158
input_dtype=dtype,
159
)
160
161
- if not current_platform.is_fp8_fnuz():
162
- kernel = self.w8a8_block_fp8_linear.kernel
163
- orig_quant = kernel.quant_fp8
164
- kernel.quant_fp8 = lambda *a, use_triton=False, **kw: orig_quant(
165
- *a, use_triton=True, **kw
166
- )
167
-
161
self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
162
163
def forward(self, x):
@@ -175,9 +168,7 @@ class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
168
def ops_in_model_before(self):
169
return [
170
SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul,
178
- rocm_aiter_ops.get_group_quant_op()
179
- if current_platform.is_fp8_fnuz()
180
- else torch.ops.vllm.triton_per_token_group_quant_fp8.default,
171
+ rocm_aiter_ops.get_group_quant_op(),
172
]
173
174
def ops_in_model_after(self):
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+1
-2
@@ -1416,8 +1416,7 @@ class AiterAllreduceFusedAddRMSNormGroupQuantWithIndexerPattern(
1416
The trailing FP8 group-quant is matched via ``MatcherQuantFP8`` (consistent
1417
with the sibling patterns above), which traces both ``QuantFP8.forward_hip``
1418
and ``forward_native`` paths and so matches whichever op the call site
1419
- lowers to (``vllm.triton_per_token_group_quant_fp8`` or
1420
- ``vllm.rocm_aiter_group_fp8_quant``).
1419
+ lowers to (``vllm.rocm_aiter_group_fp8_quant``).
1420
"""
1421
1422
def __init__(
vllm/model_executor/layers/quantization/input_quant_fp8.py
-5
@@ -139,11 +139,6 @@ class QuantFP8(CustomOp):
139
scale_ub: torch.Tensor | None = None,
140
use_triton: bool = False,
141
) -> tuple[torch.Tensor, torch.Tensor]:
142
- if self.is_group_quant and use_triton:
143
- assert scale is None, "Dynamic group quantization does not use scale"
144
-
145
- return torch.ops.vllm.triton_per_token_group_quant_fp8(x, self.group_size)
146
-
142
use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous()
143
use_aiter_per_tensor_quant = (
144
use_aiter_quant and self.group_shape.is_per_tensor()
vllm/model_executor/layers/quantization/utils/fp8_utils.py
-34
@@ -34,7 +34,6 @@ from vllm.utils.deep_gemm import (
34
transform_sf_into_required_layout,
35
)
36
from vllm.utils.platform_utils import get_device_name_as_file_name
37
-from vllm.utils.torch_utils import direct_register_custom_op
37
38
logger = init_logger(__name__)
39
@@ -45,39 +44,6 @@ def is_fp8(x: torch.dtype | torch.Tensor) -> bool:
44
return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
45
46
48
-def _triton_per_token_group_quant_fp8_impl(
49
- x: torch.Tensor,
50
- group_size: int,
51
-) -> tuple[torch.Tensor, torch.Tensor]:
52
- return per_token_group_quant_fp8(
53
- x, group_size, column_major_scales=False, use_ue8m0=False
54
- )
55
-
56
-
57
-def _triton_per_token_group_quant_fp8_fake(
58
- x: torch.Tensor,
59
- group_size: int,
60
-) -> tuple[torch.Tensor, torch.Tensor]:
61
- M, N = x.shape
62
- x_fp8 = torch.empty((M, N), dtype=current_platform.fp8_dtype(), device=x.device)
63
- out_bs = torch.empty(
64
- (
65
- M,
66
- (N + group_size - 1) // group_size,
67
- ),
68
- dtype=torch.float32,
69
- device=x.device,
70
- )
71
- return x_fp8, out_bs
72
-
73
-
74
-direct_register_custom_op(
75
- "triton_per_token_group_quant_fp8",
76
- _triton_per_token_group_quant_fp8_impl,
77
- fake_impl=_triton_per_token_group_quant_fp8_fake,
78
-)
79
-
80
-
47
def input_to_float8(
48
x: torch.Tensor, dtype: torch.dtype | None = None
49
) -> tuple[torch.Tensor, torch.Tensor]: