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 9, 17:06 UTC
[CI][AMD][BugFix] Update wvSplitK (and other skinny_gemm wrappers) to ensure tensors passed will be made contiguous for the kernel (#32831)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Jan 23, 2026 at 15:35 UTC
6cc6d92be5e64f449e31e06d43c80955f646809e
2 files changed
+23
csrc/rocm/skinny_gemms.cu
+7
@@ -13,6 +13,13 @@
13
#include "dispatch_utils.h"
14
#include "quantization/w8a8/fp8/common.cuh"
15
16
+// TODO(rasmith): The kernels in this file are susceptible to integer overflow
17
+// issues, do not take strides, and are unable to handle PyTorch tensors that
18
+// return is_contiguous() as False (the tensors may actually be contiguous
19
+// in memory).
20
+//
21
+// However, it may be possible to fix these kernels to handle both issues.
22
+
23
#if defined(__HIPCC__) && \
24
(defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
25
#define __HIP__GFX9__
vllm/_custom_ops.py
+16
@@ -2027,20 +2027,35 @@ def selective_scan_fwd(
2027
)
2028
2029
2030
+# NOTE: The wvSplitK kernel (and all of the kernels in skinny_gemms.cu)
2031
+# are unable to properly handle non-contiguous
2032
+# tensors. It might be a good TODO(rasmith) to augment these kernels
2033
+# to be able to handle non-contiguous kernels for better performance.
2034
+def rocm_enforce_contiguous_skinny_gemm_inputs(
2035
+ a: torch.Tensor, b: torch.Tensor
2036
+) -> tuple[torch.Tensor, torch.Tensor]:
2037
+ a = a.contiguous() # no-op if already contiguous, else clone
2038
+ b = b.contiguous() # no-op if already contiguous, else clone
2039
+ return a, b
2040
+
2041
+
2042
# ROCm skinny gemms
2043
def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
2044
+ a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2045
return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
2046
2047
2048
def wvSplitK(
2049
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2050
) -> torch.Tensor:
2051
+ a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2052
return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
2053
2054
2055
def wvSplitKrc(
2056
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
2057
) -> torch.Tensor:
2058
+ a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2059
return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count)
2060
2061
@@ -2053,6 +2068,7 @@ def wvSplitKQ(
2068
cu_count: int,
2069
bias: torch.Tensor = None,
2070
) -> torch.Tensor:
2071
+ a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
2072
out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device)
2073
torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count)
2074
return out