[Core] Use FlashInfer workspace sizing helper

docker/Dockerfile

@@ -792,7 +792,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.12
+ARG FLASHINFER_VERSION=0.6.13
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

docker/Dockerfile.nightly_torch

@@ -256,13 +256,13 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.
 
 
 # build flashinfer for torch nightly from source around 10 mins
-# release version: v0.6.12
+# release version: v0.6.13
 # todo(elainewy): cache flashinfer build result for faster build
 ENV CCACHE_DIR=/root/.cache/ccache
 RUN --mount=type=cache,target=/root/.cache/ccache \
     --mount=type=cache,target=/root/.cache/uv \
     echo "git clone flashinfer..." \
-    && git clone --depth 1 --branch v0.6.12 --recursive https://github.com/flashinfer-ai/flashinfer.git \
+    && git clone --depth 1 --branch v0.6.13 --recursive https://github.com/flashinfer-ai/flashinfer.git \
     && cd flashinfer \
     && git submodule update --init --recursive \
     && echo "finish git clone flashinfer..." \

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.12"
+      "default": "0.6.13"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

requirements/cuda.txt

@@ -9,8 +9,8 @@ torchaudio==2.11.0
 # These must be updated alongside torch
 torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.12
-flashinfer-cubin==0.6.12
+flashinfer-python==0.6.13
+flashinfer-cubin==0.6.13
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 nvidia-cudnn-frontend>=1.19.1

