A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[ROCm][Bugfix] Fix the case where there's bias (#24895)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Sep 15, 2025 at 22:05 UTC
2891603efdbd38ac7197d5f41b0e245fb6a53b82
2 files changed
+32
-1
tests/kernels/quantization/test_rocm_skinny_gemms.py
+31
@@ -5,6 +5,8 @@ import torch
5
6
import vllm._custom_ops as ops
7
from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
8
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
9
+ rocm_per_tensor_w8a8_scaled_mm_impl)
10
from vllm.platforms import current_platform
11
12
DTYPES = [torch.bfloat16, torch.float16]
@@ -116,3 +118,32 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
118
current_platform.get_cu_count())
119
120
assert torch.allclose(out, ref_out, rtol=0.01)
121
+
122
+
123
+@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
124
+@pytest.mark.parametrize("dtype", DTYPES)
125
+@pytest.mark.parametrize("seed", SEEDS)
126
+@pytest.mark.parametrize("use_bias", [True, False])
127
+@pytest.mark.skipif(
128
+ not (current_platform.is_rocm() and current_platform.supports_fp8()),
129
+ reason="only test for rocm fp8")
130
+def test_rocm_per_tensor_w8a8_scaled_mm_impl(n, k, m, dtype, seed, use_bias):
131
+ torch.manual_seed(seed)
132
+
133
+ A = torch.rand(n, k, device="cuda")
134
+ B = torch.rand(m, k, device="cuda")
135
+
136
+ A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
137
+ B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
138
+
139
+ bias = torch.rand(1, m, dtype=dtype, device="cuda") if use_bias else None
140
+
141
+ output = rocm_per_tensor_w8a8_scaled_mm_impl(A, B.t(), dtype, scale_a,
142
+ scale_b, bias)
143
+ ref_out = torch._scaled_mm(A,
144
+ B.t(),
145
+ out_dtype=dtype,
146
+ scale_a=scale_a,
147
+ scale_b=scale_b,
148
+ bias=bias)
149
+ assert torch.allclose(output, ref_out, rtol=0.01)
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+1
-1
@@ -179,7 +179,7 @@ def rocm_per_tensor_w8a8_scaled_mm_impl(qinput: torch.Tensor,
179
bias: torch.Tensor) -> torch.Tensor:
180
from vllm.platforms.rocm import on_mi3xx
181
if envs.VLLM_ROCM_USE_SKINNY_GEMM and on_mi3xx(
182
- ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0:
182
+ ) and qinput.shape[0] == 1 and qinput.shape[1] % 16 == 0 and bias is None:
183
output = ops.wvSplitKQ(weight.t(), qinput, out_dtype, scale_a, scale_b,
184
current_platform.get_cu_count())
185
else: