@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a5c9179e73

[Qwen3.8-Flash-Next] Compact indexer logits workspace to improve prefill efficiency (#54915)

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

Thien Tran committed Sep 4, 2026 at 15:43 UTC a5c9179e731e8d2ec7e03485a3ab0754a5a27182

4 files changed +46 -16

tests/models/qwen4_exp/test_qsa_reference.py

+20 -3

@@ -244,6 +244,7 @@ def test_qsa_side_metadata_marks_cudagraph_padding_inert() -> None:

244 num_actual_tokens=16,

245 num_reqs=4,

246 max_query_len=4,

247 + max_seq_len=68,

248 query_start_loc=query_start_loc,

249 query_start_loc_cpu=query_start_loc.cpu(),

250 seq_lens=torch.tensor([68, 68, 68, 0], dtype=torch.int32, device=device),

@@ -297,6 +298,7 @@ def test_qsa_circular_buffer_metadata_keeps_only_each_requests_suffix() -> None:

298 num_actual_tokens=16,

299 num_reqs=3,

300 max_query_len=7,

301 + max_seq_len=11,

302 query_start_loc=query_start_loc,

303 query_start_loc_cpu=query_start_loc.cpu(),

304 seq_lens=torch.tensor([9, 11, 0], dtype=torch.int32, device=device),

@@ -428,6 +430,7 @@ def test_qsa_compressed_metadata_keeps_dummy_slots_inert() -> None:

430 num_actual_tokens=8,

431 num_reqs=3,

432 max_query_len=5,

433 + max_seq_len=12,

434 query_start_loc=query_start_loc,

435 query_start_loc_cpu=query_start_loc.cpu(),

436 seq_lens=torch.tensor([7, 0, 12], dtype=torch.int32, device=device),

@@ -673,13 +676,24 @@ def test_qsa_decode_selection_correctness(

676

677

678 @requires_qsa_kernels

676 -def test_qsa_prefill_selection_correctness() -> None:

679 +@pytest.mark.parametrize("seq_len_slack", [0, 1792])

680 +@pytest.mark.parametrize("force_chunk", [False, True])

681 +def test_qsa_prefill_selection_correctness(

682 + monkeypatch: pytest.MonkeyPatch, seq_len_slack: int, force_chunk: bool

683 +) -> None:

684 + # page_size=24 (does not divide the 64-aligned clipped width) and an

685 + # oversized page table, so the clipped logits width comes from

686 + # max_seq_len, not page geometry. seq_len_slack > 0 simulates the

687 + # spec-decode case where the bound is an over-estimate. force_chunk

688 + # drives the logits budget to one row per chunk.

689 + if force_chunk:

690 + monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "0")

691 torch.manual_seed(2)

692 query_lens = [3, 33]

693 rows, heads, head_dim = sum(query_lens), 4, 128

694 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)

695 + cache = torch.randn(128, 24, 1, head_dim, device="cuda", dtype=torch.bfloat16)

696 + page_table = torch.randperm(128, device="cuda", dtype=torch.int32).reshape(2, 64)

697 token_to_req = torch.repeat_interleave(

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

699 torch.tensor(query_lens, device="cuda"),

@@ -713,6 +727,7 @@ def test_qsa_prefill_selection_correctness() -> None:

727 compress_ratio,

728 max(query_lens),

729 actual,

730 + max_seq_len=sequence_lengths.max().item() + seq_len_slack,

731 )

732 expected = _qsa_select_paged_reference(

733 q,

@@ -931,6 +946,7 @@ def test_qsa_split_selection_correctness(workspace_init, decode_query_len: int)

946 compress_ratio,

947 query_lens[-1],

948 block_indices[prefill_slice],

949 + max_seq_len=sequence_lengths.max().item(),

950 )

951 actual = torch.empty(

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

@@ -983,6 +999,7 @@ def test_qsa_selection_handles_no_complete_compressed_blocks(workspace_init) ->

999 compress_ratio=4,

1000 max_query_len=2,

1001 block_indices=block_indices,

1002 + max_seq_len=64, # clamps to the page-table capacity

1003 )

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

1005 qsa_indexer_ops.expand_qsa_block_indices(

vllm/models/qwen4_exp/common/qsa_cache.py

@@ -591,6 +591,7 @@ class QSAForwardMetadata(AttentionMetadata):

591 num_prefill_tokens: int

592 max_query_len: int

593 decode_query_len: int

594 + max_seq_len: int

595 storage_block_size: int

596 compress_ratio: int

597

@@ -717,6 +718,7 @@ class QSAMetadataBuilder(AttentionMetadataBuilder[QSAForwardMetadata]):

718 num_prefill_tokens=num_prefill_tokens,

719 max_query_len=common_attn_metadata.max_query_len,

720 decode_query_len=decode_query_len,

721 + max_seq_len=common_attn_metadata.max_seq_len,

722 storage_block_size=self.storage_block_size,

723 compress_ratio=self.compress_ratio,

724 )

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

