[ROCm][AITER] Skip AITER norm kernels when flattening to 2D would copy (#55991)
Signed-off-by: Zheng Gong <zgong@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
Zheng Gong committed
Sep 16, 2026 at 15:26 UTC
ab35354c21cc3c36439e79e18070f85cfafc3af2
2 files changed
+90
-2
tests/kernels/ir/test_aiter_norm_dispatch.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the AITER norm support predicates.
4
+
5
+The AITER norm wrappers flatten >2-D activations with ``Tensor.reshape``, which
6
+silently copies when the flattened shape is not expressible with the input's
7
+strides. The support predicates must reject those inputs so dispatch falls
8
+through to a provider that handles arbitrary strides.
9
+"""
10
+
11
+import pytest
12
+import torch
13
+
14
+from vllm._aiter_ops import is_aiter_found_and_supported
15
+from vllm.kernels.aiter_ops import flatten_to_2d_is_free
16
+
17
+pytestmark = pytest.mark.skipif(
18
+ not is_aiter_found_and_supported(),
19
+ reason="Only test on ROCm with AITER installed and supported",
20
+)
21
+
22
+
23
+def _qkv_slice_by_head(num_tokens, num_q_heads, num_kv_heads, head_dim):
24
+ """q viewed per-head from a fused QKV projection, as Qwen3-style QK-norm does."""
25
+ q_size, kv_size = num_q_heads * head_dim, num_kv_heads * head_dim
26
+ qkv = torch.empty(num_tokens, q_size + 2 * kv_size)
27
+ q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0]
28
+ return q.view(*q.shape[:-1], num_q_heads, head_dim)
29
+
30
+
31
+@pytest.mark.parametrize(
32
+ "x",
33
+ [
34
+ torch.empty(8, 16),
35
+ torch.empty(16),
36
+ torch.empty(2, 4, 16),
37
+ torch.empty(2, 1, 16),
38
+ # A contiguous tensor stays flattenable after a leading-dim slice.
39
+ torch.empty(8, 4, 16)[2:6],
40
+ ],
41
+)
42
+def test_flattenable(x):
43
+ assert flatten_to_2d_is_free(x)
44
+ assert x.reshape(-1, x.shape[-1]).data_ptr() == x.data_ptr()
45
+
46
+
47
+@pytest.mark.parametrize(
48
+ "x",
49
+ [
50
+ _qkv_slice_by_head(8, 32, 8, 128),
51
+ # Non-unit last-dim stride.
52
+ torch.empty(4, 8, 32).transpose(-1, -2),
53
+ # Leading dims cannot be merged: a slice along the middle dim.
54
+ torch.empty(4, 8, 32)[:, :4],
55
+ ],
56
+)
57
+def test_not_flattenable(x):
58
+ assert not flatten_to_2d_is_free(x)
59
+ # reshape has to copy, which is exactly what the predicate guards against.
60
+ assert x.reshape(-1, x.shape[-1]).data_ptr() != x.data_ptr()
vllm/kernels/aiter_ops.py
+30
-2
@@ -34,13 +34,38 @@ direct_register_aiter_op = functools.partial(
34
AITER_SUPPORTED = is_aiter_found()
35
"""Most kernels in this file are supported if AITER is installed."""
36
37
+
38
+def flatten_to_2d_is_free(x: Tensor) -> bool:
39
+ """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows.
40
+
41
+ The AITER norm kernels only take dense 2D inputs, so the wrappers below
42
+ flatten the leading dims with ``Tensor.reshape``, which silently falls back
43
+ to ``contiguous()`` when the flattened shape is not expressible with the
44
+ existing strides, adding a whole-tensor device copy that costs more than
45
+ the kernel saves. A strided last dim flattens without a copy but not into a
46
+ dense buffer, so it is rejected too.
47
+ """
48
+ if x.dim() <= 2:
49
+ return True
50
+ if x.stride(-1) != 1:
51
+ return False
52
+ expected_stride = x.size(-1)
53
+ for i in range(x.dim() - 2, -1, -1):
54
+ if x.size(i) != 1 and x.stride(i) != expected_stride:
55
+ return False
56
+ expected_stride *= x.size(i)
57
+ return True
58
+
59
+
60
rms_no_var_16bit_only = (
61
lambda x, weight, epsilon, variance_size=None: variance_size is None
62
and x.dtype in (torch.float16, torch.bfloat16)
63
and (weight is None or weight.dtype == x.dtype)
64
+ and flatten_to_2d_is_free(x)
65
)
66
"""AITER rms_norm only supports float16 and bfloat16 acts, no var_size override,
43
-and requires weight dtype to match x dtype."""
67
+requires weight dtype to match x dtype, and requires flattening the activation
68
+to 2D to be free."""
69
70
71
@ir.ops.rms_norm.register_impl(
@@ -79,10 +104,13 @@ rms_add_no_var_16bit_only = (
104
lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None
105
and x.dtype in (torch.float16, torch.bfloat16)
106
and (weight is None or weight.dtype == x.dtype)
107
+ and flatten_to_2d_is_free(x)
108
+ and flatten_to_2d_is_free(x_residual)
109
)
110
"""
111
AITER fused_add_rms_norm only supports 16-bit activations and no var_size override.
85
-Requires weight dtype to match x dtype.
112
+Requires weight dtype to match x dtype, and flattening both the activation and the
113
+residual to 2D to be free.
114
"""
115
116