[Bugfix][HiSparse] Read resident prefill rows by request state slot (#62739)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -449,6 +449,7 @@ def test_compressed_slot_mapping_dcp_uses_rank_local_state_positions():
 
 def test_index_conversion_warmup_uses_physical_block_stride():
     config = SimpleNamespace(
+        attention_config=SimpleNamespace(hisparse_config=None),
         cache_config=SimpleNamespace(block_size=64),
         model_config=SimpleNamespace(
             max_model_len=1024,

tests/v1/attention/test_sparse_mla_backends.py

@@ -4061,42 +4061,69 @@ def test_sparse_mqa_hisparse_lanes(
     )
 
 
-@pytest.mark.parametrize("max_swap_rows", [1, 4], ids=["over_capacity", "mixed"])
-def test_hisparse_resident_prefill_uses_attention_block_stride(
-    monkeypatch, max_swap_rows
+@pytest.mark.parametrize(
+    ("num_decode_tokens", "max_swap_rows"),
+    [(3, 1), (1, 4)],
+    ids=["over_capacity", "mixed"],
+)
+def test_hisparse_resident_prefill_reads_residency_by_state_row(
+    num_decode_tokens, max_swap_rows
 ):
-    """Resident batches with prefill rows read resident pages directly, whether
-    they exceed the swap-row capacity or fit it."""
-    expected = torch.tensor([[19]], dtype=torch.int32)
-    cache_handle = _StubHiSparseCacheHandle(
-        all_context_pages_resident=True,
-        num_decode_tokens=1,
-        num_actual_tokens=2,
-        view=SimpleNamespace(block_size=64, attention_block_stride=832),
-        block_table=torch.tensor([[3]], dtype=torch.int32),
-        batch_block_table=lambda: torch.tensor([[3]], dtype=torch.int32),
-        runtime=SimpleNamespace(max_swap_rows=max_swap_rows),
+    """Resident prefill conversion reads each request's residency row through
+    its state slot, at the attention block stride. An unset slot reads slot 0."""
+    topk = torch.zeros((3, 128), dtype=torch.int32, device=DEVICE_TYPE)
+    # [resident groups, state slots, pages]; the leader reads group 1.
+    residency = torch.tensor(
+        [[[7], [8]], [[3], [5]]], dtype=torch.int32, device=DEVICE_TYPE
+    )
+    runtime = SimpleNamespace(
+        is_group_leader=True,
+        begin_forward=lambda: None,
+        eager_host_mirror=True,
+        max_swap_rows=max_swap_rows,
+        request_state_indices=torch.tensor(
+            [1, 0, -1], dtype=torch.int32, device=DEVICE_TYPE
+        ),
+    )
+    cache = HiSparseCacheHandle(runtime)
+    cache.view = SimpleNamespace(block_size=64, attention_block_stride=832)
+    cache.block_table = residency[1]
+    group = HiSparseMLAIndexGroup(
+        logical_topk_indices=topk,
+        physical_topk_indices=torch.empty_like(topk),
+        valid_topk_counts=torch.empty(3, dtype=torch.int32, device=DEVICE_TYPE),
+        side_stream=torch.Stream(),
+        logical_topk_ready=torch.Event(),
+        physical_topk_ready=torch.Event(),
+        has_indexer=True,
+        caches=[cache],
     )
-    index_group = object.__new__(HiSparseMLAIndexGroup)
-    index_group.caches = [cache_handle]
-    index_group.physical_topk_indices = torch.empty((2, 1), dtype=torch.int32)
-    convert_once = MagicMock(return_value=expected)
-    monkeypatch.setattr(index_group, "_convert_once", convert_once)
-    topk = torch.zeros((2, 1), dtype=torch.int32)
     metadata = SimpleNamespace(
-        req_id_per_token=torch.zeros(2, dtype=torch.int32),
+        num_actual_tokens=3,
+        num_decode_tokens=num_decode_tokens,
+        max_query_len=1 if num_decode_tokens == 3 else 2,
+        num_reqs=3,
+        req_id_per_token=torch.tensor([0, 1, 2], dtype=torch.int32, device=DEVICE_TYPE),
     )
 
-    result = index_group.convert_logical_to_physical_topk(
-        0,
-        topk,
-        metadata,
-        block_stride_rows=None,
-        return_valid_counts=False,
-    )
-
-    assert result is expected
-    assert convert_once.call_args.kwargs["block_stride_rows"] == 832
+    # Repeat after compilation so its host delay cannot hide a stream race, and
+    # change the residency so a stale table cannot pass.
+    for step in range(2):
+        residency.add_(step * 10)
+        expected = torch.tensor(
+            [[5 + step * 10], [3 + step * 10], [3 + step * 10]], device=DEVICE_TYPE
+        )
+        expected *= 832
+        torch.accelerator.synchronize()
+        group.prepare_for_batch(0, metadata)
+        group.set_logical_topk_ready(0)
+        # Compute-stream work between the indexer and the layer must not delay
+        # the inputs the side-stream converter reads.
+        torch.cuda._sleep(200_000_000)
+        physical = group.convert_logical_to_physical_topk(
+            0, topk, metadata, block_stride_rows=None, return_valid_counts=False
+        )
+        torch.testing.assert_close(physical, expected.int().expand_as(topk))
 
 
 def test_hisparse_flashmla_reorders_full_speculative_window_as_decode():

tests/v1/kv_connector/unit/test_hisparse_connector.py

@@ -510,7 +510,7 @@ def test_residency_updates_persist_by_state_row():
     An update for a lost page and an appended page must leave the other pages
     and the other request's row intact across steps that reorder the batch.
     """
-    table = torch.zeros((4, 2, 4), dtype=torch.int32)
+    table = torch.zeros((2, 4, 4), dtype=torch.int32)
     update_hisparse_residency(
         table,
         {
@@ -527,5 +527,5 @@ def test_residency_updates_persist_by_state_row():
         torch.tensor([0, 2], dtype=torch.int32),
     )
 
-    assert table[[0, 2], 0].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]]
-    assert table[[0, 2], 1].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]]
+    assert table[0, [0, 2]].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]]
+    assert table[1, [0, 2]].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]]

vllm/v1/attention/backends/mla/index_group.py

@@ -116,6 +116,7 @@ class SparseMLAIndexGroup:
         *,
         block_stride_rows: int | None,
         return_valid_counts: bool,
+        idx_mapping: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         num_tokens = logical_topk_indices.shape[0]
         if num_tokens > self.physical_topk_indices.shape[0]:
@@ -127,6 +128,7 @@ class SparseMLAIndexGroup:
                 BLOCK_STRIDE_ROWS=block_stride_rows,
                 NUM_TOPK_TOKENS=logical_topk_indices.shape[1],
                 return_valid_counts=return_valid_counts,
+                idx_mapping=idx_mapping,
             )
 
         physical_topk_indices = self.physical_topk_indices[:num_tokens]
@@ -150,6 +152,7 @@ class SparseMLAIndexGroup:
                     valid_counts_out=(
                         valid_topk_counts if return_valid_counts else None
                     ),
+                    idx_mapping=idx_mapping,
                 )
                 self.physical_topk_ready.record(self.side_stream)
         compute_stream.wait_event(self.physical_topk_ready)
@@ -244,10 +247,11 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
                 layer_index,
                 logical_topk_indices,
                 req_id_per_token,
-                leader.batch_block_table(),
+                leader.block_table,
                 leader.view.block_size,
                 block_stride_rows=leader.view.attention_block_stride,
                 return_valid_counts=return_valid_counts,
+                idx_mapping=leader.runtime.request_state_indices,
             )
         source_block_table = cache.source_block_table
         assert source_block_table is not None

vllm/v1/attention/backends/mla/sparse_utils.py

@@ -81,6 +81,7 @@ class ConvertReqIndexToGlobalIndexKernel(
         block_stride_rows: int
         block_n: int
         has_prefill_workspace: bool
+        has_idx_mapping: bool
         count_valid: bool
         single_tile: bool
         compact_to_front: bool
@@ -99,6 +100,7 @@ class ConvertReqIndexToGlobalIndexKernel(
         token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
         out_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
         valid_count_ptr,  # int32 [num_tokens] - output valid count per row
+        idx_mapping_ptr,  # int32 [num_requests] - block table row per request
         prefill_request_id_ptr,  # int32 [num_tokens], -1 for decode, >=0 for prefill
         workspace_starts_ptr,  # int32 [num_prefill_reqs+1] or nullptr
         workspace_rank_stride,  # rows per DCP rank of a gathered prefill workspace
@@ -109,6 +111,7 @@ class ConvertReqIndexToGlobalIndexKernel(
         BLOCK_N: tl.constexpr,  # tile width along columns
         NUM_TOPK_TOKENS: tl.constexpr,
         HAS_PREFILL: tl.constexpr,
+        HAS_IDX_MAPPING: tl.constexpr,
         COUNT_VALID: tl.constexpr,  # whether to count valid indices
         # BLOCK_N >= NUM_TOPK_TOKENS: one program owns the query, so its count
         # is an in-register reduction and needs no atomic.
@@ -141,6 +144,8 @@ class ConvertReqIndexToGlobalIndexKernel(
 
         # Load request id for this token (no mask: grid is exact)
         req = tl.load(req_id_ptr + token_id)
+        if HAS_IDX_MAPPING:
+            req = tl.maximum(tl.load(idx_mapping_ptr + req), 0)
 
         # Load token indices for this tile
         ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
@@ -239,6 +244,7 @@ class ConvertReqIndexToGlobalIndexKernel(
         BLOCK_N: int,
         NUM_TOPK_TOKENS: int,
         HAS_PREFILL_WORKSPACE: bool,
+        HAS_IDX_MAPPING: bool,
         COUNT_VALID: bool,
         COMPACT_TO_FRONT: bool,
         DCP_SIZE: int,
@@ -252,6 +258,7 @@ class ConvertReqIndexToGlobalIndexKernel(
             block_stride_rows=BLOCK_STRIDE_ROWS,
             block_n=tiling[1],
             has_prefill_workspace=HAS_PREFILL_WORKSPACE,
+            has_idx_mapping=HAS_IDX_MAPPING,
             count_valid=COUNT_VALID,
             single_tile=tiling[0],
             compact_to_front=COMPACT_TO_FRONT,
@@ -279,6 +286,20 @@ class ConvertReqIndexToGlobalIndexKernel(
             block_size * dcp_size,
         )
         max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
+        hisparse_combos: list[dict[str, Any]] = []
+        if vllm_config.attention_config.hisparse_config is not None:
+            hisparse_combos.extend(
+                dict(
+                    HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=True,
+                    COUNT_VALID=count_valid,
+                    COMPACT_TO_FRONT=count_valid,
+                    DCP_SIZE=1,
+                    DCP_RANK=0,
+                    DCP_INTERLEAVE=1,
+                )
+                for count_valid in (False, True)
+            )
         pcp_dcp_combos = []
         if (
             vllm_config.parallel_config.prefill_context_parallel_size > 1
@@ -288,6 +309,7 @@ class ConvertReqIndexToGlobalIndexKernel(
             pcp_dcp_combos.append(
                 dict(
                     HAS_PREFILL_WORKSPACE=True,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=True,
                     DCP_SIZE=dcp_size,
@@ -299,6 +321,7 @@ class ConvertReqIndexToGlobalIndexKernel(
             zip_inputs(
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=False,
                     COMPACT_TO_FRONT=False,
                     DCP_SIZE=1,
@@ -307,6 +330,7 @@ class ConvertReqIndexToGlobalIndexKernel(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=True,
                     DCP_SIZE=1,
@@ -315,6 +339,7 @@ class ConvertReqIndexToGlobalIndexKernel(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=True,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=True,
                     DCP_SIZE=1,
@@ -323,6 +348,7 @@ class ConvertReqIndexToGlobalIndexKernel(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=False,
                     DCP_SIZE=dcp_size,
@@ -331,6 +357,7 @@ class ConvertReqIndexToGlobalIndexKernel(
                 ),
                 dict(
                     HAS_PREFILL_WORKSPACE=False,
+                    HAS_IDX_MAPPING=False,
                     COUNT_VALID=True,
                     COMPACT_TO_FRONT=dcp_size > 1,
                     DCP_SIZE=dcp_size,
@@ -338,6 +365,7 @@ class ConvertReqIndexToGlobalIndexKernel(
                     DCP_INTERLEAVE=dcp_interleave,
                 ),
                 *pcp_dcp_combos,
+                *hisparse_combos,
             ),
             BLOCK_SIZE=block_size,
             BLOCK_STRIDE_ROWS=block_stride_rows,
@@ -373,6 +401,7 @@ class ConvertReqIndexToGlobalIndexKernel(
             prefill_workspace_starts=(
                 int32_ptr if compile_key.has_prefill_workspace else None
             ),
+            idx_mapping=(int32_ptr if compile_key.has_idx_mapping else None),
             max_num_blocks_per_req=1,
             workspace_rank_stride=1,
             BLOCK_SIZE=compile_key.block_size,
@@ -380,6 +409,7 @@ class ConvertReqIndexToGlobalIndexKernel(
             BLOCK_N=compile_key.block_n,
             NUM_TOPK_TOKENS=compile_key.num_topk_tokens,
             HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace,
+            HAS_IDX_MAPPING=compile_key.has_idx_mapping,
             COUNT_VALID=compile_key.count_valid,
             COMPACT_TO_FRONT=compile_key.compact_to_front,
             DCP_SIZE=compile_key.dcp_size,
@@ -397,6 +427,7 @@ class ConvertReqIndexToGlobalIndexKernel(
         valid_counts: torch.Tensor | None,
         prefill_workspace_request_ids: torch.Tensor | None,
         prefill_workspace_starts: torch.Tensor | None,
+        idx_mapping: torch.Tensor | None,
         *,
         workspace_rank_stride: int,
         max_num_blocks_per_req: int,
@@ -405,6 +436,7 @@ class ConvertReqIndexToGlobalIndexKernel(
         BLOCK_N: int,
         NUM_TOPK_TOKENS: int,
         HAS_PREFILL_WORKSPACE: bool,
+        HAS_IDX_MAPPING: bool,
         COUNT_VALID: bool,
         COMPACT_TO_FRONT: bool,
         DCP_SIZE: int,
@@ -419,6 +451,7 @@ class ConvertReqIndexToGlobalIndexKernel(
             valid_count_ptr=valid_counts,
             prefill_request_id_ptr=prefill_workspace_request_ids,
             workspace_starts_ptr=prefill_workspace_starts,
+            idx_mapping_ptr=idx_mapping,
             BLOCK_N=block_n,
             HAS_PREFILL=HAS_PREFILL_WORKSPACE,
             SINGLE_TILE=single_tile,
@@ -475,6 +508,7 @@ def triton_convert_req_index_to_global_index(
     return_valid_counts: bool = False,
     out: torch.Tensor | None = None,
     valid_counts_out: torch.Tensor | None = None,
+    idx_mapping: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """out[token_id, indice_id] =
         block_table[req_id[token_id],
@@ -494,11 +528,17 @@ def triton_convert_req_index_to_global_index(
     prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace
         starts for each prefill request
 
+    When idx_mapping (int32 [num_requests]) is given, block_table is
+    indexed by idx_mapping[req_id[token_id]] instead of req_id[token_id].
+
     When return_valid_counts is True, also returns the count of valid (non -1)
     indices per row, computed during the same kernel pass (no extra overhead).
     """
     assert req_id.dtype == torch.int32
     assert block_table.dtype == torch.int32
+    assert idx_mapping is None or (
+        idx_mapping.dtype == torch.int32 and idx_mapping.is_contiguous()
+    )
     assert token_indices.dtype == torch.int32
     assert req_id.shape[0] == token_indices.shape[0], (
         f"req_id ({req_id.shape[0]}) and token_indices ({token_indices.shape[0]}) "
@@ -572,6 +612,7 @@ def triton_convert_req_index_to_global_index(
         valid_counts,
         prefill_workspace_request_ids,
         prefill_workspace_starts,
+        idx_mapping,
         # shapes / constexprs
         workspace_rank_stride=prefill_workspace_rank_stride or 0,
         max_num_blocks_per_req=max_num_blocks_per_req,
@@ -582,6 +623,7 @@ def triton_convert_req_index_to_global_index(
         BLOCK_N=BLOCK_N,
         NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
         HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
+        HAS_IDX_MAPPING=idx_mapping is not None,
         COUNT_VALID=return_valid_counts,
         COMPACT_TO_FRONT=return_valid_counts,
         DCP_SIZE=dcp_size,
@@ -680,6 +722,7 @@ def triton_filter_and_convert_dcp_index(
         # No prefill workspace on the DCP decode path.
         None,
         None,
+        None,
         workspace_rank_stride=0,
         max_num_blocks_per_req=max_num_blocks_per_req,
         BLOCK_SIZE=BLOCK_SIZE,
@@ -689,6 +732,7 @@ def triton_filter_and_convert_dcp_index(
         BLOCK_N=BLOCK_N,
         NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
         HAS_PREFILL_WORKSPACE=False,
+        HAS_IDX_MAPPING=False,
         COUNT_VALID=count_valid,
         COMPACT_TO_FRONT=compact_valid_to_front,
         DCP_SIZE=dcp_size,

vllm/v1/hisparse/binding.py

@@ -211,7 +211,7 @@ def bind_hisparse_kv_caches(
         assert len(table_shapes) == 1
         max_num_reqs, max_num_pages = table_shapes.pop()
         residency = torch.zeros(
-            (max_num_reqs, len(resident_groups), max_num_pages),
+            (len(resident_groups), max_num_reqs, max_num_pages),
             dtype=torch.int32,
             device=block_tables.input_block_tables[resident_groups[0][0]].device,
         )
@@ -229,7 +229,7 @@ def bind_hisparse_kv_caches(
                 block_stride=tensor_config.block_stride,
                 num_blocks=kv_cache_config.num_blocks,
                 block_size=group.kv_cache_spec.block_size,
-                block_table=residency[:, resident_source_index],
+                block_table=residency[resident_source_index],
                 slot_mapping=block_tables.slot_mappings[group_id],
             )
             cache_handle.residency = residency

vllm/v1/hisparse/runtime.py

@@ -1145,13 +1145,13 @@ def update_hisparse_residency(
 ) -> None:
     """Apply scheduled requests' residency changes.
 
-    ``residency`` is ``[state rows, resident groups, pages]``: the GPU block of
+    ``residency`` is ``[resident groups, state rows, pages]``: the GPU block of
     each page, or block 0 when the page is read from the host.
     """
     if not updates:
         return
     batch_rows = {request_id: row for row, request_id in enumerate(request_ids)}
-    num_groups = residency.shape[1]
+    num_groups = residency.shape[0]
     rows: list[int] = []
     pages: list[int] = []
     block_ids: list[list[int]] = [[] for _ in range(num_groups)]
@@ -1172,7 +1172,7 @@ def update_hisparse_residency(
         host_values, device=request_state_indices.device, dtype=torch.int32
     )
     state_rows = request_state_indices[values[:, 0]]
-    residency[state_rows, :, values[:, 1]] = values[:, 2:]
+    residency[:, state_rows, values[:, 1]] = values[:, 2:].T
 
 
 class HiSparseCacheHandle:
@@ -1197,7 +1197,6 @@ class HiSparseCacheHandle:
         # Speculator layers write their rows after the target forward.
         self.draft_layer = False
         self.index_group_caches: list[HiSparseCacheHandle] = [self]
-        self._batch_block_table: torch.Tensor | None = None
 
     def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
         assert self.runtime.is_group_leader
@@ -1206,7 +1205,6 @@ class HiSparseCacheHandle:
 
     def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
         self.dummy_batch = attn_metadata is None
-        self._batch_block_table = None
         self.runtime.begin_forward()
         self.num_actual_tokens = (
             attn_metadata.num_actual_tokens if attn_metadata is not None else 0
@@ -1222,17 +1220,6 @@ class HiSparseCacheHandle:
             not self.decode_batch or self.runtime.eager_host_mirror
         )
 
-    def batch_block_table(self) -> torch.Tensor:
-        """Resident rows by this step's batch row, for the prefill paths."""
-        if self._batch_block_table is None:
-            assert self.block_table is not None
-            indices = self.runtime.request_state_indices
-            assert indices is not None
-            self._batch_block_table = self.block_table.index_select(
-                0, indices.clamp(min=0)
-            )
-        return self._batch_block_table
-
     def write_target(
         self, num_input_rows: int, num_slot_rows: int
     ) -> tuple[torch.Tensor, torch.Tensor, int]: