[ROCm] Cut 69 wasted contiguous copies per decode step from the skinny GEMM path (#58566)

Signed-off-by: lifulu <fululi12@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
fululi12
2026-09-24 22:00:58 -05:00
committed by GitHub
co-authored by Claude Opus 5 Nick Hill
parent f55419c43c
commit 4d79d5c638
+3 -1
View File
@@ -346,12 +346,14 @@ def rocm_unquantized_gemm_impl(
if use_skinny:
# The skinny kernels assume contiguous K elements. A shape-preserving
# reshape can retain a transposed activation's non-contiguous strides.
x_view = x.reshape(-1, x.size(-1)).contiguous()
# Note: Only build that view inside the branches that consume it.
if (m == 1 or m > 8) and 0 < n <= 5:
x_view = x.reshape(-1, x.size(-1)).contiguous()
cu_count = num_compute_units()
out = ops.wvSplitK(weight, x_view, cu_count, bias)
return out.reshape(*x.shape[:-1], weight.shape[0])
elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
x_view = x.reshape(-1, x.size(-1)).contiguous()
out = ops.LLMM1(weight, x_view, 4)
return out.reshape(*x.shape[:-1], weight.shape[0])