| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from types import SimpleNamespace |
| 6 | from typing import Any |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.layers.layernorm import GemmaRMSNorm |
| 12 | from vllm.models.qwen4_exp.amd import ( |
| 13 | model as _qwen4_exp_model, # noqa: F401 |
| 14 | ) |
| 15 | from vllm.models.qwen4_exp.amd import ple_layer as ple_layer_module |
| 16 | from vllm.models.qwen4_exp.amd.indexer_qsa import ( |
| 17 | apply_qsa_rmsnorm, |
| 18 | apply_qsa_rope, |
| 19 | ) |
| 20 | from vllm.models.qwen4_exp.amd.ops import qsa as qsa_ops |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.triton_utils import HAS_TRITON |
| 23 | |
| 24 | pytestmark = pytest.mark.skipif( |
| 25 | not current_platform.is_rocm(), |
| 26 | reason="AMD QSA requires ROCm", |
| 27 | ) |
| 28 | |
| 29 | requires_qsa_kernels = pytest.mark.skipif( |
| 30 | not HAS_TRITON, |
| 31 | reason="AMD QSA kernels require Triton", |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | def test_ple_ngram_embedding_custom_op_uses_resident_weight( |
| 36 | monkeypatch: pytest.MonkeyPatch, |
| 37 | ) -> None: |
| 38 | layer_name = "model.layers.0.ple" |
| 39 | layer = ple_layer_module.Qwen4ExpPLELayer.__new__(ple_layer_module.Qwen4ExpPLELayer) |
| 40 | torch.nn.Module.__init__(layer) |
| 41 | layer.ple_embedding = torch.nn.Module() |
| 42 | layer.ple_embedding.ngram_embedding = torch.nn.Embedding(8, 3) |
| 43 | context = SimpleNamespace(no_compile_layers={layer_name: layer}) |
| 44 | monkeypatch.setattr(ple_layer_module, "get_forward_context", lambda: context) |
| 45 | |
| 46 | ngram_ids = torch.tensor([[0, 1], [2, 3]]) |
| 47 | output = torch.empty(2, 6) |
| 48 | ple_layer_module.qwen4_exp_amd_ple_ngram_embedding( |
| 49 | ngram_ids, |
| 50 | output, |
| 51 | layer_name, |
| 52 | ) |
| 53 | |
| 54 | expected = layer.ple_embedding.ngram_embedding(ngram_ids).flatten(-2) |
| 55 | torch.testing.assert_close(output, expected) |
| 56 | |
| 57 | |
| 58 | def _qsa_sparse_paged_attention_reference( |
| 59 | q: torch.Tensor, |
| 60 | k_cache: torch.Tensor, |
| 61 | v_cache: torch.Tensor, |
| 62 | logical_indices: torch.Tensor, |
| 63 | block_table: torch.Tensor, |
| 64 | token_to_req: torch.Tensor, |
| 65 | softmax_scale: float, |
| 66 | ) -> torch.Tensor: |
| 67 | output = torch.zeros_like(q) |
| 68 | repeats = q.shape[1] // k_cache.shape[2] |
| 69 | page_size = k_cache.shape[1] |
| 70 | for row in range(q.shape[0]): |
| 71 | logical = logical_indices[row] |
| 72 | logical = logical[logical >= 0].long() |
| 73 | if not logical.numel(): |
| 74 | continue |
| 75 | request = token_to_req[row].long() |
| 76 | pages = block_table[request, logical // page_size].long() |
| 77 | offsets = logical % page_size |
| 78 | keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1) |
| 79 | values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1) |
| 80 | scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float()) |
| 81 | probabilities = torch.softmax(scores * softmax_scale, dim=-1) |
| 82 | output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to( |
| 83 | q.dtype |
| 84 | ) |
| 85 | return output |
| 86 | |
| 87 | |
| 88 | def test_qsa_rope_uses_platform_dispatch() -> None: |
| 89 | tensor = torch.arange(16, dtype=torch.float32).reshape(2, 2, 4) |
| 90 | positions = torch.tensor([0, 1]) |
| 91 | calls = [] |
| 92 | |
| 93 | def apply_rotary_emb( |
| 94 | rotary_input: torch.Tensor, |
| 95 | cos: torch.Tensor, |
| 96 | sin: torch.Tensor, |
| 97 | ) -> torch.Tensor: |
| 98 | calls.append((rotary_input, cos, sin)) |
| 99 | return rotary_input + 1 |
| 100 | |
| 101 | rotary_emb = SimpleNamespace( |
| 102 | rotary_dim=2, |
| 103 | apply_rotary_emb=apply_rotary_emb, |
| 104 | _match_cos_sin_cache_dtype=lambda _: torch.zeros(2, 4), |
| 105 | ) |
| 106 | |
| 107 | output = apply_qsa_rope(rotary_emb, positions, tensor) |
| 108 | |
| 109 | assert len(calls) == 1 |
| 110 | torch.testing.assert_close(output[..., :2], tensor[..., :2] + 1) |
| 111 | torch.testing.assert_close(output[..., 2:], tensor[..., 2:]) |
| 112 | |
| 113 | |
| 114 | def test_qsa_rmsnorm_uses_portable_implementation(default_vllm_config) -> None: |
| 115 | norm = GemmaRMSNorm(4, eps=1e-6) |
| 116 | norm.weight.data.copy_(torch.tensor([0.1, -0.2, 0.3, -0.4])) |
| 117 | tensor = torch.arange(8, dtype=torch.float32).reshape(2, 4) |
| 118 | output = apply_qsa_rmsnorm(norm, tensor) |
| 119 | |
| 120 | torch.testing.assert_close(output, norm.forward_native(tensor)) |
| 121 | |
| 122 | |
| 123 | def test_qsa_selection_uses_portable_topk_on_rocm( |
| 124 | monkeypatch: pytest.MonkeyPatch, |
| 125 | ) -> None: |
| 126 | rows = 2 |
| 127 | token_topk = 8 |
| 128 | compress_ratio = 4 |
| 129 | block_topk = token_topk // compress_ratio |
| 130 | blocks = torch.empty((rows, block_topk), dtype=torch.int32) |
| 131 | visible_blocks = torch.tensor([4, 6], dtype=torch.int32) |
| 132 | logits = torch.empty((rows, 8), dtype=torch.float32) |
| 133 | selection_output = torch.empty( |
| 134 | (rows, token_topk + compress_ratio - 1), dtype=torch.int32 |
| 135 | ) |
| 136 | topk_call: dict[str, Any] = {} |
| 137 | |
| 138 | monkeypatch.setattr( |
| 139 | qsa_ops, |
| 140 | "qsa_mqa_paged", |
| 141 | lambda *args, **kwargs: (logits, visible_blocks), |
| 142 | ) |
| 143 | monkeypatch.setattr(qsa_ops.current_platform, "is_cuda", lambda: False) |
| 144 | |
| 145 | def top_k_per_row_decode( |
| 146 | topk_logits, |
| 147 | next_n, |
| 148 | seq_lens, |
| 149 | raw_topk_indices, |
| 150 | num_rows, |
| 151 | stride0, |
| 152 | stride1, |
| 153 | topk_tokens, |
| 154 | ) -> None: |
| 155 | topk_call.update( |
| 156 | logits=topk_logits, |
| 157 | next_n=next_n, |
| 158 | seq_lens=seq_lens, |
| 159 | raw_topk_indices=raw_topk_indices, |
| 160 | num_rows=num_rows, |
| 161 | strides=(stride0, stride1), |
| 162 | topk_tokens=topk_tokens, |
| 163 | ) |
| 164 | raw_topk_indices.zero_() |
| 165 | |
| 166 | def expand_qsa_block_indices(*args) -> None: |
| 167 | args[-1].fill_(-1) |
| 168 | |
| 169 | monkeypatch.setattr(qsa_ops.ops, "top_k_per_row_decode", top_k_per_row_decode) |
| 170 | monkeypatch.setattr( |
| 171 | qsa_ops, |
| 172 | "expand_qsa_block_indices_cuda", |
| 173 | expand_qsa_block_indices, |
| 174 | ) |
| 175 | |
| 176 | output = qsa_ops.qsa_select_paged_tokens( |
| 177 | torch.empty((rows, 1, 1)), |
| 178 | torch.empty((1, 4, 1, 1)), |
| 179 | torch.empty((1, 2), dtype=torch.int32), |
| 180 | torch.zeros(rows, dtype=torch.int32), |
| 181 | torch.arange(rows, dtype=torch.int32), |
| 182 | torch.tensor([8], dtype=torch.int32), |
| 183 | token_topk, |
| 184 | compress_ratio, |
| 185 | selection_output, |
| 186 | ) |
| 187 | |
| 188 | assert output is selection_output |
| 189 | assert torch.all(output == -1) |
| 190 | assert topk_call["logits"] is logits |
| 191 | assert topk_call["next_n"] == 1 |
| 192 | assert topk_call["seq_lens"] is visible_blocks |
| 193 | assert topk_call["raw_topk_indices"].shape == blocks.shape |
| 194 | assert topk_call["num_rows"] == rows |
| 195 | assert topk_call["strides"] == (logits.stride(0), logits.stride(1)) |
| 196 | assert topk_call["topk_tokens"] == block_topk |
| 197 | |
| 198 | |
| 199 | @requires_qsa_kernels |
| 200 | @pytest.mark.parametrize( |
| 201 | ("num_rows", "num_query_heads", "num_kv_heads", "page_size"), |
| 202 | [ |
| 203 | pytest.param(1, 24, 2, 1792, id="tp1_split64"), |
| 204 | pytest.param(16, 12, 1, 1792, id="tp2_split32"), |
| 205 | pytest.param(32, 6, 1, 1024, id="tp4_split8"), |
| 206 | pytest.param(257, 6, 1, 1024, id="tp4_split4"), |
| 207 | pytest.param(513, 6, 1, 1024, id="tp4_split1"), |
| 208 | ], |
| 209 | ) |
| 210 | def test_qsa_sparse_paged_attention_matches_reference( |
| 211 | num_rows: int, |
| 212 | num_query_heads: int, |
| 213 | num_kv_heads: int, |
| 214 | page_size: int, |
| 215 | ) -> None: |
| 216 | torch.manual_seed(2) |
| 217 | head_dim = 256 |
| 218 | num_requests = 2 |
| 219 | num_selected_pages = 64 |
| 220 | num_pages_per_request = num_selected_pages + 1 |
| 221 | num_cache_blocks = num_requests * num_pages_per_request |
| 222 | indexer_budget = 2048 |
| 223 | indexer_compress_ratio = 4 |
| 224 | selection_width = indexer_budget + indexer_compress_ratio - 1 |
| 225 | q = torch.randn( |
| 226 | num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 |
| 227 | ) |
| 228 | kv_cache = torch.randn( |
| 229 | num_cache_blocks, |
| 230 | page_size, |
| 231 | num_kv_heads, |
| 232 | 2 * head_dim, |
| 233 | device="cuda", |
| 234 | dtype=torch.bfloat16, |
| 235 | ) |
| 236 | k_cache, v_cache = kv_cache.split(head_dim, dim=-1) |
| 237 | block_table = ( |
| 238 | torch.randperm(num_cache_blocks, device="cuda") |
| 239 | .reshape(num_requests, num_pages_per_request) |
| 240 | .to(torch.int32) |
| 241 | ) |
| 242 | rows_per_request = math.ceil(num_rows / num_requests) |
| 243 | row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32) |
| 244 | token_to_req = row_indices // rows_per_request |
| 245 | request_row_counts = torch.tensor( |
| 246 | [rows_per_request, num_rows - rows_per_request], |
| 247 | device="cuda", |
| 248 | dtype=torch.int32, |
| 249 | ) |
| 250 | |
| 251 | context_length = num_pages_per_request * page_size - 1 |
| 252 | block_topk = indexer_budget // indexer_compress_ratio |
| 253 | compressed_blocks_per_page = page_size // indexer_compress_ratio |
| 254 | selection = torch.arange(block_topk, device="cuda") |
| 255 | selected_pages = selection % num_selected_pages |
| 256 | selected_offsets = selection // num_selected_pages |
| 257 | row_shifts = 2 * row_indices.unsqueeze(1) |
| 258 | selected_offsets = (selected_offsets + row_shifts) % compressed_blocks_per_page |
| 259 | block_indices = (selected_pages * compressed_blocks_per_page + selected_offsets).to( |
| 260 | torch.int32 |
| 261 | ) |
| 262 | rows_within_request = row_indices % rows_per_request |
| 263 | query_positions = ( |
| 264 | context_length - request_row_counts[token_to_req.long()] + rows_within_request |
| 265 | ).to(torch.int64) |
| 266 | sequence_lengths = torch.full( |
| 267 | (num_requests,), context_length, device="cuda", dtype=torch.int32 |
| 268 | ) |
| 269 | logical_indices = qsa_ops.expand_qsa_block_indices_cuda( |
| 270 | block_indices, |
| 271 | query_positions, |
| 272 | sequence_lengths, |
| 273 | token_to_req, |
| 274 | indexer_compress_ratio, |
| 275 | indexer_budget, |
| 276 | ) |
| 277 | assert logical_indices.shape == (num_rows, selection_width) |
| 278 | |
| 279 | actual = qsa_ops.qsa_sparse_paged_attention( |
| 280 | q, |
| 281 | k_cache, |
| 282 | v_cache, |
| 283 | logical_indices, |
| 284 | block_table, |
| 285 | token_to_req, |
| 286 | ) |
| 287 | expected = _qsa_sparse_paged_attention_reference( |
| 288 | q, |
| 289 | k_cache, |
| 290 | v_cache, |
| 291 | logical_indices, |
| 292 | block_table, |
| 293 | token_to_req, |
| 294 | q.shape[-1] ** -0.5, |
| 295 | ) |
| 296 | |
| 297 | torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) |