[Qwen3.8-Flash-Next] Avoid memory fragmentation in QSA indexer logits workspace (#57105)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
Thien Tran committed
Sep 27, 2026 at 10:11 UTC
7a877ae34dcf4a7d46de08eb16db4699e1e0047c
1 file changed
+13
-2
vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py
+13
-2
@@ -383,6 +383,7 @@ def _prefill_logits(
383
page_table: torch.Tensor,
384
query_start_loc: torch.Tensor,
385
visible_blocks: torch.Tensor,
386
+ logits_workspace: torch.Tensor,
387
max_query_len: int,
388
logits_width: int,
389
query_offset: int,
@@ -392,10 +393,14 @@ def _prefill_logits(
393
assert visible_blocks.shape == (q.shape[0],)
394
assert 0 <= query_offset <= query_offset + num_queries <= q.shape[0]
395
assert 0 < logits_width <= page_table.shape[1] * k_cache.shape[1]
396
+ assert logits_workspace.is_contiguous()
397
+ assert logits_workspace.numel() >= num_queries * logits_width
398
396
- logits = torch.empty(
397
- (num_queries, logits_width), dtype=torch.float32, device=q.device
399
+ # slice from pre-allocated workspace
400
+ logits = logits_workspace[: num_queries * logits_width].view(
401
+ num_queries, logits_width
402
)
403
+
404
# tuned on GB300
405
if k_cache.dtype == torch.float8_e4m3fn:
406
TILE_R, STAGES, num_warps = 32, 2, 8
@@ -610,8 +615,13 @@ def qsa_select_paged_prefill(
615
logits_width = min(max(64, logits_width), page_table.shape[1] * k_cache.shape[1])
616
617
# chunk the inputs to keep temp logits below VLLM_SPARSE_INDEXER_MAX_LOGITS_MB
618
+ # always allocate the worst case to avoid memory fragmentation.
619
max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
620
rows_per_chunk = max(1, max_logits_bytes // (logits_width * 4))
621
+
622
+ budget_bytes = max(max_logits_bytes, page_table.shape[1] * k_cache.shape[1] * 4)
623
+ logits_workspace = q.new_empty(budget_bytes // 4, dtype=torch.float32)
624
+
625
topk_workspace = torch.empty(
626
(_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device
627
)
@@ -625,6 +635,7 @@ def qsa_select_paged_prefill(
635
page_table,
636
query_start_loc,
637
visible_blocks,
638
+ logits_workspace,
639
max_query_len,
640
logits_width,
641
query_offset=query_start,