karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 2 Packages Activity

Remove triton per group quant [ROCm] [Bugfix] (#49621)

Browse Source 
Signed-off-by: Andy Friedrich <afriedri@amd.com>
This commit is contained in:
afriedri
2026-07-28 05:53:46 +00:00
committed by GitHub
parent 74587939b1
commit f472ab0a4c
6 changed files with 8 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/compile/passes/distributed/test_fusion_all_reduce.py
+6 -25
View File
@@ -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)
tests/compile/passes/test_fusion.py
-2
View File
@@ -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()]
tests/compile/passes/test_silu_mul_quant_fusion.py
+1 -10
View File
@@ -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):
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+1 -2
View File
@@ -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__(
vllm/model_executor/layers/quantization/input_quant_fp8.py
-5
View File
@@ -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()
vllm/model_executor/layers/quantization/utils/fp8_utils.py
-34
View File
@@ -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]: