From a9ec46d4b7db6ae1d0cb0f27ce1a7ddb0f749ef8 Mon Sep 17 00:00:00 2001 From: Matthias Gehre Date: Thu, 28 May 2026 18:28:21 +0200 Subject: [PATCH] [ROCm][Perf] Support N=5 in wvSplitK skinny GEMM kernels for speculative decoding (#40687) Signed-off-by: Matthias Gehre --- csrc/rocm/skinny_gemms.cu | 6 ++++++ vllm/model_executor/layers/utils.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/csrc/rocm/skinny_gemms.cu b/csrc/rocm/skinny_gemms.cu index 3342db37be9..10e3cbf2e0b 100644 --- a/csrc/rocm/skinny_gemms.cu +++ b/csrc/rocm/skinny_gemms.cu @@ -1277,6 +1277,12 @@ torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, else WVSPLIT_TILE_CFG(64, 16, sYT, 4) break; + case 5: + if (use_wave32) + WVSPLIT_TILE_CFG(32, 16, sYT, 5) + else + WVSPLIT_TILE_CFG(64, 16, sYT, 5) + break; default: throw std::runtime_error( "Unsupported N value: " + std::to_string(M_in) + "," + diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index dbc05273f30..aa40020052c 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -178,7 +178,7 @@ def rocm_unquantized_gemm_impl( if use_skinny: x_view = x.reshape(-1, x.size(-1)) - if m > 8 and 0 < n <= 4: + if m > 8 and 0 < n <= 5: cu_count = num_compute_units() out = ops.wvSplitK(weight, x_view, cu_count, bias) return out.reshape(*x.shape[:-1], weight.shape[0])