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 14, 23:45 UTC
[AMD][BugFix] Fix omission of wvSplitK kernel for small batch sizes (1-4) due to torch.compile (#21350)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Jul 28, 2025 at 14:38 UTC
b361f14e394861e31f61ab980d15dc07e5c74290
1 file changed
+28
-4
vllm/model_executor/layers/utils.py
+28
-4
@@ -8,6 +8,7 @@ import torch
8
from vllm import _custom_ops as ops
9
from vllm import envs
10
from vllm.platforms import current_platform
11
+from vllm.utils import direct_register_custom_op
12
13
14
def get_token_bin_counts_and_mask(
@@ -70,10 +71,10 @@ def default_unquantized_gemm(layer: torch.nn.Module,
71
return torch.nn.functional.linear(x, weight, bias)
72
73
73
-def rocm_unquantized_gemm(layer: torch.nn.Module,
74
- x: torch.Tensor,
75
- weight: torch.Tensor,
76
- bias: Optional[torch.Tensor] = None):
74
+def rocm_unquantized_gemm_impl(
75
+ x: torch.Tensor,
76
+ weight: torch.Tensor,
77
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
78
from vllm.platforms.rocm import on_gfx9
79
k = weight.shape[1]
80
use_skinny = (envs.VLLM_ROCM_USE_SKINNY_GEMM and on_gfx9() and \
@@ -97,6 +98,29 @@ def rocm_unquantized_gemm(layer: torch.nn.Module,
98
return torch.nn.functional.linear(x, weight, bias)
99
100
101
+def rocm_unquantized_gemm_impl_fake(
102
+ x: torch.Tensor,
103
+ weight: torch.Tensor,
104
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
105
+ return x.new_empty((*x.shape[:-1], weight.shape[0]))
106
+
107
+
108
+def rocm_unquantized_gemm(layer: torch.nn.Module,
109
+ x: torch.Tensor,
110
+ weight: torch.Tensor,
111
+ bias: Optional[torch.Tensor] = None) -> torch.Tensor:
112
+ return torch.ops.vllm.rocm_unquantized_gemm_impl(x, weight, bias)
113
+
114
+
115
+direct_register_custom_op(
116
+ op_name="rocm_unquantized_gemm_impl",
117
+ op_func=rocm_unquantized_gemm_impl,
118
+ mutates_args=[],
119
+ fake_impl=rocm_unquantized_gemm_impl_fake,
120
+ dispatch_key=current_platform.dispatch_key,
121
+)
122
+
123
+
124
def cpu_unquantized_gemm(layer: torch.nn.Module,
125
x: torch.Tensor,
126
weight: torch.Tensor,