@@ -405,6 +405,7 @@ class QSAIndexer(nn.Module):

405 self.compress_ratio,

406 compressed_metadata.max_query_len,

407 block_indices[prefill_slice],

408 + compressed_metadata.max_seq_len,

409 )

410 expand_qsa_block_indices(

411 block_indices,

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

+23 -13

@@ -97,7 +97,7 @@ def _qsa_mqa_paged_uniform_kernel(

97 scores,

98 (BLOCK_N, DECODE_QUERY_LEN_PADDED, NUM_HEADS_PADDED),

99 )

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

100 + score = tl.sum(scores, axis=2)

101 tl.store(

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

103 score,

@@ -123,8 +123,8 @@ def _qsa_mqa_paged_prefill_kernel(

123 stride_logits_row,

124 num_rows,

125 query_offset,

126 + page_table_width,

127 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,

@@ -132,7 +132,7 @@ def _qsa_mqa_paged_prefill_kernel(

132 K_TILES: tl.constexpr,

133 STAGES: tl.constexpr,

134 ) -> None:

135 - NUM_COLUMNS: tl.constexpr = PAGE_TABLE_WIDTH * PAGE_SIZE

135 + num_columns = 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))

@@ -163,7 +163,7 @@ def _qsa_mqa_paged_prefill_kernel(

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))

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)

@@ -184,7 +184,7 @@ def _qsa_mqa_paged_prefill_kernel(

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)

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,

@@ -202,11 +202,11 @@ def _qsa_mqa_paged_prefill_kernel(

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

205 + score = tl.sum(tl.maximum(scores, 0.0), axis=2)

206 store_mask = (

207 valid_rows[None, :]

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

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

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

210 )

211 tl.store(

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

@@ -372,22 +372,25 @@ def _prefill_logits(

372 query_start_loc: torch.Tensor,

373 visible_blocks: torch.Tensor,

374 max_query_len: int,

375 + logits_width: int,

376 query_offset: int,

377 num_queries: int,

378 ) -> torch.Tensor:

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

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

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

382 + assert 0 < logits_width <= page_table.shape[1] * k_cache.shape[1]

383

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

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

384 + logits = torch.empty(

385 + (num_queries, logits_width), dtype=torch.float32, device=q.device

386 + )

387 TILE_R = 64

388 BLOCK_N = 64

389 K_TILES = 16

390 grid = (

391 page_table.shape[0],

392 triton.cdiv(min(num_queries, max_query_len), TILE_R),

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

393 + triton.cdiv(logits_width, BLOCK_N * K_TILES),

394 )

395 _qsa_mqa_paged_prefill_kernel[grid](

396 q,

@@ -402,8 +405,8 @@ def _prefill_logits(

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

406 num_queries,

407 query_offset,

408 + page_table.shape[1],

409 PAGE_SIZE=k_cache.shape[1],

406 - PAGE_TABLE_WIDTH=page_table.shape[1],

410 NUM_HEADS=q.shape[1],

411 HEAD_DIM=q.shape[2],

412 TILE_R=TILE_R,

@@ -557,6 +560,7 @@ def qsa_select_paged_prefill(

560 compress_ratio: int,

561 max_query_len: int,

562 block_indices: torch.Tensor,

563 + max_seq_len: int,

564 ) -> None:

565 """Score and select compressed prefill blocks in bounded chunks.

566

@@ -572,16 +576,21 @@ def qsa_select_paged_prefill(

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

577 max_query_len: Maximum number of query tokens in one request.

578 block_indices: Compressed-index output buffer.

579 + max_seq_len: Longest context length in the batch this step.

580 """

581

582 assert token_topk % compress_ratio == 0

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

584 rows = q.shape[0]

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

585 + # No row scores beyond cdiv(max_seq_len, compress_ratio) compressed

586 + # columns. Round up to 64 to keep the logits row stride

587 + # cooperative_topk-compatible.

588 + logits_width = triton.cdiv(triton.cdiv(max_seq_len, compress_ratio), 64) * 64

589 + logits_width = min(max(64, logits_width), page_table.shape[1] * k_cache.shape[1])

590

591 # chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB

592 max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024

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

593 + rows_per_chunk = max(1, max_logits_bytes // (logits_width * 4))

594 topk_workspace = torch.empty(

595 (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device

596 )

@@ -596,6 +605,7 @@ def qsa_select_paged_prefill(

605 query_start_loc,

606 visible_blocks,

607 max_query_len,

608 + logits_width,

609 query_offset=query_start,

610 num_queries=query_end - query_start,

611 )