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 2, 13:30 UTC
[ROCm][Perf] Support N=5 in wvSplitK skinny GEMM kernels for speculative decoding (#40687)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>
Matthias Gehre committed
May 28, 2026 at 18:28 UTC
a9ec46d4b7db6ae1d0cb0f27ce1a7ddb0f749ef8
2 files changed
+7
-1
csrc/rocm/skinny_gemms.cu
+6
@@ -1277,6 +1277,12 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b,
1277
else
1278
WVSPLIT_TILE_CFG(64, 16, sYT, 4)
1279
break;
1280
+ case 5:
1281
+ if (use_wave32)
1282
+ WVSPLIT_TILE_CFG(32, 16, sYT, 5)
1283
+ else
1284
+ WVSPLIT_TILE_CFG(64, 16, sYT, 5)
1285
+ break;
1286
default:
1287
throw std::runtime_error(
1288
"Unsupported N value: " + std::to_string(M_in) + "," +
vllm/model_executor/layers/utils.py
+1
-1
@@ -178,7 +178,7 @@ def rocm_unquantized_gemm_impl(
178
179
if use_skinny:
180
x_view = x.reshape(-1, x.size(-1))
181
- if m > 8 and 0 < n <= 4:
181
+ if m > 8 and 0 < n <= 5:
182
cu_count = num_compute_units()
183
out = ops.wvSplitK(weight, x_view, cu_count, bias)
184
return out.reshape(*x.shape[:-1], weight.shape[0])