@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 003e34341a

[Qwen3.8-Flash-Next] Separate prefill and decode paths for QSA indexer (#54513)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>

Thien Tran committed Sep 2, 2026 at 09:08 UTC 003e34341a824ab8e3bec4861e92b5e84180af8b

7 files changed +1105 -523

tests/models/qwen4_exp/test_qsa_reference.py

+252 -83

@@ -14,6 +14,7 @@ from vllm.models.qwen4_exp.nvidia import (

14 model as _qwen4_exp_model, # noqa: F401

15 )

16 from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops

17 +from vllm.models.qwen4_exp.nvidia.ops import qsa_indexer as qsa_indexer_ops

18 from vllm.platforms import current_platform

19 from vllm.triton_utils import HAS_TRITON

@@ -67,8 +68,13 @@ def test_qsa_mtp_index_share_updates_cache_but_skips_selection(

68 lambda *args, **kwargs: updates.append((args, kwargs)),

69 )

70 monkeypatch.setattr(

70 - qsa_ops,

71 - "qsa_select_paged_tokens",

71 + qsa_indexer_ops,

72 + "qsa_select_paged_decode",

73 + lambda *args, **kwargs: selections.append((args, kwargs)),

74 + )

75 + monkeypatch.setattr(

76 + qsa_indexer_ops,

77 + "qsa_select_paged_prefill",

78 lambda *args, **kwargs: selections.append((args, kwargs)),

79 )

@@ -158,7 +164,7 @@ def _expand_qsa_indices_reference(

164 return result.gather(1, torch.argsort(sort_key, dim=1, stable=True)).to(torch.int32)

165

166

161 -def _qsa_select_paged_tokens_reference(

167 +def _qsa_select_paged_reference(

168 q: torch.Tensor,

169 k_cache: torch.Tensor,

170 page_table: torch.Tensor,

@@ -181,19 +187,12 @@ def _qsa_select_paged_tokens_reference(

187 visible_blocks,

188 )

189 starts = torch.zeros_like(visible_blocks)

184 - blocks = _qsa_relative_topk_reference(

190 + return _qsa_relative_topk_reference(

191 logits,

192 starts,

193 visible_blocks,

194 token_topk // compress_ratio,

195 )

190 - return _expand_qsa_indices_reference(

191 - blocks,

192 - query_positions,

193 - row_sequence_lengths,

194 - compress_ratio,

195 - token_topk,

196 - )

196

197

198 def _qsa_sparse_paged_attention_reference(

@@ -231,16 +230,20 @@ def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None:

230 device = torch.device("cuda")

231 builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)

232 builder.compress_ratio = 1

233 + builder.reorder_batch_threshold = 4

234 builder.is_circular_buffer = False

235 builder.storage_block_size = 64

236 builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device)

237 builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device)

238 builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device)

239 + builder.visible_blocks_buffer = torch.empty(16, dtype=torch.int32, device=device)

240 builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device)

241 query_start_loc = torch.tensor([0, 4, 8, 12, 12], dtype=torch.int32, device=device)

242 token_to_req = torch.tensor([0] * 4 + [1] * 4 + [2] * 4 + [0] * 4, device=device)

243 common = SimpleNamespace(

244 num_actual_tokens=16,

245 + num_reqs=4,

246 + max_query_len=4,

247 query_start_loc=query_start_loc,

248 query_start_loc_cpu=query_start_loc.cpu(),

249 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:

273 -1,

274 ]

275 assert metadata.slot_mapping.tolist() == list(range(12)) + [-1] * 4

276 + assert metadata.visible_blocks.tolist() == [65, 66, 67, 68] * 3 + [0] * 4

277

278

279 @requires_qsa_kernels

@@ -277,18 +281,22 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None:

281 device = torch.device("cuda")

282 builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)

283 builder.compress_ratio = 4

284 + builder.reorder_batch_threshold = 1

285 builder.is_circular_buffer = True

286 builder.kv_cache_spec = SimpleNamespace(block_size=4)

287 builder.storage_block_size = 4

288 builder.token_to_req_buffer = torch.empty(16, dtype=torch.int32, device=device)

289 builder.slot_mapping_buffer = torch.empty(16, dtype=torch.int64, device=device)

290 builder.logical_positions_buffer = torch.empty(16, dtype=torch.int64, device=device)

291 + builder.visible_blocks_buffer = torch.empty(16, dtype=torch.int32, device=device)

292 builder.k_work_metadata_buffer = torch.empty(0, 2, dtype=torch.int32, device=device)

293 query_start_loc = torch.tensor([0, 7, 13, 13], dtype=torch.int32, device=device)

294 token_to_req = torch.tensor([0] * 7 + [1] * 6 + [0] * 3, device=device)

295 block_table = torch.tensor([[1], [0], [2]], dtype=torch.int32, device=device)

296 common = SimpleNamespace(

297 num_actual_tokens=16,

298 + num_reqs=3,

299 + max_query_len=7,

300 query_start_loc=query_start_loc,

301 query_start_loc_cpu=query_start_loc.cpu(),

302 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:

410 device = torch.device("cuda")

411 builder = QSAMetadataBuilder.__new__(QSAMetadataBuilder)

412 builder.compress_ratio = 4

413 + builder.reorder_batch_threshold = 1

414 builder.is_circular_buffer = False

415 builder.storage_block_size = 16

416 builder.token_to_req_buffer = torch.empty(8, dtype=torch.int32, device=device)

417 builder.slot_mapping_buffer = torch.empty(8, dtype=torch.int64, device=device)

418 builder.logical_positions_buffer = torch.empty(8, dtype=torch.int64, device=device)

419 + builder.visible_blocks_buffer = torch.empty(8, dtype=torch.int32, device=device)

420 # Simulate max_num_seqs exceeding the three live requests below.

421 builder.request_capacity = 8

422 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:

426 )

427 common = SimpleNamespace(

428 num_actual_tokens=8,

429 + num_reqs=3,

430 + max_query_len=5,

431 query_start_loc=query_start_loc,

432 query_start_loc_cpu=query_start_loc.cpu(),

433 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:

439 metadata = builder.build(0, common)

440

441 assert metadata.slot_mapping.tolist() == [-1] * 8

442 + assert metadata.visible_blocks.tolist() == [1, 1, 1, 2, 2, 2, 2, 3]

443 assert metadata.k_work_metadata.tolist() == [[0, 0], [2, 0], [2, 1], [-1, -1]]

444

445

@@ -473,10 +486,11 @@ def test_qsa_triton_metadata_matches_pytorch(

486 token_to_req_indices=lambda buffer: buffer.copy_(token_to_req),

487 )

488

476 - def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

489 + def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

490 return (

491 torch.empty(num_tokens, dtype=torch.int32, device=device),

492 torch.empty(num_tokens, dtype=torch.int64, device=device),

493 + torch.empty(num_tokens, dtype=torch.int32, device=device),

494 torch.empty(num_tokens, dtype=torch.int64, device=device),

495 )

496

@@ -543,10 +557,11 @@ def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None:

557 token_to_req_indices=lambda buffer: buffer.zero_(),

558 )

559

546 - def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

560 + def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

561 return (

562 torch.empty(num_tokens, dtype=torch.int32, device=device),

563 torch.empty(num_tokens, dtype=torch.int64, device=device),

564 + torch.empty(num_tokens, dtype=torch.int32, device=device),

565 torch.empty(num_tokens, dtype=torch.int64, device=device),

566 )

567

@@ -575,70 +590,163 @@ def test_qsa_fused_metadata_matches_pytorch_for_large_padded_prefill() -> None:

590

591 @requires_qsa_kernels

592 @pytest.mark.parametrize(

578 - "num_rows",

593 + ("decode_query_len", "num_requests"),

594 [

580 - pytest.param(3, id="one_tile_per_program"),

581 - pytest.param(33, id="looped_tiles"),

595 + (1, 2),

596 + (2, 2),

597 + (3, 2),

598 + (4, 2),

599 + (4, 33),

600 ],

601 )

584 -def test_qsa_mqa_paged_matches_test_reference(num_rows: int) -> None:

602 +def test_qsa_decode_selection_correctness(

603 + decode_query_len: int, num_requests: int

604 +) -> None:

605 torch.manual_seed(1)

586 - head_dim = 128

587 - q = torch.randn(num_rows, 4, head_dim, device="cuda", dtype=torch.bfloat16)

588 - cache = torch.randn(40, 4, 1, head_dim, device="cuda", dtype=torch.bfloat16)

589 - page_table = torch.randperm(40, device="cuda", dtype=torch.int32).reshape(2, 20)

590 - request_rows = [(num_rows + 1) // 2, num_rows // 2]

606 + heads, head_dim = 4, 128

607 + rows = num_requests * decode_query_len

608 + q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)

609 + page_size, pages_per_request, max_sequence_length = (

610 + (16, 40, 2560) if num_requests > 32 else (4, 20, 320)

611 + )

612 + num_pages = num_requests * pages_per_request

613 + cache = torch.randn(

614 + num_pages,

615 + page_size,

616 + 1,

617 + head_dim,

618 + device="cuda",

619 + dtype=torch.bfloat16,

620 + )

621 + page_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(

622 + num_requests, pages_per_request

623 + )

624 token_to_req = torch.repeat_interleave(

592 - torch.arange(2, device="cuda", dtype=torch.int32),

593 - torch.tensor(request_rows, device="cuda"),

625 + torch.arange(num_requests, device="cuda", dtype=torch.int32),

626 + decode_query_len,

627 )

595 - sequence_length_values = [320, 264]

596 - sequence_lengths = torch.tensor(

597 - sequence_length_values, device="cuda", dtype=torch.int32

628 + sequence_lengths = max_sequence_length - 4 * (

629 + torch.arange(num_requests, device="cuda", dtype=torch.int32) % 8

630 )

631 query_positions = torch.cat(

632 [

601 - torch.arange(length - rows, length, device="cuda", dtype=torch.int32)

602 - for rows, length in zip(request_rows, sequence_length_values, strict=True)

633 + torch.arange(

634 + length - decode_query_len,

635 + length,

636 + device="cuda",

637 + dtype=torch.int32,

638 + )

639 + for length in sequence_lengths.tolist()

640 ]

641 )

605 - compress_ratio = 4

606 - visible_lengths = (query_positions + 1) // compress_ratio

642 + visible_blocks = torch.minimum(

643 + (query_positions + 1) // 4,

644 + sequence_lengths.index_select(0, token_to_req.long()) // 4,

645 + )

646

608 - actual, actual_visible_blocks = qsa_ops.qsa_mqa_paged(

647 + token_topk, compress_ratio = 2048, 4

648 + actual = torch.empty(

649 + (rows, token_topk // compress_ratio), device="cuda", dtype=torch.int32

650 + )

651 + qsa_indexer_ops.qsa_select_paged_decode(

652 + q,

653 + cache,

654 + page_table,

655 + visible_blocks,

656 + token_topk,

657 + compress_ratio,

658 + decode_query_len,

659 + actual,

660 + )

661 + expected = _qsa_select_paged_reference(

662 q,

663 cache,

664 page_table,

665 token_to_req,

666 query_positions,

667 sequence_lengths,

615 - compress_ratio=compress_ratio,

668 + token_topk,

669 + compress_ratio,

670 + )

671 +

672 + torch.testing.assert_close(actual.sort().values, expected.sort().values)

673 +

674 +

675 +@requires_qsa_kernels

676 +def test_qsa_prefill_selection_correctness() -> None:

677 + torch.manual_seed(2)

678 + query_lens = [3, 33]

679 + rows, heads, head_dim = sum(query_lens), 4, 128

680 + q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)

681 + cache = torch.randn(160, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16)

682 + page_table = torch.randperm(160, device="cuda", dtype=torch.int32).reshape(2, 80)

683 + token_to_req = torch.repeat_interleave(

684 + torch.arange(2, device="cuda", dtype=torch.int32),

685 + torch.tensor(query_lens, device="cuda"),

686 + )

687 + query_start_loc = torch.tensor([0, 3, 36], device="cuda", dtype=torch.int32)

688 + sequence_lengths = torch.tensor([5120, 4224], device="cuda", dtype=torch.int32)

689 + query_positions = torch.cat(

690 + [

691 + torch.arange(length - query_len, length, device="cuda", dtype=torch.int32)

692 + for query_len, length in zip(

693 + query_lens, sequence_lengths.tolist(), strict=True

694 + )

695 + ]

696 + )

697 + token_topk, compress_ratio = 2048, 4

698 + visible_blocks = torch.minimum(

699 + (query_positions + 1) // compress_ratio,

700 + sequence_lengths.index_select(0, token_to_req.long()) // compress_ratio,

701 + )

702 +

703 + actual = torch.empty(

704 + (rows, token_topk // compress_ratio), device="cuda", dtype=torch.int32

705 )

617 - expected = _qsa_mqa_paged_reference(

618 - q, cache, page_table, token_to_req, visible_lengths

706 + qsa_indexer_ops.qsa_select_paged_prefill(

707 + q,

708 + cache,

709 + page_table,

710 + query_start_loc,

711 + visible_blocks,

712 + token_topk,

713 + compress_ratio,

714 + max(query_lens),

715 + actual,

716 + )

717 + expected = _qsa_select_paged_reference(

718 + q,

719 + cache,

720 + page_table,

721 + token_to_req,

722 + query_positions,

723 + sequence_lengths,

724 + token_topk,

725 + compress_ratio,

726 )

727

621 - torch.testing.assert_close(actual_visible_blocks, visible_lengths)

622 - # Top-k is bounded by visible_blocks; later columns are intentionally unwritten.

623 - columns = torch.arange(actual.shape[1], device=actual.device)

624 - visible = columns[None, :] < actual_visible_blocks[:, None]

625 - torch.testing.assert_close(actual[visible], expected[visible], rtol=1e-3, atol=1e-3)

728 + torch.testing.assert_close(actual.sort().values, expected.sort().values)

729

730

731 @requires_qsa_kernels

629 -def test_qsa_block_expansion_matches_test_reference() -> None:

732 +def test_qsa_block_expansion_correctness() -> None:

733 blocks = torch.tensor([[0, -1], [1, 0]], device="cuda", dtype=torch.int32)

734 query_positions = torch.tensor([5, 10], device="cuda")

735 sequence_lengths = torch.tensor([6, 11], device="cuda")

736 token_to_req = torch.tensor([0, 1], device="cuda", dtype=torch.int32)

737 + visible_blocks = torch.minimum(

738 + (query_positions + 1) // 4,

739 + sequence_lengths.index_select(0, token_to_req.long()) // 4,

740 + ).to(torch.int32)

741

635 - actual = qsa_ops.expand_qsa_block_indices_cuda(

742 + actual = torch.empty((2, 11), device="cuda", dtype=torch.int32)

743 + qsa_indexer_ops.expand_qsa_block_indices(

744 blocks,

745 query_positions,

638 - sequence_lengths,

639 - token_to_req,

746 + visible_blocks,

747 compress_ratio=4,

748 token_topk=8,

749 + out=actual,

750 )

751 expected = _expand_qsa_indices_reference(

752 blocks,

@@ -663,7 +771,7 @@ def test_qsa_block_expansion_matches_test_reference() -> None:

771 pytest.param(513, 6, 1, 1024, id="tp4_split1"),

772 ],

773 )

666 -def test_qsa_sparse_paged_attention_matches_test_reference(

774 +def test_qsa_sparse_paged_attention_correctness(

775 num_rows: int,

776 num_query_heads: int,

777 num_kv_heads: int,

@@ -724,13 +832,20 @@ def test_qsa_sparse_paged_attention_matches_test_reference(

832 sequence_lengths = torch.full(

833 (num_requests,), context_length, device="cuda", dtype=torch.int32

834 )

727 - logical_indices = qsa_ops.expand_qsa_block_indices_cuda(

835 + visible_blocks = torch.minimum(

836 + (query_positions + 1) // indexer_compress_ratio,

837 + sequence_lengths.index_select(0, token_to_req.long()) // indexer_compress_ratio,

838 + ).to(torch.int32)

839 + logical_indices = torch.empty(

840 + (num_rows, selection_width), device="cuda", dtype=torch.int32

841 + )

842 + qsa_indexer_ops.expand_qsa_block_indices(

843 block_indices,

844 query_positions,

730 - sequence_lengths,

731 - token_to_req,

845 + visible_blocks,

846 indexer_compress_ratio,

847 indexer_budget,

848 + logical_indices,

849 )

850 assert logical_indices.shape == (num_rows, selection_width)

851 scale = q.shape[-1] ** -0.5

@@ -757,40 +872,78 @@ def test_qsa_sparse_paged_attention_matches_test_reference(

872

873

874 @requires_qsa_kernels

760 -def test_qsa_selection_chunks_workspace_and_matches_test_reference(

761 - monkeypatch: pytest.MonkeyPatch,

762 - workspace_init,

763 -) -> None:

764 - rows, keys, heads, head_dim = 65, 640, 4, 16

875 +@pytest.mark.parametrize("decode_query_len", [1, 2, 3, 4])

876 +def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int) -> None:

877 + query_lens = [decode_query_len, decode_query_len, 33]

878 + rows, heads, head_dim = sum(query_lens), 4, 128

879 token_topk, compress_ratio = 2048, 4

766 - torch.manual_seed(3)

880 + torch.manual_seed(13)

881 q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)

768 - cache = torch.randn(40, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16)

769 - page_table = torch.randperm(40, device="cuda", dtype=torch.int32).unsqueeze(0)

770 - token_to_req = torch.zeros(rows, device="cuda", dtype=torch.int32)

771 - query_positions = torch.full((rows,), 2559, device="cuda", dtype=torch.int32)

772 - sequence_lengths = torch.tensor([2560], device="cuda", dtype=torch.int32)

773 - monkeypatch.setattr(qsa_ops, "_LOGITS_WORKSPACE_BYTES", 32 * keys * 4)

774 - original_score = qsa_ops.qsa_mqa_paged

775 - scored_row_counts = []

776 -

777 - def record_score(query: torch.Tensor, *args, **kwargs):

778 - scored_row_counts.append(query.shape[0])

779 - return original_score(query, *args, **kwargs)

780 -

781 - monkeypatch.setattr(qsa_ops, "qsa_mqa_paged", record_score)

782 -

783 - actual = qsa_ops.qsa_select_paged_tokens(

784 - q,

882 + cache = torch.randn(120, 16, 1, head_dim, device="cuda", dtype=torch.bfloat16)

883 + page_table = torch.arange(120, device="cuda", dtype=torch.int32).view(3, 40)

884 + token_to_req = torch.repeat_interleave(

885 + torch.arange(3, device="cuda", dtype=torch.int32),

886 + torch.tensor(query_lens, device="cuda"),

887 + )

888 + query_start_loc = torch.tensor(

889 + [0, decode_query_len, 2 * decode_query_len, rows],

890 + device="cuda",

891 + dtype=torch.int32,

892 + )

893 + sequence_lengths = torch.full((3,), 2560, device="cuda", dtype=torch.int32)

894 + query_positions = torch.cat(

895 + [

896 + torch.arange(2560 - query_len, 2560, device="cuda")

897 + for query_len in query_lens

898 + ]

899 + )

900 +

901 + block_indices = torch.empty(

902 + rows,

903 + token_topk // compress_ratio,

904 + device="cuda",

905 + dtype=torch.int32,

906 + )

907 + visible_blocks = torch.minimum(

908 + (query_positions + 1) // compress_ratio,

909 + sequence_lengths.index_select(0, token_to_req.long()) // compress_ratio,

910 + ).to(torch.int32)

911 + num_decode_tokens = 2 * decode_query_len

912 + decode_slice = slice(0, num_decode_tokens)

913 + qsa_indexer_ops.qsa_select_paged_decode(

914 + q[decode_slice],

915 cache,

786 - page_table,

787 - token_to_req,

788 - query_positions,

789 - sequence_lengths,

916 + page_table[:2],

917 + visible_blocks[decode_slice],

918 token_topk,

919 compress_ratio,

920 + decode_query_len,

921 + block_indices[decode_slice],

922 + )

923 + prefill_slice = slice(num_decode_tokens, rows)

924 + qsa_indexer_ops.qsa_select_paged_prefill(

925 + q[prefill_slice],

926 + cache,

927 + page_table[2:],

928 + query_start_loc[2:],

929 + visible_blocks[prefill_slice],

930 + token_topk,

931 + compress_ratio,

932 + query_lens[-1],

933 + block_indices[prefill_slice],

934 + )

935 + actual = torch.empty(

936 + (rows, token_topk + compress_ratio - 1), device="cuda", dtype=torch.int32

937 + )

938 + qsa_indexer_ops.expand_qsa_block_indices(

939 + block_indices,

940 + query_positions,

941 + visible_blocks,

942 + compress_ratio,

943 + token_topk,

944 + actual,

945 )

793 - expected = _qsa_select_paged_tokens_reference(

946 + expected_blocks = _qsa_select_paged_reference(

947 q,

948 cache,

949 page_table,

@@ -800,9 +953,15 @@ def test_qsa_selection_chunks_workspace_and_matches_test_reference(

953 token_topk,

954 compress_ratio,

955 )

956 + expected = _expand_qsa_indices_reference(

957 + expected_blocks,

958 + query_positions,

959 + sequence_lengths.index_select(0, token_to_req.long()),

960 + compress_ratio,

961 + token_topk,

962 + )

963

964 torch.testing.assert_close(actual.sort().values, expected.sort().values)

805 - assert scored_row_counts == [32, 32, 1]

965

966

967 @requires_qsa_kernels

@@ -810,19 +969,29 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) ->

969 q = torch.zeros(2, 4, 8, device="cuda", dtype=torch.bfloat16)

970 cache = torch.zeros(1, 16, 1, 8, device="cuda", dtype=torch.bfloat16)

971 page_table = torch.zeros(1, 1, device="cuda", dtype=torch.int32)

813 - token_to_req = torch.zeros(2, device="cuda", dtype=torch.int32)

972 query_positions = torch.tensor([1, 2], device="cuda", dtype=torch.int32)

815 - sequence_lengths = torch.tensor([3], device="cuda", dtype=torch.int32)

973 + visible_blocks = torch.zeros(2, device="cuda", dtype=torch.int32)

974

817 - selected = qsa_ops.qsa_select_paged_tokens(

975 + block_indices = torch.empty((2, 512), device="cuda", dtype=torch.int32)

976 + qsa_indexer_ops.qsa_select_paged_prefill(

977 q,

978 cache,

979 page_table,

821 - token_to_req,

822 - query_positions,

823 - sequence_lengths,

980 + torch.tensor([0, 2], device="cuda", dtype=torch.int32),

981 + visible_blocks,

982 token_topk=2048,

983 compress_ratio=4,

984 + max_query_len=2,

985 + block_indices=block_indices,

986 + )

987 + selected = torch.empty((2, 2051), device="cuda", dtype=torch.int32)

988 + qsa_indexer_ops.expand_qsa_block_indices(

989 + block_indices,

990 + query_positions,

991 + visible_blocks,

992 + compress_ratio=4,

993 + token_topk=2048,

994 + out=selected,

995 )

996

997 assert selected[0, :2].tolist() == [0, 1]

vllm/model_executor/warmup/kernel_warmup.py

@@ -37,6 +37,9 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (

37 from vllm.model_executor.warmup.kimi_k3_triton_warmup import (

38 kimi_k3_triton_warmup,

39 )

40 +from vllm.model_executor.warmup.qwen4_exp_qsa_warmup import (

41 + qwen4_exp_qsa_triton_warmup,

42 +)

43 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup

44 from vllm.model_executor.warmup.replayssm_warmup import (

45 replayssm_autotune_warmup,

@@ -157,6 +160,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):

160 if worker.vllm_config.kernel_config.enable_jit_warmup:

161 kimi_k3_triton_warmup(worker)

162 fa4_cutedsl_warmup(worker)

163 + qwen4_exp_qsa_triton_warmup(worker)

164

165 if current_platform.has_device_capability(90):

166 _warmup_ll_bf16_router_gemm(worker.get_model())

vllm/model_executor/warmup/qwen4_exp_qsa_warmup.py new

+66

@@ -0,0 +1,66 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Connect loaded Qwen4Exp QSA modules to their kernel-owned warmup."""

4 +

5 +import sys

6 +from typing import TYPE_CHECKING, cast

7 +

8 +from vllm.logger import init_logger

9 +

10 +if TYPE_CHECKING:

11 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2

12 + from vllm.v1.worker.gpu_worker import Worker

13 +

14 +logger = init_logger(__name__)

15 +

16 +

17 +def qwen4_exp_qsa_triton_warmup(worker: "Worker") -> None:

18 + """Warm every reachable QSA decode-query-length specialization."""

19 +

20 + qsa_module = sys.modules.get("vllm.models.qwen4_exp.nvidia.indexer_qsa")

21 + if qsa_module is None:

22 + return

23 + indexer = next(

24 + (

25 + layer

26 + for layer in worker.get_model().modules()

27 + if isinstance(layer, qsa_module.QSAIndexer)

28 + ),

29 + None,

30 + )

31 + if indexer is None:

32 + return

33 +

34 + runner = worker.model_runner

35 + prefix = indexer.compressed_key_cache.prefix

36 + group_id = next(

37 + i

38 + for i, group in enumerate(runner.kv_cache_config.kv_cache_groups)

39 + if prefix in group.layer_names

40 + )

41 + if worker.use_v2_model_runner:

42 + runner_v2 = cast("GPUModelRunnerV2", runner)

43 + block_table = runner_v2.block_tables.input_block_tables[group_id]

44 + max_decode_query_len = runner_v2.decode_query_len

45 + else:

46 + block_table = runner.input_batch.block_table[group_id].get_device_tensor(

47 + runner.max_num_reqs

48 + )

49 + max_decode_query_len = runner.uniform_decode_query_len

50 +

51 + from vllm.models.qwen4_exp.nvidia.ops.qsa_indexer import (

52 + warmup_qsa_mqa_paged_decode,

53 + )

54 +

55 + k_cache = indexer.compressed_key_cache.kv_cache

56 + assert k_cache.numel()

57 + profiles = warmup_qsa_mqa_paged_decode(

58 + k_cache,

59 + block_table,

60 + num_heads=indexer.index_n_heads,

61 + head_dim=indexer.index_head_dim,

62 + max_decode_query_len=max_decode_query_len,

63 + max_num_reqs=runner.max_num_reqs,

64 + max_num_batched_tokens=runner.max_num_tokens,

65 + )

66 + logger.info("Warmed up Qwen4Exp QSA decode kernels: %s.", profiles)

vllm/models/qwen4_exp/common/qsa_cache.py

+86 -20

@@ -11,8 +11,6 @@ table follows the main KV-cache lifecycle. Their physical tensor storage is

11 shared by the generic cache-layout planner.

12 """

14 -from __future__ import annotations

15 -

14 import math

15 from dataclasses import dataclass

16 from functools import cache

@@ -34,7 +32,10 @@ from vllm.v1.attention.backend import (

32 AttentionMetadataBuilder,

33 CommonAttentionMetadata,

34 )

37 -from vllm.v1.attention.backends.utils import PAD_SLOT_ID

35 +from vllm.v1.attention.backends.utils import (

36 + PAD_SLOT_ID,

37 + split_decodes_and_prefills,

38 +)

39 from vllm.v1.kv_cache_interface import (

40 AttentionSpec,

41 CircularBufferSpec,

@@ -211,6 +212,7 @@ def _build_qsa_metadata_kernel(

212 block_table_ptr,

213 token_to_req_ptr,

214 logical_positions_ptr,

215 + visible_blocks_ptr,

216 slot_mapping_ptr,

217 k_work_metadata_ptr,

218 block_table_stride_0: tl.constexpr,

@@ -266,6 +268,14 @@ def _build_qsa_metadata_kernel(

268 logical_position,

269 mask=store_mask,

270 )

271 + visible_blocks = tl.maximum(

272 + 0,

273 + tl.minimum(

274 + (logical_position + 1) // compress_ratio,

275 + seq_len // compress_ratio,

276 + ),

277 + ).to(tl.int32)

278 + tl.store(visible_blocks_ptr + token_idx, visible_blocks, mask=store_mask)

279

280 # circular_buffer_size is constexpr, so each builder instance compiles out

281 # the other QSA cache owner's slot-mapping rule.

@@ -380,6 +390,7 @@ def build_qsa_metadata_triton(

390 common_attn_metadata: CommonAttentionMetadata,

391 token_to_req_buffer: torch.Tensor,

392 logical_positions_buffer: torch.Tensor,

393 + visible_blocks_buffer: torch.Tensor,

394 slot_mapping_buffer: torch.Tensor,

395 *,

396 storage_block_size: int,

@@ -387,12 +398,13 @@ def build_qsa_metadata_triton(

398 circular_buffer_size: int = 0,

399 k_work_metadata_buffer: torch.Tensor | None = None,

400 request_capacity: int | None = None,

390 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

401 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

402 """Build QSA side-cache and optional pre-indexer work metadata."""

403 num_tokens = common_attn_metadata.num_actual_tokens

404 num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])

405 token_to_req = token_to_req_buffer[:num_tokens]

406 logical_positions = logical_positions_buffer[:num_tokens]

407 + visible_blocks = visible_blocks_buffer[:num_tokens]

408 slot_mapping = slot_mapping_buffer[:num_tokens]

409 num_reqs = common_attn_metadata.query_start_loc.shape[0] - 1

410 assert num_reqs > 0

@@ -409,7 +421,7 @@ def build_qsa_metadata_triton(

421 max_num_work = 0

422

423 if num_tokens == 0 and k_work_metadata_buffer is None:

412 - return token_to_req, logical_positions, slot_mapping

424 + return token_to_req, logical_positions, visible_blocks, slot_mapping

425

426 block_table = common_attn_metadata.block_table_tensor

427 num_search_steps = int(math.ceil(math.log2(num_reqs)))

@@ -426,6 +438,7 @@ def build_qsa_metadata_triton(

438 block_table,

439 token_to_req,

440 logical_positions,

441 + visible_blocks,

442 slot_mapping,

443 k_work_metadata_buffer,

444 block_table.stride(0),

@@ -448,13 +461,14 @@ def build_qsa_metadata_triton(

461 )

462 if circular_buffer_size == 0 and compress_ratio == 1:

463 slot_mapping = common_attn_metadata.slot_mapping[:num_tokens]

451 - return token_to_req, logical_positions, slot_mapping

464 + return token_to_req, logical_positions, visible_blocks, slot_mapping

465

466

467 def _build_qsa_metadata_torch(

468 common_attn_metadata: CommonAttentionMetadata,

469 token_to_req_buffer: torch.Tensor,

470 logical_positions_buffer: torch.Tensor,

471 + visible_blocks_buffer: torch.Tensor,

472 slot_mapping_buffer: torch.Tensor,

473 *,

474 storage_block_size: int,

@@ -462,11 +476,12 @@ def _build_qsa_metadata_torch(

476 circular_buffer_size: int = 0,

477 k_work_metadata_buffer: torch.Tensor | None = None,

478 request_capacity: int | None = None,

465 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

479 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

480 del request_capacity

481 num_tokens = common_attn_metadata.num_actual_tokens

482 num_mapped_tokens = int(common_attn_metadata.query_start_loc_cpu[-1])

483 logical_positions = logical_positions_buffer[:num_tokens]

484 + visible_blocks = visible_blocks_buffer[:num_tokens]

485

486 token_to_req = common_attn_metadata.token_to_req_indices(token_to_req_buffer)[

487 :num_tokens

@@ -481,6 +496,15 @@ def _build_qsa_metadata_torch(

496 )

497 if num_mapped_tokens < num_tokens:

498 logical_positions[num_mapped_tokens:].fill_(-1)

499 + row_seq_lens = common_attn_metadata.seq_lens.index_select(0, token_to_req.long())

500 + visible_blocks.copy_(

501 + torch.minimum(

502 + (logical_positions + 1) // compress_ratio,

503 + row_seq_lens // compress_ratio,

504 + )

505 + .clamp_min_(0)

506 + .to(torch.int32)

507 + )

508 if circular_buffer_size > 0:

509 slot_mapping = circular_qsa_slot_mapping(

510 common_attn_metadata.block_table_tensor,

@@ -539,7 +563,7 @@ def _build_qsa_metadata_torch(

563 k_work_metadata_buffer[:, 1].copy_(

564 torch.where(active, work_in_request, -1).to(torch.int32)

565 )

542 - return token_to_req, logical_positions, slot_mapping

566 + return token_to_req, logical_positions, visible_blocks, slot_mapping

567

568

569 # Resolve the fallback outside the per-step metadata hot path.

@@ -558,8 +582,15 @@ class QSAForwardMetadata(AttentionMetadata):

582 query_start_loc: torch.Tensor

583 token_to_req: torch.Tensor

584 logical_positions: torch.Tensor

585 + visible_blocks: torch.Tensor

586 k_work_metadata: torch.Tensor

587 num_actual_tokens: int

588 + num_decodes: int

589 + num_decode_tokens: int

590 + num_prefills: int

591 + num_prefill_tokens: int

592 + max_query_len: int

593 + decode_query_len: int

594 storage_block_size: int

595 compress_ratio: int

596

@@ -577,6 +608,8 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

608 device: torch.device,

609 ) -> None:

610 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

611 + self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

612 + assert self.reorder_batch_threshold is not None

613 self.is_circular_buffer = isinstance(kv_cache_spec, CircularBufferSpec)

614 if isinstance(kv_cache_spec, MLAAttentionSpec):

615 compress_ratio = kv_cache_spec.tokens_per_state

@@ -597,6 +630,9 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

630 self.logical_positions_buffer = torch.empty(

631 max_tokens, dtype=torch.int64, device=device

632 )

633 + self.visible_blocks_buffer = torch.empty(

634 + max_tokens, dtype=torch.int32, device=device

635 + )

636 max_requests = vllm_config.scheduler_config.max_num_seqs

637 self.request_capacity = max_requests

638 if not self.is_circular_buffer and self.compress_ratio != 1:

@@ -619,6 +655,26 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

655 ) -> QSAForwardMetadata:

656 del common_prefix_len, fast_build

657 num_tokens = common_attn_metadata.num_actual_tokens

658 + decode_threshold = self.reorder_batch_threshold

659 + assert decode_threshold is not None

660 + num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

661 + split_decodes_and_prefills(

662 + common_attn_metadata,

663 + decode_threshold=decode_threshold,

664 + require_uniform=True,

665 + )

666 + )

667 + assert num_decodes + num_prefills == common_attn_metadata.num_reqs

668 + assert num_decode_tokens + num_prefill_tokens == num_tokens

669 + decode_query_len = 0

670 + if num_decodes > 0:

671 + query_lens_cpu = torch.diff(

672 + common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]

673 + )

674 + nonzero_query_lens = query_lens_cpu[query_lens_cpu > 0]

675 + if nonzero_query_lens.numel() > 0:

676 + decode_query_len = int(nonzero_query_lens[0].item())

677 + assert torch.all(nonzero_query_lens == decode_query_len)

678 build_k_work = not self.is_circular_buffer and self.compress_ratio != 1

679 k_work_metadata = self.k_work_metadata_buffer

680 request_capacity = None

@@ -629,18 +685,21 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

685 num_tokens + (self.compress_ratio - 1) * num_requests

686 ) // self.compress_ratio

687 k_work_metadata = self.k_work_metadata_buffer[:max_num_work]

632 - token_to_req, logical_positions, slot_mapping = build_qsa_metadata(

633 - common_attn_metadata,

634 - self.token_to_req_buffer,

635 - self.logical_positions_buffer,

636 - self.slot_mapping_buffer,

637 - storage_block_size=self.storage_block_size,

638 - compress_ratio=self.compress_ratio,

639 - circular_buffer_size=(

640 - self.kv_cache_spec.block_size if self.is_circular_buffer else 0

641 - ),

642 - k_work_metadata_buffer=k_work_metadata if build_k_work else None,

643 - request_capacity=request_capacity,

688 + token_to_req, logical_positions, visible_blocks, slot_mapping = (

689 + build_qsa_metadata(

690 + common_attn_metadata,

691 + self.token_to_req_buffer,

692 + self.logical_positions_buffer,

693 + self.visible_blocks_buffer,

694 + self.slot_mapping_buffer,

695 + storage_block_size=self.storage_block_size,

696 + compress_ratio=self.compress_ratio,

697 + circular_buffer_size=(

698 + self.kv_cache_spec.block_size if self.is_circular_buffer else 0

699 + ),

700 + k_work_metadata_buffer=k_work_metadata if build_k_work else None,

701 + request_capacity=request_capacity,

702 + )

703 )

704 return QSAForwardMetadata(

705 block_table=common_attn_metadata.block_table_tensor,

@@ -649,8 +708,15 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

708 query_start_loc=common_attn_metadata.query_start_loc,

709 token_to_req=token_to_req,

710 logical_positions=logical_positions,

711 + visible_blocks=visible_blocks,

712 k_work_metadata=k_work_metadata,

713 num_actual_tokens=num_tokens,

714 + num_decodes=num_decodes,

715 + num_decode_tokens=num_decode_tokens,

716 + num_prefills=num_prefills,

717 + num_prefill_tokens=num_prefill_tokens,

718 + max_query_len=common_attn_metadata.max_query_len,

719 + decode_query_len=decode_query_len,

720 storage_block_size=self.storage_block_size,

721 compress_ratio=self.compress_ratio,

722 )

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

+79 -15

@@ -2,8 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Qwen4Exp weight-free QSA indexer."""

5 -from __future__ import annotations

6 -

5 from typing import cast

7 import torch

@@ -185,6 +183,22 @@ class QSAIndexer(nn.Module):

183 )

184 if raw.num_actual_tokens != compressed.num_actual_tokens:

185 raise RuntimeError("QSA side-cache metadata token counts disagree")

186 + raw_split = (

187 + raw.num_decodes,

188 + raw.num_decode_tokens,

189 + raw.num_prefills,

190 + raw.num_prefill_tokens,

191 + raw.decode_query_len,

192 + )

193 + compressed_split = (

194 + compressed.num_decodes,

195 + compressed.num_decode_tokens,

196 + compressed.num_prefills,

197 + compressed.num_prefill_tokens,

198 + compressed.decode_query_len,

199 + )

200 + if raw_split != compressed_split:

201 + raise RuntimeError("QSA side-cache metadata batch splits disagree")

202 if not raw.logical_positions.is_cuda and (

203 not torch.equal(raw.logical_positions, compressed.logical_positions)

204 ):

@@ -215,10 +229,11 @@ class QSAIndexer(nn.Module):

229 return out

230 return result

231

218 - from .ops.qsa import (

219 - qsa_compress_groups_with_ratio,

220 - qsa_select_paged_tokens,

221 - qsa_store_cache_rows,

232 + from .ops.qsa import qsa_compress_groups_with_ratio, qsa_store_cache_rows

233 + from .ops.qsa_indexer import (

234 + expand_qsa_block_indices,

235 + qsa_select_paged_decode,

236 + qsa_select_paged_prefill,

237 )

238

239 raw_metadata, compressed_metadata = metadata

@@ -339,18 +354,67 @@ class QSAIndexer(nn.Module):

354 raise RuntimeError("QSA top-k reuse requires an output buffer")

355 return out

356

342 - # Score compressed keys, select blocks, then expand them to token indices.

343 - return qsa_select_paged_tokens(

344 - q,

345 - compressed_key_cache,

346 - compressed_metadata.block_table,

347 - compressed_metadata.token_to_req,

348 - compressed_metadata.logical_positions,

349 - compressed_metadata.seq_lens,

350 - self.token_topk,

357 + if out is None:

358 + out = torch.empty(

359 + num_tokens,

360 + self.output_width,

361 + dtype=torch.int32,

362 + device=q.device,

363 + )

364 + elif out.shape != (num_tokens, self.output_width):

365 + raise ValueError("QSA selection output has an invalid shape")

366 +

367 + num_decode_tokens = compressed_metadata.num_decode_tokens

368 + decode_query_len = compressed_metadata.decode_query_len

369 + visible_blocks = compressed_metadata.visible_blocks[:num_tokens]

370 + block_indices = torch.empty(

371 + num_tokens,

372 + self.token_topk // self.compress_ratio,

373 + dtype=torch.int32,

374 + device=q.device,

375 + )

376 +

377 + # Decode requests occupy the leading rows and share one query length.

378 + if num_decode_tokens:

379 + num_decodes = compressed_metadata.num_decodes

380 + if num_decodes * decode_query_len != num_decode_tokens:

381 + raise ValueError("QSA decode rows must form a uniform request batch")

382 + decode_slice = slice(0, num_decode_tokens)

383 + qsa_select_paged_decode(

384 + q[decode_slice],

385 + compressed_key_cache,

386 + compressed_metadata.block_table[:num_decodes],

387 + visible_blocks[decode_slice],

388 + self.token_topk,

389 + self.compress_ratio,

390 + decode_query_len,

391 + block_indices[decode_slice],

392 + )

393 +

394 + # Prefill requests follow the leading decode rows in the reordered batch.

395 + if num_decode_tokens < num_tokens:

396 + num_decodes = compressed_metadata.num_decodes

397 + prefill_slice = slice(num_decode_tokens, num_tokens)

398 + qsa_select_paged_prefill(

399 + q[prefill_slice],

400 + compressed_key_cache,

401 + compressed_metadata.block_table[num_decodes:],

402 + compressed_metadata.query_start_loc[num_decodes:],

403 + visible_blocks[prefill_slice],

404 + self.token_topk,

405 + self.compress_ratio,

406 + compressed_metadata.max_query_len,

407 + block_indices[prefill_slice],

408 + )

409 + expand_qsa_block_indices(

410 + block_indices,

411 + compressed_metadata.logical_positions[:num_tokens],

412 + visible_blocks,

413 self.compress_ratio,

414 + self.token_topk,

415 out,

416 )

417 + return out

418

419

420 __all__ = ["QSAIndexer", "apply_qsa_rope"]

vllm/models/qwen4_exp/nvidia/ops/qsa.py

+1 -405

@@ -1,193 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Triton kernels for the Qwen4Exp weight-free QSA path."""

3 +"""Triton kernels for Qwen4Exp QSA sparse attention and cache updates."""

5 from __future__ import annotations

7 -import math

8 -

7 import torch

11 -from vllm.platforms import current_platform

9 from vllm.triton_utils import HAS_TRITON, tl, triton

14 -_LOGITS_WORKSPACE_BYTES = 128 * 1024 * 1024

15 -_TOPK_WORKSPACE_BYTES = 1024 * 1024

16 -

17 -

18 -@triton.jit

19 -def _qsa_mqa_paged_kernel(

20 - q_ptr,

21 - k_cache_ptr,

22 - page_table_ptr,

23 - token_to_req_ptr,

24 - query_positions_ptr,

25 - sequence_lengths_ptr,

26 - visible_blocks_ptr,

27 - logits_ptr,

28 - stride_q_row,

29 - stride_q_head,

30 - stride_q_dim,

31 - stride_cache_block,

32 - stride_cache_token,

33 - stride_cache_dim,

34 - stride_table_req,

35 - stride_table_page,

36 - stride_logits_row,

37 - num_rows,

38 - num_columns,

39 - num_pages,

40 - num_requests,

41 - score_divisor,

42 - PAGE_SIZE: tl.constexpr,

43 - PAGE_TABLE_WIDTH: tl.constexpr,

44 - NUM_HEADS: tl.constexpr,

45 - HEAD_DIM: tl.constexpr,

46 - BLOCK_N: tl.constexpr,

47 - BLOCK_D: tl.constexpr,

48 - TILES_PER_PROG: tl.constexpr,

49 - STAGES: tl.constexpr,

50 - MAX_N: tl.constexpr,

51 - COMPRESS_RATIO: tl.constexpr,

52 -) -> None:

53 - row = tl.program_id(0)

54 - dims = tl.arange(0, BLOCK_D)

55 - heads = tl.arange(0, MAX_N)

56 - request = tl.load(token_to_req_ptr + row)

57 - safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)

58 - query_position = tl.load(query_positions_ptr + row)

59 - sequence_length = tl.load(

60 - sequence_lengths_ptr + safe_request,

61 - mask=(request >= 0) & (request < num_requests),

62 - other=0,

63 - )

64 - visible = tl.minimum(

65 - (query_position + 1) // COMPRESS_RATIO,

66 - sequence_length // COMPRESS_RATIO,

67 - )

68 - if tl.program_id(1) == 0:

69 - tl.store(visible_blocks_ptr + row, visible)

70 - tile_start = tl.program_id(1) * TILES_PER_PROG

71 - # Top-k is bounded by visible_blocks, so columns beyond it need no value.

72 - if tile_start * BLOCK_N >= visible:

73 - return

74 - tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(visible, BLOCK_N))

75 - tile_end = tl.minimum(tile_end, tl.cdiv(num_columns, BLOCK_N))

76 -

77 - # Pad the small head axis to a tensor-core-compatible N dimension.

78 - query = tl.load(

79 - q_ptr

80 - + row * stride_q_row

81 - + heads[None, :] * stride_q_head

82 - + dims[:, None] * stride_q_dim,

83 - mask=(heads[None, :] < NUM_HEADS) & (dims[:, None] < HEAD_DIM),

84 - other=0.0,

85 - )

86 - column_offsets = tl.arange(0, BLOCK_N)

87 - for tile in tl.range(tile_start, tile_end, num_stages=STAGES):

88 - columns = tile * BLOCK_N + column_offsets

89 - live = columns < visible

90 - logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1)

91 - page_offset = columns % PAGE_SIZE

92 - physical_page = tl.load(

93 - page_table_ptr

94 - + safe_request * stride_table_req

95 - + logical_page * stride_table_page,

96 - mask=live,

97 - other=-1,

98 - )

99 - page_valid = live & (physical_page >= 0) & (physical_page < num_pages)

100 - # physical_page * block stride can overflow int32 for large caches.

101 - safe_physical_page = tl.maximum(physical_page, 0).to(tl.int64)

102 - keys = tl.load(

103 - k_cache_ptr

104 - + safe_physical_page[:, None] * stride_cache_block

105 - + page_offset[:, None] * stride_cache_token

106 - + dims[None, :] * stride_cache_dim,

107 - mask=page_valid[:, None] & (dims[None, :] < HEAD_DIM),

108 - other=0.0,

109 - eviction_policy="evict_first",

110 - )

111 - scores = tl.dot(keys, query, out_dtype=tl.float32)

112 - scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(scores, 0.0), 0.0)

113 - score = tl.sum(scores, axis=1) / score_divisor

114 - tl.store(

115 - logits_ptr + row * stride_logits_row + columns,

116 - tl.where(page_valid, score, -float("inf")),

117 - mask=live & (columns < num_columns),

118 - )

119 -

120 -

121 -@triton.jit

122 -def _expand_qsa_indices_kernel(

123 - block_indices_ptr,

124 - query_positions_ptr,

125 - sequence_lengths_ptr,

126 - token_to_req_ptr,

127 - output_ptr,

128 - stride_blocks_row,

129 - stride_blocks_column,

130 - stride_output_row,

131 - stride_output_column,

132 - rows,

133 - num_requests,

134 - BLOCK_TOPK: tl.constexpr,

135 - COMPRESS_RATIO: tl.constexpr,

136 - TOKEN_TOPK: tl.constexpr,

137 - OUTPUT_WIDTH: tl.constexpr,

138 - COLUMN_BLOCK: tl.constexpr,

139 -) -> None:

140 - row = tl.program_id(0)

141 - columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK)

142 - query_position = tl.load(query_positions_ptr + row)

143 - request = tl.load(token_to_req_ptr + row)

144 - safe_request = tl.minimum(tl.maximum(request, 0), num_requests - 1)

145 - sequence_length = tl.load(

146 - sequence_lengths_ptr + safe_request,

147 - mask=(request >= 0) & (request < num_requests),

148 - other=0,

149 - )

150 - complete_blocks = tl.minimum(

151 - tl.minimum(

152 - (query_position + 1) // COMPRESS_RATIO,

153 - sequence_length // COMPRESS_RATIO,

154 - ),

155 - BLOCK_TOPK,

156 - )

157 - expanded_count = complete_blocks * COMPRESS_RATIO

158 - tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO

159 - tail_count = (query_position + 1) - tail_start

160 -

161 - is_expanded = columns < expanded_count

162 - block_rank = columns // COMPRESS_RATIO

163 - offset = columns % COMPRESS_RATIO

164 - safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1)

165 - block = tl.load(

166 - block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column,

167 - mask=(row < rows) & is_expanded,

168 - other=-1,

169 - )

170 - expanded = block * COMPRESS_RATIO + offset

171 - tail_offset = columns - expanded_count

172 - is_tail = (

173 - (columns >= expanded_count)

174 - & (tail_offset < tail_count)

175 - & (tail_offset < COMPRESS_RATIO - 1)

176 - )

177 - token = tl.where(is_expanded, expanded, tail_start + tail_offset)

178 - valid = (

179 - (row < rows)

180 - & (columns < OUTPUT_WIDTH)

181 - & (is_expanded | is_tail)

182 - & (token >= 0)

183 - & (token < sequence_length)

184 - )

185 - tl.store(

186 - output_ptr + row * stride_output_row + columns * stride_output_column,

187 - tl.where(valid, token, -1),

188 - mask=(row < rows) & (columns < OUTPUT_WIDTH),

189 - )

190 -

12 @triton.jit

13 def _qsa_sparse_paged_gqa_splitk_kernel(

@@ -588,227 +408,6 @@ def _compress_qsa_groups_kernel(

408 )

409

410

591 -def _validate_mqa(q: torch.Tensor) -> None:

592 - if q.ndim != 3 or q.shape[1] <= 0 or q.shape[2] <= 0:

593 - raise ValueError("QSA query must be [rows, heads, head_dim]")

594 -

595 -

596 -def qsa_mqa_paged(

597 - q: torch.Tensor,

598 - k_cache: torch.Tensor,

599 - page_table: torch.Tensor,

600 - token_to_req: torch.Tensor,

601 - query_positions: torch.Tensor,

602 - sequence_lengths: torch.Tensor,

603 - compress_ratio: int,

604 - num_columns: int | None = None,

605 - score_scale: float | None = None,

606 -) -> tuple[torch.Tensor, torch.Tensor]:

607 - """Compute QSA scores directly from a paged compressed-key cache."""

608 -

609 - _validate_mqa(q)

610 - if not q.is_cuda or not HAS_TRITON:

611 - raise RuntimeError("paged QSA scoring requires CUDA and Triton")

612 - if k_cache.ndim != 4 or k_cache.shape[2] != 1:

613 - raise ValueError("QSA cache must be [pages, page_size, 1, head_dim]")

614 - if k_cache.shape[3] != q.shape[2]:

615 - raise ValueError("QSA query and cache dimensions must match")

616 - if page_table.ndim != 2:

617 - raise ValueError("QSA page table must be two-dimensional")

618 - if q.shape[0] and (not all(k_cache.shape[:2]) or not all(page_table.shape)):

619 - raise ValueError("QSA paged scoring cache and page table must be nonempty")

620 - if token_to_req.shape != (q.shape[0],):

621 - raise ValueError("QSA request mapping must match query rows")

622 - if query_positions.shape != (q.shape[0],):

623 - raise ValueError("QSA query positions must match query rows")

624 - if sequence_lengths.shape != (page_table.shape[0],):

625 - raise ValueError("QSA sequence lengths must match page-table requests")

626 - if compress_ratio <= 0:

627 - raise ValueError("QSA compression ratio must be positive")

628 - score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale

629 - if score_divisor <= 0:

630 - raise ValueError("QSA score scale must be positive")

631 -

632 - capacity = page_table.shape[1] * k_cache.shape[1]

633 - columns = capacity if num_columns is None else num_columns

634 - if columns < 0:

635 - raise ValueError("QSA score width must be non-negative")

636 - logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device)

637 - visible_blocks = torch.empty(q.shape[0], dtype=torch.int32, device=q.device)

638 - if not q.shape[0] or not columns:

639 - return logits, visible_blocks

640 - BLOCK_N = 64

641 - BLOCK_D = max(16, triton.next_power_of_2(q.shape[2]))

642 - MAX_N = max(16, triton.next_power_of_2(q.shape[1]))

643 - # Tuned on GB300: larger row batches provide enough parallelism to reuse Q.

644 - tiles_per_program = 1 if q.shape[0] <= 32 else 8

645 - _qsa_mqa_paged_kernel[

646 - (q.shape[0], triton.cdiv(columns, BLOCK_N * tiles_per_program))

647 - ](

648 - q,

649 - k_cache,

650 - page_table,

651 - token_to_req,

652 - query_positions,

653 - sequence_lengths,

654 - visible_blocks,

655 - logits,

656 - q.stride(0),

657 - q.stride(1),

658 - q.stride(2),

659 - k_cache.stride(0),

660 - k_cache.stride(1),

661 - k_cache.stride(3),

662 - page_table.stride(0),

663 - page_table.stride(1),

664 - logits.stride(0),

665 - q.shape[0],

666 - columns,

667 - k_cache.shape[0],

668 - page_table.shape[0],

669 - float(score_divisor),

670 - PAGE_SIZE=k_cache.shape[1],

671 - PAGE_TABLE_WIDTH=page_table.shape[1],

672 - NUM_HEADS=q.shape[1],

673 - HEAD_DIM=q.shape[2],

674 - BLOCK_N=BLOCK_N,

675 - BLOCK_D=BLOCK_D,

676 - TILES_PER_PROG=tiles_per_program,

677 - STAGES=2,

678 - MAX_N=MAX_N,

679 - COMPRESS_RATIO=compress_ratio,

680 - num_warps=2,

681 - )

682 - return logits, visible_blocks

683 -

684 -

685 -def expand_qsa_block_indices_cuda(

686 - block_indices: torch.Tensor,

687 - query_positions: torch.Tensor,

688 - sequence_lengths: torch.Tensor,

689 - token_to_req: torch.Tensor,

690 - compress_ratio: int,

691 - token_topk: int,

692 - out: torch.Tensor | None = None,

693 -) -> torch.Tensor:

694 - """Expand compressed blocks and compact the causal tail of the open group."""

695 -

696 - if not block_indices.is_cuda or not HAS_TRITON:

697 - raise RuntimeError("QSA CUDA expansion requires Triton")

698 - if token_topk % compress_ratio:

699 - raise ValueError("QSA token top-k must be divisible by compression ratio")

700 - block_topk = token_topk // compress_ratio

701 - output_width = token_topk + compress_ratio - 1

702 - if block_indices.shape != (query_positions.numel(), block_topk):

703 - raise ValueError("QSA compressed top-k has an invalid shape")

704 - if token_to_req.shape != query_positions.shape:

705 - raise ValueError("QSA request mapping must match query positions")

706 - if sequence_lengths.ndim != 1 or not sequence_lengths.shape[0]:

707 - raise ValueError("QSA request sequence lengths must be nonempty")

708 - if out is None:

709 - out = torch.empty(

710 - (block_indices.shape[0], output_width),

711 - dtype=torch.int32,

712 - device=block_indices.device,

713 - )

714 - elif out.shape != (block_indices.shape[0], output_width):

715 - raise ValueError("QSA expansion output has an invalid shape")

716 - if not block_indices.shape[0]:

717 - return out

718 - column_block = 256

719 - _expand_qsa_indices_kernel[

720 - (block_indices.shape[0], triton.cdiv(output_width, column_block))

721 - ](

722 - block_indices,

723 - query_positions,

724 - sequence_lengths,

725 - token_to_req,

726 - out,

727 - block_indices.stride(0),

728 - block_indices.stride(1),

729 - out.stride(0),

730 - out.stride(1),

731 - block_indices.shape[0],

732 - sequence_lengths.shape[0],

733 - BLOCK_TOPK=block_topk,

734 - COMPRESS_RATIO=compress_ratio,

735 - TOKEN_TOPK=token_topk,

736 - OUTPUT_WIDTH=output_width,

737 - COLUMN_BLOCK=column_block,

738 - num_warps=4,

739 - )

740 - return out

741 -

742 -

743 -def qsa_select_paged_tokens(

744 - q: torch.Tensor,

745 - k_cache: torch.Tensor,

746 - page_table: torch.Tensor,

747 - token_to_req: torch.Tensor,

748 - query_positions: torch.Tensor,

749 - sequence_lengths: torch.Tensor,

750 - token_topk: int,

751 - compress_ratio: int,

752 - out: torch.Tensor | None = None,

753 -) -> torch.Tensor:

754 - """Score, select, and expand QSA indices without host synchronization."""

755 -

756 - rows = q.shape[0]

757 - output_width = token_topk + compress_ratio - 1

758 - if out is None:

759 - out = torch.empty((rows, output_width), dtype=torch.int32, device=q.device)

760 - if out.shape != (rows, output_width):

761 - raise ValueError("QSA selection output has an invalid shape")

762 - if not rows:

763 - return out

764 -

765 - columns = page_table.shape[1] * k_cache.shape[1]

766 - block_topk = token_topk // compress_ratio

767 - rows_per_chunk = max(1, _LOGITS_WORKSPACE_BYTES // max(columns * 4, 1))

768 - chunk_rows = min(rows, rows_per_chunk)

769 - blocks_buffer = torch.empty(

770 - (chunk_rows, block_topk), dtype=torch.int32, device=q.device

771 - )

772 - topk_workspace = torch.empty(

773 - (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device

774 - )

775 - for row_start in range(0, rows, rows_per_chunk):

776 - row_end = min(row_start + rows_per_chunk, rows)

777 - row_slice = slice(row_start, row_end)

778 - logits, visible_blocks = qsa_mqa_paged(

779 - q[row_slice],

780 - k_cache,

781 - page_table,

782 - token_to_req[row_slice],

783 - query_positions[row_slice],

784 - sequence_lengths,

785 - compress_ratio,

786 - )

787 - blocks = blocks_buffer[: row_end - row_start]

788 - use_cooperative_topk = (

789 - blocks.shape[0] <= 32

790 - and logits.stride(0) % 4 == 0

791 - and current_platform.has_device_capability(90)

792 - and not current_platform.is_device_capability_family(120)

793 - )

794 - topk_op = (

795 - torch.ops._C.cooperative_topk

796 - if use_cooperative_topk

797 - else torch.ops._C.persistent_topk

798 - )

799 - topk_op(logits, visible_blocks, blocks, topk_workspace, block_topk, columns)

800 - expand_qsa_block_indices_cuda(

801 - blocks,

802 - query_positions[row_slice],

803 - sequence_lengths,

804 - token_to_req[row_slice],

805 - compress_ratio,

806 - token_topk,

807 - out[row_slice],

808 - )

809 - return out

810 -

811 -

411 def qsa_sparse_paged_attention(

412 q: torch.Tensor,

413 k_cache: torch.Tensor,

@@ -1106,10 +705,7 @@ def qsa_compress_groups_with_ratio(

705

706

707 __all__ = [

1109 - "expand_qsa_block_indices_cuda",

708 "qsa_compress_groups_with_ratio",

1111 - "qsa_mqa_paged",

1112 - "qsa_select_paged_tokens",

709 "qsa_sparse_paged_attention",

710 "qsa_store_cache_rows",

711 ]

vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py new

+617

@@ -0,0 +1,617 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Triton kernels for Qwen4Exp QSA index selection."""

4 +

5 +import torch

6 +

7 +import vllm.envs as envs

8 +from vllm.model_executor.warmup.jit_warmup_triton_helper import (

9 + TritonWarmupTensor,

10 +)

11 +from vllm.platforms import current_platform

12 +from vllm.triton_utils import tl, triton

13 +

14 +_TOPK_WORKSPACE_BYTES = 1024 * 1024

15 +_DECODE_BLOCK_N = 64

16 +

17 +

18 +@triton.jit

19 +def _qsa_mqa_paged_uniform_kernel(

20 + q_ptr,

21 + k_cache_ptr,

22 + page_table_ptr,

23 + visible_blocks_ptr,

24 + logits_ptr,

25 + stride_q_row,

26 + stride_q_head,

27 + stride_cache_block,

28 + stride_cache_token,

29 + stride_table_req,

30 + stride_logits_row,

31 + PAGE_SIZE: tl.constexpr,

32 + PAGE_TABLE_WIDTH: tl.constexpr,

33 + NUM_HEADS: tl.constexpr,

34 + HEAD_DIM: tl.constexpr,

35 + DECODE_QUERY_LEN: tl.constexpr,

36 + BLOCK_N: tl.constexpr,

37 + TILES_PER_PROG: tl.constexpr,

38 + STAGES: tl.constexpr,

39 +) -> None:

40 + NUM_COLUMNS: tl.constexpr = PAGE_TABLE_WIDTH * PAGE_SIZE

41 + DECODE_QUERY_LEN_PADDED: tl.constexpr = triton.next_power_of_2(DECODE_QUERY_LEN)

42 + NUM_HEADS_PADDED: tl.constexpr = triton.next_power_of_2(NUM_HEADS)

43 + # tl.dot requires a reduction dimension of at least 16.

44 + BLOCK_D: tl.constexpr = max(16, triton.next_power_of_2(HEAD_DIM))

45 + request = tl.program_id(0)

46 + tile_start = tl.program_id(1) * TILES_PER_PROG

47 + query_offsets = tl.arange(0, DECODE_QUERY_LEN_PADDED)

48 + valid_query_offsets = query_offsets < DECODE_QUERY_LEN

49 + rows = request * DECODE_QUERY_LEN + query_offsets

50 + visible = tl.load(

51 + visible_blocks_ptr + rows,

52 + mask=valid_query_offsets,

53 + other=0,

54 + )

55 + max_visible = tl.max(visible, axis=0)

56 + if tile_start * BLOCK_N >= max_visible:

57 + return

58 + tile_end = tl.minimum(tile_start + TILES_PER_PROG, tl.cdiv(max_visible, BLOCK_N))

59 + tile_end = tl.minimum(tile_end, tl.cdiv(NUM_COLUMNS, BLOCK_N))

60 +

61 + dims = tl.arange(0, BLOCK_D)

62 + n = tl.arange(0, DECODE_QUERY_LEN_PADDED * NUM_HEADS_PADDED)

63 + query_offset = n // NUM_HEADS_PADDED

64 + head = n % NUM_HEADS_PADDED

65 + valid_query = (query_offset < DECODE_QUERY_LEN) & (head < NUM_HEADS)

66 + query = tl.load(

67 + q_ptr

68 + + (request * DECODE_QUERY_LEN + query_offset)[None, :] * stride_q_row

69 + + head[None, :] * stride_q_head

70 + + dims[:, None],

71 + mask=valid_query[None, :] & (dims[:, None] < HEAD_DIM),

72 + other=0.0,

73 + )

74 + column_offsets = tl.arange(0, BLOCK_N)

75 + for tile in tl.range(tile_start, tile_end, num_stages=STAGES):

76 + columns = tile * BLOCK_N + column_offsets

77 + live = columns < max_visible

78 + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1)

79 + page_offset = columns % PAGE_SIZE

80 + physical_page = tl.load(

81 + page_table_ptr + request * stride_table_req + logical_page,

82 + mask=live,

83 + other=0,

84 + )

85 + keys = tl.load(

86 + k_cache_ptr

87 + + physical_page[:, None].to(tl.int64) * stride_cache_block

88 + + page_offset[:, None] * stride_cache_token

89 + + dims[None, :],

90 + mask=live[:, None] & (dims[None, :] < HEAD_DIM),

91 + other=0.0,

92 + eviction_policy="evict_first",

93 + )

94 + scores = tl.dot(keys, query, out_dtype=tl.float32)

95 + scores = tl.where(valid_query[None, :], tl.maximum(scores, 0.0), 0.0)

96 + scores = tl.reshape(

97 + scores,

98 + (BLOCK_N, DECODE_QUERY_LEN_PADDED, NUM_HEADS_PADDED),

99 + )

100 + score = tl.sum(scores, axis=2) / HEAD_DIM**0.5

101 + tl.store(

102 + logits_ptr + rows[None, :] * stride_logits_row + columns[:, None],

103 + score,

104 + mask=valid_query_offsets[None, :]

105 + & (columns[:, None] < NUM_COLUMNS)

106 + & (columns[:, None] < visible[None, :]),

107 + )

108 +

109 +

110 +@triton.jit(do_not_specialize=["num_rows", "query_offset"])

111 +def _qsa_mqa_paged_prefill_kernel(

112 + q_ptr,

113 + k_cache_ptr,

114 + page_table_ptr,

115 + query_start_loc_ptr,

116 + visible_blocks_ptr,

117 + logits_ptr,

118 + stride_q_row,

119 + stride_q_head,

120 + stride_cache_block,

121 + stride_cache_token,

122 + stride_table_req,

123 + stride_logits_row,

124 + num_rows,

125 + query_offset,

126 + PAGE_SIZE: tl.constexpr,

127 + PAGE_TABLE_WIDTH: tl.constexpr,

128 + NUM_HEADS: tl.constexpr,

129 + HEAD_DIM: tl.constexpr,

130 + TILE_R: tl.constexpr,

131 + BLOCK_N: tl.constexpr,

132 + K_TILES: tl.constexpr,

133 + STAGES: tl.constexpr,

134 +) -> None:

135 + NUM_COLUMNS: tl.constexpr = PAGE_TABLE_WIDTH * PAGE_SIZE

136 + NUM_HEADS_PADDED: tl.constexpr = triton.next_power_of_2(NUM_HEADS)

137 + # tl.dot requires a reduction dimension of at least 16.

138 + BLOCK_D: tl.constexpr = max(16, triton.next_power_of_2(HEAD_DIM))

139 + request = tl.program_id(0)

140 + query_base = tl.load(query_start_loc_ptr)

141 + query_end = query_offset + num_rows

142 + request_start = tl.maximum(

143 + tl.load(query_start_loc_ptr + request) - query_base, query_offset

144 + )

145 + request_end = tl.minimum(

146 + tl.load(query_start_loc_ptr + request + 1) - query_base, query_end

147 + )

148 + absolute_row_start = request_start + tl.program_id(1) * TILE_R

149 + if absolute_row_start >= request_end:

150 + return

151 +

152 + lanes = tl.arange(0, TILE_R)

153 + absolute_rows = absolute_row_start + lanes

154 + rows = absolute_rows - query_offset

155 + valid_rows = absolute_rows < request_end

156 + visible = tl.load(

157 + visible_blocks_ptr + absolute_rows,

158 + mask=valid_rows,

159 + other=0,

160 + )

161 + max_visible = tl.max(visible, axis=0)

162 + k_tile_start = tl.program_id(2) * K_TILES

163 + if k_tile_start * BLOCK_N >= max_visible:

164 + return

165 + k_tile_end = tl.minimum(k_tile_start + K_TILES, tl.cdiv(max_visible, BLOCK_N))

166 + k_tile_end = tl.minimum(k_tile_end, tl.cdiv(NUM_COLUMNS, BLOCK_N))

167 +

168 + dims = tl.arange(0, BLOCK_D)

169 + m = tl.arange(0, TILE_R * NUM_HEADS_PADDED)

170 + q_row_offsets = m // NUM_HEADS_PADDED

171 + q_rows = absolute_row_start + q_row_offsets

172 + heads = m % NUM_HEADS_PADDED

173 + query = tl.load(

174 + q_ptr

175 + + q_rows[None, :] * stride_q_row

176 + + heads[None, :] * stride_q_head

177 + + dims[:, None],

178 + mask=(heads[None, :] < NUM_HEADS)

179 + & (absolute_row_start + q_row_offsets[None, :] < request_end)

180 + & (dims[:, None] < HEAD_DIM),

181 + other=0.0,

182 + )

183 + column_offsets = tl.arange(0, BLOCK_N)

184 + for tile in tl.range(k_tile_start, k_tile_end, num_stages=STAGES):

185 + columns = tile * BLOCK_N + column_offsets

186 + live = columns < max_visible

187 + logical_page = tl.minimum(columns // PAGE_SIZE, PAGE_TABLE_WIDTH - 1)

188 + page_offset = columns % PAGE_SIZE

189 + physical_page = tl.load(

190 + page_table_ptr + request * stride_table_req + logical_page,

191 + mask=live,

192 + other=0,

193 + )

194 + keys = tl.load(

195 + k_cache_ptr

196 + + physical_page[:, None].to(tl.int64) * stride_cache_block

197 + + page_offset[:, None] * stride_cache_token

198 + + dims[None, :],

199 + mask=live[:, None] & (dims[None, :] < HEAD_DIM),

200 + other=0.0,

201 + eviction_policy="evict_first",

202 + )

203 + scores = tl.dot(keys, query, out_dtype=tl.float32)

204 + scores = tl.reshape(scores, (BLOCK_N, TILE_R, NUM_HEADS_PADDED))

205 + score = tl.sum(tl.maximum(scores, 0.0), axis=2) / HEAD_DIM**0.5

206 + store_mask = (

207 + valid_rows[None, :]

208 + & (columns[:, None] < visible[None, :])

209 + & (columns[:, None] < NUM_COLUMNS)

210 + )

211 + tl.store(

212 + logits_ptr + rows[None, :] * stride_logits_row + columns[:, None],

213 + score,

214 + mask=store_mask,

215 + )

216 +

217 +

218 +@triton.jit

219 +def _expand_qsa_indices_kernel(

220 + block_indices_ptr,

221 + query_positions_ptr,

222 + visible_blocks_ptr,

223 + output_ptr,

224 + stride_blocks_row,

225 + stride_blocks_column,

226 + stride_output_row,

227 + stride_output_column,

228 + BLOCK_TOPK: tl.constexpr,

229 + COMPRESS_RATIO: tl.constexpr,

230 + COLUMN_BLOCK: tl.constexpr,

231 +) -> None:

232 + OUTPUT_WIDTH: tl.constexpr = BLOCK_TOPK * COMPRESS_RATIO + COMPRESS_RATIO - 1

233 + row = tl.program_id(0)

234 + columns = tl.program_id(1) * COLUMN_BLOCK + tl.arange(0, COLUMN_BLOCK)

235 + query_position = tl.load(query_positions_ptr + row)

236 + visible_blocks = tl.load(visible_blocks_ptr + row)

237 + complete_blocks = tl.minimum(visible_blocks, BLOCK_TOPK)

238 + expanded_count = complete_blocks * COMPRESS_RATIO

239 + tail_start = ((query_position + 1) // COMPRESS_RATIO) * COMPRESS_RATIO

240 + tail_count = (query_position + 1) - tail_start

241 +

242 + is_expanded = columns < expanded_count

243 + block_rank = columns // COMPRESS_RATIO

244 + offset = columns % COMPRESS_RATIO

245 + safe_rank = tl.minimum(block_rank, BLOCK_TOPK - 1)

246 + block = tl.load(

247 + block_indices_ptr + row * stride_blocks_row + safe_rank * stride_blocks_column,

248 + mask=is_expanded,

249 + other=-1,

250 + )

251 + expanded = block * COMPRESS_RATIO + offset

252 + tail_offset = columns - expanded_count

253 + is_tail = (

254 + (columns >= expanded_count)

255 + & (tail_offset < tail_count)

256 + & (tail_offset < COMPRESS_RATIO - 1)

257 + )

258 + token = tl.where(is_expanded, expanded, tail_start + tail_offset)

259 + valid = (columns < OUTPUT_WIDTH) & (is_expanded | is_tail) & (token >= 0)

260 + tl.store(

261 + output_ptr + row * stride_output_row + columns * stride_output_column,

262 + tl.where(valid, token, -1),

263 + mask=columns < OUTPUT_WIDTH,

264 + )

265 +

266 +

267 +def _decode_tiles_per_program(num_requests: int, columns: int) -> int:

268 + programs = num_requests * triton.cdiv(columns, _DECODE_BLOCK_N)

269 + if programs < 16384:

270 + return 1

271 + if programs < 32768:

272 + return 2

273 + if programs < 131072:

274 + return 4

275 + return 8

276 +

277 +

278 +def _qsa_decode_warmup_profiles(

279 + max_dql: int,

280 + max_num_reqs: int,

281 + max_num_batched_tokens: int,

282 + columns: int,

283 +) -> tuple[tuple[int, int], ...]:

284 + profiles: list[tuple[int, int]] = []

285 + for dql in range(1, max_dql + 1):

286 + max_reqs = min(max_num_reqs, max_num_batched_tokens // dql)

287 + requests_by_grouping: dict[int, int] = {}

288 + for num_requests in range(1, max_reqs + 1):

289 + tiles_per_program = _decode_tiles_per_program(num_requests, columns)

290 + requests_by_grouping.setdefault(tiles_per_program, num_requests)

291 + profiles.extend(

292 + (dql, num_requests) for num_requests in requests_by_grouping.values()

293 + )

294 + return tuple(profiles)

295 +

296 +

297 +def warmup_qsa_mqa_paged_decode(

298 + k_cache: torch.Tensor,

299 + page_table: torch.Tensor,

300 + *,

301 + num_heads: int,

302 + head_dim: int,

303 + max_decode_query_len: int,

304 + max_num_reqs: int,

305 + max_num_batched_tokens: int,

306 +) -> tuple[tuple[int, int], ...]:

307 + """Compile every reachable decode specialization without launching it."""

308 +

309 + page_size = k_cache.shape[1]

310 + page_table_width = page_table.shape[1]

311 + columns = page_table_width * page_size

312 + profiles = _qsa_decode_warmup_profiles(

313 + max_decode_query_len,

314 + max_num_reqs,

315 + max_num_batched_tokens,

316 + columns,

317 + )

318 + if not profiles:

319 + return ()

320 +

321 + k_cache_ptr = TritonWarmupTensor(k_cache.dtype, shape=tuple(k_cache.shape))

322 + page_table_ptr = TritonWarmupTensor(

323 + page_table.dtype,

324 + shape=(max_num_reqs, page_table_width),

325 + )

326 + visible_blocks_ptr = TritonWarmupTensor(torch.int32)

327 +

328 + for decode_query_len, num_requests in profiles:

329 + num_rows = decode_query_len * num_requests

330 + q_ptr = TritonWarmupTensor(

331 + torch.bfloat16,

332 + shape=(num_rows, num_heads, head_dim),

333 + )

334 + logits_ptr = TritonWarmupTensor(

335 + torch.float32,

336 + shape=(num_rows, columns),

337 + )

338 + tiles_per_program = _decode_tiles_per_program(num_requests, columns)

339 + _qsa_mqa_paged_uniform_kernel.warmup(

340 + q_ptr,

341 + k_cache_ptr,

342 + page_table_ptr,

343 + visible_blocks_ptr,

344 + logits_ptr,

345 + num_heads * head_dim,

346 + head_dim,

347 + k_cache.stride(0),

348 + k_cache.stride(1),

349 + page_table.stride(0),

350 + columns,

351 + PAGE_SIZE=page_size,

352 + PAGE_TABLE_WIDTH=page_table_width,

353 + NUM_HEADS=num_heads,

354 + HEAD_DIM=head_dim,

355 + DECODE_QUERY_LEN=decode_query_len,

356 + BLOCK_N=_DECODE_BLOCK_N,

357 + TILES_PER_PROG=tiles_per_program,

358 + STAGES=2,

359 + num_warps=2,

360 + grid=(

361 + num_requests,

362 + triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program),

363 + ),

364 + )

365 + return profiles

366 +

367 +

368 +def _prefill_logits(

369 + q: torch.Tensor,

370 + k_cache: torch.Tensor,

371 + page_table: torch.Tensor,

372 + query_start_loc: torch.Tensor,

373 + visible_blocks: torch.Tensor,

374 + max_query_len: int,

375 + query_offset: int,

376 + num_queries: int,

377 +) -> torch.Tensor:

378 + assert query_start_loc.shape == (page_table.shape[0] + 1,)

379 + assert visible_blocks.shape == (q.shape[0],)

380 + assert 0 <= query_offset <= query_offset + num_queries <= q.shape[0]

381 +

382 + columns = page_table.shape[1] * k_cache.shape[1]

383 + logits = torch.empty((num_queries, columns), dtype=torch.float32, device=q.device)

384 + TILE_R = 64

385 + BLOCK_N = 64

386 + K_TILES = 16

387 + grid = (

388 + page_table.shape[0],

389 + triton.cdiv(min(num_queries, max_query_len), TILE_R),

390 + triton.cdiv(columns, BLOCK_N * K_TILES),

391 + )

392 + _qsa_mqa_paged_prefill_kernel[grid](

393 + q,

394 + k_cache,

395 + page_table,

396 + query_start_loc,

397 + visible_blocks,

398 + logits,

399 + *q.stride()[:-1],

400 + *k_cache.stride()[:2],

401 + *page_table.stride()[:-1],

402 + *logits.stride()[:-1],

403 + num_queries,

404 + query_offset,

405 + PAGE_SIZE=k_cache.shape[1],

406 + PAGE_TABLE_WIDTH=page_table.shape[1],

407 + NUM_HEADS=q.shape[1],

408 + HEAD_DIM=q.shape[2],

409 + TILE_R=TILE_R,

410 + BLOCK_N=BLOCK_N,

411 + K_TILES=K_TILES,

412 + STAGES=2,

413 + num_warps=4,

414 + )

415 + return logits

416 +

417 +

418 +def expand_qsa_block_indices(

419 + block_indices: torch.Tensor,

420 + query_positions: torch.Tensor,

421 + visible_blocks: torch.Tensor,

422 + compress_ratio: int,

423 + token_topk: int,

424 + out: torch.Tensor,

425 +) -> None:

426 + """Expand compressed blocks and compact the causal tail of the open group."""

427 +

428 + assert token_topk % compress_ratio == 0

429 + block_topk = token_topk // compress_ratio

430 + output_width = token_topk + compress_ratio - 1

431 + assert block_indices.shape == (query_positions.numel(), block_topk)

432 + assert visible_blocks.shape == query_positions.shape

433 + assert out.shape == (block_indices.shape[0], output_width)

434 + column_block = 256

435 + grid = (block_indices.shape[0], triton.cdiv(output_width, column_block))

436 + _expand_qsa_indices_kernel[grid](

437 + block_indices,

438 + query_positions,

439 + visible_blocks,

440 + out,

441 + *block_indices.stride(),

442 + *out.stride(),

443 + BLOCK_TOPK=block_topk,

444 + COMPRESS_RATIO=compress_ratio,

445 + COLUMN_BLOCK=column_block,

446 + num_warps=4,

447 + )

448 +

449 +

450 +def _topk(

451 + logits: torch.Tensor,

452 + visible_blocks: torch.Tensor,

453 + token_topk: int,

454 + compress_ratio: int,

455 + block_indices: torch.Tensor,

456 + topk_workspace: torch.Tensor,

457 +) -> None:

458 + # similar dispatch logic as DeepSeek indexer

459 + block_topk = token_topk // compress_ratio

460 + use_cooperative_topk = (

461 + logits.shape[0] <= 64

462 + and logits.stride(0) % 4 == 0

463 + and current_platform.has_device_capability(90)

464 + and not current_platform.is_device_capability_family(120)

465 + )

466 + topk_op = (

467 + torch.ops._C.cooperative_topk

468 + if use_cooperative_topk

469 + else torch.ops._C.persistent_topk

470 + )

471 + topk_op(

472 + logits,

473 + visible_blocks,

474 + block_indices,

475 + topk_workspace,

476 + block_topk,

477 + logits.shape[1],

478 + )

479 +

480 +

481 +def qsa_select_paged_decode(

482 + q: torch.Tensor,

483 + k_cache: torch.Tensor,

484 + page_table: torch.Tensor,

485 + visible_blocks: torch.Tensor,

486 + token_topk: int,

487 + compress_ratio: int,

488 + decode_query_len: int,

489 + block_indices: torch.Tensor,

490 +) -> None:

491 + """Score and select compressed blocks for a request-major decode batch.

492 +

493 + Args:

494 + q: Query tensor shaped ``[num_requests * decode_query_len, heads,

495 + head_dim]``.

496 + k_cache: Compressed key cache shaped ``[blocks, page_size, 1,

497 + head_dim]``.

498 + page_table: Request block table shaped ``[num_requests, max_pages]``.

499 + visible_blocks: Number of visible compressed blocks per query.

500 + token_topk: Number of logical tokens selected per query.

501 + compress_ratio: Number of logical tokens represented by a cache row.

502 + decode_query_len: Number of query tokens per request.

503 + block_indices: Compressed-index output buffer.

504 + """

505 +

506 + assert token_topk % compress_ratio == 0

507 + assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)

508 + assert decode_query_len > 0 and q.shape[0] % decode_query_len == 0

509 + num_requests = q.shape[0] // decode_query_len

510 + assert page_table.shape[0] == num_requests

511 + assert visible_blocks.shape == (q.shape[0],)

512 +

513 + columns = page_table.shape[1] * k_cache.shape[1]

514 + logits = torch.empty((q.shape[0], columns), dtype=torch.float32, device=q.device)

515 + tiles_per_program = _decode_tiles_per_program(num_requests, columns)

516 + grid = (

517 + num_requests,

518 + triton.cdiv(columns, _DECODE_BLOCK_N * tiles_per_program),

519 + )

520 + _qsa_mqa_paged_uniform_kernel[grid](

521 + q,

522 + k_cache,

523 + page_table,

524 + visible_blocks,

525 + logits,

526 + *q.stride()[:-1],

527 + *k_cache.stride()[:2],

528 + *page_table.stride()[:-1],

529 + *logits.stride()[:-1],

530 + PAGE_SIZE=k_cache.shape[1],

531 + PAGE_TABLE_WIDTH=page_table.shape[1],

532 + NUM_HEADS=q.shape[1],

533 + HEAD_DIM=q.shape[2],

534 + DECODE_QUERY_LEN=decode_query_len,

535 + BLOCK_N=_DECODE_BLOCK_N,

536 + TILES_PER_PROG=tiles_per_program,

537 + STAGES=2,

538 + num_warps=2,

539 + )

540 + _topk(

541 + logits,

542 + visible_blocks,

543 + token_topk,

544 + compress_ratio,

545 + block_indices,

546 + torch.empty((_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device),

547 + )

548 +

549 +

550 +def qsa_select_paged_prefill(

551 + q: torch.Tensor,

552 + k_cache: torch.Tensor,

553 + page_table: torch.Tensor,

554 + query_start_loc: torch.Tensor,

555 + visible_blocks: torch.Tensor,

556 + token_topk: int,

557 + compress_ratio: int,

558 + max_query_len: int,

559 + block_indices: torch.Tensor,

560 +) -> None:

561 + """Score and select compressed prefill blocks in bounded chunks.

562 +

563 + Args:

564 + q: Packed prefill query tensor shaped ``[num_tokens, heads, head_dim]``.

565 + k_cache: Compressed key cache shaped ``[blocks, page_size, 1,

566 + head_dim]``.

567 + page_table: Block table shaped ``[num_requests, max_pages]``.

568 + query_start_loc: Packed prefill query offsets with a terminal offset.

569 + Offsets may share the base of a larger backing tensor.

570 + visible_blocks: Number of visible compressed blocks per query.

571 + token_topk: Number of logical tokens selected per query.

572 + compress_ratio: Number of logical tokens represented by a cache row.

573 + max_query_len: Maximum number of query tokens in one request.

574 + block_indices: Compressed-index output buffer.

575 + """

576 +

577 + assert token_topk % compress_ratio == 0

578 + assert block_indices.shape == (q.shape[0], token_topk // compress_ratio)

579 + rows = q.shape[0]

580 + columns = page_table.shape[1] * k_cache.shape[1]

581 +

582 + # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB

583 + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024

584 + rows_per_chunk = max(1, max_logits_bytes // (columns * 4))

585 + topk_workspace = torch.empty(

586 + (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device

587 + )

588 +

589 + for query_start in range(0, rows, rows_per_chunk):

590 + query_end = min(query_start + rows_per_chunk, rows)

591 + query_slice = slice(query_start, query_end)

592 + logits = _prefill_logits(

593 + q,

594 + k_cache,

595 + page_table,

596 + query_start_loc,

597 + visible_blocks,

598 + max_query_len,

599 + query_offset=query_start,

600 + num_queries=query_end - query_start,

601 + )

602 + _topk(

603 + logits,

604 + visible_blocks[query_slice],

605 + token_topk,

606 + compress_ratio,

607 + block_indices[query_slice],

608 + topk_workspace,

609 + )

610 +

611 +

612 +__all__ = [

613 + "expand_qsa_block_indices",

614 + "qsa_select_paged_decode",

615 + "qsa_select_paged_prefill",

616 + "warmup_qsa_mqa_paged_decode",

617 +]