diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 276eb31b1759..163f3a4f9c1e 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/tests/models/qwen4_exp/test_qsa_reference.py @@ -14,6 +14,7 @@ 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 @@ -67,8 +68,13 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection( lambda *args, **kwargs: updates.append((args, kwargs)), ) monkeypatch.setattr( - qsa_ops, - "qsa_select_paged_tokens", + 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)), ) @@ -158,7 +164,7 @@ def _expand_qsa_indices_reference( return result.gather(1, torch.argsort(sort_key, dim=1, stable=True)).to(torch.int32) -def _qsa_select_paged_tokens_reference( +def _qsa_select_paged_reference( q: torch.Tensor, k_cache: torch.Tensor, page_table: torch.Tensor, @@ -181,19 +187,12 @@ def _qsa_select_paged_tokens_reference( visible_blocks, ) starts = torch.zeros_like(visible_blocks) - blocks = _qsa_relative_topk_reference( + return _qsa_relative_topk_reference( logits, starts, visible_blocks, token_topk // compress_ratio, ) - return _expand_qsa_indices_reference( - blocks, - query_positions, - row_sequence_lengths, - compress_ratio, - token_topk, - ) def _qsa_sparse_paged_attention_reference( @@ -231,16 +230,20 @@ 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), @@ -270,6 +273,7 @@ def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None: -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 @@ -277,18 +281,22 @@ 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), @@ -402,11 +410,13 @@ 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) @@ -416,6 +426,8 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None: ) 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), @@ -427,6 +439,7 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None: 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]] @@ -473,10 +486,11 @@ def test_qsa_triton_metadata_matches_pytorch( token_to_req_indices=lambda buffer: buffer.copy_(token_to_req), ) - def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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), ) @@ -543,10 +557,11 @@ def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None: token_to_req_indices=lambda buffer: buffer.zero_(), ) - def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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), ) @@ -575,70 +590,163 @@ def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @requires_qsa_kernels @pytest.mark.parametrize( - "num_rows", + ("decode_query_len", "num_requests"), [ - pytest.param(3, id="one_tile_per_program"), - pytest.param(33, id="looped_tiles"), + (1, 2), + (2, 2), + (3, 2), + (4, 2), + (4, 33), ], ) -def test_qsa_mqa_paged_matches_test_reference(num_rows: int) -> None: +def test_qsa_decode_selection_correctness( + decode_query_len: int, num_requests: int +) -> None: torch.manual_seed(1) - head_dim = 128 - q = torch.randn(num_rows, 4, head_dim, device="cuda", dtype=torch.bfloat16) - cache = torch.randn(40, 4, 1, head_dim, device="cuda", dtype=torch.bfloat16) - page_table = torch.randperm(40, device="cuda", dtype=torch.int32).reshape(2, 20) - request_rows = [(num_rows + 1) // 2, num_rows // 2] + 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(2, device="cuda", dtype=torch.int32), - torch.tensor(request_rows, device="cuda"), + torch.arange(num_requests, device="cuda", dtype=torch.int32), + decode_query_len, ) - sequence_length_values = [320, 264] - sequence_lengths = torch.tensor( - sequence_length_values, device="cuda", dtype=torch.int32 + sequence_lengths = max_sequence_length - 4 * ( + torch.arange(num_requests, device="cuda", dtype=torch.int32) % 8 ) query_positions = torch.cat( [ - torch.arange(length - rows, length, device="cuda", dtype=torch.int32) - for rows, length in zip(request_rows, sequence_length_values, strict=True) + torch.arange( + length - decode_query_len, + length, + device="cuda", + dtype=torch.int32, + ) + for length in sequence_lengths.tolist() ] ) - compress_ratio = 4 - visible_lengths = (query_positions + 1) // compress_ratio + visible_blocks = torch.minimum( + (query_positions + 1) // 4, + sequence_lengths.index_select(0, token_to_req.long()) // 4, + ) - actual, actual_visible_blocks = qsa_ops.qsa_mqa_paged( + 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, - compress_ratio=compress_ratio, + 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 ) - expected = _qsa_mqa_paged_reference( - q, cache, page_table, token_to_req, visible_lengths + 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_visible_blocks, visible_lengths) - # Top-k is bounded by visible_blocks; later columns are intentionally unwritten. - columns = torch.arange(actual.shape[1], device=actual.device) - visible = columns[None, :] < actual_visible_blocks[:, None] - torch.testing.assert_close(actual[visible], expected[visible], rtol=1e-3, atol=1e-3) + torch.testing.assert_close(actual.sort().values, expected.sort().values) @requires_qsa_kernels -def test_qsa_block_expansion_matches_test_reference() -> None: +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 = qsa_ops.expand_qsa_block_indices_cuda( + actual = torch.empty((2, 11), device="cuda", dtype=torch.int32) + qsa_indexer_ops.expand_qsa_block_indices( blocks, query_positions, - sequence_lengths, - token_to_req, + visible_blocks, compress_ratio=4, token_topk=8, + out=actual, ) expected = _expand_qsa_indices_reference( blocks, @@ -663,7 +771,7 @@ def test_qsa_block_expansion_matches_test_reference() -> None: pytest.param(513, 6, 1, 1024, id="tp4_split1"), ], ) -def test_qsa_sparse_paged_attention_matches_test_reference( +def test_qsa_sparse_paged_attention_correctness( num_rows: int, num_query_heads: int, num_kv_heads: int, @@ -724,13 +832,20 @@ def test_qsa_sparse_paged_attention_matches_test_reference( sequence_lengths = torch.full( (num_requests,), context_length, device="cuda", dtype=torch.int32 ) - logical_indices = qsa_ops.expand_qsa_block_indices_cuda( + 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, - sequence_lengths, - token_to_req, + visible_blocks, indexer_compress_ratio, indexer_budget, + logical_indices, ) assert logical_indices.shape == (num_rows, selection_width) scale = q.shape[-1] ** -0.5 @@ -757,40 +872,78 @@ def test_qsa_sparse_paged_attention_matches_test_reference( @requires_qsa_kernels -def test_qsa_selection_chunks_workspace_and_matches_test_reference( - monkeypatch: pytest.MonkeyPatch, - workspace_init, -) -> None: - rows, keys, heads, head_dim = 65, 640, 4, 16 +@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(3) + torch.manual_seed(13) q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16) - cache = torch.randn(40, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16) - page_table = torch.randperm(40, device="cuda", dtype=torch.int32).unsqueeze(0) - token_to_req = torch.zeros(rows, device="cuda", dtype=torch.int32) - query_positions = torch.full((rows,), 2559, device="cuda", dtype=torch.int32) - sequence_lengths = torch.tensor([2560], device="cuda", dtype=torch.int32) - monkeypatch.setattr(qsa_ops, "_LOGITS_WORKSPACE_BYTES", 32 * keys * 4) - original_score = qsa_ops.qsa_mqa_paged - scored_row_counts = [] - - def record_score(query: torch.Tensor, *args, **kwargs): - scored_row_counts.append(query.shape[0]) - return original_score(query, *args, **kwargs) - - monkeypatch.setattr(qsa_ops, "qsa_mqa_paged", record_score) - - actual = qsa_ops.qsa_select_paged_tokens( - q, + 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, - token_to_req, - query_positions, - sequence_lengths, + 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 = _qsa_select_paged_tokens_reference( + expected_blocks = _qsa_select_paged_reference( q, cache, page_table, @@ -800,9 +953,15 @@ def record_score(query: torch.Tensor, *args, **kwargs): 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) - assert scored_row_counts == [32, 32, 1] @requires_qsa_kernels @@ -810,19 +969,29 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) -> 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) - token_to_req = torch.zeros(2, device="cuda", dtype=torch.int32) query_positions = torch.tensor([1, 2], device="cuda", dtype=torch.int32) - sequence_lengths = torch.tensor([3], device="cuda", dtype=torch.int32) + visible_blocks = torch.zeros(2, device="cuda", dtype=torch.int32) - selected = qsa_ops.qsa_select_paged_tokens( + block_indices = torch.empty((2, 512), device="cuda", dtype=torch.int32) + qsa_indexer_ops.qsa_select_paged_prefill( q, cache, page_table, - token_to_req, - query_positions, - sequence_lengths, + 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] diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index ca816cdb44e8..ad11976650b0 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -37,6 +37,9 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import ( kimi_k3_triton_warmup, ) +from vllm.model_executor.warmup.qwen4_exp_qsa_warmup import ( + qwen4_exp_qsa_triton_warmup, +) from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup from vllm.model_executor.warmup.replayssm_warmup import ( replayssm_autotune_warmup, @@ -157,6 +160,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): if worker.vllm_config.kernel_config.enable_jit_warmup: kimi_k3_triton_warmup(worker) fa4_cutedsl_warmup(worker) + qwen4_exp_qsa_triton_warmup(worker) if current_platform.has_device_capability(90): _warmup_ll_bf16_router_gemm(worker.get_model()) diff --git a/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py new file mode 100644 index 000000000000..7f2c5a20b6d1 --- /dev/null +++ b/vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Connect loaded Qwen4Exp QSA modules to their kernel-owned warmup.""" + +import sys +from typing import TYPE_CHECKING, cast + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2 + from vllm.v1.worker.gpu_worker import Worker + +logger = init_logger(__name__) + + +def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None: + """Warm every reachable QSA decode-query-length specialization.""" + + qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa") + if qsa_module is None: + return + indexer = next( + ( + layer + for layer in worker.get_model().modules() + if isinstance(layer, qsa_module.QSAIndexer) + ), + None, + ) + if indexer is None: + return + + runner = worker.model_runner + prefix = indexer.compressed_key_cache.prefix + group_id = next( + i + for i, group in enumerate(runner.kv_cache_config.kv_cache_groups) + if prefix in group.layer_names + ) + if worker.use_v2_model_runner: + runner_v2 = cast("GPUModelRunnerV2", runner) + block_table = runner_v2.block_tables.input_block_tables[group_id] + max_decode_query_len = runner_v2.decode_query_len + else: + block_table = runner.input_batch.block_table[group_id].get_device_tensor( + runner.max_num_reqs + ) + max_decode_query_len = runner.uniform_decode_query_len + + from vllm.models.qwen4_exp.nvidia.ops.qsa_indexer import ( + warmup_qsa_mqa_paged_decode, + ) + + k_cache = indexer.compressed_key_cache.kv_cache + assert k_cache.numel() + profiles = warmup_qsa_mqa_paged_decode( + k_cache, + block_table, + num_heads=indexer.index_n_heads, + head_dim=indexer.index_head_dim, + max_decode_query_len=max_decode_query_len, + max_num_reqs=runner.max_num_reqs, + max_num_batched_tokens=runner.max_num_tokens, + ) + logger.info("Warmed up Qwen4Exp QSA decode kernels: %s.", profiles) diff --git a/vllm/models/qwen4_exp/common/qsa_cache.py b/vllm/models/qwen4_exp/common/qsa_cache.py index cffb961def14..539b7c112130 100644 --- a/vllm/models/qwen4_exp/common/qsa_cache.py +++ b/vllm/models/qwen4_exp/common/qsa_cache.py @@ -11,8 +11,6 @@ shared by the generic cache-layout planner. """ -from __future__ import annotations - import math from dataclasses import dataclass from functools import cache @@ -34,7 +32,10 @@ AttentionMetadataBuilder, CommonAttentionMetadata, ) -from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.attention.backends.utils import ( + PAD_SLOT_ID, + split_decodes_and_prefills, +) from vllm.v1.kv_cache_interface import ( AttentionSpec, CircularBufferSpec, @@ -211,6 +212,7 @@ def _build_qsa_metadata_kernel( block_table_ptr, token_to_req_ptr, logical_positions_ptr, + visible_blocks_ptr, slot_mapping_ptr, k_work_metadata_ptr, block_table_stride_0: tl.constexpr, @@ -266,6 +268,14 @@ def _build_qsa_metadata_kernel( logical_position, mask=store_mask, ) + visible_blocks = tl.maximum( + 0, + tl.minimum( + (logical_position + 1) // compress_ratio, + seq_len // compress_ratio, + ), + ).to(tl.int32) + tl.store(visible_blocks_ptr + token_idx, visible_blocks, mask=store_mask) # circular_buffer_size is constexpr, so each builder instance compiles out # the other QSA cache owner's slot-mapping rule. @@ -380,6 +390,7 @@ def build_qsa_metadata_triton( common_attn_metadata: CommonAttentionMetadata, token_to_req_buffer: torch.Tensor, logical_positions_buffer: torch.Tensor, + visible_blocks_buffer: torch.Tensor, slot_mapping_buffer: torch.Tensor, *, storage_block_size: int, @@ -387,12 +398,13 @@ def build_qsa_metadata_triton( circular_buffer_size: int = 0, k_work_metadata_buffer: torch.Tensor | None = None, request_capacity: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Build QSA side-cache and optional pre-indexer work metadata.""" num_tokens = common_attn_metadata.num_actual_tokens num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) token_to_req = token_to_req_buffer[:num_tokens] logical_positions = logical_positions_buffer[:num_tokens] + visible_blocks = visible_blocks_buffer[:num_tokens] slot_mapping = slot_mapping_buffer[:num_tokens] num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1 assert num_reqs > 0 @@ -409,7 +421,7 @@ def build_qsa_metadata_triton( max_num_work = 0 if num_tokens == 0 and k_work_metadata_buffer is None: - return token_to_req, logical_positions, slot_mapping + return token_to_req, logical_positions, visible_blocks, slot_mapping block_table = common_attn_metadata.block_table_tensor num_search_steps = int(math.ceil(math.log2(num_reqs))) @@ -426,6 +438,7 @@ def build_qsa_metadata_triton( block_table, token_to_req, logical_positions, + visible_blocks, slot_mapping, k_work_metadata_buffer, block_table.stride(0), @@ -448,13 +461,14 @@ def build_qsa_metadata_triton( ) if circular_buffer_size == 0 and compress_ratio == 1: slot_mapping = common_attn_metadata.slot_mapping[:num_tokens] - return token_to_req, logical_positions, slot_mapping + return token_to_req, logical_positions, visible_blocks, slot_mapping def _build_qsa_metadata_torch( common_attn_metadata: CommonAttentionMetadata, token_to_req_buffer: torch.Tensor, logical_positions_buffer: torch.Tensor, + visible_blocks_buffer: torch.Tensor, slot_mapping_buffer: torch.Tensor, *, storage_block_size: int, @@ -462,11 +476,12 @@ def _build_qsa_metadata_torch( circular_buffer_size: int = 0, k_work_metadata_buffer: torch.Tensor | None = None, request_capacity: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: del request_capacity num_tokens = common_attn_metadata.num_actual_tokens num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) logical_positions = logical_positions_buffer[:num_tokens] + visible_blocks = visible_blocks_buffer[:num_tokens] token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[ :num_tokens @@ -481,6 +496,15 @@ def _build_qsa_metadata_torch( ) if num_mapped_tokens < num_tokens: logical_positions[num_mapped_tokens:].fill_(-1) + row_seq_lens = common_attn_metadata.seq_lens.index_select(0, token_to_req.long()) + visible_blocks.copy_( + torch.minimum( + (logical_positions + 1) // compress_ratio, + row_seq_lens // compress_ratio, + ) + .clamp_min_(0) + .to(torch.int32) + ) if circular_buffer_size > 0: slot_mapping = circular_qsa_slot_mapping( common_attn_metadata.block_table_tensor, @@ -539,7 +563,7 @@ def _build_qsa_metadata_torch( k_work_metadata_buffer[:, 1].copy_( torch.where(active, work_in_request, -1).to(torch.int32) ) - return token_to_req, logical_positions, slot_mapping + return token_to_req, logical_positions, visible_blocks, slot_mapping # Resolve the fallback outside the per-step metadata hot path. @@ -558,8 +582,15 @@ class QSAForwardMetadata(AttentionMetadata): query_start_loc: torch.Tensor token_to_req: torch.Tensor logical_positions: torch.Tensor + visible_blocks: torch.Tensor k_work_metadata: torch.Tensor num_actual_tokens: int + num_decodes: int + num_decode_tokens: int + num_prefills: int + num_prefill_tokens: int + max_query_len: int + decode_query_len: int storage_block_size: int compress_ratio: int @@ -577,6 +608,8 @@ def __init__( device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) + assert self.reorder_batch_threshold is not None self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec) if isinstance(kv_cache_spec, MLAAttentionSpec): compress_ratio = kv_cache_spec.tokens_per_state @@ -597,6 +630,9 @@ def __init__( self.logical_positions_buffer = torch.empty( max_tokens, dtype=torch.int64, device=device ) + self.visible_blocks_buffer = torch.empty( + max_tokens, dtype=torch.int32, device=device + ) max_requests = vllm_config.scheduler_config.max_num_seqs self.request_capacity = max_requests if not self.is_circular_buffer and self.compress_ratio != 1: @@ -619,6 +655,26 @@ def build( ) -> QSAForwardMetadata: del common_prefix_len, fast_build num_tokens = common_attn_metadata.num_actual_tokens + decode_threshold = self.reorder_batch_threshold + assert decode_threshold is not None + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( + split_decodes_and_prefills( + common_attn_metadata, + decode_threshold=decode_threshold, + require_uniform=True, + ) + ) + assert num_decodes + num_prefills == common_attn_metadata.num_reqs + assert num_decode_tokens + num_prefill_tokens == num_tokens + decode_query_len = 0 + if num_decodes > 0: + query_lens_cpu = torch.diff( + common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] + ) + nonzero_query_lens = query_lens_cpu[query_lens_cpu > 0] + if nonzero_query_lens.numel() > 0: + decode_query_len = int(nonzero_query_lens[0].item()) + assert torch.all(nonzero_query_lens == decode_query_len) build_k_work = not self.is_circular_buffer and self.compress_ratio != 1 k_work_metadata = self.k_work_metadata_buffer request_capacity = None @@ -629,18 +685,21 @@ def build( num_tokens + (self.compress_ratio - 1) * num_requests ) // self.compress_ratio k_work_metadata = self.k_work_metadata_buffer[:max_num_work] - token_to_req, logical_positions, slot_mapping = build_qsa_metadata( - common_attn_metadata, - self.token_to_req_buffer, - self.logical_positions_buffer, - self.slot_mapping_buffer, - storage_block_size=self.storage_block_size, - compress_ratio=self.compress_ratio, - circular_buffer_size=( - self.kv_cache_spec.block_size if self.is_circular_buffer else 0 - ), - k_work_metadata_buffer=k_work_metadata if build_k_work else None, - request_capacity=request_capacity, + token_to_req, logical_positions, visible_blocks, slot_mapping = ( + build_qsa_metadata( + common_attn_metadata, + self.token_to_req_buffer, + self.logical_positions_buffer, + self.visible_blocks_buffer, + self.slot_mapping_buffer, + storage_block_size=self.storage_block_size, + compress_ratio=self.compress_ratio, + circular_buffer_size=( + self.kv_cache_spec.block_size if self.is_circular_buffer else 0 + ), + k_work_metadata_buffer=k_work_metadata if build_k_work else None, + request_capacity=request_capacity, + ) ) return QSAForwardMetadata( block_table=common_attn_metadata.block_table_tensor, @@ -649,8 +708,15 @@ def build( query_start_loc=common_attn_metadata.query_start_loc, token_to_req=token_to_req, logical_positions=logical_positions, + visible_blocks=visible_blocks, k_work_metadata=k_work_metadata, num_actual_tokens=num_tokens, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + num_prefills=num_prefills, + num_prefill_tokens=num_prefill_tokens, + max_query_len=common_attn_metadata.max_query_len, + decode_query_len=decode_query_len, storage_block_size=self.storage_block_size, compress_ratio=self.compress_ratio, ) diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index 1e6c308718f1..ffaac6434c60 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py @@ -2,8 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Qwen4Exp weight-free QSA indexer.""" -from __future__ import annotations - from typing import cast import torch @@ -185,6 +183,22 @@ def _metadata( ) if raw.num_actual_tokens != compressed.num_actual_tokens: raise RuntimeError("QSA side-cache metadata token counts disagree") + raw_split = ( + raw.num_decodes, + raw.num_decode_tokens, + raw.num_prefills, + raw.num_prefill_tokens, + raw.decode_query_len, + ) + compressed_split = ( + compressed.num_decodes, + compressed.num_decode_tokens, + compressed.num_prefills, + compressed.num_prefill_tokens, + compressed.decode_query_len, + ) + if raw_split != compressed_split: + raise RuntimeError("QSA side-cache metadata batch splits disagree") if not raw.logical_positions.is_cuda and ( not torch.equal(raw.logical_positions, compressed.logical_positions) ): @@ -215,10 +229,11 @@ def forward( return out return result - from .ops.qsa import ( - qsa_compress_groups_with_ratio, - qsa_select_paged_tokens, - qsa_store_cache_rows, + from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows + from .ops.qsa_indexer import ( + expand_qsa_block_indices, + qsa_select_paged_decode, + qsa_select_paged_prefill, ) raw_metadata, compressed_metadata = metadata @@ -339,18 +354,67 @@ def forward( raise RuntimeError("QSA top-k reuse requires an output buffer") return out - # Score compressed keys, select blocks, then expand them to token indices. - return qsa_select_paged_tokens( - q, - compressed_key_cache, - compressed_metadata.block_table, - compressed_metadata.token_to_req, - compressed_metadata.logical_positions, - compressed_metadata.seq_lens, - self.token_topk, + if out is None: + out = torch.empty( + num_tokens, + self.output_width, + dtype=torch.int32, + device=q.device, + ) + elif out.shape != (num_tokens, self.output_width): + raise ValueError("QSA selection output has an invalid shape") + + num_decode_tokens = compressed_metadata.num_decode_tokens + decode_query_len = compressed_metadata.decode_query_len + visible_blocks = compressed_metadata.visible_blocks[:num_tokens] + block_indices = torch.empty( + num_tokens, + self.token_topk // self.compress_ratio, + dtype=torch.int32, + device=q.device, + ) + + # Decode requests occupy the leading rows and share one query length. + if num_decode_tokens: + num_decodes = compressed_metadata.num_decodes + if num_decodes * decode_query_len != num_decode_tokens: + raise ValueError("QSA decode rows must form a uniform request batch") + decode_slice = slice(0, num_decode_tokens) + qsa_select_paged_decode( + q[decode_slice], + compressed_key_cache, + compressed_metadata.block_table[:num_decodes], + visible_blocks[decode_slice], + self.token_topk, + self.compress_ratio, + decode_query_len, + block_indices[decode_slice], + ) + + # Prefill requests follow the leading decode rows in the reordered batch. + if num_decode_tokens < num_tokens: + num_decodes = compressed_metadata.num_decodes + prefill_slice = slice(num_decode_tokens, num_tokens) + qsa_select_paged_prefill( + q[prefill_slice], + compressed_key_cache, + compressed_metadata.block_table[num_decodes:], + compressed_metadata.query_start_loc[num_decodes:], + visible_blocks[prefill_slice], + self.token_topk, + self.compress_ratio, + compressed_metadata.max_query_len, + block_indices[prefill_slice], + ) + expand_qsa_block_indices( + block_indices, + compressed_metadata.logical_positions[:num_tokens], + visible_blocks, self.compress_ratio, + self.token_topk, out, ) + return out __all__ = ["QSAIndexer", "apply_qsa_rope"] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index d6aa06f3a3d5..312160210cab 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -1,193 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Triton kernels for the Qwen4Exp weight-free QSA path.""" +"""Triton kernels for Qwen4Exp QSA sparse attention and cache updates.""" from __future__ import annotations -import math - import torch -from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton -_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024 -_TOPK_WORKSPACE_BYTES = 1024 * 1024 - - -@triton.jit -def _qsa_mqa_paged_kernel( - q_ptr, - k_cache_ptr, - page_table_ptr, - token_to_req_ptr, - query_positions_ptr, - sequence_lengths_ptr, - visible_blocks_ptr, - logits_ptr, - stride_q_row, - stride_q_head, - stride_q_dim, - stride_cache_block, - stride_cache_token, - stride_cache_dim, - stride_table_req, - stride_table_page, - stride_logits_row, - num_rows, - num_columns, - num_pages, - num_requests, - score_divisor, - PAGE_SIZE: tl.constexpr, - PAGE_TABLE_WIDTH: tl.constexpr, - NUM_HEADS: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_D: tl.constexpr, - TILES_PER_PROG: tl.constexpr, - STAGES: tl.constexpr, - MAX_N: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, -) -> None: - row = tl.program_id(0) - dims = tl.arange(0, BLOCK_D) - heads = tl.arange(0, MAX_N) - request = tl.load(token_to_req_ptr + row) - safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) - query_position = tl.load(query_positions_ptr + row) - sequence_length = tl.load( - sequence_lengths_ptr + safe_request, - mask=(request >= 0) & (request < num_requests), - other=0, - ) - visible = tl.minimum( - (query_position + 1) // COMPRESS_RATIO, - sequence_length // COMPRESS_RATIO, - ) - if tl.program_id(1) == 0: - tl.store(visible_blocks_ptr + row, visible) - tile_start = tl.program_id(1) * TILES_PER_PROG - # Top-k is bounded by visible_blocks, so columns beyond it need no value. - if tile_start * BLOCK_N >= visible: - return - tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible, BLOCK_N)) - tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N)) - - # Pad the small head axis to a tensor-core-compatible N dimension. - query = tl.load( - q_ptr - + row * stride_q_row - + heads[None, :] * stride_q_head - + dims[:, None] * stride_q_dim, - mask=(heads[None, :] < NUM_HEADS) & (dims[:, None] < HEAD_DIM), - other=0.0, - ) - column_offsets = tl.arange(0, BLOCK_N) - for tile in tl.range(tile_start, tile_end, num_stages=STAGES): - columns = tile * BLOCK_N + column_offsets - live = columns < visible - logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) - page_offset = columns % PAGE_SIZE - physical_page = tl.load( - page_table_ptr - + safe_request * stride_table_req - + logical_page * stride_table_page, - mask=live, - other=-1, - ) - page_valid = live & (physical_page >= 0) & (physical_page < num_pages) - # physical_page * block stride can overflow int32 for large caches. - safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64) - keys = tl.load( - k_cache_ptr - + safe_physical_page[:, None] * stride_cache_block - + page_offset[:, None] * stride_cache_token - + dims[None, :] * stride_cache_dim, - mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM), - other=0.0, - eviction_policy="evict_first", - ) - scores = tl.dot(keys, query, out_dtype=tl.float32) - scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(scores, 0.0), 0.0) - score = tl.sum(scores, axis=1) / score_divisor - tl.store( - logits_ptr + row * stride_logits_row + columns, - tl.where(page_valid, score, -float("inf")), - mask=live & (columns < num_columns), - ) - - -@triton.jit -def _expand_qsa_indices_kernel( - block_indices_ptr, - query_positions_ptr, - sequence_lengths_ptr, - token_to_req_ptr, - output_ptr, - stride_blocks_row, - stride_blocks_column, - stride_output_row, - stride_output_column, - rows, - num_requests, - BLOCK_TOPK: tl.constexpr, - COMPRESS_RATIO: tl.constexpr, - TOKEN_TOPK: tl.constexpr, - OUTPUT_WIDTH: tl.constexpr, - COLUMN_BLOCK: tl.constexpr, -) -> None: - row = tl.program_id(0) - columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK) - query_position = tl.load(query_positions_ptr + row) - request = tl.load(token_to_req_ptr + row) - safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1) - sequence_length = tl.load( - sequence_lengths_ptr + safe_request, - mask=(request >= 0) & (request < num_requests), - other=0, - ) - complete_blocks = tl.minimum( - tl.minimum( - (query_position + 1) // COMPRESS_RATIO, - sequence_length // COMPRESS_RATIO, - ), - BLOCK_TOPK, - ) - expanded_count = complete_blocks * COMPRESS_RATIO - tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO - tail_count = (query_position + 1) - tail_start - - is_expanded = columns < expanded_count - block_rank = columns // COMPRESS_RATIO - offset = columns % COMPRESS_RATIO - safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1) - block = tl.load( - block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column, - mask=(row < rows) & is_expanded, - other=-1, - ) - expanded = block * COMPRESS_RATIO + offset - tail_offset = columns - expanded_count - is_tail = ( - (columns >= expanded_count) - & (tail_offset < tail_count) - & (tail_offset < COMPRESS_RATIO - 1) - ) - token = tl.where(is_expanded, expanded, tail_start + tail_offset) - valid = ( - (row < rows) - & (columns < OUTPUT_WIDTH) - & (is_expanded | is_tail) - & (token >= 0) - & (token < sequence_length) - ) - tl.store( - output_ptr + row * stride_output_row + columns * stride_output_column, - tl.where(valid, token, -1), - mask=(row < rows) & (columns < OUTPUT_WIDTH), - ) - @triton.jit def _qsa_sparse_paged_gqa_splitk_kernel( @@ -588,227 +408,6 @@ def _compress_qsa_groups_kernel( ) -def _validate_mqa(q: torch.Tensor) -> None: - if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0: - raise ValueError("QSA query must be [rows, heads, head_dim]") - - -def qsa_mqa_paged( - q: torch.Tensor, - k_cache: torch.Tensor, - page_table: torch.Tensor, - token_to_req: torch.Tensor, - query_positions: torch.Tensor, - sequence_lengths: torch.Tensor, - compress_ratio: int, - num_columns: int | None = None, - score_scale: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Compute QSA scores directly from a paged compressed-key cache.""" - - _validate_mqa(q) - if not q.is_cuda or not HAS_TRITON: - raise RuntimeError("paged QSA scoring requires CUDA and Triton") - if k_cache.ndim != 4 or k_cache.shape[2] != 1: - raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]") - if k_cache.shape[3] != q.shape[2]: - raise ValueError("QSA query and cache dimensions must match") - if page_table.ndim != 2: - raise ValueError("QSA page table must be two-dimensional") - if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)): - raise ValueError("QSA paged scoring cache and page table must be nonempty") - if token_to_req.shape != (q.shape[0],): - raise ValueError("QSA request mapping must match query rows") - if query_positions.shape != (q.shape[0],): - raise ValueError("QSA query positions must match query rows") - if sequence_lengths.shape != (page_table.shape[0],): - raise ValueError("QSA sequence lengths must match page-table requests") - if compress_ratio <= 0: - raise ValueError("QSA compression ratio must be positive") - score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale - if score_divisor <= 0: - raise ValueError("QSA score scale must be positive") - - capacity = page_table.shape[1] * k_cache.shape[1] - columns = capacity if num_columns is None else num_columns - if columns < 0: - raise ValueError("QSA score width must be non-negative") - logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) - visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device) - if not q.shape[0] or not columns: - return logits, visible_blocks - BLOCK_N = 64 - BLOCK_D = max(16, triton.next_power_of_2(q.shape[2])) - MAX_N = max(16, triton.next_power_of_2(q.shape[1])) - # Tuned on GB300: larger row batches provide enough parallelism to reuse Q. - tiles_per_program = 1 if q.shape[0] <= 32 else 8 - _qsa_mqa_paged_kernel[ - (q.shape[0], triton.cdiv(columns, BLOCK_N * tiles_per_program)) - ]( - q, - k_cache, - page_table, - token_to_req, - query_positions, - sequence_lengths, - visible_blocks, - logits, - q.stride(0), - q.stride(1), - q.stride(2), - k_cache.stride(0), - k_cache.stride(1), - k_cache.stride(3), - page_table.stride(0), - page_table.stride(1), - logits.stride(0), - q.shape[0], - columns, - k_cache.shape[0], - page_table.shape[0], - float(score_divisor), - PAGE_SIZE=k_cache.shape[1], - PAGE_TABLE_WIDTH=page_table.shape[1], - NUM_HEADS=q.shape[1], - HEAD_DIM=q.shape[2], - BLOCK_N=BLOCK_N, - BLOCK_D=BLOCK_D, - TILES_PER_PROG=tiles_per_program, - STAGES=2, - MAX_N=MAX_N, - COMPRESS_RATIO=compress_ratio, - num_warps=2, - ) - return logits, visible_blocks - - -def expand_qsa_block_indices_cuda( - block_indices: torch.Tensor, - query_positions: torch.Tensor, - sequence_lengths: torch.Tensor, - token_to_req: torch.Tensor, - compress_ratio: int, - token_topk: int, - out: torch.Tensor | None = None, -) -> torch.Tensor: - """Expand compressed blocks and compact the causal tail of the open group.""" - - if not block_indices.is_cuda or not HAS_TRITON: - raise RuntimeError("QSA CUDA expansion requires Triton") - if token_topk % compress_ratio: - raise ValueError("QSA token top-k must be divisible by compression ratio") - block_topk = token_topk // compress_ratio - output_width = token_topk + compress_ratio - 1 - if block_indices.shape != (query_positions.numel(), block_topk): - raise ValueError("QSA compressed top-k has an invalid shape") - if token_to_req.shape != query_positions.shape: - raise ValueError("QSA request mapping must match query positions") - if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]: - raise ValueError("QSA request sequence lengths must be nonempty") - if out is None: - out = torch.empty( - (block_indices.shape[0], output_width), - dtype=torch.int32, - device=block_indices.device, - ) - elif out.shape != (block_indices.shape[0], output_width): - raise ValueError("QSA expansion output has an invalid shape") - if not block_indices.shape[0]: - return out - column_block = 256 - _expand_qsa_indices_kernel[ - (block_indices.shape[0], triton.cdiv(output_width, column_block)) - ]( - block_indices, - query_positions, - sequence_lengths, - token_to_req, - out, - block_indices.stride(0), - block_indices.stride(1), - out.stride(0), - out.stride(1), - block_indices.shape[0], - sequence_lengths.shape[0], - BLOCK_TOPK=block_topk, - COMPRESS_RATIO=compress_ratio, - TOKEN_TOPK=token_topk, - OUTPUT_WIDTH=output_width, - COLUMN_BLOCK=column_block, - num_warps=4, - ) - return out - - -def qsa_select_paged_tokens( - 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, - out: torch.Tensor | None = None, -) -> torch.Tensor: - """Score, select, and expand QSA indices without host synchronization.""" - - rows = q.shape[0] - output_width = token_topk + compress_ratio - 1 - if out is None: - out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device) - if out.shape != (rows, output_width): - raise ValueError("QSA selection output has an invalid shape") - if not rows: - return out - - columns = page_table.shape[1] * k_cache.shape[1] - block_topk = token_topk // compress_ratio - rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1)) - chunk_rows = min(rows, rows_per_chunk) - blocks_buffer = torch.empty( - (chunk_rows, block_topk), dtype=torch.int32, device=q.device - ) - topk_workspace = torch.empty( - (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device - ) - for row_start in range(0, rows, rows_per_chunk): - row_end = min(row_start + rows_per_chunk, rows) - row_slice = slice(row_start, row_end) - logits, visible_blocks = qsa_mqa_paged( - q[row_slice], - k_cache, - page_table, - token_to_req[row_slice], - query_positions[row_slice], - sequence_lengths, - compress_ratio, - ) - blocks = blocks_buffer[: row_end - row_start] - use_cooperative_topk = ( - blocks.shape[0] <= 32 - and logits.stride(0) % 4 == 0 - and current_platform.has_device_capability(90) - and not current_platform.is_device_capability_family(120) - ) - topk_op = ( - torch.ops._C.cooperative_topk - if use_cooperative_topk - else torch.ops._C.persistent_topk - ) - topk_op(logits, visible_blocks, blocks, topk_workspace, block_topk, columns) - expand_qsa_block_indices_cuda( - blocks, - query_positions[row_slice], - sequence_lengths, - token_to_req[row_slice], - compress_ratio, - token_topk, - out[row_slice], - ) - return out - - def qsa_sparse_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -1106,10 +705,7 @@ def qsa_compress_groups_with_ratio( __all__ = [ - "expand_qsa_block_indices_cuda", "qsa_compress_groups_with_ratio", - "qsa_mqa_paged", - "qsa_select_paged_tokens", "qsa_sparse_paged_attention", "qsa_store_cache_rows", ] diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py new file mode 100644 index 000000000000..d58f28078b0e --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -0,0 +1,617 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Triton kernels for Qwen4Exp QSA index selection.""" + +import torch + +import vllm.envs as envs +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + TritonWarmupTensor, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton + +_TOPK_WORKSPACE_BYTES = 1024 * 1024 +_DECODE_BLOCK_N = 64 + + +@triton.jit +def _qsa_mqa_paged_uniform_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_cache_block, + stride_cache_token, + stride_table_req, + stride_logits_row, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + DECODE_QUERY_LEN: tl.constexpr, + BLOCK_N: tl.constexpr, + TILES_PER_PROG: tl.constexpr, + STAGES: tl.constexpr, +) -> None: + NUM_COLUMNS: tl.constexpr = PAGE_TABLE_WIDTH * PAGE_SIZE + DECODE_QUERY_LEN_PADDED: tl.constexpr = triton.next_power_of_2(DECODE_QUERY_LEN) + NUM_HEADS_PADDED: tl.constexpr = triton.next_power_of_2(NUM_HEADS) + # tl.dot requires a reduction dimension of at least 16. + BLOCK_D: tl.constexpr = max(16, triton.next_power_of_2(HEAD_DIM)) + request = tl.program_id(0) + tile_start = tl.program_id(1) * TILES_PER_PROG + query_offsets = tl.arange(0, DECODE_QUERY_LEN_PADDED) + valid_query_offsets = query_offsets < DECODE_QUERY_LEN + rows = request * DECODE_QUERY_LEN + query_offsets + visible = tl.load( + visible_blocks_ptr + rows, + mask=valid_query_offsets, + other=0, + ) + max_visible = tl.max(visible, axis=0) + if tile_start * BLOCK_N >= max_visible: + return + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(max_visible, BLOCK_N)) + tile_end = tl.minimum(tile_end, tl.cdiv(NUM_COLUMNS, BLOCK_N)) + + dims = tl.arange(0, BLOCK_D) + n = tl.arange(0, DECODE_QUERY_LEN_PADDED * NUM_HEADS_PADDED) + query_offset = n // NUM_HEADS_PADDED + head = n % NUM_HEADS_PADDED + valid_query = (query_offset < DECODE_QUERY_LEN) & (head < NUM_HEADS) + query = tl.load( + q_ptr + + (request * DECODE_QUERY_LEN + query_offset)[None, :] * stride_q_row + + head[None, :] * stride_q_head + + dims[:, None], + mask=valid_query[None, :] & (dims[:, None] < HEAD_DIM), + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < max_visible + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + request * stride_table_req + logical_page, + mask=live, + other=0, + ) + keys = tl.load( + k_cache_ptr + + physical_page[:, None].to(tl.int64) * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + scores = tl.dot(keys, query, out_dtype=tl.float32) + scores = tl.where(valid_query[None, :], tl.maximum(scores, 0.0), 0.0) + scores = tl.reshape( + scores, + (BLOCK_N, DECODE_QUERY_LEN_PADDED, NUM_HEADS_PADDED), + ) + score = tl.sum(scores, axis=2) / HEAD_DIM**0.5 + tl.store( + logits_ptr + rows[None, :] * stride_logits_row + columns[:, None], + score, + mask=valid_query_offsets[None, :] + & (columns[:, None] < NUM_COLUMNS) + & (columns[:, None] < visible[None, :]), + ) + + +@triton.jit(do_not_specialize=["num_rows", "query_offset"]) +def _qsa_mqa_paged_prefill_kernel( + q_ptr, + k_cache_ptr, + page_table_ptr, + query_start_loc_ptr, + visible_blocks_ptr, + logits_ptr, + stride_q_row, + stride_q_head, + stride_cache_block, + stride_cache_token, + stride_table_req, + stride_logits_row, + num_rows, + query_offset, + PAGE_SIZE: tl.constexpr, + PAGE_TABLE_WIDTH: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + TILE_R: tl.constexpr, + BLOCK_N: tl.constexpr, + K_TILES: tl.constexpr, + STAGES: tl.constexpr, +) -> None: + NUM_COLUMNS: tl.constexpr = PAGE_TABLE_WIDTH * PAGE_SIZE + NUM_HEADS_PADDED: tl.constexpr = triton.next_power_of_2(NUM_HEADS) + # tl.dot requires a reduction dimension of at least 16. + BLOCK_D: tl.constexpr = max(16, triton.next_power_of_2(HEAD_DIM)) + request = tl.program_id(0) + query_base = tl.load(query_start_loc_ptr) + query_end = query_offset + num_rows + request_start = tl.maximum( + tl.load(query_start_loc_ptr + request) - query_base, query_offset + ) + request_end = tl.minimum( + tl.load(query_start_loc_ptr + request + 1) - query_base, query_end + ) + absolute_row_start = request_start + tl.program_id(1) * TILE_R + if absolute_row_start >= request_end: + return + + lanes = tl.arange(0, TILE_R) + absolute_rows = absolute_row_start + lanes + rows = absolute_rows - query_offset + valid_rows = absolute_rows < request_end + visible = tl.load( + visible_blocks_ptr + absolute_rows, + mask=valid_rows, + other=0, + ) + max_visible = tl.max(visible, axis=0) + k_tile_start = tl.program_id(2) * K_TILES + if k_tile_start * BLOCK_N >= max_visible: + return + k_tile_end = tl.minimum(k_tile_start + K_TILES, tl.cdiv(max_visible, BLOCK_N)) + k_tile_end = tl.minimum(k_tile_end, tl.cdiv(NUM_COLUMNS, BLOCK_N)) + + dims = tl.arange(0, BLOCK_D) + m = tl.arange(0, TILE_R * NUM_HEADS_PADDED) + q_row_offsets = m // NUM_HEADS_PADDED + q_rows = absolute_row_start + q_row_offsets + heads = m % NUM_HEADS_PADDED + query = tl.load( + q_ptr + + q_rows[None, :] * stride_q_row + + heads[None, :] * stride_q_head + + dims[:, None], + mask=(heads[None, :] < NUM_HEADS) + & (absolute_row_start + q_row_offsets[None, :] < request_end) + & (dims[:, None] < HEAD_DIM), + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(k_tile_start, k_tile_end, num_stages=STAGES): + columns = tile * BLOCK_N + column_offsets + live = columns < max_visible + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1) + page_offset = columns % PAGE_SIZE + physical_page = tl.load( + page_table_ptr + request * stride_table_req + logical_page, + mask=live, + other=0, + ) + keys = tl.load( + k_cache_ptr + + physical_page[:, None].to(tl.int64) * stride_cache_block + + page_offset[:, None] * stride_cache_token + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + scores = tl.dot(keys, query, out_dtype=tl.float32) + scores = tl.reshape(scores, (BLOCK_N, TILE_R, NUM_HEADS_PADDED)) + score = tl.sum(tl.maximum(scores, 0.0), axis=2) / HEAD_DIM**0.5 + store_mask = ( + valid_rows[None, :] + & (columns[:, None] < visible[None, :]) + & (columns[:, None] < NUM_COLUMNS) + ) + tl.store( + logits_ptr + rows[None, :] * stride_logits_row + columns[:, None], + score, + mask=store_mask, + ) + + +@triton.jit +def _expand_qsa_indices_kernel( + block_indices_ptr, + query_positions_ptr, + visible_blocks_ptr, + output_ptr, + stride_blocks_row, + stride_blocks_column, + stride_output_row, + stride_output_column, + BLOCK_TOPK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + COLUMN_BLOCK: tl.constexpr, +) -> None: + OUTPUT_WIDTH: tl.constexpr = BLOCK_TOPK * COMPRESS_RATIO + COMPRESS_RATIO - 1 + row = tl.program_id(0) + columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK) + query_position = tl.load(query_positions_ptr + row) + visible_blocks = tl.load(visible_blocks_ptr + row) + complete_blocks = tl.minimum(visible_blocks, BLOCK_TOPK) + expanded_count = complete_blocks * COMPRESS_RATIO + tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO + tail_count = (query_position + 1) - tail_start + + is_expanded = columns < expanded_count + block_rank = columns // COMPRESS_RATIO + offset = columns % COMPRESS_RATIO + safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1) + block = tl.load( + block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column, + mask=is_expanded, + other=-1, + ) + expanded = block * COMPRESS_RATIO + offset + tail_offset = columns - expanded_count + is_tail = ( + (columns >= expanded_count) + & (tail_offset < tail_count) + & (tail_offset < COMPRESS_RATIO - 1) + ) + token = tl.where(is_expanded, expanded, tail_start + tail_offset) + valid = (columns < OUTPUT_WIDTH) & (is_expanded | is_tail) & (token >= 0) + tl.store( + output_ptr + row * stride_output_row + columns * stride_output_column, + tl.where(valid, token, -1), + mask=columns < OUTPUT_WIDTH, + ) + + +def _decode_tiles_per_program(num_requests: int, columns: int) -> int: + programs = num_requests * triton.cdiv(columns, _DECODE_BLOCK_N) + if programs < 16384: + return 1 + if programs < 32768: + return 2 + if programs < 131072: + return 4 + return 8 + + +def _qsa_decode_warmup_profiles( + max_dql: int, + max_num_reqs: int, + max_num_batched_tokens: int, + columns: int, +) -> tuple[tuple[int, int], ...]: + profiles: list[tuple[int, int]] = [] + for dql in range(1, max_dql + 1): + max_reqs = min(max_num_reqs, max_num_batched_tokens // dql) + requests_by_grouping: dict[int, int] = {} + for num_requests in range(1, max_reqs + 1): + tiles_per_program = _decode_tiles_per_program(num_requests, columns) + requests_by_grouping.setdefault(tiles_per_program, num_requests) + profiles.extend( + (dql, num_requests) for num_requests in requests_by_grouping.values() + ) + return tuple(profiles) + + +def warmup_qsa_mqa_paged_decode( + k_cache: torch.Tensor, + page_table: torch.Tensor, + *, + num_heads: int, + head_dim: int, + max_decode_query_len: int, + max_num_reqs: int, + max_num_batched_tokens: int, +) -> tuple[tuple[int, int], ...]: + """Compile every reachable decode specialization without launching it.""" + + page_size = k_cache.shape[1] + page_table_width = page_table.shape[1] + columns = page_table_width * page_size + profiles = _qsa_decode_warmup_profiles( + max_decode_query_len, + max_num_reqs, + max_num_batched_tokens, + columns, + ) + if not profiles: + return () + + k_cache_ptr = TritonWarmupTensor(k_cache.dtype, shape=tuple(k_cache.shape)) + page_table_ptr = TritonWarmupTensor( + page_table.dtype, + shape=(max_num_reqs, page_table_width), + ) + visible_blocks_ptr = TritonWarmupTensor(torch.int32) + + for decode_query_len, num_requests in profiles: + num_rows = decode_query_len * num_requests + q_ptr = TritonWarmupTensor( + torch.bfloat16, + shape=(num_rows, num_heads, head_dim), + ) + logits_ptr = TritonWarmupTensor( + torch.float32, + shape=(num_rows, columns), + ) + tiles_per_program = _decode_tiles_per_program(num_requests, columns) + _qsa_mqa_paged_uniform_kernel.warmup( + q_ptr, + k_cache_ptr, + page_table_ptr, + visible_blocks_ptr, + logits_ptr, + num_heads * head_dim, + head_dim, + k_cache.stride(0), + k_cache.stride(1), + page_table.stride(0), + columns, + PAGE_SIZE=page_size, + PAGE_TABLE_WIDTH=page_table_width, + NUM_HEADS=num_heads, + HEAD_DIM=head_dim, + DECODE_QUERY_LEN=decode_query_len, + BLOCK_N=_DECODE_BLOCK_N, + TILES_PER_PROG=tiles_per_program, + STAGES=2, + num_warps=2, + grid=( + num_requests, + triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program), + ), + ) + return profiles + + +def _prefill_logits( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + query_start_loc: torch.Tensor, + visible_blocks: torch.Tensor, + max_query_len: int, + query_offset: int, + num_queries: int, +) -> torch.Tensor: + assert query_start_loc.shape == (page_table.shape[0] + 1,) + assert visible_blocks.shape == (q.shape[0],) + assert 0 <= query_offset <= query_offset + num_queries <= q.shape[0] + + columns = page_table.shape[1] * k_cache.shape[1] + logits = torch.empty((num_queries, columns), dtype=torch.float32, device=q.device) + TILE_R = 64 + BLOCK_N = 64 + K_TILES = 16 + grid = ( + page_table.shape[0], + triton.cdiv(min(num_queries, max_query_len), TILE_R), + triton.cdiv(columns, BLOCK_N * K_TILES), + ) + _qsa_mqa_paged_prefill_kernel[grid]( + q, + k_cache, + page_table, + query_start_loc, + visible_blocks, + logits, + *q.stride()[:-1], + *k_cache.stride()[:2], + *page_table.stride()[:-1], + *logits.stride()[:-1], + num_queries, + query_offset, + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + NUM_HEADS=q.shape[1], + HEAD_DIM=q.shape[2], + TILE_R=TILE_R, + BLOCK_N=BLOCK_N, + K_TILES=K_TILES, + STAGES=2, + num_warps=4, + ) + return logits + + +def expand_qsa_block_indices( + block_indices: torch.Tensor, + query_positions: torch.Tensor, + visible_blocks: torch.Tensor, + compress_ratio: int, + token_topk: int, + out: torch.Tensor, +) -> None: + """Expand compressed blocks and compact the causal tail of the open group.""" + + assert token_topk % compress_ratio == 0 + block_topk = token_topk // compress_ratio + output_width = token_topk + compress_ratio - 1 + assert block_indices.shape == (query_positions.numel(), block_topk) + assert visible_blocks.shape == query_positions.shape + assert out.shape == (block_indices.shape[0], output_width) + column_block = 256 + grid = (block_indices.shape[0], triton.cdiv(output_width, column_block)) + _expand_qsa_indices_kernel[grid]( + block_indices, + query_positions, + visible_blocks, + out, + *block_indices.stride(), + *out.stride(), + BLOCK_TOPK=block_topk, + COMPRESS_RATIO=compress_ratio, + COLUMN_BLOCK=column_block, + num_warps=4, + ) + + +def _topk( + logits: torch.Tensor, + visible_blocks: torch.Tensor, + token_topk: int, + compress_ratio: int, + block_indices: torch.Tensor, + topk_workspace: torch.Tensor, +) -> None: + # similar dispatch logic as DeepSeek indexer + block_topk = token_topk // compress_ratio + use_cooperative_topk = ( + logits.shape[0] <= 64 + and logits.stride(0) % 4 == 0 + and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) + ) + topk_op = ( + torch.ops._C.cooperative_topk + if use_cooperative_topk + else torch.ops._C.persistent_topk + ) + topk_op( + logits, + visible_blocks, + block_indices, + topk_workspace, + block_topk, + logits.shape[1], + ) + + +def qsa_select_paged_decode( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + visible_blocks: torch.Tensor, + token_topk: int, + compress_ratio: int, + decode_query_len: int, + block_indices: torch.Tensor, +) -> None: + """Score and select compressed blocks for a request-major decode batch. + + Args: + q: Query tensor shaped ``[num_requests * decode_query_len, heads, + head_dim]``. + k_cache: Compressed key cache shaped ``[blocks, page_size, 1, + head_dim]``. + page_table: Request block table shaped ``[num_requests, max_pages]``. + visible_blocks: Number of visible compressed blocks per query. + token_topk: Number of logical tokens selected per query. + compress_ratio: Number of logical tokens represented by a cache row. + decode_query_len: Number of query tokens per request. + block_indices: Compressed-index output buffer. + """ + + assert token_topk % compress_ratio == 0 + assert block_indices.shape == (q.shape[0], token_topk // compress_ratio) + assert decode_query_len > 0 and q.shape[0] % decode_query_len == 0 + num_requests = q.shape[0] // decode_query_len + assert page_table.shape[0] == num_requests + assert visible_blocks.shape == (q.shape[0],) + + columns = page_table.shape[1] * k_cache.shape[1] + logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device) + tiles_per_program = _decode_tiles_per_program(num_requests, columns) + grid = ( + num_requests, + triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program), + ) + _qsa_mqa_paged_uniform_kernel[grid]( + q, + k_cache, + page_table, + visible_blocks, + logits, + *q.stride()[:-1], + *k_cache.stride()[:2], + *page_table.stride()[:-1], + *logits.stride()[:-1], + PAGE_SIZE=k_cache.shape[1], + PAGE_TABLE_WIDTH=page_table.shape[1], + NUM_HEADS=q.shape[1], + HEAD_DIM=q.shape[2], + DECODE_QUERY_LEN=decode_query_len, + BLOCK_N=_DECODE_BLOCK_N, + TILES_PER_PROG=tiles_per_program, + STAGES=2, + num_warps=2, + ) + _topk( + logits, + visible_blocks, + token_topk, + compress_ratio, + block_indices, + torch.empty((_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device), + ) + + +def qsa_select_paged_prefill( + q: torch.Tensor, + k_cache: torch.Tensor, + page_table: torch.Tensor, + query_start_loc: torch.Tensor, + visible_blocks: torch.Tensor, + token_topk: int, + compress_ratio: int, + max_query_len: int, + block_indices: torch.Tensor, +) -> None: + """Score and select compressed prefill blocks in bounded chunks. + + Args: + q: Packed prefill query tensor shaped ``[num_tokens, heads, head_dim]``. + k_cache: Compressed key cache shaped ``[blocks, page_size, 1, + head_dim]``. + page_table: Block table shaped ``[num_requests, max_pages]``. + query_start_loc: Packed prefill query offsets with a terminal offset. + Offsets may share the base of a larger backing tensor. + visible_blocks: Number of visible compressed blocks per query. + token_topk: Number of logical tokens selected per query. + compress_ratio: Number of logical tokens represented by a cache row. + max_query_len: Maximum number of query tokens in one request. + block_indices: Compressed-index output buffer. + """ + + assert token_topk % compress_ratio == 0 + assert block_indices.shape == (q.shape[0], token_topk // compress_ratio) + rows = q.shape[0] + columns = page_table.shape[1] * k_cache.shape[1] + + # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + rows_per_chunk = max(1, max_logits_bytes // (columns * 4)) + topk_workspace = torch.empty( + (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device + ) + + for query_start in range(0, rows, rows_per_chunk): + query_end = min(query_start + rows_per_chunk, rows) + query_slice = slice(query_start, query_end) + logits = _prefill_logits( + q, + k_cache, + page_table, + query_start_loc, + visible_blocks, + max_query_len, + query_offset=query_start, + num_queries=query_end - query_start, + ) + _topk( + logits, + visible_blocks[query_slice], + token_topk, + compress_ratio, + block_indices[query_slice], + topk_workspace, + ) + + +__all__ = [ + "expand_qsa_block_indices", + "qsa_select_paged_decode", + "qsa_select_paged_prefill", + "warmup_qsa_mqa_paged_decode", +]