| 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 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.models.qwen4_exp.common import qsa_cache |
| 11 | from vllm.models.qwen4_exp.common.qsa_cache import QSAMetadataBuilder |
| 12 | from vllm.models.qwen4_exp.nvidia import indexer_qsa |
| 13 | from vllm.models.qwen4_exp.nvidia import ( |
| 14 | model as _qwen4_exp_model, # noqa: F401 |
| 15 | ) |
| 16 | from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops |
| 17 | from vllm.models.qwen4_exp.nvidia.ops import qsa_indexer as qsa_indexer_ops |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.triton_utils import HAS_TRITON |
| 20 | from vllm.v1.worker.utils import clear_layer_kv_caches |
| 21 | |
| 22 | requires_qsa_kernels = pytest.mark.skipif( |
| 23 | not current_platform.is_cuda() or not HAS_TRITON, |
| 24 | reason="QSA kernels require CUDA and Triton", |
| 25 | ) |
| 26 | |
| 27 | |
| 28 | def test_qsa_mtp_index_share_updates_cache_but_skips_selection( |
| 29 | monkeypatch: pytest.MonkeyPatch, |
| 30 | ) -> None: |
| 31 | rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) |
| 32 | raw_metadata = SimpleNamespace( |
| 33 | num_actual_tokens=2, |
| 34 | slot_mapping=torch.arange(2), |
| 35 | block_table=torch.empty(0), |
| 36 | query_start_loc=torch.arange(3), |
| 37 | logical_positions=torch.arange(2), |
| 38 | ) |
| 39 | compressed_metadata = SimpleNamespace( |
| 40 | num_actual_tokens=2, |
| 41 | slot_mapping=torch.arange(2), |
| 42 | k_work_metadata=torch.empty(0), |
| 43 | ) |
| 44 | updates = [] |
| 45 | selections = [] |
| 46 | indexer = SimpleNamespace( |
| 47 | skip_topk=True, |
| 48 | _metadata=lambda: (raw_metadata, compressed_metadata), |
| 49 | index_n_heads=1, |
| 50 | index_kv_heads=1, |
| 51 | index_head_dim=1, |
| 52 | indexer_dtype=torch.bfloat16, |
| 53 | raw_key_cache=SimpleNamespace( |
| 54 | kv_cache=torch.empty(0), |
| 55 | rope_position_cache=None, |
| 56 | rope_position_offset=0, |
| 57 | ), |
| 58 | compressed_key_cache=SimpleNamespace(kv_cache=torch.empty(0)), |
| 59 | rotary_emb=SimpleNamespace(cos_sin_cache=torch.empty(0)), |
| 60 | q_layernorm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6), |
| 61 | k_layernorm=SimpleNamespace(weight=torch.ones(1)), |
| 62 | compress_ratio=2, |
| 63 | ) |
| 64 | attn = SimpleNamespace( |
| 65 | use_fused_qsa_prepare=True, |
| 66 | kv_cache=torch.empty(0, 1, 1, 2), |
| 67 | kv_cache_dtype="auto", |
| 68 | q_norm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6), |
| 69 | k_norm=SimpleNamespace(weight=torch.ones(1)), |
| 70 | _k_scale_float=1.0, |
| 71 | _v_scale_float=1.0, |
| 72 | ) |
| 73 | |
| 74 | monkeypatch.setattr( |
| 75 | indexer_qsa, |
| 76 | "qsa_prepare", |
| 77 | lambda *args, **kwargs: updates.append((args, kwargs)), |
| 78 | ) |
| 79 | monkeypatch.setattr( |
| 80 | qsa_indexer_ops, |
| 81 | "qsa_select_paged_decode", |
| 82 | lambda *args, **kwargs: selections.append((args, kwargs)), |
| 83 | ) |
| 84 | monkeypatch.setattr( |
| 85 | qsa_indexer_ops, |
| 86 | "qsa_select_paged_prefill", |
| 87 | lambda *args, **kwargs: selections.append((args, kwargs)), |
| 88 | ) |
| 89 | |
| 90 | actual, _ = indexer_qsa.QSAIndexer.forward( |
| 91 | indexer, |
| 92 | torch.zeros(2, 2), |
| 93 | torch.tensor([7, 8]), |
| 94 | rows, |
| 95 | attn=attn, |
| 96 | qkv=torch.zeros(2, 4), |
| 97 | slot_mapping=torch.arange(2), |
| 98 | ) |
| 99 | |
| 100 | assert actual is rows |
| 101 | assert len(updates) == 1 |
| 102 | assert not selections |
| 103 | |
| 104 | |
| 105 | def _qsa_mqa_paged_reference( |
| 106 | q: torch.Tensor, |
| 107 | k_cache: torch.Tensor, |
| 108 | page_table: torch.Tensor, |
| 109 | token_to_req: torch.Tensor, |
| 110 | visible_lengths: torch.Tensor, |
| 111 | ) -> torch.Tensor: |
| 112 | pages = page_table.index_select(0, token_to_req.long()).long() |
| 113 | keys = k_cache[pages, :, 0, :].flatten(1, 2) |
| 114 | scores = torch.einsum("rhd,rnd->rnh", q.float(), keys.float()) |
| 115 | logits = torch.relu(scores).sum(dim=-1) / math.sqrt(q.shape[-1]) |
| 116 | positions = torch.arange(keys.shape[1], device=q.device).unsqueeze(0) |
| 117 | return logits.masked_fill(positions >= visible_lengths.unsqueeze(1), -torch.inf) |
| 118 | |
| 119 | |
| 120 | def _qsa_relative_topk_reference( |
| 121 | logits: torch.Tensor, |
| 122 | row_starts: torch.Tensor, |
| 123 | row_ends: torch.Tensor, |
| 124 | topk: int, |
| 125 | ) -> torch.Tensor: |
| 126 | output = torch.full( |
| 127 | (logits.shape[0], topk), -1, dtype=torch.int32, device=logits.device |
| 128 | ) |
| 129 | for row in range(logits.shape[0]): |
| 130 | start = int(row_starts[row].item()) |
| 131 | length = int((row_ends[row] - row_starts[row]).item()) |
| 132 | width = min(length, topk) |
| 133 | if width: |
| 134 | output[row, :width] = torch.topk( |
| 135 | logits[row, start : start + length], width |
| 136 | ).indices.to(torch.int32) |
| 137 | return output |
| 138 | |
| 139 | |
| 140 | def _expand_qsa_indices_reference( |
| 141 | block_indices: torch.Tensor, |
| 142 | query_positions: torch.Tensor, |
| 143 | sequence_lengths: torch.Tensor, |
| 144 | compress_ratio: int, |
| 145 | token_topk: int, |
| 146 | ) -> torch.Tensor: |
| 147 | rows = block_indices.shape[0] |
| 148 | block_topk = token_topk // compress_ratio |
| 149 | output_width = token_topk + compress_ratio - 1 |
| 150 | offsets = torch.arange(compress_ratio, device=block_indices.device) |
| 151 | blocks = block_indices.long() |
| 152 | expanded = blocks.unsqueeze(-1) * compress_ratio + offsets |
| 153 | expanded = torch.where( |
| 154 | blocks.unsqueeze(-1) >= 0, expanded, torch.full_like(expanded, -1) |
| 155 | ).reshape(rows, block_topk * compress_ratio) |
| 156 | expanded = expanded[:, :token_topk] |
| 157 | expanded = torch.where( |
| 158 | (expanded >= 0) & (expanded < sequence_lengths.unsqueeze(1)), |
| 159 | expanded, |
| 160 | torch.full_like(expanded, -1), |
| 161 | ) |
| 162 | |
| 163 | tail_offsets = torch.arange(compress_ratio - 1, device=block_indices.device) |
| 164 | visible_tokens = query_positions + 1 |
| 165 | tail_start = visible_tokens // compress_ratio * compress_ratio |
| 166 | tail = tail_start.unsqueeze(1) + tail_offsets.unsqueeze(0) |
| 167 | tail_count = (visible_tokens - tail_start).unsqueeze(1) |
| 168 | tail_valid = (tail_offsets.unsqueeze(0) < tail_count) & ( |
| 169 | tail < sequence_lengths.unsqueeze(1) |
| 170 | ) |
| 171 | tail = torch.where(tail_valid, tail, torch.full_like(tail, -1)) |
| 172 | |
| 173 | result = torch.cat((expanded, tail), dim=1) |
| 174 | order = torch.arange(output_width, device=result.device).expand(rows, -1) |
| 175 | sort_key = torch.where(result >= 0, order, order + output_width) |
| 176 | return result.gather(1, torch.argsort(sort_key, dim=1, stable=True)).to(torch.int32) |
| 177 | |
| 178 | |
| 179 | def _qsa_select_paged_reference( |
| 180 | q: torch.Tensor, |
| 181 | k_cache: torch.Tensor, |
| 182 | page_table: torch.Tensor, |
| 183 | token_to_req: torch.Tensor, |
| 184 | query_positions: torch.Tensor, |
| 185 | sequence_lengths: torch.Tensor, |
| 186 | token_topk: int, |
| 187 | compress_ratio: int, |
| 188 | ) -> torch.Tensor: |
| 189 | row_sequence_lengths = sequence_lengths.index_select(0, token_to_req.long()) |
| 190 | visible_blocks = torch.minimum( |
| 191 | (query_positions + 1) // compress_ratio, |
| 192 | row_sequence_lengths // compress_ratio, |
| 193 | ).to(torch.int32) |
| 194 | logits = _qsa_mqa_paged_reference( |
| 195 | q, |
| 196 | k_cache, |
| 197 | page_table, |
| 198 | token_to_req, |
| 199 | visible_blocks, |
| 200 | ) |
| 201 | starts = torch.zeros_like(visible_blocks) |
| 202 | return _qsa_relative_topk_reference( |
| 203 | logits, |
| 204 | starts, |
| 205 | visible_blocks, |
| 206 | token_topk // compress_ratio, |
| 207 | ) |
| 208 | |
| 209 | |
| 210 | def _qsa_sparse_paged_attention_reference( |
| 211 | q: torch.Tensor, |
| 212 | k_cache: torch.Tensor, |
| 213 | v_cache: torch.Tensor, |
| 214 | logical_indices: torch.Tensor, |
| 215 | block_table: torch.Tensor, |
| 216 | token_to_req: torch.Tensor, |
| 217 | softmax_scale: float, |
| 218 | k_scale: float = 1.0, |
| 219 | v_scale: float = 1.0, |
| 220 | ) -> torch.Tensor: |
| 221 | """Dense reference for QSA sparse paged attention. |
| 222 | |
| 223 | Mirrors the kernel's dequant: fp8-e4m3 K/V caches are dequantized with the |
| 224 | per-tensor k_scale/v_scale host floats; bf16 caches use unit scales. |
| 225 | """ |
| 226 | output = torch.zeros_like(q) |
| 227 | repeats = q.shape[1] // k_cache.shape[2] |
| 228 | page_size = k_cache.shape[1] |
| 229 | for row in range(q.shape[0]): |
| 230 | logical = logical_indices[row] |
| 231 | logical = logical[logical >= 0].long() |
| 232 | if not logical.numel(): |
| 233 | continue |
| 234 | request = token_to_req[row].long() |
| 235 | pages = block_table[request, logical // page_size].long() |
| 236 | offsets = logical % page_size |
| 237 | keys = (k_cache[pages, offsets].float() * k_scale).repeat_interleave( |
| 238 | repeats, dim=1 |
| 239 | ) |
| 240 | values = (v_cache[pages, offsets].float() * v_scale).repeat_interleave( |
| 241 | repeats, dim=1 |
| 242 | ) |
| 243 | scores = torch.einsum("hd,khd->hk", q[row].float(), keys) |
| 244 | probabilities = torch.softmax(scores * softmax_scale, dim=-1) |
| 245 | output[row] = torch.einsum("hk,khd->hd", probabilities, values).to(q.dtype) |
| 246 | return output |
| 247 | |
| 248 | |
| 249 | @requires_qsa_kernels |
| 250 | def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None: |
| 251 | device = torch.device("cuda") |
| 252 | builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder) |
| 253 | builder.compress_ratio = 1 |
| 254 | builder.reorder_batch_threshold = 4 |
| 255 | builder.is_circular_buffer = False |
| 256 | builder.storage_block_size = 64 |
| 257 | builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device) |
| 258 | builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device) |
| 259 | builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device) |
| 260 | builder.visible_blocks_buffer = torch.empty(16, dtype=torch.int32, device=device) |
| 261 | builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device) |
| 262 | query_start_loc = torch.tensor([0, 4, 8, 12, 12], dtype=torch.int32, device=device) |
| 263 | token_to_req = torch.tensor([0] * 4 + [1] * 4 + [2] * 4 + [0] * 4, device=device) |
| 264 | common = SimpleNamespace( |
| 265 | num_actual_tokens=16, |
| 266 | num_reqs=4, |
| 267 | max_query_len=4, |
| 268 | max_seq_len=68, |
| 269 | query_start_loc=query_start_loc, |
| 270 | query_start_loc_cpu=query_start_loc.cpu(), |
| 271 | seq_lens=torch.tensor([68, 68, 68, 0], dtype=torch.int32, device=device), |
| 272 | slot_mapping=torch.tensor(list(range(12)) + [-1] * 4, device=device), |
| 273 | block_table_tensor=torch.empty((4, 0), dtype=torch.int32, device=device), |
| 274 | token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), |
| 275 | ) |
| 276 | |
| 277 | metadata = builder.build(0, common) |
| 278 | |
| 279 | assert metadata.logical_positions.tolist() == [ |
| 280 | 64, |
| 281 | 65, |
| 282 | 66, |
| 283 | 67, |
| 284 | 64, |
| 285 | 65, |
| 286 | 66, |
| 287 | 67, |
| 288 | 64, |
| 289 | 65, |
| 290 | 66, |
| 291 | 67, |
| 292 | -1, |
| 293 | -1, |
| 294 | -1, |
| 295 | -1, |
| 296 | ] |
| 297 | assert metadata.slot_mapping.tolist() == list(range(12)) + [-1] * 4 |
| 298 | assert metadata.visible_blocks.tolist() == [65, 66, 67, 68] * 3 + [0] * 4 |
| 299 | |
| 300 | |
| 301 | @requires_qsa_kernels |
| 302 | def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: |
| 303 | device = torch.device("cuda") |
| 304 | builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder) |
| 305 | builder.compress_ratio = 4 |
| 306 | builder.reorder_batch_threshold = 1 |
| 307 | builder.is_circular_buffer = True |
| 308 | builder.kv_cache_spec = SimpleNamespace(block_size=4) |
| 309 | builder.storage_block_size = 4 |
| 310 | builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device) |
| 311 | builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device) |
| 312 | builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device) |
| 313 | builder.visible_blocks_buffer = torch.empty(16, dtype=torch.int32, device=device) |
| 314 | builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device) |
| 315 | query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device) |
| 316 | token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device) |
| 317 | block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device) |
| 318 | common = SimpleNamespace( |
| 319 | num_actual_tokens=16, |
| 320 | num_reqs=3, |
| 321 | max_query_len=7, |
| 322 | max_seq_len=11, |
| 323 | query_start_loc=query_start_loc, |
| 324 | query_start_loc_cpu=query_start_loc.cpu(), |
| 325 | seq_lens=torch.tensor([9, 11, 0], dtype=torch.int32, device=device), |
| 326 | slot_mapping=torch.full((16,), -1, dtype=torch.int64, device=device), |
| 327 | block_table_tensor=block_table, |
| 328 | token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), |
| 329 | ) |
| 330 | |
| 331 | metadata = builder.build(0, common) |
| 332 | expected = [ |
| 333 | -1, |
| 334 | -1, |
| 335 | -1, |
| 336 | 5, |
| 337 | 6, |
| 338 | 7, |
| 339 | 4, |
| 340 | -1, |
| 341 | -1, |
| 342 | 3, |
| 343 | 0, |
| 344 | 1, |
| 345 | 2, |
| 346 | -1, |
| 347 | -1, |
| 348 | -1, |
| 349 | ] |
| 350 | |
| 351 | assert metadata.slot_mapping.tolist() == expected |
| 352 | |
| 353 | |
| 354 | @pytest.mark.parametrize("chunk_start", list(range(8))) |
| 355 | def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> None: |
| 356 | """A speculative step must not overwrite the open group's committed keys. |
| 357 | |
| 358 | The step stores every row it computes, drafts included, before acceptance |
| 359 | is known, while the next step still reads the earlier members of the group |
| 360 | being compressed from the ring. A ring sized at the compression ratio makes |
| 361 | those rows alias, so a rejected draft silently replaces a committed key. |
| 362 | """ |
| 363 | compress_ratio = 4 |
| 364 | num_spec = 3 |
| 365 | capacity = compress_ratio * -(-(compress_ratio + num_spec) // compress_ratio) |
| 366 | query_len = num_spec + 1 |
| 367 | |
| 368 | slots = qsa_cache.circular_qsa_slot_mapping( |
| 369 | torch.tensor([[0]], dtype=torch.int32), |
| 370 | torch.zeros(query_len, dtype=torch.int32), |
| 371 | torch.arange(chunk_start, chunk_start + query_len), |
| 372 | capacity, |
| 373 | query_start_loc=torch.tensor([0, query_len], dtype=torch.int32), |
| 374 | ) |
| 375 | |
| 376 | committed = torch.arange(chunk_start - chunk_start % compress_ratio, chunk_start) |
| 377 | assert set(slots.tolist()).isdisjoint((committed % capacity).tolist()) |
| 378 | |
| 379 | |
| 380 | def _qsa_key_cache( |
| 381 | block_size: int, compress_ratio: int, **kwargs |
| 382 | ) -> qsa_cache.QSAKeyStateCache: |
| 383 | return qsa_cache.QSAKeyStateCache( |
| 384 | head_size=64, |
| 385 | dtype=torch.bfloat16, |
| 386 | cache_config=SimpleNamespace(block_size=block_size), |
| 387 | prefix=f"raw.{block_size}.{compress_ratio}", |
| 388 | vllm_config=SimpleNamespace( |
| 389 | compilation_config=SimpleNamespace(static_forward_context={}) |
| 390 | ), |
| 391 | compress_ratio=compress_ratio, |
| 392 | **kwargs, |
| 393 | ) |
| 394 | |
| 395 | |
| 396 | def test_qsa_state_caches_adapt_the_unified_logical_layout() -> None: |
| 397 | raw_cache = _qsa_key_cache(block_size=32, compress_ratio=4) |
| 398 | compressed_cache = qsa_cache.QSACompressedKeyCache( |
| 399 | head_size=64, |
| 400 | dtype=torch.bfloat16, |
| 401 | cache_config=SimpleNamespace(block_size=32), |
| 402 | prefix="compressed.bind", |
| 403 | vllm_config=SimpleNamespace( |
| 404 | compilation_config=SimpleNamespace(static_forward_context={}) |
| 405 | ), |
| 406 | compress_ratio=4, |
| 407 | ) |
| 408 | raw_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16) |
| 409 | compressed_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16) |
| 410 | |
| 411 | raw_cache.bind_kv_cache(raw_view) |
| 412 | compressed_cache.bind_kv_cache(compressed_view) |
| 413 | |
| 414 | assert raw_cache.kv_cache.shape == (2, 8, 1, 64) |
| 415 | assert compressed_cache.kv_cache.shape == (2, 8, 1, 64) |
| 416 | assert raw_cache.kv_cache.data_ptr() == raw_view.data_ptr() |
| 417 | assert compressed_cache.kv_cache.data_ptr() == compressed_view.data_ptr() |
| 418 | |
| 419 | |
| 420 | def test_clearing_qsa_key_cache_releases_its_storage() -> None: |
| 421 | """Profiling teardown clears `kv_cache`; no derived view may keep it alive.""" |
| 422 | raw_cache = _qsa_key_cache( |
| 423 | block_size=32, compress_ratio=4, cache_rope_positions=True |
| 424 | ) |
| 425 | kv = torch.zeros(2, 1, 8, raw_cache.head_size, dtype=torch.bfloat16) |
| 426 | references_before_bind = torch._C._storage_Use_Count(kv.untyped_storage()._cdata) |
| 427 | raw_cache.bind_kv_cache(kv) |
| 428 | assert raw_cache.key_cache.shape[-1] == 64 |
| 429 | assert raw_cache.rope_position_cache.dtype == torch.int64 |
| 430 | |
| 431 | clear_layer_kv_caches([raw_cache]) |
| 432 | |
| 433 | references = torch._C._storage_Use_Count(kv.untyped_storage()._cdata) |
| 434 | assert references == references_before_bind |
| 435 | |
| 436 | |
| 437 | @pytest.mark.parametrize( |
| 438 | ("compress_ratio", "num_spec", "expected"), |
| 439 | [(4, 0, 4), (4, 1, 8), (4, 3, 8), (4, 4, 8), (4, 5, 12), (2, 3, 6)], |
| 440 | ) |
| 441 | def test_qsa_ring_capacity_covers_one_speculative_step( |
| 442 | compress_ratio: int, num_spec: int, expected: int |
| 443 | ) -> None: |
| 444 | """Capacity spans the open group plus one speculative step, in whole groups.""" |
| 445 | spec = _qsa_key_cache( |
| 446 | block_size=48, compress_ratio=compress_ratio |
| 447 | ).get_kv_cache_spec(SimpleNamespace(num_speculative_tokens=num_spec)) |
| 448 | assert spec.block_size == expected |
| 449 | |
| 450 | |
| 451 | @requires_qsa_kernels |
| 452 | def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None: |
| 453 | device = torch.device("cuda") |
| 454 | builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder) |
| 455 | builder.compress_ratio = 4 |
| 456 | builder.reorder_batch_threshold = 1 |
| 457 | builder.is_circular_buffer = False |
| 458 | builder.storage_block_size = 16 |
| 459 | builder.token_to_req_buffer = torch.empty(8, dtype=torch.int32, device=device) |
| 460 | builder.slot_mapping_buffer = torch.empty(8, dtype=torch.int64, device=device) |
| 461 | builder.logical_positions_buffer = torch.empty(8, dtype=torch.int64, device=device) |
| 462 | builder.visible_blocks_buffer = torch.empty(8, dtype=torch.int32, device=device) |
| 463 | # Simulate max_num_seqs exceeding the three live requests below. |
| 464 | builder.request_capacity = 8 |
| 465 | builder.k_work_metadata_buffer = torch.empty(4, 2, dtype=torch.int32, device=device) |
| 466 | query_start_loc = torch.tensor([0, 3, 3, 8], dtype=torch.int32, device=device) |
| 467 | token_to_req = torch.tensor( |
| 468 | [0, 0, 0, 2, 2, 2, 2, 2], dtype=torch.int32, device=device |
| 469 | ) |
| 470 | common = SimpleNamespace( |
| 471 | num_actual_tokens=8, |
| 472 | num_reqs=3, |
| 473 | max_query_len=5, |
| 474 | max_seq_len=12, |
| 475 | query_start_loc=query_start_loc, |
| 476 | query_start_loc_cpu=query_start_loc.cpu(), |
| 477 | seq_lens=torch.tensor([7, 0, 12], dtype=torch.int32, device=device), |
| 478 | slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device), |
| 479 | block_table_tensor=torch.zeros((3, 1), dtype=torch.int32, device=device), |
| 480 | token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), |
| 481 | ) |
| 482 | |
| 483 | metadata = builder.build(0, common) |
| 484 | |
| 485 | assert metadata.slot_mapping.tolist() == [-1] * 8 |
| 486 | assert metadata.visible_blocks.tolist() == [1, 1, 1, 2, 2, 2, 2, 3] |
| 487 | assert metadata.k_work_metadata.tolist() == [[0, 0], [2, 0], [2, 1], [-1, -1]] |
| 488 | |
| 489 | |
| 490 | @requires_qsa_kernels |
| 491 | @pytest.mark.usefixtures("default_vllm_config") |
| 492 | def test_qsa_unfused_cache_update_ignores_padded_qk() -> None: |
| 493 | """Padded projected Q/K rows must not affect either side cache.""" |
| 494 | from vllm.model_executor.layers.rotary_embedding import get_rope |
| 495 | |
| 496 | device = torch.device("cuda") |
| 497 | # Five tokens complete one compressed group and retain four keys in the ring. |
| 498 | raw_metadata = SimpleNamespace( |
| 499 | num_actual_tokens=5, |
| 500 | slot_mapping=torch.tensor([-1, 1, 2, 3, 0], device=device), |
| 501 | block_table=torch.zeros((1, 1), dtype=torch.int32, device=device), |
| 502 | token_to_req=torch.zeros(5, dtype=torch.int32, device=device), |
| 503 | query_start_loc=torch.tensor([0, 5], dtype=torch.int32, device=device), |
| 504 | logical_positions=torch.arange(5, device=device), |
| 505 | ) |
| 506 | compressed_metadata = SimpleNamespace( |
| 507 | slot_mapping=torch.tensor([-1, -1, -1, 0, -1], device=device), |
| 508 | ) |
| 509 | with torch.device(device): |
| 510 | rope = get_rope( |
| 511 | head_size=128, |
| 512 | max_position=32, |
| 513 | rope_parameters={"rope_type": "default", "partial_rotary_factor": 0.5}, |
| 514 | dtype=torch.bfloat16, |
| 515 | ) |
| 516 | raw_cache = torch.zeros((1, 4, 1, 64), dtype=torch.bfloat16, device=device) |
| 517 | compressed_cache = torch.zeros((1, 2, 1, 64), dtype=torch.bfloat16, device=device) |
| 518 | norm = SimpleNamespace( |
| 519 | weight=torch.zeros(64, dtype=torch.bfloat16, device=device), |
| 520 | variance_epsilon=1e-6, |
| 521 | ) |
| 522 | indexer = SimpleNamespace( |
| 523 | _metadata=lambda: (raw_metadata, compressed_metadata), |
| 524 | skip_topk=True, |
| 525 | index_kv_heads=1, |
| 526 | index_n_heads=1, |
| 527 | index_head_dim=64, |
| 528 | indexer_dtype=torch.bfloat16, |
| 529 | q_layernorm=norm, |
| 530 | k_layernorm=norm, |
| 531 | rotary_emb=rope, |
| 532 | compress_ratio=4, |
| 533 | raw_key_cache=SimpleNamespace( |
| 534 | kv_cache=raw_cache, key_cache=raw_cache, rope_position_cache=None |
| 535 | ), |
| 536 | compressed_key_cache=SimpleNamespace(kv_cache=compressed_cache), |
| 537 | ) |
| 538 | keys = torch.arange(1, 6, dtype=torch.bfloat16, device=device)[:, None].expand( |
| 539 | 5, 64 |
| 540 | ) |
| 541 | padded_keys = torch.full((8, 64), torch.nan, dtype=torch.bfloat16, device=device) |
| 542 | padded_keys[:5].copy_(keys) |
| 543 | indexer_qsa.QSAIndexer.forward( |
| 544 | indexer, |
| 545 | torch.cat((torch.ones_like(padded_keys), padded_keys), dim=-1), |
| 546 | torch.zeros(8, dtype=torch.long, device=device), |
| 547 | torch.full((5, 5), -1, dtype=torch.int32, device=device), |
| 548 | attn=SimpleNamespace(use_fused_qsa_prepare=False), |
| 549 | ) |
| 550 | torch.testing.assert_close(raw_cache[0, :, 0], keys[[4, 1, 2, 3]]) |
| 551 | expected_compressed = torch.zeros_like(compressed_cache) |
| 552 | expected_compressed[0, 0] = 1 |
| 553 | torch.testing.assert_close(compressed_cache, expected_compressed) |
| 554 | |
| 555 | |
| 556 | @requires_qsa_kernels |
| 557 | @pytest.mark.parametrize("compress_ratio", [1, 4]) |
| 558 | @pytest.mark.parametrize("num_reqs", [2, 3, 4, 7, 8, 9]) |
| 559 | def test_qsa_triton_metadata_matches_pytorch( |
| 560 | compress_ratio: int, num_reqs: int |
| 561 | ) -> None: |
| 562 | device = torch.device("cuda") |
| 563 | num_tokens = 8 |
| 564 | query_start_loc = torch.tensor( |
| 565 | [0, 3, *([3] * (num_reqs - 2)), 8], dtype=torch.int32, device=device |
| 566 | ) |
| 567 | token_to_req = torch.tensor( |
| 568 | [0, 0, 0, *([num_reqs - 1] * 5)], |
| 569 | dtype=torch.int32, |
| 570 | device=device, |
| 571 | ) |
| 572 | block_table_rows = torch.tensor( |
| 573 | [ |
| 574 | [4, -1, 8, -1, 12, -1], |
| 575 | [1, -1, 2, -1, 3, -1], |
| 576 | [7, -1, 9, -1, 11, -1], |
| 577 | ], |
| 578 | dtype=torch.int32, |
| 579 | device=device, |
| 580 | ) |
| 581 | block_table_storage = block_table_rows[ |
| 582 | torch.arange(num_reqs, device=device) % block_table_rows.shape[0] |
| 583 | ] |
| 584 | seq_lens = torch.zeros(num_reqs, dtype=torch.int32, device=device) |
| 585 | seq_lens[0] = 10 |
| 586 | seq_lens[-1] = 20 |
| 587 | common = SimpleNamespace( |
| 588 | num_actual_tokens=num_tokens, |
| 589 | query_start_loc=query_start_loc, |
| 590 | query_start_loc_cpu=query_start_loc.cpu(), |
| 591 | seq_lens=seq_lens, |
| 592 | slot_mapping=torch.tensor( |
| 593 | [0, 1, -1, 3, 4, -1, -1, -1], dtype=torch.int64, device=device |
| 594 | ), |
| 595 | block_table_tensor=block_table_storage[:, ::2], |
| 596 | token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), |
| 597 | ) |
| 598 | |
| 599 | def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 600 | return ( |
| 601 | torch.empty(num_tokens, dtype=torch.int32, device=device), |
| 602 | torch.empty(num_tokens, dtype=torch.int64, device=device), |
| 603 | torch.empty(num_tokens, dtype=torch.int32, device=device), |
| 604 | torch.empty(num_tokens, dtype=torch.int64, device=device), |
| 605 | ) |
| 606 | |
| 607 | max_num_work = ( |
| 608 | (num_tokens + (compress_ratio - 1) * num_reqs) // compress_ratio |
| 609 | if compress_ratio != 1 |
| 610 | else 0 |
| 611 | ) |
| 612 | actual_k_work = ( |
| 613 | torch.empty(max_num_work, 2, dtype=torch.int32, device=device) |
| 614 | if max_num_work |
| 615 | else None |
| 616 | ) |
| 617 | actual_buffers = make_buffers() |
| 618 | actual = qsa_cache.build_qsa_metadata_triton( |
| 619 | common, |
| 620 | *actual_buffers, |
| 621 | storage_block_size=2, |
| 622 | compress_ratio=compress_ratio, |
| 623 | k_work_metadata_buffer=actual_k_work, |
| 624 | request_capacity=num_reqs, |
| 625 | ) |
| 626 | |
| 627 | expected_k_work = ( |
| 628 | torch.empty_like(actual_k_work) if actual_k_work is not None else None |
| 629 | ) |
| 630 | expected_buffers = make_buffers() |
| 631 | expected = qsa_cache._build_qsa_metadata_torch( |
| 632 | common, |
| 633 | *expected_buffers, |
| 634 | storage_block_size=2, |
| 635 | compress_ratio=compress_ratio, |
| 636 | k_work_metadata_buffer=expected_k_work, |
| 637 | request_capacity=num_reqs, |
| 638 | ) |
| 639 | |
| 640 | for actual_tensor, expected_tensor in zip(actual, expected): |
| 641 | torch.testing.assert_close(actual_tensor, expected_tensor) |
| 642 | if actual_k_work is not None: |
| 643 | torch.testing.assert_close(actual_k_work, expected_k_work) |
| 644 | |
| 645 | |
| 646 | @requires_qsa_kernels |
| 647 | def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None: |
| 648 | device = torch.device("cuda") |
| 649 | num_mapped_tokens = 4096 |
| 650 | num_tokens = 4224 |
| 651 | query_start_loc = torch.tensor( |
| 652 | [0, num_mapped_tokens], dtype=torch.int32, device=device |
| 653 | ) |
| 654 | common = SimpleNamespace( |
| 655 | num_actual_tokens=num_tokens, |
| 656 | query_start_loc=query_start_loc, |
| 657 | query_start_loc_cpu=query_start_loc.cpu(), |
| 658 | seq_lens=torch.tensor( |
| 659 | [num_mapped_tokens + 32], dtype=torch.int32, device=device |
| 660 | ), |
| 661 | block_table_tensor=torch.arange(256, dtype=torch.int32, device=device)[None], |
| 662 | slot_mapping=torch.tensor( |
| 663 | [0] * num_mapped_tokens + [-1] * (num_tokens - num_mapped_tokens), |
| 664 | dtype=torch.int64, |
| 665 | device=device, |
| 666 | ), |
| 667 | token_to_req_indices=lambda buffer: buffer.zero_(), |
| 668 | ) |
| 669 | |
| 670 | def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 671 | return ( |
| 672 | torch.empty(num_tokens, dtype=torch.int32, device=device), |
| 673 | torch.empty(num_tokens, dtype=torch.int64, device=device), |
| 674 | torch.empty(num_tokens, dtype=torch.int32, device=device), |
| 675 | torch.empty(num_tokens, dtype=torch.int64, device=device), |
| 676 | ) |
| 677 | |
| 678 | max_num_work = (num_tokens + 3) // 4 |
| 679 | actual_k_work = torch.empty(max_num_work, 2, dtype=torch.int32, device=device) |
| 680 | expected_k_work = torch.empty_like(actual_k_work) |
| 681 | actual = qsa_cache.build_qsa_metadata_triton( |
| 682 | common, |
| 683 | *make_buffers(), |
| 684 | storage_block_size=8, |
| 685 | compress_ratio=4, |
| 686 | k_work_metadata_buffer=actual_k_work, |
| 687 | ) |
| 688 | expected = qsa_cache._build_qsa_metadata_torch( |
| 689 | common, |
| 690 | *make_buffers(), |
| 691 | storage_block_size=8, |
| 692 | compress_ratio=4, |
| 693 | k_work_metadata_buffer=expected_k_work, |
| 694 | ) |
| 695 | |
| 696 | for actual_tensor, expected_tensor in zip(actual, expected): |
| 697 | torch.testing.assert_close(actual_tensor, expected_tensor) |
| 698 | torch.testing.assert_close(actual_k_work, expected_k_work) |
| 699 | |
| 700 | |
| 701 | @requires_qsa_kernels |
| 702 | @pytest.mark.parametrize( |
| 703 | ("decode_query_len", "num_requests"), |
| 704 | [ |
| 705 | (1, 2), |
| 706 | (2, 2), |
| 707 | (3, 2), |
| 708 | (4, 2), |
| 709 | (4, 33), |
| 710 | ], |
| 711 | ) |
| 712 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) |
| 713 | def test_qsa_decode_selection_correctness( |
| 714 | decode_query_len: int, num_requests: int, dtype: torch.dtype |
| 715 | ) -> None: |
| 716 | torch.manual_seed(1) |
| 717 | heads, head_dim = 4, 128 |
| 718 | rows = num_requests * decode_query_len |
| 719 | q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16).to( |
| 720 | dtype |
| 721 | ) |
| 722 | page_size, pages_per_request, max_sequence_length = ( |
| 723 | (16, 40, 2560) if num_requests > 32 else (4, 20, 320) |
| 724 | ) |
| 725 | num_pages = num_requests * pages_per_request |
| 726 | cache = torch.randn( |
| 727 | num_pages, |
| 728 | page_size, |
| 729 | 1, |
| 730 | head_dim, |
| 731 | device="cuda", |
| 732 | dtype=torch.bfloat16, |
| 733 | ).to(dtype) |
| 734 | page_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape( |
| 735 | num_requests, pages_per_request |
| 736 | ) |
| 737 | token_to_req = torch.repeat_interleave( |
| 738 | torch.arange(num_requests, device="cuda", dtype=torch.int32), |
| 739 | decode_query_len, |
| 740 | ) |
| 741 | sequence_lengths = max_sequence_length - 4 * ( |
| 742 | torch.arange(num_requests, device="cuda", dtype=torch.int32) % 8 |
| 743 | ) |
| 744 | query_positions = torch.cat( |
| 745 | [ |
| 746 | torch.arange( |
| 747 | length - decode_query_len, |
| 748 | length, |
| 749 | device="cuda", |
| 750 | dtype=torch.int32, |
| 751 | ) |
| 752 | for length in sequence_lengths.tolist() |
| 753 | ] |
| 754 | ) |
| 755 | visible_blocks = torch.minimum( |
| 756 | (query_positions + 1) // 4, |
| 757 | sequence_lengths.index_select(0, token_to_req.long()) // 4, |
| 758 | ) |
| 759 | |
| 760 | token_topk, compress_ratio = 2048, 4 |
| 761 | actual = torch.empty( |
| 762 | (rows, token_topk // compress_ratio), device="cuda", dtype=torch.int32 |
| 763 | ) |
| 764 | qsa_indexer_ops.qsa_select_paged_decode( |
| 765 | q, |
| 766 | cache, |
| 767 | page_table, |
| 768 | visible_blocks, |
| 769 | token_topk, |
| 770 | compress_ratio, |
| 771 | decode_query_len, |
| 772 | actual, |
| 773 | ) |
| 774 | expected = _qsa_select_paged_reference( |
| 775 | q, |
| 776 | cache, |
| 777 | page_table, |
| 778 | token_to_req, |
| 779 | query_positions, |
| 780 | sequence_lengths, |
| 781 | token_topk, |
| 782 | compress_ratio, |
| 783 | ) |
| 784 | |
| 785 | if dtype == torch.float8_e4m3fn: |
| 786 | # fp8 logits tie at the top-k boundary more often than bf16, so index |
| 787 | # identity is not stable; compare the selected value multisets. |
| 788 | # SM90 wgmma accumulates fp8 in reduced precision (~3e-4 abs |
| 789 | # observed); SM100 tcgen05 is exact fp32. |
| 790 | rtol = atol = 1e-3 if current_platform.is_device_capability(90) else None |
| 791 | logits = _qsa_mqa_paged_reference( |
| 792 | q, cache, page_table, token_to_req, visible_blocks |
| 793 | ) |
| 794 | for row in range(rows): |
| 795 | selected = actual[row][actual[row] >= 0] |
| 796 | wanted = expected[row][expected[row] >= 0] |
| 797 | assert selected.numel() == wanted.numel() |
| 798 | torch.testing.assert_close( |
| 799 | logits[row, selected.long()].sort().values, |
| 800 | logits[row, wanted.long()].sort().values, |
| 801 | rtol=rtol, |
| 802 | atol=atol, |
| 803 | ) |
| 804 | return |
| 805 | |
| 806 | torch.testing.assert_close(actual.sort().values, expected.sort().values) |
| 807 | |
| 808 | |
| 809 | @requires_qsa_kernels |
| 810 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) |
| 811 | @pytest.mark.parametrize("seq_len_slack", [0, 1792]) |
| 812 | @pytest.mark.parametrize("force_chunk", [False, True]) |
| 813 | def test_qsa_prefill_selection_correctness( |
| 814 | monkeypatch: pytest.MonkeyPatch, |
| 815 | seq_len_slack: int, |
| 816 | force_chunk: bool, |
| 817 | dtype: torch.dtype, |
| 818 | ) -> None: |
| 819 | # page_size=24 (does not divide the 64-aligned clipped width) and an |
| 820 | # oversized page table, so the clipped logits width comes from |
| 821 | # max_seq_len, not page geometry. seq_len_slack > 0 simulates the |
| 822 | # spec-decode case where the bound is an over-estimate. force_chunk |
| 823 | # drives the logits budget to one row per chunk. |
| 824 | if force_chunk: |
| 825 | monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "0") |
| 826 | torch.manual_seed(2) |
| 827 | query_lens = [3, 33] |
| 828 | rows, heads, head_dim = sum(query_lens), 4, 128 |
| 829 | q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16).to( |
| 830 | dtype |
| 831 | ) |
| 832 | cache = torch.randn(128, 24, 1, head_dim, device="cuda", dtype=torch.bfloat16).to( |
| 833 | dtype |
| 834 | ) |
| 835 | page_table = torch.randperm(128, device="cuda", dtype=torch.int32).reshape(2, 64) |
| 836 | token_to_req = torch.repeat_interleave( |
| 837 | torch.arange(2, device="cuda", dtype=torch.int32), |
| 838 | torch.tensor(query_lens, device="cuda"), |
| 839 | ) |
| 840 | query_start_loc = torch.tensor([0, 3, 36], device="cuda", dtype=torch.int32) |
| 841 | sequence_lengths = torch.tensor([5120, 4224], device="cuda", dtype=torch.int32) |
| 842 | query_positions = torch.cat( |
| 843 | [ |
| 844 | torch.arange(length - query_len, length, device="cuda", dtype=torch.int32) |
| 845 | for query_len, length in zip( |
| 846 | query_lens, sequence_lengths.tolist(), strict=True |
| 847 | ) |
| 848 | ] |
| 849 | ) |
| 850 | token_topk, compress_ratio = 2048, 4 |
| 851 | visible_blocks = torch.minimum( |
| 852 | (query_positions + 1) // compress_ratio, |
| 853 | sequence_lengths.index_select(0, token_to_req.long()) // compress_ratio, |
| 854 | ) |
| 855 | |
| 856 | actual = torch.empty( |
| 857 | (rows, token_topk // compress_ratio), device="cuda", dtype=torch.int32 |
| 858 | ) |
| 859 | qsa_indexer_ops.qsa_select_paged_prefill( |
| 860 | q, |
| 861 | cache, |
| 862 | page_table, |
| 863 | query_start_loc, |
| 864 | visible_blocks, |
| 865 | token_topk, |
| 866 | compress_ratio, |
| 867 | max(query_lens), |
| 868 | actual, |
| 869 | max_seq_len=sequence_lengths.max().item() + seq_len_slack, |
| 870 | ) |
| 871 | expected = _qsa_select_paged_reference( |
| 872 | q, |
| 873 | cache, |
| 874 | page_table, |
| 875 | token_to_req, |
| 876 | query_positions, |
| 877 | sequence_lengths, |
| 878 | token_topk, |
| 879 | compress_ratio, |
| 880 | ) |
| 881 | |
| 882 | if dtype == torch.float8_e4m3fn: |
| 883 | # fp8 logits tie at the top-k boundary more often than bf16, so index |
| 884 | # identity is not stable; compare the selected value multisets. |
| 885 | # SM90 wgmma accumulates fp8 in reduced precision (~3e-4 abs |
| 886 | # observed); SM100 tcgen05 is exact fp32. |
| 887 | rtol = atol = 1e-3 if current_platform.is_device_capability(90) else None |
| 888 | logits = _qsa_mqa_paged_reference( |
| 889 | q, cache, page_table, token_to_req, visible_blocks |
| 890 | ) |
| 891 | for row in range(rows): |
| 892 | selected = actual[row][actual[row] >= 0] |
| 893 | wanted = expected[row][expected[row] >= 0] |
| 894 | assert selected.numel() == wanted.numel() |
| 895 | torch.testing.assert_close( |
| 896 | logits[row, selected.long()].sort().values, |
| 897 | logits[row, wanted.long()].sort().values, |
| 898 | rtol=rtol, |
| 899 | atol=atol, |
| 900 | ) |
| 901 | return |
| 902 | |
| 903 | torch.testing.assert_close(actual.sort().values, expected.sort().values) |
| 904 | |
| 905 | |
| 906 | @requires_qsa_kernels |
| 907 | def test_qsa_block_expansion_correctness() -> None: |
| 908 | blocks = torch.tensor([[0, -1], [1, 0]], device="cuda", dtype=torch.int32) |
| 909 | query_positions = torch.tensor([5, 10], device="cuda") |
| 910 | sequence_lengths = torch.tensor([6, 11], device="cuda") |
| 911 | token_to_req = torch.tensor([0, 1], device="cuda", dtype=torch.int32) |
| 912 | visible_blocks = torch.minimum( |
| 913 | (query_positions + 1) // 4, |
| 914 | sequence_lengths.index_select(0, token_to_req.long()) // 4, |
| 915 | ).to(torch.int32) |
| 916 | |
| 917 | # Packed layout: one trailing column per row holds the valid-entry count |
| 918 | # (never a token index). Row 0: 1 visible block + 2 tail; row 1: 2 blocks |
| 919 | # + 3 tail. |
| 920 | actual = torch.empty((2, 12), device="cuda", dtype=torch.int32) |
| 921 | qsa_indexer_ops.expand_qsa_block_indices( |
| 922 | blocks, |
| 923 | query_positions, |
| 924 | visible_blocks, |
| 925 | compress_ratio=4, |
| 926 | token_topk=8, |
| 927 | out=actual, |
| 928 | ) |
| 929 | expected = _expand_qsa_indices_reference( |
| 930 | blocks, |
| 931 | query_positions, |
| 932 | sequence_lengths, |
| 933 | compress_ratio=4, |
| 934 | token_topk=8, |
| 935 | ) |
| 936 | |
| 937 | torch.testing.assert_close(actual[:, :11], expected) |
| 938 | assert actual[:, 11].tolist() == [6, 11] |
| 939 | |
| 940 | |
| 941 | @requires_qsa_kernels |
| 942 | @pytest.mark.parametrize( |
| 943 | ( |
| 944 | "num_rows", |
| 945 | "num_query_heads", |
| 946 | "num_kv_heads", |
| 947 | "page_size", |
| 948 | "use_prefill_config", |
| 949 | "num_requests", |
| 950 | "fp8", |
| 951 | ), |
| 952 | [ |
| 953 | # Production page sizes from hybrid-cache block alignment: 784/800 |
| 954 | # at TP4 and 1568/1600 at TP1/TP2 (no-MTP / MTP num_spec=3). Head |
| 955 | # splits are per-rank TP1/TP2/TP4; the largest batch runs both |
| 956 | # use_prefill_config variants. |
| 957 | pytest.param(1, 24, 2, 1600, True, 2, False, id="tp1_r1"), |
| 958 | pytest.param(16, 12, 1, 1600, True, 3, False, id="tp2_r16"), |
| 959 | pytest.param(32, 6, 1, 800, True, 5, False, id="tp4_r32"), |
| 960 | pytest.param(128, 24, 2, 1568, True, 7, False, id="tp1_r128"), |
| 961 | pytest.param(257, 6, 1, 800, True, 13, False, id="tp4_r257"), |
| 962 | pytest.param(513, 6, 1, 784, True, 17, False, id="tp4_r513"), |
| 963 | pytest.param(700, 6, 1, 800, True, 23, False, id="tp4_r700"), |
| 964 | pytest.param(1024, 24, 2, 1600, True, 33, False, id="tp1_r1024"), |
| 965 | pytest.param(2048, 24, 2, 1600, True, 63, False, id="tp1_r2048_prefill"), |
| 966 | pytest.param(2048, 24, 2, 1600, False, 63, False, id="tp1_r2048_uniform"), |
| 967 | # fp8_e4m3 K/V caches on the TP1 head split. |
| 968 | pytest.param(1, 24, 2, 1600, True, 2, True, id="tp1_r1_fp8"), |
| 969 | pytest.param(128, 24, 2, 1568, True, 7, True, id="tp1_r128_fp8"), |
| 970 | pytest.param(2048, 24, 2, 1600, True, 63, True, id="tp1_r2048_prefill_fp8"), |
| 971 | pytest.param(2048, 24, 2, 1600, False, 63, True, id="tp1_r2048_uniform_fp8"), |
| 972 | ], |
| 973 | ) |
| 974 | def test_qsa_sparse_paged_attention_correctness( |
| 975 | num_rows: int, |
| 976 | num_query_heads: int, |
| 977 | num_kv_heads: int, |
| 978 | page_size: int, |
| 979 | use_prefill_config: bool, |
| 980 | num_requests: int, |
| 981 | fp8: bool, |
| 982 | ) -> None: |
| 983 | """QSA sparse paged attention matches the dense reference. |
| 984 | |
| 985 | fp8 only changes the K/V cache dtype (e4m3 with a per-tensor scale pair) and |
| 986 | the scales; the reference dequantizes the same cache with those scales, so |
| 987 | both paths compare the production kernel against the reference on identical |
| 988 | inputs. fp8=True additionally covers the host-side scale folding. |
| 989 | """ |
| 990 | torch.manual_seed(2) |
| 991 | # One QSA attention problem: bf16 Q and paged K/V, a packed selection with |
| 992 | # the trailing count column, block table and row-to-request map. |
| 993 | head_dim = 256 |
| 994 | num_selected_pages = 64 |
| 995 | # Keep the newest page outside the synthetic top-k as causal headroom. |
| 996 | num_pages_per_request = num_selected_pages + 1 |
| 997 | num_cache_blocks = num_requests * num_pages_per_request |
| 998 | indexer_budget = 2048 |
| 999 | indexer_compress_ratio = 4 |
| 1000 | selection_width = indexer_budget + indexer_compress_ratio - 1 |
| 1001 | q = torch.randn( |
| 1002 | num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 |
| 1003 | ) |
| 1004 | output_gate = torch.randn_like(q) |
| 1005 | kv_cache = torch.randn( |
| 1006 | num_cache_blocks, |
| 1007 | page_size, |
| 1008 | num_kv_heads, |
| 1009 | 2 * head_dim, |
| 1010 | device="cuda", |
| 1011 | dtype=torch.bfloat16, |
| 1012 | ) |
| 1013 | k_cache, v_cache = kv_cache.split(head_dim, dim=-1) |
| 1014 | block_table = ( |
| 1015 | torch.randperm(num_cache_blocks, device="cuda") |
| 1016 | .reshape(num_requests, num_pages_per_request) |
| 1017 | .to(torch.int32) |
| 1018 | ) |
| 1019 | rows_per_request = math.ceil(num_rows / num_requests) |
| 1020 | row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32) |
| 1021 | token_to_req = row_indices // rows_per_request |
| 1022 | # Uniform row split; the last request takes the remainder (possibly 0). |
| 1023 | request_row_counts = torch.full( |
| 1024 | (num_requests,), rows_per_request, device="cuda", dtype=torch.int32 |
| 1025 | ) |
| 1026 | request_row_counts[-1] = num_rows - rows_per_request * (num_requests - 1) |
| 1027 | |
| 1028 | # Mix context lengths: every third request is short-context, attending |
| 1029 | # to only its first few pages; the rest fill their cache. |
| 1030 | context_lengths = torch.full( |
| 1031 | (num_requests,), |
| 1032 | num_pages_per_request * page_size - 1, |
| 1033 | device="cuda", |
| 1034 | dtype=torch.int32, |
| 1035 | ) |
| 1036 | short_requests = torch.arange(num_requests, device="cuda") % 3 == 1 |
| 1037 | context_lengths[short_requests] = request_row_counts[short_requests] + 8 |
| 1038 | block_topk = indexer_budget // indexer_compress_ratio |
| 1039 | compressed_blocks_per_page = page_size // indexer_compress_ratio |
| 1040 | selection = torch.arange(block_topk, device="cuda") |
| 1041 | selected_pages = selection % num_selected_pages |
| 1042 | selected_offsets = selection // num_selected_pages |
| 1043 | row_shifts = 2 * row_indices.unsqueeze(1) |
| 1044 | # Eight blocks per page; adjacent rows overlap by six of those eight. |
| 1045 | selected_offsets = (selected_offsets + row_shifts) % compressed_blocks_per_page |
| 1046 | block_indices = (selected_pages * compressed_blocks_per_page + selected_offsets).to( |
| 1047 | torch.int32 |
| 1048 | ) |
| 1049 | rows_within_request = row_indices % rows_per_request |
| 1050 | query_positions = ( |
| 1051 | context_lengths[token_to_req.long()] |
| 1052 | - request_row_counts[token_to_req.long()] |
| 1053 | + rows_within_request |
| 1054 | ).to(torch.int64) |
| 1055 | sequence_lengths = context_lengths |
| 1056 | visible_blocks = torch.minimum( |
| 1057 | (query_positions + 1) // indexer_compress_ratio, |
| 1058 | sequence_lengths.index_select(0, token_to_req.long()) // indexer_compress_ratio, |
| 1059 | ).to(torch.int32) |
| 1060 | # +1: the packed trailing column holds each row's valid-entry count |
| 1061 | # (never a token index); the reference reads only the selection region. |
| 1062 | logical_indices = torch.empty( |
| 1063 | (num_rows, selection_width + 1), device="cuda", dtype=torch.int32 |
| 1064 | ) |
| 1065 | qsa_indexer_ops.expand_qsa_block_indices( |
| 1066 | block_indices, |
| 1067 | query_positions, |
| 1068 | visible_blocks, |
| 1069 | indexer_compress_ratio, |
| 1070 | indexer_budget, |
| 1071 | logical_indices, |
| 1072 | ) |
| 1073 | |
| 1074 | scale = head_dim**-0.5 |
| 1075 | |
| 1076 | if fp8: |
| 1077 | # A fixed non-unit pair (k != v) exercises the host-side scale folding |
| 1078 | # and catches a k/v swap; scales are host floats, as the layer exposes |
| 1079 | # them. Stored values are the scaled ones, as reshape_and_cache does. |
| 1080 | k_scale, v_scale = 0.5, 2.0 |
| 1081 | k_cache = (k_cache / k_scale).to(torch.float8_e4m3fn) |
| 1082 | v_cache = (v_cache / v_scale).to(torch.float8_e4m3fn) |
| 1083 | else: |
| 1084 | k_scale, v_scale = 1.0, 1.0 |
| 1085 | |
| 1086 | actual = qsa_ops.qsa_sparse_paged_attention( |
| 1087 | q, |
| 1088 | k_cache, |
| 1089 | v_cache, |
| 1090 | logical_indices, |
| 1091 | block_table, |
| 1092 | token_to_req, |
| 1093 | use_prefill_config=use_prefill_config, |
| 1094 | k_scale=k_scale, |
| 1095 | v_scale=v_scale, |
| 1096 | output_gate=output_gate, |
| 1097 | ) |
| 1098 | expected = _qsa_sparse_paged_attention_reference( |
| 1099 | q, |
| 1100 | k_cache, |
| 1101 | v_cache, |
| 1102 | logical_indices[:, :selection_width], |
| 1103 | block_table, |
| 1104 | token_to_req, |
| 1105 | scale, |
| 1106 | k_scale=k_scale, |
| 1107 | v_scale=v_scale, |
| 1108 | ) |
| 1109 | expected = expected * torch.sigmoid(output_gate) |
| 1110 | |
| 1111 | torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) |
| 1112 | |
| 1113 | |
| 1114 | @requires_qsa_kernels |
| 1115 | @pytest.mark.parametrize("decode_query_len", [1, 2, 3, 4]) |
| 1116 | def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int) -> None: |
| 1117 | query_lens = [decode_query_len, decode_query_len, 33] |
| 1118 | rows, heads, head_dim = sum(query_lens), 4, 128 |
| 1119 | token_topk, compress_ratio = 2048, 4 |
| 1120 | torch.manual_seed(13) |
| 1121 | q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16) |
| 1122 | cache = torch.randn(120, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16) |
| 1123 | page_table = torch.arange(120, device="cuda", dtype=torch.int32).view(3, 40) |
| 1124 | token_to_req = torch.repeat_interleave( |
| 1125 | torch.arange(3, device="cuda", dtype=torch.int32), |
| 1126 | torch.tensor(query_lens, device="cuda"), |
| 1127 | ) |
| 1128 | query_start_loc = torch.tensor( |
| 1129 | [0, decode_query_len, 2 * decode_query_len, rows], |
| 1130 | device="cuda", |
| 1131 | dtype=torch.int32, |
| 1132 | ) |
| 1133 | sequence_lengths = torch.full((3,), 2560, device="cuda", dtype=torch.int32) |
| 1134 | query_positions = torch.cat( |
| 1135 | [ |
| 1136 | torch.arange(2560 - query_len, 2560, device="cuda") |
| 1137 | for query_len in query_lens |
| 1138 | ] |
| 1139 | ) |
| 1140 | |
| 1141 | block_indices = torch.empty( |
| 1142 | rows, |
| 1143 | token_topk // compress_ratio, |
| 1144 | device="cuda", |
| 1145 | dtype=torch.int32, |
| 1146 | ) |
| 1147 | visible_blocks = torch.minimum( |
| 1148 | (query_positions + 1) // compress_ratio, |
| 1149 | sequence_lengths.index_select(0, token_to_req.long()) // compress_ratio, |
| 1150 | ).to(torch.int32) |
| 1151 | num_decode_tokens = 2 * decode_query_len |
| 1152 | decode_slice = slice(0, num_decode_tokens) |
| 1153 | qsa_indexer_ops.qsa_select_paged_decode( |
| 1154 | q[decode_slice], |
| 1155 | cache, |
| 1156 | page_table[:2], |
| 1157 | visible_blocks[decode_slice], |
| 1158 | token_topk, |
| 1159 | compress_ratio, |
| 1160 | decode_query_len, |
| 1161 | block_indices[decode_slice], |
| 1162 | ) |
| 1163 | prefill_slice = slice(num_decode_tokens, rows) |
| 1164 | qsa_indexer_ops.qsa_select_paged_prefill( |
| 1165 | q[prefill_slice], |
| 1166 | cache, |
| 1167 | page_table[2:], |
| 1168 | query_start_loc[2:], |
| 1169 | visible_blocks[prefill_slice], |
| 1170 | token_topk, |
| 1171 | compress_ratio, |
| 1172 | query_lens[-1], |
| 1173 | block_indices[prefill_slice], |
| 1174 | max_seq_len=sequence_lengths.max().item(), |
| 1175 | ) |
| 1176 | # +1: the packed trailing count column (never a token index; excluded |
| 1177 | # from the comparison). |
| 1178 | actual = torch.empty( |
| 1179 | (rows, token_topk + compress_ratio), device="cuda", dtype=torch.int32 |
| 1180 | ) |
| 1181 | qsa_indexer_ops.expand_qsa_block_indices( |
| 1182 | block_indices, |
| 1183 | query_positions, |
| 1184 | visible_blocks, |
| 1185 | compress_ratio, |
| 1186 | token_topk, |
| 1187 | actual, |
| 1188 | ) |
| 1189 | expected_blocks = _qsa_select_paged_reference( |
| 1190 | q, |
| 1191 | cache, |
| 1192 | page_table, |
| 1193 | token_to_req, |
| 1194 | query_positions, |
| 1195 | sequence_lengths, |
| 1196 | token_topk, |
| 1197 | compress_ratio, |
| 1198 | ) |
| 1199 | expected = _expand_qsa_indices_reference( |
| 1200 | expected_blocks, |
| 1201 | query_positions, |
| 1202 | sequence_lengths.index_select(0, token_to_req.long()), |
| 1203 | compress_ratio, |
| 1204 | token_topk, |
| 1205 | ) |
| 1206 | |
| 1207 | torch.testing.assert_close( |
| 1208 | actual[:, : token_topk + compress_ratio - 1].sort().values, |
| 1209 | expected.sort().values, |
| 1210 | ) |
| 1211 | |
| 1212 | |
| 1213 | @requires_qsa_kernels |
| 1214 | def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) -> None: |
| 1215 | q = torch.zeros(2, 4, 8, device="cuda", dtype=torch.bfloat16) |
| 1216 | cache = torch.zeros(1, 16, 1, 8, device="cuda", dtype=torch.bfloat16) |
| 1217 | page_table = torch.zeros(1, 1, device="cuda", dtype=torch.int32) |
| 1218 | query_positions = torch.tensor([1, 2], device="cuda", dtype=torch.int32) |
| 1219 | visible_blocks = torch.zeros(2, device="cuda", dtype=torch.int32) |
| 1220 | |
| 1221 | block_indices = torch.empty((2, 512), device="cuda", dtype=torch.int32) |
| 1222 | qsa_indexer_ops.qsa_select_paged_prefill( |
| 1223 | q, |
| 1224 | cache, |
| 1225 | page_table, |
| 1226 | torch.tensor([0, 2], device="cuda", dtype=torch.int32), |
| 1227 | visible_blocks, |
| 1228 | token_topk=2048, |
| 1229 | compress_ratio=4, |
| 1230 | max_query_len=2, |
| 1231 | block_indices=block_indices, |
| 1232 | max_seq_len=64, # clamps to the page-table capacity |
| 1233 | ) |
| 1234 | selected = torch.empty((2, 2052), device="cuda", dtype=torch.int32) |
| 1235 | qsa_indexer_ops.expand_qsa_block_indices( |
| 1236 | block_indices, |
| 1237 | query_positions, |
| 1238 | visible_blocks, |
| 1239 | compress_ratio=4, |
| 1240 | token_topk=2048, |
| 1241 | out=selected, |
| 1242 | ) |
| 1243 | |
| 1244 | assert selected[0, :2].tolist() == [0, 1] |
| 1245 | assert selected[1, :3].tolist() == [0, 1, 2] |
| 1246 | assert torch.all(selected[0, 2:2051] == -1) |
| 1247 | assert torch.all(selected[1, 3:2051] == -1) |
| 1248 | # The packed trailing column holds each row's valid-entry count. |
| 1249 | assert selected[:, 2051].tolist() == [2, 3] |
| 1250 | |
| 1251 | |
| 1252 | @requires_qsa_kernels |
| 1253 | def test_qsa_streaming_compression_and_compressor_state_store_match_reference() -> None: |
| 1254 | head_dim = 8 |
| 1255 | current_pairs = [ |
| 1256 | *((0, position) for position in range(2, 9)), |
| 1257 | *((1, position) for position in range(5, 11)), |
| 1258 | ] |
| 1259 | |
| 1260 | def key_row(request: int, position: int) -> torch.Tensor: |
| 1261 | return ( |
| 1262 | torch.arange(head_dim, dtype=torch.float32) + request * 1000 + position * 10 |
| 1263 | ) |
| 1264 | |
| 1265 | def position_row(request: int, position: int) -> torch.Tensor: |
| 1266 | return torch.tensor( |
| 1267 | [ |
| 1268 | request * 1000 + position, |
| 1269 | request * 1000 + position + 100, |
| 1270 | request * 1000 + position + 200, |
| 1271 | ], |
| 1272 | dtype=torch.int64, |
| 1273 | ) |
| 1274 | |
| 1275 | raw_keys = ( |
| 1276 | torch.stack([key_row(request, position) for request, position in current_pairs]) |
| 1277 | .unsqueeze(1) |
| 1278 | .to(device="cuda", dtype=torch.bfloat16) |
| 1279 | ) |
| 1280 | raw_positions = ( |
| 1281 | torch.stack( |
| 1282 | [position_row(request, position) for request, position in current_pairs] |
| 1283 | ) |
| 1284 | .unsqueeze(1) |
| 1285 | .to(device="cuda") |
| 1286 | ) |
| 1287 | token_to_req = torch.tensor( |
| 1288 | [request for request, _ in current_pairs], |
| 1289 | dtype=torch.int32, |
| 1290 | device="cuda", |
| 1291 | ) |
| 1292 | logical_positions = torch.tensor( |
| 1293 | [position for _, position in current_pairs], |
| 1294 | dtype=torch.int64, |
| 1295 | device="cuda", |
| 1296 | ) |
| 1297 | query_start_loc = torch.tensor([0, 7, 13], dtype=torch.int32, device="cuda") |
| 1298 | compressor_state_block_table = torch.tensor( |
| 1299 | [[1], [0]], dtype=torch.int32, device="cuda" |
| 1300 | ) |
| 1301 | compressor_state_cache = torch.zeros( |
| 1302 | 2, 4, 1, head_dim, dtype=torch.bfloat16, device="cuda" |
| 1303 | ) |
| 1304 | rope_cache = torch.zeros(2, 4, 1, 3, dtype=torch.int64, device="cuda") |
| 1305 | for request, position, block in ((0, 0, 1), (0, 1, 1), (1, 4, 0)): |
| 1306 | compressor_state_cache[block, position % 4, 0] = key_row(request, position).to( |
| 1307 | device="cuda", dtype=torch.bfloat16 |
| 1308 | ) |
| 1309 | rope_cache[block, position % 4, 0] = position_row(request, position).to("cuda") |
| 1310 | |
| 1311 | compressed_slots = torch.full( |
| 1312 | (len(current_pairs),), -1, dtype=torch.int64, device="cuda" |
| 1313 | ) |
| 1314 | valid_rows = torch.tensor([1, 5, 9], dtype=torch.int64, device="cuda") |
| 1315 | compressed_slots[valid_rows] = torch.arange(3, device="cuda") |
| 1316 | pooled, first_positions = qsa_ops.qsa_compress_groups_with_ratio( |
| 1317 | raw_keys, |
| 1318 | raw_positions, |
| 1319 | compressor_state_cache, |
| 1320 | compressor_state_block_table, |
| 1321 | token_to_req, |
| 1322 | query_start_loc, |
| 1323 | logical_positions, |
| 1324 | compressed_slots, |
| 1325 | compress_ratio=4, |
| 1326 | rope_cache=rope_cache, |
| 1327 | ) |
| 1328 | pooled_without_rope, scalar_first_positions = ( |
| 1329 | qsa_ops.qsa_compress_groups_with_ratio( |
| 1330 | raw_keys, |
| 1331 | raw_positions, |
| 1332 | compressor_state_cache, |
| 1333 | compressor_state_block_table, |
| 1334 | token_to_req, |
| 1335 | query_start_loc, |
| 1336 | logical_positions, |
| 1337 | compressed_slots, |
| 1338 | compress_ratio=4, |
| 1339 | ) |
| 1340 | ) |
| 1341 | |
| 1342 | groups = [ |
| 1343 | [(0, position) for position in range(0, 4)], |
| 1344 | [(0, position) for position in range(4, 8)], |
| 1345 | [(1, position) for position in range(4, 8)], |
| 1346 | ] |
| 1347 | expected_pooled = ( |
| 1348 | torch.stack( |
| 1349 | [ |
| 1350 | torch.stack([key_row(*pair) for pair in group]).mean(dim=0) |
| 1351 | for group in groups |
| 1352 | ] |
| 1353 | ) |
| 1354 | .unsqueeze(1) |
| 1355 | .to(device="cuda", dtype=torch.bfloat16) |
| 1356 | ) |
| 1357 | expected_positions = torch.stack( |
| 1358 | [position_row(0, 0), position_row(0, 4), position_row(1, 4)] |
| 1359 | ).to("cuda") |
| 1360 | expected_scalar_positions = torch.tensor( |
| 1361 | [[0, 0, 0], [4, 4, 4], [4, 4, 4]], |
| 1362 | dtype=torch.int64, |
| 1363 | device="cuda", |
| 1364 | ) |
| 1365 | |
| 1366 | torch.testing.assert_close(pooled[valid_rows], expected_pooled) |
| 1367 | torch.testing.assert_close(pooled_without_rope[valid_rows], expected_pooled) |
| 1368 | torch.testing.assert_close(first_positions[valid_rows], expected_positions) |
| 1369 | torch.testing.assert_close( |
| 1370 | scalar_first_positions[valid_rows], expected_scalar_positions |
| 1371 | ) |
| 1372 | |
| 1373 | compressor_state_slots = torch.tensor( |
| 1374 | [-1, -1, -1, 5, 6, 7, 4, -1, -1, 3, 0, 1, 2], |
| 1375 | dtype=torch.int64, |
| 1376 | device="cuda", |
| 1377 | ) |
| 1378 | qsa_ops.qsa_store_cache_rows( |
| 1379 | compressor_state_cache, compressor_state_slots, raw_keys |
| 1380 | ) |
| 1381 | qsa_ops.qsa_store_cache_rows(rope_cache, compressor_state_slots, raw_positions) |
| 1382 | for request, positions, block in ((0, range(5, 9), 1), (1, range(7, 11), 0)): |
| 1383 | for position in positions: |
| 1384 | torch.testing.assert_close( |
| 1385 | compressor_state_cache[block, position % 4, 0], |
| 1386 | key_row(request, position).to(device="cuda", dtype=torch.bfloat16), |
| 1387 | ) |
| 1388 | torch.testing.assert_close( |
| 1389 | rope_cache[block, position % 4, 0], |
| 1390 | position_row(request, position).to("cuda"), |
| 1391 | ) |