[Bugfix][HiSparse] Drop the unaccounted, write-only prefill mirror staging buffer (#60429)

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

tests/v1/attention/test_sparse_mla_backends.py

@@ -2578,7 +2578,7 @@ def test_hisparse_multi_step_writes_request_major_output():
 
 
 @requires_hisparse_ops
-def test_hisparse_kv_update_writes_resident_and_staging_caches():
+def test_hisparse_kv_update_writes_resident_cache():
     device = torch.device(DEVICE_TYPE)
     block_size = 4
     row_width = 8
@@ -2597,18 +2597,11 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches():
         block_table=torch.tensor([[1]], dtype=torch.int32, device=device),
         slot_mapping=resident_slots,
     )
-    cache_handle.mirror_staging_cache = torch.empty(
-        (1, block_size, row_width), dtype=torch.float32, device=device
-    )
-    cache_handle.mirror_staging_slots = torch.arange(
-        block_size, dtype=torch.int64, device=device
-    )
     slots = torch.tensor([3, 7, -1], dtype=torch.int64, device=device)
     kv_c = torch.randn(8, row_width - 2, device=device)
     k_pe = torch.randn(8, 1, 2, device=device)
     cache_handle.num_actual_tokens = slots.numel()
     cache_handle.decode_batch = False
-    cache_handle.host_mirror_required = True
     source_cache = torch.zeros_like(cache_handle.view.cache)
     impl = object.__new__(FlashMLASparseImpl)
     layer = SimpleNamespace(
@@ -2635,9 +2628,6 @@ def test_hisparse_kv_update_writes_resident_and_staging_caches():
         expected.to(device),
     )
     torch.testing.assert_close(source_cache, torch.zeros_like(source_cache))
-    staged = cache_handle.mirror_staging_cache.view(-1, row_width)
-    staged_expected = torch.cat([kv_c[:3], k_pe[:3, 0]], dim=-1)
-    torch.testing.assert_close(staged[:3], staged_expected)
 
 
 @requires_hisparse_ops

tests/v1/worker/test_attn_utils.py

@@ -434,7 +434,7 @@ def test_init_hisparse_rolls_back_shared_region(monkeypatch, failure_phase):
         cache_config=SimpleNamespace(
             get_resolved_kv_cache_layout=lambda: KVCacheLayout.BLHNC
         ),
-        scheduler_config=SimpleNamespace(max_num_seqs=1, max_num_batched_tokens=1),
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
     )
 
     def allocate(*args):

vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py

@@ -417,7 +417,6 @@ class HiSparseConnectorWorker:
         self._clear_forward_mirror_state()
         for handle in self.cache_handles:
             handle.all_context_pages_resident = metadata.all_context_pages_resident
-            handle.mirror_from_resident = True
         self._copy_host_blocks(metadata.host_block_copies, previous_host_write_event)
         transfers = (
             metadata.command.page_transfers if metadata.command is not None else []

vllm/model_executor/layers/attention/mla_attention.py

@@ -802,18 +802,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             kv_cache_dtype,
             k_scale,
         )
-        if cache is not None:
-            mirror_target = cache.mirror_write_target(kv_c_normed.shape[0])
-            if mirror_target is not None:
-                mirror_cache, mirror_slots = mirror_target
-                self.impl.do_kv_cache_update(  # type: ignore[attr-defined]
-                    kv_c_normed,
-                    k_pe,
-                    mirror_cache,
-                    mirror_slots,
-                    kv_cache_dtype,
-                    k_scale,
-                )
 
     def prepare_kv_cache_update(
         self, attn_metadata: "MLACommonMetadata | None"

vllm/v1/hisparse/binding.py

@@ -147,7 +147,6 @@ def init_hisparse_kv_cache(
         initialize_hisparse_runtime_buffers(
             cache_handles,
             max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
-            max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
         )
         return kv_caches
     except Exception:
@@ -197,9 +196,6 @@ def release_hisparse_profiling_cache(forward_context: dict[str, Any]) -> None:
         )
     )
     release_pinned_state(list(runtimes.values()), registered_pools, shared_region)
-    for cache in cache_handles:
-        cache.mirror_staging_cache = None
-        cache.mirror_staging_slots = None
 
 
 def bind_hisparse_kv_caches(

vllm/v1/hisparse/runtime.py

@@ -1120,10 +1120,7 @@ class HiSparseCacheHandle:
         self.num_decode_tokens = 0
         self.req_id_per_token: torch.Tensor | None = None
         self.host_mirror_required = False
-        self.mirror_from_resident = False
         self.mirror_slot_mapping: torch.Tensor | None = None
-        self.mirror_staging_cache: torch.Tensor | None = None
-        self.mirror_staging_slots: torch.Tensor | None = None
         self.submit_layer_mirror: Callable[[], None] | None = None
         # Speculator layers write their rows after the target forward.
         self.draft_layer = False
@@ -1160,26 +1157,6 @@ class HiSparseCacheHandle:
             num_rows = min(num_rows, self.runtime.max_num_reqs)
         return self.view.cache, self.slot_mapping[:num_rows], num_rows
 
-    def mirror_write_target(
-        self, num_rows: int
-    ) -> tuple[torch.Tensor, torch.Tensor] | None:
-        if (
-            not self.host_mirror_required
-            or self.decode_batch
-            or self.mirror_from_resident
-        ):
-            return None
-        cache = self.mirror_staging_cache
-        slots = self.mirror_staging_slots
-        if cache is None or slots is None:
-            raise RuntimeError("HiSparse prefill mirror staging is not bound.")
-        if num_rows > slots.numel():
-            raise RuntimeError(
-                "HiSparse prefill mirror exceeds staging capacity: "
-                f"{num_rows} > {slots.numel()}."
-            )
-        return cache, slots[:num_rows]
-
     def finish_kv_update(self) -> None:
         if self.dummy_batch:
             return
@@ -1242,9 +1219,8 @@ def initialize_hisparse_runtime_buffers(
     cache_handles: list[HiSparseCacheHandle],
     *,
     max_num_reqs: int,
-    max_num_batched_tokens: int,
 ) -> None:
-    """Allocate shared request state and per-layer prefill staging buffers."""
+    """Allocate the request state shared by every HiSparse runtime."""
     assert cache_handles
     resident = cache_handles[0].view
     assert resident is not None
@@ -1252,26 +1228,8 @@ def initialize_hisparse_runtime_buffers(
     request_state_indices = torch.full(
         (max_num_reqs,), -1, dtype=torch.int32, device=device
     )
-    staging_blocks = (
-        max_num_batched_tokens + resident.block_size - 1
-    ) // resident.block_size
-    mirror_staging_caches = torch.empty(
-        (
-            len(cache_handles),
-            staging_blocks,
-            resident.block_size,
-            resident.cache.shape[-1],
-        ),
-        dtype=resident.cache.dtype,
-        device=device,
-    )
-    mirror_staging_slots = torch.arange(
-        max_num_batched_tokens, dtype=torch.int64, device=device
-    )
-    for layer_index, cache_handle in enumerate(cache_handles):
+    for cache_handle in cache_handles:
         cache_handle.runtime.request_state_indices = request_state_indices
-        cache_handle.mirror_staging_cache = mirror_staging_caches[layer_index]
-        cache_handle.mirror_staging_slots = mirror_staging_slots
 
 
 def create_hisparse_cache_handle(
@@ -1325,11 +1283,4 @@ def create_hisparse_cache_handle(
         config.device_buffer_size,
         max_num_reqs,
     )
-    handle = HiSparseCacheHandle(runtime)
-    speculative_config = vllm_config.speculative_config
-    handle.mirror_from_resident = bool(
-        vllm_config.scheduler_config.async_scheduling
-        and speculative_config is not None
-        and speculative_config.uses_draft_kv_cache()
-    )
-    return handle
+    return HiSparseCacheHandle(runtime)