[ROCm][Perf] Support N=5 in wvSplitK skinny GEMM kernels for speculative decoding (#40687)

Signed-off-by: Matthias Gehre <matthias.gehre@amd.com>

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) + "," +

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])