[Bugfix][QSA] Reuse bounded prefill logits workspace

tests/models/qwen4_exp/test_qsa_reference.py

@@ -765,11 +765,86 @@ def test_qsa_decode_selection_correctness(
     torch.testing.assert_close(actual.sort().values, expected.sort().values)
 
 
+@requires_qsa_kernels
+@pytest.mark.parametrize("logits_mb", [0, 1])
+def test_qsa_prefill_reuses_profiled_workspace(
+    workspace_init, monkeypatch: pytest.MonkeyPatch, logits_mb: int
+) -> None:
+    """Growing contexts reuse profiled storage without corrupting top-k selection."""
+    from vllm.models.qwen4_exp.nvidia import qsa as qsa_layer
+    from vllm.v1.worker.workspace import current_workspace_manager
+
+    monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", str(logits_mb))
+    max_model_len, compress_ratio, token_topk = 8192, 4, 2048
+    owner = SimpleNamespace(
+        indexer=SimpleNamespace(max_logits_width=max_model_len // compress_ratio)
+    )
+    monkeypatch.setattr(
+        qsa_layer, "get_forward_context", lambda: SimpleNamespace(attn_metadata=None)
+    )
+    dummy = torch.empty(5, 128, device="cuda")
+    qsa_layer.Qwen4ExpQSAAttention._run_qsa(
+        owner, dummy, torch.arange(5, device="cuda"), dummy, dummy, dummy, dummy
+    )
+    manager = current_workspace_manager()
+    manager.lock()
+    workspace, scratch = qsa_indexer_ops.get_qsa_prefill_workspace(
+        max_model_len // compress_ratio
+    )
+    assert workspace.data_ptr() + workspace.numel() * 4 <= scratch.data_ptr()
+    storage_ptrs = []
+    topk = qsa_indexer_ops._topk
+
+    def record_topk(logits, *args, **kwargs):
+        storage_ptrs.append(logits.data_ptr())
+        return topk(logits, *args, **kwargs)
+
+    monkeypatch.setattr(qsa_indexer_ops, "_topk", record_topk)
+    torch.manual_seed(2)
+    q = torch.randn(5, 4, 128, device="cuda", dtype=torch.bfloat16)
+    cache = torch.randn(32, 64, 1, 128, device="cuda", dtype=torch.bfloat16)
+    page_table = torch.arange(32, device="cuda", dtype=torch.int32).view(1, -1)
+    query_start_loc = torch.tensor([0, 5], device="cuda", dtype=torch.int32)
+    token_to_req = torch.zeros(5, device="cuda", dtype=torch.int32)
+    actual = torch.empty(
+        5, token_topk // compress_ratio, device="cuda", dtype=torch.int32
+    )
+    for context_len in (2048, 4096, 8192):
+        positions = torch.arange(context_len - 5, context_len, device="cuda")
+        visible = ((positions + 1) // compress_ratio).to(torch.int32)
+        workspace.fill_(float("nan"))
+        qsa_indexer_ops.qsa_select_paged_prefill(
+            q,
+            cache,
+            page_table,
+            query_start_loc,
+            visible,
+            token_topk,
+            compress_ratio,
+            5,
+            actual,
+            context_len,
+        )
+        expected = _qsa_select_paged_reference(
+            q,
+            cache,
+            page_table,
+            token_to_req,
+            positions,
+            torch.tensor([context_len], device="cuda"),
+            token_topk,
+            compress_ratio,
+        )
+        torch.testing.assert_close(actual.sort().values, expected.sort().values)
+    assert set(storage_ptrs) == {workspace.data_ptr()}
+
+
 @requires_qsa_kernels
 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
 @pytest.mark.parametrize("seq_len_slack", [0, 1792])
 @pytest.mark.parametrize("force_chunk", [False, True])
 def test_qsa_prefill_selection_correctness(
+    workspace_init,
     monkeypatch: pytest.MonkeyPatch,
     seq_len_slack: int,
     force_chunk: bool,

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

@@ -16,6 +16,7 @@
 from vllm.transformers_utils.configs.qwen4_exp import (
     Qwen4ExpTextConfig,
 )
+from vllm.utils.math_utils import cdiv, round_up
 
 from ..common.qsa_cache import (
     QSACompressedKeyCache,
@@ -117,6 +118,9 @@ def __init__(
         self.index_head_dim = int(config.indexer_head_dim)
         self.token_topk = int(config.indexer_budget)
         self.compress_ratio = int(config.indexer_compress_ratio)
+        self.max_logits_width = round_up(
+            cdiv(vllm_config.model_config.max_model_len, self.compress_ratio), 64
+        )
         self.rotary_emb = rotary_emb
         self.use_fused_pre_indexer = _supports_fused_pre_indexer(
             rotary_emb,

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

@@ -10,6 +10,7 @@
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
+from vllm.v1.worker.workspace import current_workspace_manager
 
 _TOPK_WORKSPACE_BYTES = 1024 * 1024
 _DECODE_BLOCK_N = 64
@@ -385,18 +386,15 @@ def _prefill_logits(
     query_start_loc: torch.Tensor,
     visible_blocks: torch.Tensor,
     max_query_len: int,
-    logits_width: int,
+    logits: torch.Tensor,
     query_offset: int,
-    num_queries: int,
-) -> torch.Tensor:
+) -> None:
+    num_queries, logits_width = logits.shape
     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]
     assert 0 < logits_width <= page_table.shape[1] * k_cache.shape[1]
 
-    logits = torch.empty(
-        (num_queries, logits_width), dtype=torch.float32, device=q.device
-    )
     # tuned on GB300
     if k_cache.dtype == torch.float8_e4m3fn:
         TILE_R, STAGES, num_warps = 32, 2, 8
@@ -432,7 +430,6 @@ def _prefill_logits(
         STAGES=STAGES,
         num_warps=num_warps,
     )
-    return logits
 
 
 def expand_qsa_block_indices(
@@ -572,6 +569,18 @@ def qsa_select_paged_decode(
     )
 
 
+def get_qsa_prefill_workspace(logits_width: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Reserve the logits budget and disjoint top-k scratch, also during profiling."""
+    max_logits_elems = max(
+        envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 // 4, logits_width
+    )
+    logits, topk = current_workspace_manager().get_simultaneous(
+        ((max_logits_elems,), torch.float32),
+        ((_TOPK_WORKSPACE_BYTES,), torch.uint8),
+    )
+    return logits, topk
+
+
 def qsa_select_paged_prefill(
     q: torch.Tensor,
     k_cache: torch.Tensor,
@@ -612,25 +621,24 @@ def qsa_select_paged_prefill(
     logits_width = min(max(64, logits_width), 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 // (logits_width * 4))
-    topk_workspace = torch.empty(
-        (_TOPK_WORKSPACE_BYTES,), dtype=torch.uint8, device=q.device
-    )
+    logits_workspace, topk_workspace = get_qsa_prefill_workspace(logits_width)
+    rows_per_chunk = logits_workspace.numel() // logits_width
 
     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(
+        logits = logits_workspace[: (query_end - query_start) * logits_width].view(
+            -1, logits_width
+        )
+        _prefill_logits(
             q,
             k_cache,
             page_table,
             query_start_loc,
             visible_blocks,
             max_query_len,
-            logits_width,
+            logits,
             query_offset=query_start,
-            num_queries=query_end - query_start,
         )
         _topk(
             logits,

vllm/models/qwen4_exp/nvidia/qsa.py

@@ -52,6 +52,7 @@
 from ..common.qsa_cache import QSAForwardMetadata
 from . import model
 from .indexer_qsa import QSAIndexer
+from .ops.qsa_indexer import get_qsa_prefill_workspace
 
 
 class Qwen4ExpQSAMetadataBuilder(FlashAttentionMetadataBuilder):
@@ -358,6 +359,7 @@ def _run_qsa(
         if isinstance(metadata, list):
             metadata = metadata[0]
         if not isinstance(metadata, dict):
+            get_qsa_prefill_workspace(self.indexer.max_logits_width)
             output.zero_()
             return
         main_metadata = cast(FlashAttentionMetadata, metadata[self.layer_name])