[ROCm] Make vllm_c RMSNorm output contiguous (#49913)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 26, 2026 at 23:59 UTC
49f31d7cee425a6d38f8c5bc76877986daf832ed
2 files changed
+26
-1
tests/kernels/ir/test_layernorm.py
+23
@@ -176,6 +176,29 @@ def test_vllm_c_rms_norm_accepts_nd_input():
176
assert_close(ir.ops.rms_norm, output, ref_output)
177
178
179
+@pytest.mark.skipif(
180
+ not current_platform.is_rocm(),
181
+ reason="ROCm vllm_c RMSNorm needs a contiguous output for strided inputs",
182
+)
183
+def test_vllm_c_rms_norm_accepts_transposed_input():
184
+ impl = ir.ops.rms_norm.impls["vllm_c"]
185
+ if not impl.supported:
186
+ pytest.skip("vllm_c impl not supported on this platform")
187
+
188
+ x = torch.randn(
189
+ 1, 320, 120, dtype=torch.float16, device=current_platform.device_type
190
+ ).transpose(1, 2)
191
+ assert x.reshape(-1, x.shape[-1]).stride(-1) != 1
192
+ weight = torch.randn(320, dtype=torch.float16, device=current_platform.device_type)
193
+ epsilon = 1e-5
194
+
195
+ output = impl.impl_fn(x, weight, epsilon)
196
+ ref_output = rms_norm_native(x, weight, epsilon)
197
+
198
+ assert output.shape == x.shape
199
+ assert_close(ir.ops.rms_norm, output, ref_output)
200
+
201
+
202
fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn
203
204
vllm/kernels/vllm_c.py
+3
-1
@@ -32,7 +32,9 @@ def rms_norm(
32
if IS_ROCM and (x.dim() > 2 or not x.is_contiguous()):
33
original_shape = x.shape
34
x = x.reshape(-1, original_shape[-1])
35
- output = torch.empty_like(x)
35
+ # empty_like preserves the strides of transposed inputs, but the
36
+ # libtorch-stable kernel requires a contiguous output tensor.
37
+ output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
38
torch.ops._C.rms_norm(output, x, weight, epsilon)
39
return output.reshape(original_shape)
40