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 9, 07:30 UTC
[Bugfix][ROCm] Fixing the skinny gemm dispatch logic from #32831 (#33366)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Jan 30, 2026 at 19:05 UTC
31aedfe7d62adea6b5a7644f3c323f585b042ce3
4 files changed
+18
-18
tests/kernels/quantization/test_rocm_skinny_gemms.py
+15
-2
@@ -87,6 +87,13 @@ NKM_FACTORS_WVSPLITK_FP8 = [
87
SEEDS = [0]
88
89
90
+def pad_weights_fp8(weight):
91
+ num_pad = 256 // weight.element_size()
92
+ import torch.nn.functional as F
93
+
94
+ return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
95
+
96
+
97
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC)
98
@pytest.mark.parametrize("dtype", DTYPES)
99
@pytest.mark.parametrize("seed", SEEDS)
@@ -191,11 +198,12 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
198
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
199
@pytest.mark.parametrize("dtype", DTYPES)
200
@pytest.mark.parametrize("seed", SEEDS)
201
+@pytest.mark.parametrize("padded", [False, True])
202
@pytest.mark.skipif(
203
not (current_platform.is_rocm() and current_platform.supports_fp8()),
204
reason="only test for rocm fp8",
205
)
198
-def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
206
+def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
207
torch.manual_seed(seed)
208
209
A = torch.rand(n, k, device="cuda") - 0.5
@@ -203,6 +211,8 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
211
212
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
213
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
214
+ if padded:
215
+ B = pad_weights_fp8(B)
216
217
ref_out = torch._scaled_mm(
218
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b
@@ -222,11 +232,12 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
232
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
233
@pytest.mark.parametrize("dtype", DTYPES)
234
@pytest.mark.parametrize("seed", SEEDS)
235
+@pytest.mark.parametrize("padded", [False, True])
236
@pytest.mark.skipif(
237
not (current_platform.is_rocm() and current_platform.supports_fp8()),
238
reason="only test for rocm fp8",
239
)
229
-def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
240
+def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded):
241
torch.manual_seed(seed)
242
243
xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
@@ -236,6 +247,8 @@ def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
247
248
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
249
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
250
+ if padded:
251
+ B = pad_weights_fp8(B)
252
253
ref_out = torch._scaled_mm(
254
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
vllm/_custom_ops.py
-16
@@ -2027,35 +2027,20 @@ def selective_scan_fwd(
2027
)
2028
2029
2030
-# NOTE: The wvSplitK kernel (and all of the kernels in skinny_gemms.cu)
2031
-# are unable to properly handle non-contiguous
2032
-# tensors. It might be a good TODO(rasmith) to augment these kernels
2033
-# to be able to handle non-contiguous kernels for better performance.
2034
-def rocm_enforce_contiguous_skinny_gemm_inputs(
2035
- a: torch.Tensor, b: torch.Tensor
2036
-) -> tuple[torch.Tensor, torch.Tensor]:
2037
- a = a.contiguous() # no-op if already contiguous, else clone
2038
- b = b.contiguous() # no-op if already contiguous, else clone
2039
- return a, b
2040
-
2041
-
2030
# ROCm skinny gemms
2031
def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
2044
- a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2032
return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
2033
2034
2035
def wvSplitK(
2036
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2037
) -> torch.Tensor:
2051
- a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2038
return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
2039
2040
2041
def wvSplitKrc(
2042
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2043
) -> torch.Tensor:
2058
- a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2044
return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count)
2045
2046
@@ -2068,7 +2053,6 @@ def wvSplitKQ(
2053
cu_count: int,
2054
bias: torch.Tensor = None,
2055
) -> torch.Tensor:
2071
- a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2056
out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device)
2057
torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count)
2058
return out
vllm/model_executor/layers/quantization/kernels/scaled_mm/rocm.py
+1
@@ -28,6 +28,7 @@ def rocm_per_tensor_float_w8a8_scaled_mm_impl(
28
A.shape[0] == 1
29
and B.shape[1] % 16 == 0
30
and ((bias is None) or (bias.dtype == out_dtype))
31
+ and A.is_contiguous()
32
):
33
output = ops.wvSplitKQ(
34
B.t(),
vllm/model_executor/layers/utils.py
+2
@@ -165,6 +165,7 @@ def rocm_unquantized_gemm_impl(
165
and n <= 128
166
and k > 512
167
and math.ceil(k / 512) * math.ceil(m / 16) < get_cu_count()
168
+ and x.is_contiguous()
169
)
170
# k == 2880 and (m == 640 or m == 128))
171
)
@@ -179,6 +180,7 @@ def rocm_unquantized_gemm_impl(
180
and on_gfx9()
181
and x.dtype in [torch.float16, torch.bfloat16]
182
and k % 8 == 0
183
+ and x.is_contiguous()
184
)
185
186
if use_skinny is not True: