[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:
co-authored by
Claude Opus 5
Nick Hill
parent
f55419c43c
commit
4d79d5c638
@@ -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])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user