[Attention] Fetch MLA prefill workspace at use instead of init (#60654)

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

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."""

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

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]
         )