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 2, 23:15 UTC
[ROCm][CI] Fix ROCm LoRA Transformers fallback with full CUDA graphs (#41577)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 22, 2026 at 23:31 UTC
d28bdf9344aec8c393f8f1184243a6b159dbba27
4 files changed
+88
-19
tests/kernels/ir/test_layernorm.py
+51
@@ -141,6 +141,29 @@ def test_aiter_rejects_unsupported_dtypes():
141
assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"
142
143
144
+@pytest.mark.skipif(
145
+ not current_platform.is_rocm(),
146
+ reason="ROCm vllm_c RMSNorm needs explicit ND input handling",
147
+)
148
+def test_vllm_c_rms_norm_accepts_nd_input():
149
+ torch.set_default_device(current_platform.device_type)
150
+ impl = ir.ops.rms_norm.impls["vllm_c"]
151
+ if not impl.supported:
152
+ pytest.skip("vllm_c impl not supported on this platform")
153
+
154
+ base = torch.randn(3, 8, 192, dtype=torch.float16)
155
+ x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
156
+ assert not x.is_contiguous()
157
+ weight = torch.randn(16, dtype=torch.float16)
158
+ epsilon = 1e-5
159
+
160
+ output = impl.impl_fn(x, weight, epsilon)
161
+ ref_output = rms_norm_native(x, weight, epsilon)
162
+
163
+ assert output.shape == x.shape
164
+ assert_close(ir.ops.rms_norm, output, ref_output)
165
+
166
+
167
fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn
168
169
@@ -167,6 +190,34 @@ def test_fused_add_rms_norm_registration():
190
assert actual == expected
191
192
193
+@pytest.mark.skipif(
194
+ not current_platform.is_rocm(),
195
+ reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling",
196
+)
197
+def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
198
+ torch.set_default_device(current_platform.device_type)
199
+ impl = ir.ops.fused_add_rms_norm.impls["vllm_c"]
200
+ if not impl.supported:
201
+ pytest.skip("vllm_c impl not supported on this platform")
202
+
203
+ base = torch.randn(3, 8, 192, dtype=torch.float16)
204
+ residual_base = torch.randn(3, 8, 192, dtype=torch.float16)
205
+ x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
206
+ x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16)
207
+ assert not x.is_contiguous()
208
+ assert not x_residual.is_contiguous()
209
+ weight = torch.randn(16, dtype=torch.float16)
210
+ epsilon = 1e-5
211
+
212
+ output, residual = impl.impl_fn(x.clone(), x_residual.clone(), weight, epsilon)
213
+ ref_output, ref_residual = fused_add_rms_norm_native(x, x_residual, weight, epsilon)
214
+
215
+ assert output.shape == x.shape
216
+ assert residual.shape == x_residual.shape
217
+ assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
218
+ assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)
219
+
220
+
221
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
222
@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
223
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
vllm/kernels/vllm_c.py
+34
-6
@@ -10,10 +10,11 @@ current_platform.import_kernels()
10
11
CUDA_ALIKE = current_platform.is_cuda_alike()
12
"""Most kernels in this file are supported on all CUDA-alike platforms."""
13
+IS_ROCM = current_platform.is_rocm()
14
+"""ROCm needs shape normalization before calling some vLLM C kernels."""
15
14
-rms_no_var_size = (
15
- lambda x, weight, epsilon, variance_size=None: variance_size is None
16
- and (weight is None or weight.dtype == x.dtype)
16
+rms_no_var_size = lambda x, weight, epsilon, variance_size=None: (
17
+ variance_size is None and (weight is None or weight.dtype == x.dtype)
18
)
19
"""vLLM kernel requires no variance_size override and matching input/weight dtype."""
20
@@ -28,14 +29,23 @@ def rms_norm(
29
# Kernel requires weight tensor, pass ones
30
weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
31
assert variance_size is None
32
+ # ROCm's vLLM C RMSNorm kernel operates on contiguous 2D tensors.
33
+ # Higher-rank callers still normalize over the last dimension, so flatten
34
+ # all leading dims. reshape handles strided views from q/k/v splits.
35
+ if IS_ROCM and (x.dim() > 2 or not x.is_contiguous()):
36
+ original_shape = x.shape
37
+ x = x.reshape(-1, original_shape[-1])
38
+ output = torch.empty_like(x)
39
+ torch.ops._C.rms_norm(output, x, weight, epsilon)
40
+ return output.reshape(original_shape)
41
+
42
output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
43
torch.ops._C.rms_norm(output, x, weight, epsilon)
44
return output
45
46
36
-rms_add_no_var_size = (
37
- lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
38
- and (weight is None or weight.dtype == x.dtype)
47
+rms_add_no_var_size = lambda x, x_residual, weight, epsilon, variance_size=None: (
48
+ variance_size is None and (weight is None or weight.dtype == x.dtype)
49
)
50
"""vLLM Kernel does not support variance_size parameter and requires
51
matching input/weight dtype."""
@@ -59,5 +69,23 @@ def fused_add_rms_norm(
69
weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
70
71
assert variance_size is None
72
+ if IS_ROCM and (not x.is_contiguous() or not x_residual.is_contiguous()):
73
+ output, residual = ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
74
+ x, x_residual, weight, epsilon
75
+ )
76
+ x.copy_(output)
77
+ x_residual.copy_(residual)
78
+ return x, x_residual
79
+
80
+ # ROCm's vLLM C RMSNorm kernel operates on contiguous 2D tensors.
81
+ # Higher-rank callers still normalize over the last dimension, so flatten
82
+ # all leading dims.
83
+ if IS_ROCM and x.dim() > 2:
84
+ original_shape = x.shape
85
+ x = x.view(-1, original_shape[-1])
86
+ x_residual = x_residual.view(-1, original_shape[-1])
87
+ torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
88
+ return x.view(original_shape), x_residual.view(original_shape)
89
+
90
torch.ops._C.fused_add_rms_norm(x, x_residual, weight, epsilon)
91
return x, x_residual
vllm/model_executor/layers/utils.py
+1
@@ -150,6 +150,7 @@ def rocm_unquantized_gemm_impl(
150
envs.VLLM_ROCM_USE_SKINNY_GEMM
151
and on_gfx950()
152
and x.dtype in [torch.float16, torch.bfloat16]
153
+ and x.dim() == 2
154
and (
155
10 <= n <= 128
156
and k % 8 == 0
vllm/model_executor/models/transformers/base.py
+2
-13
@@ -16,7 +16,6 @@
16
# limitations under the License.
17
"""Transformers modeling backend base class."""
18
19
-import sys
19
from collections.abc import Callable, Iterable
20
from itertools import chain
21
from operator import attrgetter
@@ -268,21 +267,11 @@ class Base(
267
dynamic_arg_dims,
268
)
269
271
- @support_torch_compile(
270
+ support_torch_compile(
271
dynamic_arg_dims=dynamic_arg_dims,
272
enable_if=enable_if,
273
is_encoder=is_encoder,
275
- )
276
- class SupportTorchCompileWrapper(cls): ...
277
-
278
- # Preserve __module__ so transformers v5's source-file checks
279
- # (e.g. _can_set_experts_implementation) read the original
280
- # model's module instead of this file.
281
- SupportTorchCompileWrapper.__module__ = cls.__module__
282
-
283
- # Patch the class in its module
284
- module = sys.modules[cls.__module__]
285
- setattr(module, cls.__name__, SupportTorchCompileWrapper)
274
+ )(cls)
275
276
def _decorate_for_torch_compile(self, **kwargs: dict):
277
"""