| 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() |