# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from types import SimpleNamespace import pytest import torch from vllm.models.qwen4_exp.common import qsa_cache from vllm.models.qwen4_exp.common.qsa_cache import QSAMetadataBuilder from vllm.models.qwen4_exp.nvidia import indexer_qsa from vllm.models.qwen4_exp.nvidia import ( model as _qwen4_exp_model, # noqa: F401 ) from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops from vllm.models.qwen4_exp.nvidia.ops import qsa_indexer as qsa_indexer_ops from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON requires_qsa_kernels = pytest.mark.skipif( not current_platform.is_cuda() or not HAS_TRITON, reason="QSA kernels require CUDA and Triton", ) def test_qsa_mtp_index_share_updates_cache_but_skips_selection( monkeypatch: pytest.MonkeyPatch, ) -> None: rows = torch.tensor([[3, 1, -1], [5, 2, 0]], dtype=torch.int32) raw_metadata = SimpleNamespace( num_actual_tokens=2, slot_mapping=torch.arange(2), block_table=torch.empty(0), query_start_loc=torch.arange(3), logical_positions=torch.arange(2), ) compressed_metadata = SimpleNamespace( num_actual_tokens=2, slot_mapping=torch.arange(2), k_work_metadata=torch.empty(0), ) updates = [] selections = [] indexer = SimpleNamespace( skip_topk=True, _metadata=lambda: (raw_metadata, compressed_metadata), index_qk_proj=lambda hidden: (torch.zeros(2, 2), None), index_n_heads=1, index_kv_heads=1, index_head_dim=1, raw_key_cache=SimpleNamespace( kv_cache=torch.empty(0), rope_position_cache=None, rope_position_offset=0, ), compressed_key_cache=SimpleNamespace(kv_cache=torch.empty(0)), use_fused_pre_indexer=True, rotary_emb=SimpleNamespace(cos_sin_cache=torch.empty(0)), q_layernorm=SimpleNamespace(weight=torch.ones(1), variance_epsilon=1e-6), k_layernorm=SimpleNamespace(weight=torch.ones(1)), compress_ratio=2, ) monkeypatch.setattr( indexer_qsa, "qsa_pre_indexer", lambda *args, **kwargs: updates.append((args, kwargs)), ) monkeypatch.setattr( qsa_indexer_ops, "qsa_select_paged_decode", lambda *args, **kwargs: selections.append((args, kwargs)), ) monkeypatch.setattr( qsa_indexer_ops, "qsa_select_paged_prefill", lambda *args, **kwargs: selections.append((args, kwargs)), ) actual = indexer_qsa.QSAIndexer.forward( indexer, torch.zeros(2, 4), torch.tensor([7, 8]), rows, ) assert actual is rows assert len(updates) == 1 assert not selections def _qsa_mqa_paged_reference( q: torch.Tensor, k_cache: torch.Tensor, page_table: torch.Tensor, token_to_req: torch.Tensor, visible_lengths: torch.Tensor, ) -> torch.Tensor: pages = page_table.index_select(0, token_to_req.long()).long() keys = k_cache[pages, :, 0, :].flatten(1, 2) scores = torch.einsum("rhd,rnd->rnh", q.float(), keys.float()) logits = torch.relu(scores).sum(dim=-1) / math.sqrt(q.shape[-1]) positions = torch.arange(keys.shape[1], device=q.device).unsqueeze(0) return logits.masked_fill(positions >= visible_lengths.unsqueeze(1), -torch.inf) def _qsa_relative_topk_reference( logits: torch.Tensor, row_starts: torch.Tensor, row_ends: torch.Tensor, topk: int, ) -> torch.Tensor: output = torch.full( (logits.shape[0], topk), -1, dtype=torch.int32, device=logits.device ) for row in range(logits.shape[0]): start = int(row_starts[row].item()) length = int((row_ends[row] - row_starts[row]).item()) width = min(length, topk) if width: output[row, :width] = torch.topk( logits[row, start : start + length], width ).indices.to(torch.int32) return output def _expand_qsa_indices_reference( block_indices: torch.Tensor, query_positions: torch.Tensor, sequence_lengths: torch.Tensor, compress_ratio: int, token_topk: int, ) -> torch.Tensor: rows = block_indices.shape[0] block_topk = token_topk // compress_ratio output_width = token_topk + compress_ratio - 1 offsets = torch.arange(compress_ratio, device=block_indices.device) blocks = block_indices.long() expanded = blocks.unsqueeze(-1) * compress_ratio + offsets expanded = torch.where( blocks.unsqueeze(-1) >= 0, expanded, torch.full_like(expanded, -1) ).reshape(rows, block_topk * compress_ratio) expanded = expanded[:, :token_topk] expanded = torch.where( (expanded >= 0) & (expanded < sequence_lengths.unsqueeze(1)), expanded, torch.full_like(expanded, -1), ) tail_offsets = torch.arange(compress_ratio - 1, device=block_indices.device) visible_tokens = query_positions + 1 tail_start = visible_tokens // compress_ratio * compress_ratio tail = tail_start.unsqueeze(1) + tail_offsets.unsqueeze(0) tail_count = (visible_tokens - tail_start).unsqueeze(1) tail_valid = (tail_offsets.unsqueeze(0) < tail_count) & ( tail < sequence_lengths.unsqueeze(1) ) tail = torch.where(tail_valid, tail, torch.full_like(tail, -1)) result = torch.cat((expanded, tail), dim=1) order = torch.arange(output_width, device=result.device).expand(rows, -1) sort_key = torch.where(result >= 0, order, order + output_width) return result.gather(1, torch.argsort(sort_key, dim=1, stable=True)).to(torch.int32) def _qsa_select_paged_reference( q: torch.Tensor, k_cache: torch.Tensor, page_table: torch.Tensor, token_to_req: torch.Tensor, query_positions: torch.Tensor, sequence_lengths: torch.Tensor, token_topk: int, compress_ratio: int, ) -> torch.Tensor: row_sequence_lengths = sequence_lengths.index_select(0, token_to_req.long()) visible_blocks = torch.minimum( (query_positions + 1) // compress_ratio, row_sequence_lengths // compress_ratio, ).to(torch.int32) logits = _qsa_mqa_paged_reference( q, k_cache, page_table, token_to_req, visible_blocks, ) starts = torch.zeros_like(visible_blocks) return _qsa_relative_topk_reference( logits, starts, visible_blocks, token_topk // compress_ratio, ) def _qsa_sparse_paged_attention_reference( q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, logical_indices: torch.Tensor, block_table: torch.Tensor, token_to_req: torch.Tensor, softmax_scale: float, ) -> torch.Tensor: output = torch.zeros_like(q) repeats = q.shape[1] // k_cache.shape[2] page_size = k_cache.shape[1] for row in range(q.shape[0]): logical = logical_indices[row] logical = logical[logical >= 0].long() if not logical.numel(): continue request = token_to_req[row].long() pages = block_table[request, logical // page_size].long() offsets = logical % page_size keys = k_cache[pages, offsets].repeat_interleave(repeats, dim=1) values = v_cache[pages, offsets].repeat_interleave(repeats, dim=1) scores = torch.einsum("hd,khd->hk", q[row].float(), keys.float()) probabilities = torch.softmax(scores * softmax_scale, dim=-1) output[row] = torch.einsum("hk,khd->hd", probabilities, values.float()).to( q.dtype ) return output @requires_qsa_kernels def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None: device = torch.device("cuda") builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder) builder.compress_ratio = 1 builder.reorder_batch_threshold = 4 builder.is_circular_buffer = False builder.storage_block_size = 64 builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device) builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device) builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device) builder.visible_blocks_buffer = torch.empty(16, dtype=torch.int32, device=device) builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device) query_start_loc = torch.tensor([0, 4, 8, 12, 12], dtype=torch.int32, device=device) token_to_req = torch.tensor([0] * 4 + [1] * 4 + [2] * 4 + [0] * 4, device=device) common = SimpleNamespace( num_actual_tokens=16, num_reqs=4, max_query_len=4, query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc.cpu(), seq_lens=torch.tensor([68, 68, 68, 0], dtype=torch.int32, device=device), slot_mapping=torch.tensor(list(range(12)) + [-1] * 4, device=device), block_table_tensor=torch.empty((4, 0), dtype=torch.int32, device=device), token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), ) metadata = builder.build(0, common) assert metadata.logical_positions.tolist() == [ 64, 65, 66, 67, 64, 65, 66, 67, 64, 65, 66, 67, -1, -1, -1, -1, ] assert metadata.slot_mapping.tolist() == list(range(12)) + [-1] * 4 assert metadata.visible_blocks.tolist() == [65, 66, 67, 68] * 3 + [0] * 4 @requires_qsa_kernels def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None: device = torch.device("cuda") builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder) builder.compress_ratio = 4 builder.reorder_batch_threshold = 1 builder.is_circular_buffer = True builder.kv_cache_spec = SimpleNamespace(block_size=4) builder.storage_block_size = 4 builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device) builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device) builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device) builder.visible_blocks_buffer = torch.empty(16, dtype=torch.int32, device=device) builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device) query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device) token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device) block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device) common = SimpleNamespace( num_actual_tokens=16, num_reqs=3, max_query_len=7, query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc.cpu(), seq_lens=torch.tensor([9, 11, 0], dtype=torch.int32, device=device), slot_mapping=torch.full((16,), -1, dtype=torch.int64, device=device), block_table_tensor=block_table, token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), ) metadata = builder.build(0, common) expected = [ -1, -1, -1, 5, 6, 7, 4, -1, -1, 3, 0, 1, 2, -1, -1, -1, ] assert metadata.slot_mapping.tolist() == expected @pytest.mark.parametrize("chunk_start", list(range(8))) def test_qsa_circular_buffer_survives_one_speculative_step(chunk_start: int) -> None: """A speculative step must not overwrite the open group's committed keys. The step stores every row it computes, drafts included, before acceptance is known, while the next step still reads the earlier members of the group being compressed from the ring. A ring sized at the compression ratio makes those rows alias, so a rejected draft silently replaces a committed key. """ compress_ratio = 4 num_spec = 3 capacity = compress_ratio * -(-(compress_ratio + num_spec) // compress_ratio) query_len = num_spec + 1 slots = qsa_cache.circular_qsa_slot_mapping( torch.tensor([[0]], dtype=torch.int32), torch.zeros(query_len, dtype=torch.int32), torch.arange(chunk_start, chunk_start + query_len), capacity, query_start_loc=torch.tensor([0, query_len], dtype=torch.int32), ) committed = torch.arange(chunk_start - chunk_start % compress_ratio, chunk_start) assert set(slots.tolist()).isdisjoint((committed % capacity).tolist()) def _qsa_key_cache(block_size: int, compress_ratio: int) -> qsa_cache.QSAKeyStateCache: return qsa_cache.QSAKeyStateCache( head_size=64, dtype=torch.bfloat16, cache_config=SimpleNamespace(block_size=block_size), prefix=f"raw.{block_size}.{compress_ratio}", vllm_config=SimpleNamespace( compilation_config=SimpleNamespace(static_forward_context={}) ), compress_ratio=compress_ratio, ) def test_qsa_state_caches_adapt_the_unified_logical_layout() -> None: raw_cache = _qsa_key_cache(block_size=32, compress_ratio=4) compressed_cache = qsa_cache.QSACompressedKeyCache( head_size=64, dtype=torch.bfloat16, cache_config=SimpleNamespace(block_size=32), prefix="compressed.bind", vllm_config=SimpleNamespace( compilation_config=SimpleNamespace(static_forward_context={}) ), compress_ratio=4, ) raw_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16) compressed_view = torch.empty(2, 1, 8, 64, dtype=torch.bfloat16) raw_cache.bind_kv_cache(raw_view) compressed_cache.bind_kv_cache(compressed_view) assert raw_cache.kv_cache.shape == (2, 8, 1, 64) assert compressed_cache.kv_cache.shape == (2, 8, 1, 64) assert raw_cache.kv_cache.data_ptr() == raw_view.data_ptr() assert compressed_cache.kv_cache.data_ptr() == compressed_view.data_ptr() @pytest.mark.parametrize( ("compress_ratio", "num_spec", "expected"), [(4, 0, 4), (4, 1, 8), (4, 3, 8), (4, 4, 8), (4, 5, 12), (2, 3, 6)], ) def test_qsa_ring_capacity_covers_one_speculative_step( compress_ratio: int, num_spec: int, expected: int ) -> None: """Capacity spans the open group plus one speculative step, in whole groups.""" spec = _qsa_key_cache( block_size=48, compress_ratio=compress_ratio ).get_kv_cache_spec(SimpleNamespace(num_speculative_tokens=num_spec)) assert spec.block_size == expected @requires_qsa_kernels def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None: device = torch.device("cuda") builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder) builder.compress_ratio = 4 builder.reorder_batch_threshold = 1 builder.is_circular_buffer = False builder.storage_block_size = 16 builder.token_to_req_buffer = torch.empty(8, dtype=torch.int32, device=device) builder.slot_mapping_buffer = torch.empty(8, dtype=torch.int64, device=device) builder.logical_positions_buffer = torch.empty(8, dtype=torch.int64, device=device) builder.visible_blocks_buffer = torch.empty(8, dtype=torch.int32, device=device) # Simulate max_num_seqs exceeding the three live requests below. builder.request_capacity = 8 builder.k_work_metadata_buffer = torch.empty(4, 2, dtype=torch.int32, device=device) query_start_loc = torch.tensor([0, 3, 3, 8], dtype=torch.int32, device=device) token_to_req = torch.tensor( [0, 0, 0, 2, 2, 2, 2, 2], dtype=torch.int32, device=device ) common = SimpleNamespace( num_actual_tokens=8, num_reqs=3, max_query_len=5, query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc.cpu(), seq_lens=torch.tensor([7, 0, 12], dtype=torch.int32, device=device), slot_mapping=torch.full((8,), -1, dtype=torch.int64, device=device), block_table_tensor=torch.zeros((3, 1), dtype=torch.int32, device=device), token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), ) metadata = builder.build(0, common) assert metadata.slot_mapping.tolist() == [-1] * 8 assert metadata.visible_blocks.tolist() == [1, 1, 1, 2, 2, 2, 2, 3] assert metadata.k_work_metadata.tolist() == [[0, 0], [2, 0], [2, 1], [-1, -1]] @requires_qsa_kernels @pytest.mark.parametrize("compress_ratio", [1, 4]) @pytest.mark.parametrize("num_reqs", [2, 3, 4, 7, 8, 9]) def test_qsa_triton_metadata_matches_pytorch( compress_ratio: int, num_reqs: int ) -> None: device = torch.device("cuda") num_tokens = 8 query_start_loc = torch.tensor( [0, 3, *([3] * (num_reqs - 2)), 8], dtype=torch.int32, device=device ) token_to_req = torch.tensor( [0, 0, 0, *([num_reqs - 1] * 5)], dtype=torch.int32, device=device, ) block_table_rows = torch.tensor( [ [4, -1, 8, -1, 12, -1], [1, -1, 2, -1, 3, -1], [7, -1, 9, -1, 11, -1], ], dtype=torch.int32, device=device, ) block_table_storage = block_table_rows[ torch.arange(num_reqs, device=device) % block_table_rows.shape[0] ] seq_lens = torch.zeros(num_reqs, dtype=torch.int32, device=device) seq_lens[0] = 10 seq_lens[-1] = 20 common = SimpleNamespace( num_actual_tokens=num_tokens, query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc.cpu(), seq_lens=seq_lens, slot_mapping=torch.tensor( [0, 1, -1, 3, 4, -1, -1, -1], dtype=torch.int64, device=device ), block_table_tensor=block_table_storage[:, ::2], token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), ) def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: return ( torch.empty(num_tokens, dtype=torch.int32, device=device), torch.empty(num_tokens, dtype=torch.int64, device=device), torch.empty(num_tokens, dtype=torch.int32, device=device), torch.empty(num_tokens, dtype=torch.int64, device=device), ) max_num_work = ( (num_tokens + (compress_ratio - 1) * num_reqs) // compress_ratio if compress_ratio != 1 else 0 ) actual_k_work = ( torch.empty(max_num_work, 2, dtype=torch.int32, device=device) if max_num_work else None ) actual_buffers = make_buffers() actual = qsa_cache.build_qsa_metadata_triton( common, *actual_buffers, storage_block_size=2, compress_ratio=compress_ratio, k_work_metadata_buffer=actual_k_work, request_capacity=num_reqs, ) expected_k_work = ( torch.empty_like(actual_k_work) if actual_k_work is not None else None ) expected_buffers = make_buffers() expected = qsa_cache._build_qsa_metadata_torch( common, *expected_buffers, storage_block_size=2, compress_ratio=compress_ratio, k_work_metadata_buffer=expected_k_work, request_capacity=num_reqs, ) for actual_tensor, expected_tensor in zip(actual, expected): torch.testing.assert_close(actual_tensor, expected_tensor) if actual_k_work is not None: torch.testing.assert_close(actual_k_work, expected_k_work) @requires_qsa_kernels def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None: device = torch.device("cuda") num_mapped_tokens = 4096 num_tokens = 4224 query_start_loc = torch.tensor( [0, num_mapped_tokens], dtype=torch.int32, device=device ) common = SimpleNamespace( num_actual_tokens=num_tokens, query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc.cpu(), seq_lens=torch.tensor( [num_mapped_tokens + 32], dtype=torch.int32, device=device ), block_table_tensor=torch.arange(256, dtype=torch.int32, device=device)[None], slot_mapping=torch.tensor( [0] * num_mapped_tokens + [-1] * (num_tokens - num_mapped_tokens), dtype=torch.int64, device=device, ), token_to_req_indices=lambda buffer: buffer.zero_(), ) def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: return ( torch.empty(num_tokens, dtype=torch.int32, device=device), torch.empty(num_tokens, dtype=torch.int64, device=device), torch.empty(num_tokens, dtype=torch.int32, device=device), torch.empty(num_tokens, dtype=torch.int64, device=device), ) max_num_work = (num_tokens + 3) // 4 actual_k_work = torch.empty(max_num_work, 2, dtype=torch.int32, device=device) expected_k_work = torch.empty_like(actual_k_work) actual = qsa_cache.build_qsa_metadata_triton( common, *make_buffers(), storage_block_size=8, compress_ratio=4, k_work_metadata_buffer=actual_k_work, ) expected = qsa_cache._build_qsa_metadata_torch( common, *make_buffers(), storage_block_size=8, compress_ratio=4, k_work_metadata_buffer=expected_k_work, ) for actual_tensor, expected_tensor in zip(actual, expected): torch.testing.assert_close(actual_tensor, expected_tensor) torch.testing.assert_close(actual_k_work, expected_k_work) @requires_qsa_kernels @pytest.mark.parametrize( ("decode_query_len", "num_requests"), [ (1, 2), (2, 2), (3, 2), (4, 2), (4, 33), ], ) def test_qsa_decode_selection_correctness( decode_query_len: int, num_requests: int ) -> None: torch.manual_seed(1) heads, head_dim = 4, 128 rows = num_requests * decode_query_len q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16) page_size, pages_per_request, max_sequence_length = ( (16, 40, 2560) if num_requests > 32 else (4, 20, 320) ) num_pages = num_requests * pages_per_request cache = torch.randn( num_pages, page_size, 1, head_dim, device="cuda", dtype=torch.bfloat16, ) page_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape( num_requests, pages_per_request ) token_to_req = torch.repeat_interleave( torch.arange(num_requests, device="cuda", dtype=torch.int32), decode_query_len, ) sequence_lengths = max_sequence_length - 4 * ( torch.arange(num_requests, device="cuda", dtype=torch.int32) % 8 ) query_positions = torch.cat( [ torch.arange( length - decode_query_len, length, device="cuda", dtype=torch.int32, ) for length in sequence_lengths.tolist() ] ) visible_blocks = torch.minimum( (query_positions + 1) // 4, sequence_lengths.index_select(0, token_to_req.long()) // 4, ) token_topk, compress_ratio = 2048, 4 actual = torch.empty( (rows, token_topk // compress_ratio), device="cuda", dtype=torch.int32 ) qsa_indexer_ops.qsa_select_paged_decode( q, cache, page_table, visible_blocks, token_topk, compress_ratio, decode_query_len, actual, ) expected = _qsa_select_paged_reference( q, cache, page_table, token_to_req, query_positions, sequence_lengths, token_topk, compress_ratio, ) torch.testing.assert_close(actual.sort().values, expected.sort().values) @requires_qsa_kernels def test_qsa_prefill_selection_correctness() -> None: torch.manual_seed(2) query_lens = [3, 33] rows, heads, head_dim = sum(query_lens), 4, 128 q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16) cache = torch.randn(160, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16) page_table = torch.randperm(160, device="cuda", dtype=torch.int32).reshape(2, 80) token_to_req = torch.repeat_interleave( torch.arange(2, device="cuda", dtype=torch.int32), torch.tensor(query_lens, device="cuda"), ) query_start_loc = torch.tensor([0, 3, 36], device="cuda", dtype=torch.int32) sequence_lengths = torch.tensor([5120, 4224], device="cuda", dtype=torch.int32) query_positions = torch.cat( [ torch.arange(length - query_len, length, device="cuda", dtype=torch.int32) for query_len, length in zip( query_lens, sequence_lengths.tolist(), strict=True ) ] ) token_topk, compress_ratio = 2048, 4 visible_blocks = torch.minimum( (query_positions + 1) // compress_ratio, sequence_lengths.index_select(0, token_to_req.long()) // compress_ratio, ) actual = torch.empty( (rows, token_topk // compress_ratio), device="cuda", dtype=torch.int32 ) qsa_indexer_ops.qsa_select_paged_prefill( q, cache, page_table, query_start_loc, visible_blocks, token_topk, compress_ratio, max(query_lens), actual, ) expected = _qsa_select_paged_reference( q, cache, page_table, token_to_req, query_positions, sequence_lengths, token_topk, compress_ratio, ) torch.testing.assert_close(actual.sort().values, expected.sort().values) @requires_qsa_kernels def test_qsa_block_expansion_correctness() -> None: blocks = torch.tensor([[0, -1], [1, 0]], device="cuda", dtype=torch.int32) query_positions = torch.tensor([5, 10], device="cuda") sequence_lengths = torch.tensor([6, 11], device="cuda") token_to_req = torch.tensor([0, 1], device="cuda", dtype=torch.int32) visible_blocks = torch.minimum( (query_positions + 1) // 4, sequence_lengths.index_select(0, token_to_req.long()) // 4, ).to(torch.int32) actual = torch.empty((2, 11), device="cuda", dtype=torch.int32) qsa_indexer_ops.expand_qsa_block_indices( blocks, query_positions, visible_blocks, compress_ratio=4, token_topk=8, out=actual, ) expected = _expand_qsa_indices_reference( blocks, query_positions, sequence_lengths, compress_ratio=4, token_topk=8, ) torch.testing.assert_close(actual, expected) @requires_qsa_kernels @pytest.mark.parametrize( ("num_rows", "num_query_heads", "num_kv_heads", "page_size"), [ # Kernel-visible pages with --block-size 256 and hybrid-cache alignment. pytest.param(1, 24, 2, 1792, id="tp1_split64"), pytest.param(16, 12, 1, 1792, id="tp2_split32"), pytest.param(32, 6, 1, 1024, id="tp4_split8"), pytest.param(257, 6, 1, 1024, id="tp4_split4"), pytest.param(513, 6, 1, 1024, id="tp4_split1"), ], ) def test_qsa_sparse_paged_attention_correctness( num_rows: int, num_query_heads: int, num_kv_heads: int, page_size: int, ) -> None: torch.manual_seed(2) head_dim = 256 num_requests = 2 num_selected_pages = 64 # Keep the newest page outside the synthetic top-k as causal headroom. num_pages_per_request = num_selected_pages + 1 num_cache_blocks = num_requests * num_pages_per_request indexer_budget = 2048 indexer_compress_ratio = 4 selection_width = indexer_budget + indexer_compress_ratio - 1 q = torch.randn( num_rows, num_query_heads, head_dim, device="cuda", dtype=torch.bfloat16 ) kv_cache = torch.randn( num_cache_blocks, page_size, num_kv_heads, 2 * head_dim, device="cuda", dtype=torch.bfloat16, ) k_cache, v_cache = kv_cache.split(head_dim, dim=-1) block_table = ( torch.randperm(num_cache_blocks, device="cuda") .reshape(num_requests, num_pages_per_request) .to(torch.int32) ) rows_per_request = math.ceil(num_rows / num_requests) row_indices = torch.arange(num_rows, device="cuda", dtype=torch.int32) token_to_req = row_indices // rows_per_request request_row_counts = torch.tensor( [rows_per_request, num_rows - rows_per_request], device="cuda", dtype=torch.int32, ) context_length = num_pages_per_request * page_size - 1 block_topk = indexer_budget // indexer_compress_ratio compressed_blocks_per_page = page_size // indexer_compress_ratio selection = torch.arange(block_topk, device="cuda") selected_pages = selection % num_selected_pages selected_offsets = selection // num_selected_pages row_shifts = 2 * row_indices.unsqueeze(1) # Eight blocks per page; adjacent rows overlap by six of those eight. selected_offsets = (selected_offsets + row_shifts) % compressed_blocks_per_page block_indices = (selected_pages * compressed_blocks_per_page + selected_offsets).to( torch.int32 ) rows_within_request = row_indices % rows_per_request query_positions = ( context_length - request_row_counts[token_to_req.long()] + rows_within_request ).to(torch.int64) sequence_lengths = torch.full( (num_requests,), context_length, device="cuda", dtype=torch.int32 ) visible_blocks = torch.minimum( (query_positions + 1) // indexer_compress_ratio, sequence_lengths.index_select(0, token_to_req.long()) // indexer_compress_ratio, ).to(torch.int32) logical_indices = torch.empty( (num_rows, selection_width), device="cuda", dtype=torch.int32 ) qsa_indexer_ops.expand_qsa_block_indices( block_indices, query_positions, visible_blocks, indexer_compress_ratio, indexer_budget, logical_indices, ) assert logical_indices.shape == (num_rows, selection_width) scale = q.shape[-1] ** -0.5 actual = qsa_ops.qsa_sparse_paged_attention( q, k_cache, v_cache, logical_indices, block_table, token_to_req, ) expected = _qsa_sparse_paged_attention_reference( q, k_cache, v_cache, logical_indices, block_table, token_to_req, scale, ) torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) @requires_qsa_kernels @pytest.mark.parametrize("decode_query_len", [1, 2, 3, 4]) def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int) -> None: query_lens = [decode_query_len, decode_query_len, 33] rows, heads, head_dim = sum(query_lens), 4, 128 token_topk, compress_ratio = 2048, 4 torch.manual_seed(13) q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16) cache = torch.randn(120, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16) page_table = torch.arange(120, device="cuda", dtype=torch.int32).view(3, 40) token_to_req = torch.repeat_interleave( torch.arange(3, device="cuda", dtype=torch.int32), torch.tensor(query_lens, device="cuda"), ) query_start_loc = torch.tensor( [0, decode_query_len, 2 * decode_query_len, rows], device="cuda", dtype=torch.int32, ) sequence_lengths = torch.full((3,), 2560, device="cuda", dtype=torch.int32) query_positions = torch.cat( [ torch.arange(2560 - query_len, 2560, device="cuda") for query_len in query_lens ] ) block_indices = torch.empty( rows, token_topk // compress_ratio, device="cuda", dtype=torch.int32, ) visible_blocks = torch.minimum( (query_positions + 1) // compress_ratio, sequence_lengths.index_select(0, token_to_req.long()) // compress_ratio, ).to(torch.int32) num_decode_tokens = 2 * decode_query_len decode_slice = slice(0, num_decode_tokens) qsa_indexer_ops.qsa_select_paged_decode( q[decode_slice], cache, page_table[:2], visible_blocks[decode_slice], token_topk, compress_ratio, decode_query_len, block_indices[decode_slice], ) prefill_slice = slice(num_decode_tokens, rows) qsa_indexer_ops.qsa_select_paged_prefill( q[prefill_slice], cache, page_table[2:], query_start_loc[2:], visible_blocks[prefill_slice], token_topk, compress_ratio, query_lens[-1], block_indices[prefill_slice], ) actual = torch.empty( (rows, token_topk + compress_ratio - 1), device="cuda", dtype=torch.int32 ) qsa_indexer_ops.expand_qsa_block_indices( block_indices, query_positions, visible_blocks, compress_ratio, token_topk, actual, ) expected_blocks = _qsa_select_paged_reference( q, cache, page_table, token_to_req, query_positions, sequence_lengths, token_topk, compress_ratio, ) expected = _expand_qsa_indices_reference( expected_blocks, query_positions, sequence_lengths.index_select(0, token_to_req.long()), compress_ratio, token_topk, ) torch.testing.assert_close(actual.sort().values, expected.sort().values) @requires_qsa_kernels def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) -> None: q = torch.zeros(2, 4, 8, device="cuda", dtype=torch.bfloat16) cache = torch.zeros(1, 16, 1, 8, device="cuda", dtype=torch.bfloat16) page_table = torch.zeros(1, 1, device="cuda", dtype=torch.int32) query_positions = torch.tensor([1, 2], device="cuda", dtype=torch.int32) visible_blocks = torch.zeros(2, device="cuda", dtype=torch.int32) block_indices = torch.empty((2, 512), device="cuda", dtype=torch.int32) qsa_indexer_ops.qsa_select_paged_prefill( q, cache, page_table, torch.tensor([0, 2], device="cuda", dtype=torch.int32), visible_blocks, token_topk=2048, compress_ratio=4, max_query_len=2, block_indices=block_indices, ) selected = torch.empty((2, 2051), device="cuda", dtype=torch.int32) qsa_indexer_ops.expand_qsa_block_indices( block_indices, query_positions, visible_blocks, compress_ratio=4, token_topk=2048, out=selected, ) assert selected[0, :2].tolist() == [0, 1] assert selected[1, :3].tolist() == [0, 1, 2] assert torch.all(selected[0, 2:] == -1) assert torch.all(selected[1, 3:] == -1) @requires_qsa_kernels def test_qsa_streaming_compression_and_compressor_state_store_match_reference() -> None: head_dim = 8 current_pairs = [ *((0, position) for position in range(2, 9)), *((1, position) for position in range(5, 11)), ] def key_row(request: int, position: int) -> torch.Tensor: return ( torch.arange(head_dim, dtype=torch.float32) + request * 1000 + position * 10 ) def position_row(request: int, position: int) -> torch.Tensor: return torch.tensor( [ request * 1000 + position, request * 1000 + position + 100, request * 1000 + position + 200, ], dtype=torch.int64, ) raw_keys = ( torch.stack([key_row(request, position) for request, position in current_pairs]) .unsqueeze(1) .to(device="cuda", dtype=torch.bfloat16) ) raw_positions = ( torch.stack( [position_row(request, position) for request, position in current_pairs] ) .unsqueeze(1) .to(device="cuda") ) token_to_req = torch.tensor( [request for request, _ in current_pairs], dtype=torch.int32, device="cuda", ) logical_positions = torch.tensor( [position for _, position in current_pairs], dtype=torch.int64, device="cuda", ) query_start_loc = torch.tensor([0, 7, 13], dtype=torch.int32, device="cuda") compressor_state_block_table = torch.tensor( [[1], [0]], dtype=torch.int32, device="cuda" ) compressor_state_cache = torch.zeros( 2, 4, 1, head_dim, dtype=torch.bfloat16, device="cuda" ) rope_cache = torch.zeros(2, 4, 1, 3, dtype=torch.int64, device="cuda") for request, position, block in ((0, 0, 1), (0, 1, 1), (1, 4, 0)): compressor_state_cache[block, position % 4, 0] = key_row(request, position).to( device="cuda", dtype=torch.bfloat16 ) rope_cache[block, position % 4, 0] = position_row(request, position).to("cuda") compressed_slots = torch.full( (len(current_pairs),), -1, dtype=torch.int64, device="cuda" ) valid_rows = torch.tensor([1, 5, 9], dtype=torch.int64, device="cuda") compressed_slots[valid_rows] = torch.arange(3, device="cuda") pooled, first_positions = qsa_ops.qsa_compress_groups_with_ratio( raw_keys, raw_positions, compressor_state_cache, compressor_state_block_table, token_to_req, query_start_loc, logical_positions, compressed_slots, compress_ratio=4, rope_cache=rope_cache, ) pooled_without_rope, scalar_first_positions = ( qsa_ops.qsa_compress_groups_with_ratio( raw_keys, raw_positions, compressor_state_cache, compressor_state_block_table, token_to_req, query_start_loc, logical_positions, compressed_slots, compress_ratio=4, ) ) groups = [ [(0, position) for position in range(0, 4)], [(0, position) for position in range(4, 8)], [(1, position) for position in range(4, 8)], ] expected_pooled = ( torch.stack( [ torch.stack([key_row(*pair) for pair in group]).mean(dim=0) for group in groups ] ) .unsqueeze(1) .to(device="cuda", dtype=torch.bfloat16) ) expected_positions = torch.stack( [position_row(0, 0), position_row(0, 4), position_row(1, 4)] ).to("cuda") expected_scalar_positions = torch.tensor( [[0, 0, 0], [4, 4, 4], [4, 4, 4]], dtype=torch.int64, device="cuda", ) torch.testing.assert_close(pooled[valid_rows], expected_pooled) torch.testing.assert_close(pooled_without_rope[valid_rows], expected_pooled) torch.testing.assert_close(first_positions[valid_rows], expected_positions) torch.testing.assert_close( scalar_first_positions[valid_rows], expected_scalar_positions ) compressor_state_slots = torch.tensor( [-1, -1, -1, 5, 6, 7, 4, -1, -1, 3, 0, 1, 2], dtype=torch.int64, device="cuda", ) qsa_ops.qsa_store_cache_rows( compressor_state_cache, compressor_state_slots, raw_keys ) qsa_ops.qsa_store_cache_rows(rope_cache, compressor_state_slots, raw_positions) for request, positions, block in ((0, range(5, 9), 1), (1, range(7, 11), 0)): for position in positions: torch.testing.assert_close( compressor_state_cache[block, position % 4, 0], key_row(request, position).to(device="cuda", dtype=torch.bfloat16), ) torch.testing.assert_close( rope_cache[block, position % 4, 0], position_row(request, position).to("cuda"), )