From 624a5ddafd96df23fe2b9f6d4b8f52a08b9fe9ea Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 13 Oct 2026 13:25:03 -0400 Subject: [PATCH] [Attention] Fetch MLA prefill workspace at use instead of init (#60654) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- .../v1/attention/test_mla_prefill_registry.py | 44 +++++++++++++++++++ .../backends/mla/prefill/flashinfer.py | 25 ++++++++--- .../backends/mla/prefill/trtllm_ragged.py | 14 +++--- 3 files changed, 72 insertions(+), 11 deletions(-) diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index 0068d53b88..62835b3409 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/tests/v1/attention/test_mla_prefill_registry.py @@ -2,14 +2,26 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for MLA prefill backend registry.""" +import gc +from unittest.mock import MagicMock + import pytest import torch from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend +from vllm.v1.attention.backends.mla.prefill.flashinfer import FlashInferPrefillBackend from vllm.v1.attention.backends.mla.prefill.registry import ( MLAPrefillBackendEnum, register_mla_prefill_backend, ) +from vllm.v1.attention.backends.mla.prefill.trtllm_ragged import ( + TrtllmRaggedPrefillBackend, +) +from vllm.v1.worker.workspace import ( + current_workspace_manager, + init_workspace_manager, + reset_workspace_manager, +) class CustomMLAPrefillBackend(MLAPrefillBackend): @@ -50,6 +62,38 @@ def test_prefill_backend_clone_has_isolated_metadata(): assert clone._prefill_metadata is not backend._prefill_metadata +@pytest.mark.parametrize( + "backend_cls", [TrtllmRaggedPrefillBackend, FlashInferPrefillBackend] +) +def test_prefill_backend_does_not_pin_grown_workspace(monkeypatch, backend_cls): + """Reserving workspace at init must not keep a view of it: when another user + grows the shared workspace, the old buffer must be freed.""" + monkeypatch.setenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", "4096") + init_workspace_manager(torch.device("cpu")) + try: + manager = current_workspace_manager() + backend = backend_cls( + num_heads=4, + scale=0.5, + kv_lora_rank=8, + qk_nope_head_dim=16, + qk_rope_head_dim=8, + v_head_dim=32, + vllm_config=MagicMock(), + ) + (view,) = manager.get_simultaneous(((1,), torch.uint8)) + old_storage = view.untyped_storage()._weak_ref() + del view + + manager.get_simultaneous(((1 << 20,), torch.uint8)) + gc.collect() + + assert torch.UntypedStorage._expired(old_storage) + del backend + finally: + reset_workspace_manager() + + @pytest.fixture(autouse=True) def cleanup_overrides(): """Clear any overrides after each test.""" diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 5cf47ea2ca..0846a737f0 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -91,9 +91,14 @@ class FlashInferPrefillBackend(MLAPrefillBackend): self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = [] self._global_hyperparameters: PerLayerParameters | None = None - (self._workspace_buffer,) = current_workspace_manager().get_simultaneous( - ((envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8), + self._wrapper_workspace_ptr = 0 + self._workspace_spec = ( + (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), + torch.uint8, ) + # Reserve capacity without retaining a view: a retained view would pin + # the old buffer when another user grows the shared workspace. + current_workspace_manager().get_simultaneous(self._workspace_spec) def _ensure_chunks( self, @@ -146,16 +151,24 @@ class FlashInferPrefillBackend(MLAPrefillBackend): global_hyperparameters = self._resolve_global_hyperparameters() qo_indptr = prefill_metadata.query_start_loc has_context = prefill_metadata.chunked_context is not None - if self._prefill_main is None: + (workspace_buffer,) = current_workspace_manager().get_simultaneous( + self._workspace_spec + ) + if ( + self._prefill_main is None + or self._wrapper_workspace_ptr != workspace_buffer.data_ptr() + ): self._prefill_main = BatchPrefillWithRaggedKVCacheWrapper( - self._workspace_buffer, "NHD", backend="cutlass" + workspace_buffer, "NHD", backend="cutlass" ) - self._ensure_chunks(_DEFAULT_NUM_CHUNKS, self._workspace_buffer) + self._prefill_chunks = [] + self._ensure_chunks(_DEFAULT_NUM_CHUNKS, workspace_buffer) + self._wrapper_workspace_ptr = workspace_buffer.data_ptr() if has_context: chunked_context = prefill_metadata.chunked_context assert chunked_context is not None - self._ensure_chunks(len(chunked_context.chunks), self._workspace_buffer) + self._ensure_chunks(len(chunked_context.chunks), workspace_buffer) num_qo_heads = self.num_heads num_kv_heads = num_qo_heads diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 3cb498be2b..266fd83df2 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -79,18 +79,22 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): v_head_dim=v_head_dim, vllm_config=vllm_config, ) - (self._workspace_buffer,) = current_workspace_manager().get_simultaneous( - ( - (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), - torch.uint8, - ), + self._workspace_spec = ( + (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), + torch.uint8, ) + # Reserve capacity without retaining a view: a retained view would pin + # the old buffer when another user grows the shared workspace. + current_workspace_manager().get_simultaneous(self._workspace_spec) def prepare_metadata( self, prefill_metadata: "MLACommonPrefillMetadata", ) -> None: super().prepare_metadata(prefill_metadata) + (self._workspace_buffer,) = current_workspace_manager().get_simultaneous( + self._workspace_spec + ) self._query_seq_lens = ( prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1] )