From 4ed4c4433b87e45d72f401f27f1b56c932c90071 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Thu, 8 Oct 2026 10:00:49 -0400 Subject: [PATCH] [Bugfix][HiSparse] Drop the unaccounted, write-only prefill mirror staging buffer (#60429) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../v1/attention/test_sparse_mla_backends.py | 12 +--- tests/v1/worker/test_attn_utils.py | 2 +- .../kv_connector/v1/hisparse/worker.py | 1 - .../layers/attention/mla_attention.py | 12 ---- vllm/v1/hisparse/binding.py | 4 -- vllm/v1/hisparse/runtime.py | 55 +------------------ 6 files changed, 5 insertions(+), 81 deletions(-) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 3d32aed5d6..4d33075c93 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/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 diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 4114807052..9c9c26c3c2 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/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): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index bf6329b3bf..b9dd01f054 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/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 [] diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4286502ea9..2a3048096c 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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" diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index d021d8b7be..b8c0d5dff2 100644 --- a/vllm/v1/hisparse/binding.py +++ b/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( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index fc415fa2ec..a6378ca914 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/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)