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 13, 19:30 UTC
[ROCm] gemm_a16w16 upstreaming (#26969)
Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com>
Aleksandr Malyshev committed
Nov 4, 2025 at 13:01 UTC
2d977a7a9ead3179fde9ed55d69393ef7b6cec47
2 files changed
+43
-9
vllm/model_executor/layers/utils.py
+34
-8
@@ -103,12 +103,41 @@ def default_unquantized_gemm(
103
return torch.nn.functional.linear(x, weight, bias)
104
105
106
+def use_aiter_triton_gemm(n, m, k, dtype):
107
+ if (
108
+ envs.VLLM_ROCM_USE_AITER == 0
109
+ # MI300's - fp8nuz=True
110
+ or current_platform.is_fp8_fnuz()
111
+ or dtype not in [torch.float16, torch.bfloat16]
112
+ ):
113
+ return False
114
+
115
+ # use hipblaslt for the larger GEMMs
116
+ if n > 2048 and m > 512:
117
+ return False
118
+ return (
119
+ (m == 5120 and k == 2880)
120
+ or (m == 2880 and k == 4096)
121
+ or (m == 128 and k == 2880)
122
+ or (m == 640 and k == 2880)
123
+ or (m == 2880 and k == 512)
124
+ )
125
+
126
+
127
def rocm_unquantized_gemm_impl(
128
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
129
) -> torch.Tensor:
130
from vllm.platforms.rocm import on_gfx9
131
132
+ n = x.numel() / x.size(-1)
133
+ m = weight.shape[0]
134
k = weight.shape[1]
135
+
136
+ if use_aiter_triton_gemm(n, m, k, x.dtype):
137
+ from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
138
+
139
+ return gemm_a16w16(x, weight, bias)
140
+
141
use_skinny = (
142
envs.VLLM_ROCM_USE_SKINNY_GEMM
143
and on_gfx9()
@@ -120,11 +149,8 @@ def rocm_unquantized_gemm_impl(
149
return torch.nn.functional.linear(x, weight, bias)
150
151
x_view = x.reshape(-1, x.size(-1))
123
- n = x_view.shape[0]
124
- m = weight.shape[0]
125
- cu_count = current_platform.get_cu_count()
126
-
152
if m > 8 and 0 < n <= 4:
153
+ cu_count = current_platform.get_cu_count()
154
out = ops.wvSplitK(weight, x_view, cu_count, bias)
155
return out.reshape(*x.shape[:-1], weight.shape[0])
156
elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
@@ -133,7 +159,7 @@ def rocm_unquantized_gemm_impl(
159
return torch.nn.functional.linear(x, weight, bias)
160
161
136
-def rocm_unquantized_gemm_impl_fake(
162
+def rocm_unquantized_gemm_fake(
163
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
164
) -> torch.Tensor:
165
return x.new_empty((*x.shape[:-1], weight.shape[0]))
@@ -145,13 +171,13 @@ def rocm_unquantized_gemm(
171
weight: torch.Tensor,
172
bias: torch.Tensor | None = None,
173
) -> torch.Tensor:
148
- return torch.ops.vllm.rocm_unquantized_gemm_impl(x, weight, bias)
174
+ return torch.ops.vllm.rocm_unquantized_gemm(x, weight, bias)
175
176
177
direct_register_custom_op(
152
- op_name="rocm_unquantized_gemm_impl",
178
+ op_name="rocm_unquantized_gemm",
179
op_func=rocm_unquantized_gemm_impl,
154
- fake_impl=rocm_unquantized_gemm_impl_fake,
180
+ fake_impl=rocm_unquantized_gemm_fake,
181
)
182
183
vllm/model_executor/models/gpt_oss.py
+9
-1
@@ -25,12 +25,14 @@ from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLine
25
from vllm.model_executor.layers.logits_processor import LogitsProcessor
26
from vllm.model_executor.layers.quantization import QuantizationConfig
27
from vllm.model_executor.layers.rotary_embedding import get_rope
28
+from vllm.model_executor.layers.utils import rocm_unquantized_gemm
29
from vllm.model_executor.layers.vocab_parallel_embedding import (
30
ParallelLMHead,
31
VocabParallelEmbedding,
32
)
33
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
34
from vllm.model_executor.models.utils import sequence_parallel_chunk
35
+from vllm.platforms import current_platform
36
from vllm.sequence import IntermediateTensors
37
from vllm.utils.math_utils import cdiv
38
@@ -153,6 +155,7 @@ class MLPBlock(torch.nn.Module):
155
156
self.layer_idx = layer_idx
157
self.num_experts = config.num_local_experts
158
+ self.hidden_size = config.hidden_size
159
self.experts_per_token = config.num_experts_per_tok
160
self.world_size = dist.get_world_size() if dist.is_initialized() else 1
161
self.router = torch.nn.Linear(config.hidden_size, config.num_local_experts)
@@ -177,7 +180,12 @@ class MLPBlock(torch.nn.Module):
180
if self.is_sequence_parallel:
181
x = sequence_parallel_chunk(x)
182
180
- g = self.router(x)
183
+ if current_platform.is_rocm():
184
+ g = rocm_unquantized_gemm(
185
+ self, x[:, : self.hidden_size], self.router.weight, self.router.bias
186
+ )
187
+ else:
188
+ g = self.router(x)
189
x = self.experts(hidden_states=x, router_logits=g)
190
191
if self.is_sequence_parallel: