diff --git a/tests/models/qwen4_exp/test_qsa_reference.py b/tests/models/qwen4_exp/test_qsa_reference.py index 9259972ac9eb..313dacce6456 100644 --- a/tests/models/qwen4_exp/test_qsa_reference.py +++ b/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, diff --git a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py b/vllm/models/qwen4_exp/nvidia/indexer_qsa.py index bbdc5e4feba9..5bfc002b9cc5 100644 --- a/vllm/models/qwen4_exp/nvidia/indexer_qsa.py +++ b/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, diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index 9ddd65d15935..bd73f3747091 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/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, diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index 90449d5706fb..e9b109cb411d 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/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])