tests/v1/worker/test_cudagraph_memory_profiling.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+import gc
+import weakref
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs
+from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager
+
+
+def _attention_spec(head_size: int, head_size_v: int | None = None):
+    return FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=head_size,
+        head_size_v=head_size_v,
+        dtype=torch.float16,
+    )
+
+
+def test_flashinfer_separate_cudagraph_memory_profile_gate():
+    pytest.importorskip("flashinfer")
+    from vllm.v1.attention.backends.flashinfer import FlashInferMetadataBuilder
+
+    assert not FlashInferMetadataBuilder.requires_separate_cudagraph_memory_profiling(
+        None, _attention_spec(256)
+    )
+    assert FlashInferMetadataBuilder.requires_separate_cudagraph_memory_profiling(
+        None, _attention_spec(512)
+    )
+    assert FlashInferMetadataBuilder.requires_separate_cudagraph_memory_profiling(
+        None, _attention_spec(256, head_size_v=512)
+    )
+
+    uniform_spec = UniformTypeKVCacheSpecs(
+        block_size=16,
+        kv_cache_specs={
+            "layer.0": _attention_spec(256),
+            "layer.1": _attention_spec(512),
+        },
+    )
+    assert FlashInferMetadataBuilder.requires_separate_cudagraph_memory_profiling(
+        None, uniform_spec
+    )
+
+
+def test_flashinfer_workspace_buffer_uses_workspace_manager():
+    pytest.importorskip("flashinfer")
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    FlashInferMetadataBuilder = flashinfer_backend.FlashInferMetadataBuilder
+
+    def make_builder():
+        builder = FlashInferMetadataBuilder.__new__(FlashInferMetadataBuilder)
+        builder._workspace_buffer = None
+        builder._workspace_state = flashinfer_backend._FlashInferWorkspaceState()
+        builder.device = torch.device("cpu")
+        builder.use_dcp = False
+        return builder
+
+    reset_workspace_manager()
+    init_workspace_manager(torch.device("cpu"))
+    try:
+        first_builder = make_builder()
+        first_state = first_builder.get_workspace_buffer_state()
+        first = first_builder._get_workspace_buffer(
+            first_builder._native_initial_workspace_buffer_size()
+        )
+
+        second_builder = make_builder()
+        second_builder.set_workspace_buffer_state(first_state)
+        second = second_builder._get_workspace_buffer(
+            second_builder._native_initial_workspace_buffer_size()
+        )
+
+        assert first.device.type == "cpu"
+        assert first.dtype == torch.uint8
+        assert first.numel() == 1
+        assert first.data_ptr() == second.data_ptr()
+    finally:
+        reset_workspace_manager()
+
+
+def test_flashinfer_workspace_buffer_growth_resets_registered_wrappers():
+    pytest.importorskip("flashinfer")
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    class FakeWrapper:
+        def __init__(self, float_workspace_buffer):
+            self._float_workspace_buffer = float_workspace_buffer
+            self._int_workspace_buffer = torch.empty(1, dtype=torch.uint8)
+            self.reset_calls = 0
+
+        def reset_workspace_buffer(self, float_workspace_buffer, int_workspace_buffer):
+            self._float_workspace_buffer = float_workspace_buffer
+            self._int_workspace_buffer = int_workspace_buffer
+            self.reset_calls += 1
+
+    FlashInferMetadataBuilder = flashinfer_backend.FlashInferMetadataBuilder
+    builder = FlashInferMetadataBuilder.__new__(FlashInferMetadataBuilder)
+    builder._workspace_buffer = None
+    builder._workspace_state = flashinfer_backend._FlashInferWorkspaceState()
+    builder.device = torch.device("cpu")
+    builder.use_dcp = False
+
+    reset_workspace_manager()
+    init_workspace_manager(torch.device("cpu"))
+    try:
+        wrapper = FakeWrapper(
+            builder._get_workspace_buffer(
+                builder._native_initial_workspace_buffer_size()
+            )
+        )
+        builder._register_workspace_wrapper(wrapper)
+        builder._ensure_flashinfer_wrapper_workspace(wrapper, 1024)
+
+        assert builder._workspace_buffer.numel() == 1024
+        assert wrapper._float_workspace_buffer.data_ptr() == (
+            builder._workspace_buffer.data_ptr()
+        )
+        assert wrapper._float_workspace_buffer.numel() == 1024
+        assert wrapper.reset_calls >= 1
+
+        wrapper_ref = weakref.ref(wrapper)
+        del wrapper
+        gc.collect()
+
+        builder._workspace_state.set_buffer(torch.empty(2048, dtype=torch.uint8))
+        assert wrapper_ref() is None
+        assert builder._workspace_state.wrappers == []
+    finally:
+        reset_workspace_manager()
+
+
+def test_flashinfer_nvfp4_slot_mapping_symbol_available():
+    flashinfer = pytest.importorskip("flashinfer")
+    assert hasattr(
+        flashinfer,
+        "nvfp4_quantize_append_paged_kv_cache_with_slot_mapping",
+    )
+
+
+def test_separate_profile_accounts_persistent_and_graph_pool(monkeypatch):
+    from vllm.v1.worker import gpu_model_runner
+    from vllm.v1.worker.gpu_model_runner import CUDAGraphMode, GPUModelRunner
+
+    class FakeWrapper:
+        _all_instances = []
+
+        @staticmethod
+        def clear_all_graphs():
+            pass
+
+    @contextlib.contextmanager
+    def null_context(*args, **kwargs):
+        yield
+
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.lora_config = None
+    runner.cudagraph_dispatcher = SimpleNamespace(
+        get_capture_descs=lambda: [
+            (
+                CUDAGraphMode.PIECEWISE,
+                [
+                    SimpleNamespace(
+                        num_tokens=128,
+                        uniform=False,
+                        num_active_loras=0,
+                    )
+                ],
+            )
+        ],
+        cudagraph_keys={},
+        keys_initialized=True,
+    )
+
+    warmup_calls = []
+    capture_calls = []
+    cleanup_calls = []
+
+    runner._init_minimal_kv_cache_for_profiling = lambda: None
+    runner._requires_separate_cudagraph_memory_profiling = lambda: True
+    runner._create_encoder_cudagraph_manager = lambda: None
+    runner._freeze_gc = null_context
+    runner._cleanup_profiling_kv_cache = lambda: cleanup_calls.append("cleanup")
+    runner.maybe_remove_all_loras = lambda lora_config: None
+    runner._warmup_before_cudagraph_capture = lambda *args, **kwargs: (
+        warmup_calls.append(kwargs)
+    )
+    runner._warmup_and_capture = lambda *args, **kwargs: capture_calls.append(kwargs)
+
+    memory_reserved_values = iter([1_000, 1_600])
+    mem_get_info_values = iter([(5_000, 0), (4_300, 0)])
+
+    monkeypatch.setattr(gpu_model_runner, "CUDAGraphWrapper", FakeWrapper)
+    monkeypatch.setattr(gpu_model_runner, "BreakableCUDAGraphWrapper", FakeWrapper)
+    monkeypatch.setattr(
+        gpu_model_runner,
+        "set_current_vllm_config",
+        lambda *args, **kwargs: null_context(),
+    )
+    monkeypatch.setattr(
+        gpu_model_runner, "graph_capture", lambda *args, **kwargs: null_context()
+    )
+    monkeypatch.setattr(
+        gpu_model_runner,
+        "set_cudagraph_capturing_enabled",
+        lambda enabled: None,
+    )
+    monkeypatch.setattr(
+        gpu_model_runner.current_platform,
+        "graph_pool_handle",
+        lambda: object(),
+    )
+    monkeypatch.setattr(gpu_model_runner.torch.accelerator, "synchronize", lambda: None)
+    monkeypatch.setattr(gpu_model_runner.torch.accelerator, "empty_cache", lambda: None)
+    monkeypatch.setattr(
+        gpu_model_runner.torch.accelerator,
+        "memory_reserved",
+        lambda device: next(memory_reserved_values),
+    )
+    monkeypatch.setattr(
+        gpu_model_runner.torch.cuda,
+        "mem_get_info",
+        lambda: next(mem_get_info_values),
+    )
+
+    estimate = runner.profile_cudagraph_memory()
+
+    assert estimate == 1_300
+    assert runner.cudagraph_memory_persistent_estimate == 600
+    assert runner.cudagraph_memory_graph_pool_estimate == 700
+    assert len(warmup_calls) == 1
+    assert len(capture_calls) == 1
+    assert capture_calls[0]["num_warmups"] == 0
+    assert cleanup_calls == ["cleanup"]

vllm/v1/attention/backend.py

@@ -596,6 +596,20 @@ def get_cudagraph_support(
         """Get the cudagraph support level of this builder class."""
         return cls._cudagraph_support
 
+    @classmethod
+    def requires_separate_cudagraph_memory_profiling(
+        cls,
+        vllm_config: "VllmConfig",
+        kv_cache_spec: Any,
+    ) -> bool:
+        """Return whether CUDA graph memory profiling should separate warmup
+        allocations from graph-pool allocations for this builder class.
+
+        Most backends have small or already-accounted persistent warmup
+        allocations, so the default sampled estimator is retained.
+        """
+        return False
+
     def _init_reorder_batch_threshold(
         self,
         reorder_batch_threshold: int | None = 1,