[Bugfix] Keep GEMM workspace addresses stable once CUDA graphs may have captured them

vllm/model_executor/warmup/kernel_warmup.py

@@ -162,6 +162,15 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
 
     minimax_m3_msa_warmup(worker)
 
+    # Grow FlashInfer's shared GEMM workspace to its high-water mark now,
+    # while no CUDA graph has been captured: growing it later moves the
+    # buffer and leaves captured graphs replaying against freed memory
+    # (see presize_flashinfer_gemm_workspaces).
+    if has_flashinfer() and current_platform.has_device_capability(90):
+        from vllm.utils.flashinfer import presize_flashinfer_gemm_workspaces
+
+        presize_flashinfer_gemm_workspaces(worker.device)
+
     enable_flashinfer_autotune = (
         worker.vllm_config.kernel_config.enable_flashinfer_autotune
     )

vllm/utils/flashinfer.py

@@ -884,6 +884,48 @@ def flashinfer_scaled_fp4_mm_out(
     return out
 
 
+def presize_flashinfer_gemm_workspaces(
+    device: torch.device, size: int = 64 * 1024 * 1024
+) -> None:
+    """Grow FlashInfer's shared GEMM workspace before any CUDA graph capture.
+
+    FlashInfer's cuDNN GEMM execute paths grow their shared cache workspace
+    with ``tensor.resize_()`` whenever an execution plan asks for more than
+    the current buffer (32 MiB by default). ``resize_()`` frees the old
+    storage and moves the buffer to a new device address, but launches
+    already captured into CUDA graphs keep the old raw pointer, so a
+    post-capture grow turns every subsequent replay into a use-after-free:
+    cuDNN kernels read scratch from, and write scratch over, whatever
+    tensors the allocator hands the freed block to. Observed as illegal
+    memory access / misaligned address faults and GPU wedges minutes into
+    serving (see issue #52540; a 256-byte overflow past the 32 MiB default
+    during capture was enough).
+
+    Pre-growing the shared buffer past any plan's realistic demand keeps
+    the address stable for the lifetime of the process. This guard is
+    needed until a FlashInfer release ships the allocation fix
+    (flashinfer-ai/flashinfer#4553); it is cheap (one buffer per device)
+    and a no-op when FlashInfer is absent.
+    """
+    if not has_flashinfer():
+        return
+    try:
+        from flashinfer.utils import _get_cache_buf
+    except ImportError:
+        return
+    # Workspace names used by the GEMM lanes vLLM routes through
+    # FlashInfer's autotuner (which may select the cuDNN runner).
+    # zero_init matters: cuDNN split-K plans keep completion semaphores in
+    # the workspace and spin forever on garbage values; a fresh cudaMalloc
+    # happens to hand back zeroed pages, but that is luck, not a contract.
+    for name in ("bmm_fp8_workspace",):
+        try:
+            _get_cache_buf(name, size, device, zero_init=True)
+        except TypeError:
+            # older FlashInfer without the zero_init parameter
+            _get_cache_buf(name, size, device)
+
+
 def flashinfer_scaled_fp8_mm(
     a: torch.Tensor,
     b: torch.Tensor,

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

@@ -79,6 +79,11 @@ def __init__(self, initial_workspace_size):
         self._workspace_buf = torch.empty(
             initial_workspace_size, device="cuda", dtype=torch.uint8
         )
+        # Retired workspace generations. CUDA graphs capture the raw device
+        # pointer of whichever buffer was current at capture time, so a
+        # buffer may never be freed once a graph could have captured it;
+        # growing the workspace must keep every older generation alive.
+        self._retired_bufs: list[torch.Tensor] = []
 
         self._block_size = 128  # Forced to 128
 
@@ -101,7 +106,21 @@ def ensure_size(self, attn_metadata: MLACommonMetadata, num_kv_splits: int):
         )
 
         if self._workspace_buf.shape[0] < workspace_size:
-            self._workspace_buf.resize_(workspace_size)
+            # Do NOT resize_() in place: resize_() frees the old storage,
+            # while launches already enqueued and CUDA graphs captured
+            # against the old device address keep dereferencing it. Replays
+            # would then read/write memory the allocator has reused for
+            # unrelated tensors. Allocate a new buffer and retire the old
+            # one instead; graphs captured against a retired buffer stay
+            # correct because its capture-time size was sufficient for the
+            # shapes baked into those graphs.
+            self._retired_bufs.append(self._workspace_buf)
+            # zeros, not empty: split-KV workspaces carry semaphore/accum
+            # regions that kernels expect to start zeroed; a recycled dirty
+            # block from the caching allocator can hang or corrupt them.
+            self._workspace_buf = torch.zeros(
+                workspace_size, device="cuda", dtype=torch.uint8
+            )
 
 
 g_sm100_workspace = SM100Workspace(128 * 1024 * 1024)  # 128MB