File size: 2,515 Bytes
0693474 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | --- a/vllm/model_executor/layers/sparse_attn_indexer.py
+++ b/vllm/model_executor/layers/sparse_attn_indexer.py
@@ -1,6 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Custom Sparse Attention Indexer layers."""
+
+import os
import torch
import torch.distributed as dist
@@ -286,6 +288,29 @@
return q_fp8, weights_out
+# PF1: keep the legacy layout unless explicitly enabled for a measured A/B.
+_DSV41_INDEXER_REQUEST_WORKSPACE = (
+ os.environ.get("VLLM_DSV41_INDEXER_REQUEST_WORKSPACE", "0") == "1"
+)
+
+
+def _prefill_workspace_rows(reserved_rows: int, chunk_rows: list[int]) -> int:
+ """Host-known maximum gather extent, including query-continuation chunks.
+
+ The profile reservation stays unchanged. All views are made once before
+ the loop so skip_kv_gather continuations reuse the same values AND scales.
+ No device scalar read, reallocation, altered split or changed kernel shape.
+ """
+ if reserved_rows < 0 or any(rows < 0 for rows in chunk_rows):
+ raise ValueError("Indexer prefill workspace row counts must be nonnegative")
+ rows = max(chunk_rows, default=0)
+ if rows > reserved_rows:
+ raise ValueError(
+ f"Indexer prefill gather needs {rows} rows, reservation is {reserved_rows}"
+ )
+ return rows
+
+
def _gather_workspace_shapes(
total_seq_lens: int,
head_dim: int,
@@ -573,8 +598,17 @@
# MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0
# scales) based on use_fp4_cache.
workspace_manager = current_workspace_manager()
+ # PF1: on the served TP-only path, use the actual largest gather
+ # for the workspace views and scale offset, independent of model cap.
+ # Keep PCP/DCP and the dummy/profile reservation on the legacy layout.
+ gather_rows = total_seq_lens
+ if _DSV41_INDEXER_REQUEST_WORKSPACE and not use_pcp and dcp_world_size == 1:
+ gather_rows = _prefill_workspace_rows(
+ total_seq_lens,
+ [chunk.max_local_total_seq_lens for chunk in prefill_metadata.chunks],
+ )
values_spec, scales_spec = _gather_workspace_shapes(
- total_seq_lens, head_dim, fp8_dtype, use_fp4_cache
+ gather_rows, head_dim, fp8_dtype, use_fp4_cache
)
# PCP + DCP needs two more pairs: the rank-major all-gather destination
# and the de-interleaved result.
|