diff --git a/tests/v1/worker/test_cudagraph_memory_profiling.py b/tests/v1/worker/test_cudagraph_memory_profiling.py new file mode 100644 index 000000000000..984bc7b0d27d --- /dev/null +++ b/tests/v1/worker/test_cudagraph_memory_profiling.py @@ -0,0 +1,979 @@ +# 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, + ) + + +class _FakeFlashInferWrapper: + def __init__( + self, + float_workspace_buffer: torch.Tensor | None = None, + int_workspace_bytes: int = 1, + ) -> None: + self._float_workspace_buffer = ( + float_workspace_buffer + if float_workspace_buffer is not None + else torch.empty(1, dtype=torch.uint8) + ) + self._int_workspace_buffer = torch.empty( + max(int_workspace_bytes, 1), dtype=torch.uint8 + ) + self._vllm_flashinfer_int_workspace_finalized = False + self.is_cuda_graph_enabled = False + self.reset_calls = 0 + + def reset_workspace_buffer( + self, + float_workspace_buffer: torch.Tensor, + int_workspace_buffer: torch.Tensor, + ) -> None: + self._float_workspace_buffer = float_workspace_buffer + self._int_workspace_buffer = int_workspace_buffer + self.reset_calls += 1 + + +class _FakeWorkspaceSizeArray: + def __init__(self, values: list[int]) -> None: + self.values = values + + def __len__(self) -> int: + return len(self.values) + + def __getitem__(self, index: int) -> int: + return self.values[index] + + +def _make_flashinfer_builder(flashinfer_backend): + 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 + return builder + + +def test_flashinfer_default_workspace_covers_prefill_head_footprint(monkeypatch): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + builder = _make_flashinfer_builder(flashinfer_backend) + builder.max_num_batched_tokens = 8 + builder.num_qo_heads = 4 + builder.head_dim = 16 + estimated_prefill_size = ( + builder.max_num_batched_tokens + * builder.num_qo_heads + * builder.head_dim + * flashinfer_backend.FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM + ) + + monkeypatch.setattr(flashinfer_backend.envs, "VLLM_BATCH_INVARIANT", False) + monkeypatch.setattr( + flashinfer_backend.envs, + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", + 1, + ) + assert builder._default_workspace_buffer_size() == estimated_prefill_size + + configured_size = estimated_prefill_size + 1 + monkeypatch.setattr( + flashinfer_backend.envs, + "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", + configured_size, + ) + assert builder._default_workspace_buffer_size() == configured_size + + +def test_flashinfer_workspace_size_parses_sequence_like_array(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + + assert flashinfer_backend._parse_workspace_sizes( + _FakeWorkspaceSizeArray([1024, 64]) + ) == WorkspaceSizes(1024, 64, True) + + +def test_flashinfer_workspace_size_float_only_keeps_default_int_workspace(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + builder = _make_flashinfer_builder(flashinfer_backend) + wrapper = _FakeFlashInferWrapper(int_workspace_bytes=64) + sizes = flashinfer_backend._parse_workspace_sizes(_FakeWorkspaceSizeArray([1024])) + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + builder._ensure_flashinfer_wrapper_workspace(wrapper, sizes) + finally: + reset_workspace_manager() + + assert wrapper._float_workspace_buffer.numel() == 1024 + assert wrapper._int_workspace_buffer.numel() == 64 + + +def test_flashinfer_workspace_size_explicit_int_shrinks_cudagraph_workspace(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + mib = 1 << 20 + builder = _make_flashinfer_builder(flashinfer_backend) + wrapper = _FakeFlashInferWrapper(int_workspace_bytes=8 * mib) + wrapper.is_cuda_graph_enabled = True + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 64 * 1024, True) + ) + finally: + reset_workspace_manager() + + assert wrapper._float_workspace_buffer.numel() == 1024 + assert wrapper._int_workspace_buffer.numel() == 64 * 1024 + + +def test_flashinfer_workspace_size_explicit_int_rounds_non_cudagraph_workspace(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + mib = 1 << 20 + builder = _make_flashinfer_builder(flashinfer_backend) + wrapper = _FakeFlashInferWrapper(int_workspace_bytes=8 * mib) + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 64 * 1024, True) + ) + finally: + reset_workspace_manager() + + assert wrapper._float_workspace_buffer.numel() == 1024 + assert wrapper._int_workspace_buffer.numel() == mib + + +def test_flashinfer_workspace_size_rejects_invalid_sequence_length(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + with pytest.raises(ValueError, match="workspace_size"): + flashinfer_backend._parse_workspace_sizes(_FakeWorkspaceSizeArray([1, 2, 3])) + + +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 + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + first_builder = _make_flashinfer_builder(flashinfer_backend) + first_state = first_builder.get_workspace_buffer_state() + first = first_builder._get_workspace_buffer( + first_builder._native_initial_workspace_buffer_size() + ) + + second_builder = _make_flashinfer_builder(flashinfer_backend) + 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 + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + builder = _make_flashinfer_builder(flashinfer_backend) + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + wrapper = _FakeFlashInferWrapper( + builder._get_workspace_buffer( + builder._native_initial_workspace_buffer_size() + ) + ) + builder._register_workspace_wrapper(wrapper) + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 16, True) + ) + + 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._int_workspace_buffer.numel() == 1 << 20 + reset_calls = wrapper.reset_calls + assert reset_calls >= 1 + + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 16, True) + ) + assert wrapper.reset_calls == reset_calls + + 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_int_workspace_is_per_wrapper(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + builder = _make_flashinfer_builder(flashinfer_backend) + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + first = _FakeFlashInferWrapper() + second = _FakeFlashInferWrapper() + + builder._ensure_flashinfer_wrapper_workspace( + first, WorkspaceSizes(1024, 32, True) + ) + builder._ensure_flashinfer_wrapper_workspace( + second, WorkspaceSizes(1024, 32, True) + ) + + assert first._float_workspace_buffer.data_ptr() == ( + second._float_workspace_buffer.data_ptr() + ) + assert first._int_workspace_buffer.data_ptr() != ( + second._int_workspace_buffer.data_ptr() + ) + assert first._int_workspace_buffer.numel() == 1 << 20 + assert second._int_workspace_buffer.numel() == 1 << 20 + finally: + reset_workspace_manager() + + +def test_flashinfer_finalized_int_workspace_cannot_grow(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + builder = _make_flashinfer_builder(flashinfer_backend) + wrapper = _FakeFlashInferWrapper(int_workspace_bytes=8) + wrapper._vllm_flashinfer_int_workspace_finalized = True + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + with pytest.raises(AssertionError, match="int workspace is finalized"): + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 16, True) + ) + finally: + reset_workspace_manager() + + +def test_flashinfer_non_cudagraph_int_workspace_can_grow(monkeypatch): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + builder = _make_flashinfer_builder(flashinfer_backend) + wrapper = _FakeFlashInferWrapper(int_workspace_bytes=8) + warnings = [] + + monkeypatch.setattr( + flashinfer_backend.logger, + "warning", + lambda msg, *args: warnings.append(msg), + ) + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 8, True) + ) + builder._ensure_flashinfer_wrapper_workspace( + wrapper, WorkspaceSizes(1024, 2 << 20, True) + ) + + assert wrapper._int_workspace_buffer.numel() == 2 << 20 + assert wrapper.reset_calls == 2 + assert any("Growing FlashInfer int workspace" in msg for msg in warnings) + finally: + reset_workspace_manager() + + +def test_flashinfer_reserves_prefill_tail_workspace(monkeypatch): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + 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 + builder.model_config = SimpleNamespace(max_model_len=1024, dtype=torch.float16) + builder.vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace( + max_num_batched_tokens=8, + max_num_seqs=4, + ), + speculative_config=None, + ) + builder.q_data_type_prefill = torch.float16 + builder.q_data_type_decode = torch.float16 + builder.kv_cache_dtype = torch.uint8 + builder.page_size = 16 + builder.window_left = -1 + builder.prefill_fixed_split_size = -1 + builder.disable_split_kv = False + + class FakeWrapper: + pass + + ensured = [] + observed_query_lens = [] + + def fake_workspace_size(**kwargs): + qo_indptr = kwargs["qo_indptr_cpu"] + query_lens = torch.diff(qo_indptr).tolist() + observed_query_lens.extend(query_lens) + return ( + WorkspaceSizes(4096, 64, True) + if query_lens == [3] + else WorkspaceSizes(0, 0) + ) + + monkeypatch.setattr( + builder, + "_get_prefill_workspace_size_func", + lambda **kwargs: ("fa2", object()), + ) + monkeypatch.setattr(builder, "_call_prefill_workspace_size", fake_workspace_size) + monkeypatch.setattr( + builder, "_get_prefill_wrapper", lambda causal=True: FakeWrapper() + ) + monkeypatch.setattr( + builder, + "_ensure_flashinfer_wrapper_workspace", + lambda wrapper, size: ensured.append(size), + ) + monkeypatch.setattr( + builder, + "_reserve_decode_wrapper_workspace", + lambda **kwargs: WorkspaceSizes(0, 0), + ) + builder.enable_cuda_graph = False + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + assert builder.reserve_workspace_for_cudagraph_capture() == 4160 + finally: + reset_workspace_manager() + + assert ensured == [WorkspaceSizes(4096, 64, True)] + assert 3 in observed_query_lens + assert 8 in observed_query_lens + + +def test_flashinfer_reserves_decode_cudagraph_int_workspace(monkeypatch): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + builder = _make_flashinfer_builder(flashinfer_backend) + builder.decode_fixed_split_size = -1 + builder.disable_split_kv = False + + wrapper = _FakeFlashInferWrapper() + wrapper.is_cuda_graph_enabled = True + + monkeypatch.setattr(builder, "_get_decode_wrapper", lambda *args: wrapper) + monkeypatch.setattr( + builder, + "_get_decode_workspace_size", + lambda **kwargs: WorkspaceSizes(128, 32, True), + ) + + reset_workspace_manager() + init_workspace_manager(torch.device("cpu")) + try: + sizes = builder._reserve_decode_wrapper_workspace( + batch_size=4, + num_pages=8, + last_page_len=16, + use_cudagraph=True, + ) + finally: + reset_workspace_manager() + + assert sizes == WorkspaceSizes(128, 32, True) + assert wrapper._float_workspace_buffer.numel() == 128 + assert wrapper._int_workspace_buffer.numel() == 32 + assert wrapper.reset_calls == 1 + assert wrapper._vllm_flashinfer_int_workspace_finalized + + +def test_flashinfer_workspace_debug_info_reports_retained_default_int_workspace(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + WorkspaceSizes = flashinfer_backend.WorkspaceSizes + builder = _make_flashinfer_builder(flashinfer_backend) + mib = 1 << 20 + builder._prefill_wrapper = _FakeFlashInferWrapper(int_workspace_bytes=8 * mib) + builder._decode_wrapper = _FakeFlashInferWrapper(int_workspace_bytes=4 * mib) + builder._decode_wrappers_cudagraph = { + 1: _FakeFlashInferWrapper(int_workspace_bytes=8 * mib), + 2: _FakeFlashInferWrapper(int_workspace_bytes=8 * mib), + } + builder._last_reserved_workspace_sizes = WorkspaceSizes( + float_bytes=128 * mib, + int_bytes=256 * 1024, + ) + + info = builder.get_workspace_reserve_debug_info() + + assert info["workspace_wrapper_count"] == 4 + assert info["prefill_wrappers"] == 1 + assert info["decode_wrappers"] == 1 + assert info["decode_cudagraph_wrappers"] == 2 + assert info["actual_int_workspace_bytes"] == 28 * mib + assert info["reserved_int_workspace_bytes"] == 256 * 1024 + assert info["int_workspace_over_reserved_bytes"] == 28 * mib - 256 * 1024 + assert info["default_int_workspace_wrappers"] == 3 + assert info["unique_int_workspace_buffers"] == 4 + + +def test_flashinfer_workspace_query_len_candidates(): + pytest.importorskip("flashinfer") + from vllm.v1.attention.backends import flashinfer as flashinfer_backend + + candidates = ( + flashinfer_backend.FlashInferMetadataBuilder._get_workspace_query_len_candidates + ) + + assert candidates(8) == list(range(1, 9)) + + large_candidates = candidates(1024) + assert 1 in large_candidates + assert 256 in large_candidates + assert 512 in large_candidates + assert 1024 in large_candidates + assert 257 not in large_candidates + + +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, + ), + SimpleNamespace( + num_tokens=64, + uniform=False, + num_active_loras=0, + ), + SimpleNamespace( + num_tokens=32, + uniform=False, + num_active_loras=0, + ), + ], + ), + ( + CUDAGraphMode.FULL, + [ + SimpleNamespace( + num_tokens=80, + uniform=False, + num_active_loras=0, + ), + SimpleNamespace( + num_tokens=40, + uniform=False, + num_active_loras=0, + ), + ], + ), + ], + cudagraph_keys={}, + keys_initialized=True, + ) + + warmup_calls = [] + capture_calls = [] + cleanup_calls = [] + + runner.max_model_len = 4096 + runner.max_num_tokens = 128 + 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._reserve_attention_workspace_for_cudagraph_capture = lambda: 200 + runner._warmup_before_cudagraph_capture = lambda *args, **kwargs: ( + warmup_calls.append((args[0], kwargs)) + ) + runner._warmup_and_capture = lambda *args, **kwargs: capture_calls.append( + (args[0], kwargs) + ) + + memory_reserved_values = iter([1_000, 1_600]) + get_memory_info_values = iter( + [ + (10_000_000, 0), + (8_000_000, 0), + (8_000_000, 0), + (6_500_000, 0), + (6_500_000, 0), + (4_100_000, 0), + (4_100_000, 0), + (3_000_000, 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.accelerator, + "get_memory_info", + lambda: next(get_memory_info_values), + ) + + estimate = runner.profile_cudagraph_memory() + + assert estimate == 6_500_800 + assert runner.cudagraph_memory_persistent_estimate == 800 + assert runner.cudagraph_memory_graph_pool_estimate == 6_500_000 + assert [call[0].num_tokens for call in warmup_calls] == [128, 64, 80, 40] + assert [call[0].num_tokens for call in capture_calls] == [128, 64, 80, 40] + assert all(call[1]["num_warmups"] == 0 for call in capture_calls) + assert warmup_calls[2][1]["profile_seq_lens"] == 1 + assert warmup_calls[3][1]["profile_seq_lens"] is None + assert cleanup_calls == ["cleanup"] + + +def test_v2_profile_accounts_attention_workspace(monkeypatch): + from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + @contextlib.contextmanager + def null_context(*args, **kwargs): + yield + + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.vllm_config = object() + + events = [] + + monkeypatch.setattr( + gpu_model_runner_v2, + "set_current_vllm_config", + lambda *args, **kwargs: null_context(), + ) + + def reserve_attention_workspace(): + events.append("reserve") + return 4096 + + runner._init_minimal_kv_cache_for_profiling = lambda: events.append("init") + runner._reserve_attention_workspace_for_cudagraph_capture = ( + reserve_attention_workspace + ) + runner._cleanup_profiling_kv_cache = lambda: events.append("cleanup") + + estimate = runner.profile_cudagraph_memory() + + assert estimate == 4096 + assert runner.cudagraph_memory_persistent_estimate == 4096 + assert runner.cudagraph_memory_graph_pool_estimate == 0 + assert events == ["init", "reserve", "cleanup"] + + +def test_v2_cleanup_profiling_kv_cache_releases_builder_refs(monkeypatch): + from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + class Builder: + pass + + builder = Builder() + builder_ref = weakref.ref(builder) + layer = SimpleNamespace( + kv_cache=torch.empty(1), + impl=SimpleNamespace(_k_scale_cache=object(), _v_scale_cache=object()), + ) + + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.cache_config = SimpleNamespace(num_gpu_blocks=4) + runner.kv_caches = [torch.empty(1)] + runner.attn_groups = [[SimpleNamespace(metadata_builders=[builder])]] + runner.kv_cache_config = object() + runner.block_tables = object() + runner.kernel_block_sizes = [16] + runner.cudagraph_manager = object() + runner.compilation_config = SimpleNamespace(static_forward_context={"layer": layer}) + del builder + + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, "synchronize", lambda: None + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, "empty_cache", lambda: None + ) + + runner._cleanup_profiling_kv_cache() + gc.collect() + + assert runner.kv_caches == [] + assert not hasattr(runner, "attn_groups") + assert not hasattr(runner, "kv_cache_config") + assert not hasattr(runner, "block_tables") + assert not hasattr(runner, "kernel_block_sizes") + assert not hasattr(runner, "cudagraph_manager") + assert runner.cache_config.num_gpu_blocks is None + assert isinstance(layer.kv_cache, torch.Tensor) + assert layer.kv_cache.numel() == 0 + assert layer.impl._k_scale_cache is None + assert layer.impl._v_scale_cache is None + assert builder_ref() is None + + +def test_v2_capture_reserves_workspace_before_measurement_and_locks(monkeypatch): + from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + @contextlib.contextmanager + def null_context(*args, **kwargs): + yield + + class Builder: + reserved = False + + def reserve_workspace_for_cudagraph_capture(self): + events.append("builder_reserve") + self.reserved = True + return 128 + + class FakeCudaGraphManager: + def needs_capture(self): + return True + + def capture( + self, + model, + model_state, + input_buffers, + intermediate_tensors, + block_tables, + attn_groups, + kv_cache_config, + **kwargs, + ): + events.append("capture") + assert attn_groups[0][0].metadata_builders[0].reserved + return {} + + events = [] + builder = Builder() + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.device = torch.device("cpu") + runner.cudagraph_manager = FakeCudaGraphManager() + runner.lora_config = None + runner.maybe_setup_dummy_loras = lambda lora_config: null_context() + runner.model = object() + runner.model_state = object() + runner.input_buffers = object() + runner.intermediate_tensors = None + runner.block_tables = object() + runner.attn_groups = [[SimpleNamespace(metadata_builders=[builder])]] + runner.kv_cache_config = object() + runner.use_aux_hidden_state_outputs = False + runner.speculator = None + + memory_reserved_values = iter([1_000, 1_000, 1_128, 1_128]) + memory_allocated_values = iter([500, 500, 628, 628]) + get_memory_info_values = iter([(10_000, 0), (9_000, 0)]) + + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, "synchronize", lambda: None + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, "empty_cache", lambda: None + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, + "memory_reserved", + lambda device: next(memory_reserved_values), + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, + "memory_allocated", + lambda device: next(memory_allocated_values), + ) + + def get_memory_info(): + events.append("memory_info") + return next(get_memory_info_values) + + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, + "get_memory_info", + get_memory_info, + ) + monkeypatch.setattr( + gpu_model_runner_v2, + "lock_workspace", + lambda: events.append("lock"), + ) + + assert runner.capture_model() == 1_000 + assert events == [ + "builder_reserve", + "memory_info", + "capture", + "memory_info", + "lock", + ] + + +def test_v2_attention_workspace_reserve_logs_breakdown(monkeypatch): + from vllm.v1.worker.gpu import model_runner as gpu_model_runner_v2 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + mib = 1 << 20 + workspace_buffer = torch.empty(mib, dtype=torch.uint8) + + class Builder: + def __init__(self, requested_bytes, debug_info): + self.requested_bytes = requested_bytes + self.debug_info = debug_info + + def reserve_workspace_for_cudagraph_capture(self): + return self.requested_bytes + + def get_workspace_buffer_state(self): + return SimpleNamespace(buffer=workspace_buffer) + + def get_workspace_reserve_debug_info(self): + return self.debug_info + + runner = GPUModelRunner.__new__(GPUModelRunner) + runner.device = torch.device("cpu") + runner.attn_groups = [ + [ + SimpleNamespace( + metadata_builders=[ + Builder( + 64 * mib, + { + "workspace_wrapper_count": 4, + "decode_cudagraph_wrappers": 2, + "default_int_workspace_wrappers": 3, + "actual_int_workspace_bytes": 24 * mib, + "reserved_int_workspace_bytes": 1 * mib, + "int_workspace_over_reserved_bytes": 23 * mib, + "unique_int_workspace_buffers": 4, + "unique_float_workspace_buffers": 1, + "unique_float_workspace_bytes": 24 * mib, + "workspace_state_live_wrappers": 4, + }, + ), + Builder( + 32 * mib, + { + "workspace_wrapper_count": 2, + "decode_cudagraph_wrappers": 1, + "default_int_workspace_wrappers": 1, + "actual_int_workspace_bytes": 8 * mib, + "reserved_int_workspace_bytes": 1 * mib, + "int_workspace_over_reserved_bytes": 7 * mib, + "unique_int_workspace_buffers": 2, + "unique_float_workspace_buffers": 1, + "unique_float_workspace_bytes": 8 * mib, + "workspace_state_live_wrappers": 2, + }, + ), + ] + ) + ] + ] + + memory_reserved_values = iter( + [1000 * mib, 1000 * mib, 1064 * mib, 1064 * mib, 1160 * mib, 1160 * mib] + ) + memory_allocated_values = iter( + [500 * mib, 500 * mib, 564 * mib, 564 * mib, 596 * mib, 596 * mib] + ) + debug_logs = [] + + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, + "memory_reserved", + lambda device: next(memory_reserved_values), + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, + "memory_allocated", + lambda device: next(memory_allocated_values), + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, "synchronize", lambda: None + ) + monkeypatch.setattr( + gpu_model_runner_v2.torch.accelerator, "empty_cache", lambda: None + ) + monkeypatch.setattr( + gpu_model_runner_v2.logger, + "debug", + lambda msg, *args: debug_logs.append(msg % args if args else msg), + ) + + assert runner._reserve_attention_workspace_for_cudagraph_capture() == 160 * mib + + assert any("96.00 MiB requested by builders" in log for log in debug_logs) + assert any("64.00 MiB unexplained" in log for log in debug_logs) + assert any("1 unique workspace buffers" in log for log in debug_logs) + assert any("1.00 MiB unique workspace bytes" in log for log in debug_logs) + assert any("6 wrappers" in log for log in debug_logs) + assert any("32.00 MiB actual int workspace" in log for log in debug_logs) + assert any("2.00 MiB requested int workspace" in log for log in debug_logs) + assert any("30.00 MiB int workspace over request" in log for log in debug_logs) + assert any("24.00 MiB max unique float workspace" in log for log in debug_logs) + assert any("default_int_workspaces=3" in log for log in debug_logs) + assert any("unique_float=24.00 MiB" in log for log in debug_logs) + assert ( + sum( + "Reserved attention workspace builder=Builder requested=" in log + for log in debug_logs + ) + == 2 + ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 2708e0ab1f0e..4b45550ea1e8 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -638,6 +638,28 @@ 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 reserve_workspace_for_cudagraph_capture(self) -> int: + """Reserve backend workspace that must not grow after CUDA graph capture. + + Backends that use the global WorkspaceManager can override this to size + their worst-case execution workspace before the manager is locked. + """ + return 0 + def _init_reorder_batch_threshold( self, reorder_batch_threshold: int | None = 1, diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index da4e59d4eaaf..66b097d0aafc 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -2,10 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Attention layer with FlashInfer.""" -from dataclasses import dataclass +import weakref +from dataclasses import dataclass, field from enum import Enum from functools import partial -from typing import ClassVar +from typing import Any, ClassVar, NamedTuple import numpy as np import torch @@ -16,8 +17,12 @@ MultiLevelCascadeAttentionWrapper, ) from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache -from flashinfer.prefill import trtllm_batch_context_with_kv_cache -from flashinfer.utils import FP4Tensor +from flashinfer.page import get_seq_lens +from flashinfer.prefill import ( + get_batch_prefill_module, + trtllm_batch_context_with_kv_cache, +) +from flashinfer.utils import FP4Tensor, PosEncodingMode, determine_attention_backend from typing_extensions import override from vllm import _custom_ops as custom_ops @@ -76,12 +81,19 @@ from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.kv_cache_interface import ( AttentionSpec, + KVCacheSpec, KVQuantMode, UniformTypeKVCacheSpecs, ) from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.workspace import ( + current_workspace_manager, + is_workspace_manager_initialized, +) FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024 +FLASHINFER_DEFAULT_INT_WORKSPACE_BYTES = 8 * 1024 * 1024 +FLASHINFER_INT_WORKSPACE_GRANULARITY_BYTES = 1 << 20 FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM = 16 FP8_DTYPE = current_platform.fp8_dtype() @@ -92,6 +104,110 @@ trtllm_workspace_buffer = None +def _buffer_nbytes(buffer: torch.Tensor | None) -> int: + if buffer is None: + return 0 + return buffer.numel() * buffer.element_size() + + +class WorkspaceSizes(NamedTuple): + float_bytes: int + int_bytes: int = 0 + has_int_bytes: bool = False + + @property + def total_bytes(self) -> int: + return self.float_bytes + self.int_bytes + + +def _int_workspace_allocation_bytes( + required_bytes: int, *, round_up: bool = True +) -> int: + required_bytes = max(int(required_bytes), 0) + if required_bytes == 0: + return 1 + if not round_up: + return required_bytes + return cdiv(required_bytes, FLASHINFER_INT_WORKSPACE_GRANULARITY_BYTES) * ( + FLASHINFER_INT_WORKSPACE_GRANULARITY_BYTES + ) + + +def _parse_workspace_sizes(workspace_size: Any) -> WorkspaceSizes: + if not isinstance(workspace_size, (str, bytes)): + try: + workspace_size_len = len(workspace_size) + except TypeError: + pass + else: + if workspace_size_len == 0: + raise ValueError("FlashInfer workspace_size returned an empty result") + if workspace_size_len > 2: + raise ValueError( + "FlashInfer workspace_size must return a scalar or a " + "(float_bytes, int_bytes) pair" + ) + float_bytes = int(workspace_size[0]) + if workspace_size_len == 2: + return WorkspaceSizes(float_bytes, int(workspace_size[1]), True) + return WorkspaceSizes(float_bytes, 0, False) + + return WorkspaceSizes(int(workspace_size), 0, False) + + +def _is_float8_dtype(dtype: torch.dtype) -> bool: + return dtype in {torch.float8_e4m3fn, torch.float8_e5m2} + + +@dataclass +class _FlashInferWorkspaceState: + buffer: torch.Tensor | None = None + wrappers: list[weakref.ReferenceType[object]] = field(default_factory=list) + + def register_wrapper(self, wrapper: object) -> None: + if not any(registered is wrapper for registered in self._live_wrappers()): + self.wrappers.append(weakref.ref(wrapper)) + self._reset_wrapper(wrapper) + + def set_buffer(self, buffer: torch.Tensor) -> None: + if self.buffer is buffer: + return + self.buffer = buffer + self._reset_wrappers() + + def _reset_wrappers(self) -> None: + for wrapper in self._live_wrappers(): + self._reset_wrapper(wrapper) + + def _live_wrappers(self) -> list[object]: + live_refs: list[weakref.ReferenceType[object]] = [] + live_wrappers: list[object] = [] + for wrapper_ref in self.wrappers: + wrapper = wrapper_ref() + if wrapper is None: + continue + live_refs.append(wrapper_ref) + live_wrappers.append(wrapper) + if len(live_refs) != len(self.wrappers): + self.wrappers = live_refs + return live_wrappers + + def _reset_wrapper(self, wrapper: object) -> None: + if self.buffer is None or not hasattr(wrapper, "reset_workspace_buffer"): + return + int_workspace = getattr(wrapper, "_int_workspace_buffer", None) + if int_workspace is None: + return + current_buffer = getattr(wrapper, "_float_workspace_buffer", None) + if ( + current_buffer is not None + and current_buffer.data_ptr() == self.buffer.data_ptr() + and _buffer_nbytes(current_buffer) == _buffer_nbytes(self.buffer) + ): + return + wrapper.reset_workspace_buffer(self.buffer, int_workspace) + + def _get_trtllm_workspace_buffer(): global trtllm_workspace_buffer if trtllm_workspace_buffer is None: @@ -631,6 +747,8 @@ def __init__( self.model_config = vllm_config.model_config self.attention_config = vllm_config.attention_config self._workspace_buffer = None + self._workspace_state = _FlashInferWorkspaceState() + self._last_reserved_workspace_sizes = WorkspaceSizes(0, 0) self._prefill_wrapper: ( BatchPrefillWithPagedKVCacheWrapper | BatchDCPPrefillWrapper | None ) = None # Wrapper for prefill/append @@ -932,34 +1050,286 @@ def get_cudagraph_support( else: return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE - def _get_workspace_buffer(self): - if self._workspace_buffer is None: - buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE - if envs.VLLM_BATCH_INVARIANT: - buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT - else: - # FlashInfer prefill temp buffers (batch_prefill_tmp_v, ...) - # scale with the prefill chunk and query-head footprint, NOT - # context length. The fixed ~394 MiB default is too small for - # wide-head models at the default 8192-token chunk on some - # archs (e.g. sm_120), where FlashInfer hard-errors instead of - # growing. Size to the batch's head footprint; never shrink - # below the configured default. - est = ( - self.max_num_batched_tokens - * self.num_qo_heads - * self.head_dim - * FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM - ) - buffer_size = max(buffer_size, est) - self._workspace_buffer = torch.zeros( - buffer_size, dtype=torch.uint8, device=self.device + @classmethod + def requires_separate_cudagraph_memory_profiling( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> bool: + kv_specs = ( + kv_cache_spec.kv_cache_specs.values() + if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs) + else [kv_cache_spec] + ) + for spec in kv_specs: + if not isinstance(spec, AttentionSpec): + continue + head_size_v = getattr(spec, "head_size_v", None) or spec.head_size + if max(spec.head_size, head_size_v) >= 512: + return True + return False + + def _default_workspace_buffer_size(self) -> int: + if envs.VLLM_BATCH_INVARIANT: + return FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT + # FlashInfer prefill temp buffers scale with the prefill chunk and + # query-head footprint, rather than the context length. + estimated_prefill_size = ( + self.max_num_batched_tokens + * self.num_qo_heads + * self.head_dim + * FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM + ) + return max( + envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, + estimated_prefill_size, + ) + + def _native_initial_workspace_buffer_size(self) -> int: + if envs.VLLM_BATCH_INVARIANT or self.use_dcp: + return self._default_workspace_buffer_size() + return 1 + + def _allocate_workspace_buffer(self, buffer_size: int) -> torch.Tensor: + buffer_size = max(int(buffer_size), 1) + if self.use_vllm_workspace_manager_for_workspace_buffer(): + manager = current_workspace_manager() + (workspace_buffer,) = manager.get_simultaneous( + ((buffer_size,), torch.uint8), + ) + return workspace_buffer + return torch.zeros(buffer_size, dtype=torch.uint8, device=self.device) + + def _get_workspace_buffer(self, buffer_size: int | None = None): + if buffer_size is None: + buffer_size = self._default_workspace_buffer_size() + if _buffer_nbytes(self._workspace_state.buffer) < buffer_size: + self._workspace_state.set_buffer( + self._allocate_workspace_buffer(buffer_size) ) + assert self._workspace_state.buffer is not None + self._workspace_buffer = self._workspace_state.buffer return self._workspace_buffer + @staticmethod + def use_vllm_workspace_manager_for_workspace_buffer() -> bool: + return is_workspace_manager_initialized() + + def get_workspace_buffer_state(self) -> _FlashInferWorkspaceState: + return self._workspace_state + + def set_workspace_buffer_state(self, workspace_state: _FlashInferWorkspaceState): + self._workspace_state = workspace_state + self._workspace_buffer = workspace_state.buffer + def set_workspace_buffer(self, workspace_buffer: torch.Tensor): + self._workspace_state.set_buffer(workspace_buffer) self._workspace_buffer = workspace_buffer + def _register_workspace_wrapper(self, wrapper: object) -> None: + self._workspace_state.register_wrapper(wrapper) + + def _normalize_workspace_sizes( + self, workspace_size: WorkspaceSizes | int | None + ) -> WorkspaceSizes: + if workspace_size is None: + return WorkspaceSizes(self._default_workspace_buffer_size(), 0, False) + if isinstance(workspace_size, WorkspaceSizes): + return workspace_size + return WorkspaceSizes(int(workspace_size), 0, False) + + def _ensure_flashinfer_wrapper_int_workspace( + self, + wrapper: object, + required_bytes: int, + has_required_bytes: bool, + ) -> tuple[torch.Tensor | None, bool]: + int_workspace = getattr(wrapper, "_int_workspace_buffer", None) + if int_workspace is None: + return None, False + + required_bytes = max(int(required_bytes), 0) + current_bytes = _buffer_nbytes(int_workspace) + finalized = getattr(wrapper, "_vllm_flashinfer_int_workspace_finalized", False) + prepared = getattr(wrapper, "_vllm_flashinfer_int_workspace_prepared", False) + # Reserve-time CUDA graph wrappers are still unplanned, so the + # constructor default int workspace can be replaced with the exact + # workspace_size() result. Once finalized, keep the buffer stable + # because plan data and graph replay addresses live in it. + exact_size = finalized or bool(getattr(wrapper, "is_cuda_graph_enabled", False)) + + if not has_required_bytes: + if current_bytes >= required_bytes: + return int_workspace, False + target_bytes = max(required_bytes, 1) + else: + target_bytes = _int_workspace_allocation_bytes( + required_bytes, round_up=not exact_size + ) + if finalized: + if current_bytes < required_bytes: + raise AssertionError( + "FlashInfer CUDA graph int workspace is finalized but a " + f"larger buffer is required: {current_bytes} bytes " + f"allocated, {required_bytes} bytes required." + ) + return int_workspace, False + if prepared and not exact_size and current_bytes >= required_bytes: + return int_workspace, False + if current_bytes == target_bytes: + return int_workspace, False + + if finalized: + raise AssertionError( + "FlashInfer CUDA graph int workspace is finalized but a larger " + f"buffer is required: {current_bytes} bytes " + f"allocated, {required_bytes} bytes required." + ) + + if prepared and target_bytes > current_bytes: + logger.warning( + "Growing FlashInfer int workspace after initial preparation: " + "%.2f MiB -> %.2f MiB. This is allowed for non-captured " + "wrappers, but frequent growth means workspace reserve " + "candidates are too small.", + current_bytes / (1 << 20), + target_bytes / (1 << 20), + ) + + int_workspace = torch.empty( + (target_bytes,), dtype=torch.uint8, device=self.device + ) + object.__setattr__(wrapper, "_int_workspace_buffer", int_workspace) + return int_workspace, True + + def _flashinfer_wrapper_workspace_matches( + self, + wrapper: object, + float_workspace: torch.Tensor, + int_workspace: torch.Tensor | None, + ) -> bool: + current_float = getattr(wrapper, "_float_workspace_buffer", None) + if ( + current_float is None + or current_float.data_ptr() != float_workspace.data_ptr() + or _buffer_nbytes(current_float) != _buffer_nbytes(float_workspace) + ): + return False + + current_int = getattr(wrapper, "_int_workspace_buffer", None) + if current_int is None or int_workspace is None: + return current_int is int_workspace + return current_int.data_ptr() == int_workspace.data_ptr() and _buffer_nbytes( + current_int + ) == _buffer_nbytes(int_workspace) + + def _ensure_flashinfer_wrapper_workspace( + self, + wrapper: object, + workspace_size: WorkspaceSizes | int | None, + ) -> None: + sizes = self._normalize_workspace_sizes(workspace_size) + int_workspace, int_workspace_changed = ( + self._ensure_flashinfer_wrapper_int_workspace( + wrapper, sizes.int_bytes, sizes.has_int_bytes + ) + ) + float_workspace = self._get_workspace_buffer(sizes.float_bytes) + if ( + hasattr(wrapper, "reset_workspace_buffer") + and ( + int_workspace_changed + or not self._flashinfer_wrapper_workspace_matches( + wrapper, float_workspace, int_workspace + ) + ) + and int_workspace is not None + ): + wrapper.reset_workspace_buffer(float_workspace, int_workspace) + object.__setattr__(wrapper, "_vllm_flashinfer_int_workspace_prepared", True) + self._register_workspace_wrapper(wrapper) + + def _iter_workspace_wrappers(self) -> list[tuple[str, object]]: + wrappers: list[tuple[str, object]] = [] + + def add_wrapper(kind: str, wrapper: object | None) -> None: + if wrapper is not None: + wrappers.append((kind, wrapper)) + + add_wrapper("prefill", getattr(self, "_prefill_wrapper", None)) + add_wrapper( + "noncausal_prefill", getattr(self, "_noncausal_prefill_wrapper", None) + ) + add_wrapper("decode", getattr(self, "_decode_wrapper", None)) + add_wrapper("cascade", getattr(self, "_cascade_wrapper", None)) + + decode_wrappers_cudagraph = getattr(self, "_decode_wrappers_cudagraph", {}) + for wrapper in decode_wrappers_cudagraph.values(): + add_wrapper("decode_cudagraph", wrapper) + + return wrappers + + def get_workspace_reserve_debug_info(self) -> dict[str, int]: + wrappers = self._iter_workspace_wrappers() + unique_wrappers: dict[int, tuple[str, object]] = {} + for kind, wrapper in wrappers: + unique_wrappers.setdefault(id(wrapper), (kind, wrapper)) + + kind_counts: dict[str, int] = {} + actual_int_workspace_bytes = 0 + default_int_workspace_wrappers = 0 + unique_int_buffers: set[int] = set() + unique_float_buffers: dict[int, int] = {} + + for kind, wrapper in unique_wrappers.values(): + kind_counts[kind] = kind_counts.get(kind, 0) + 1 + + int_workspace = getattr(wrapper, "_int_workspace_buffer", None) + int_workspace_bytes = ( + _buffer_nbytes(int_workspace) + if isinstance(int_workspace, torch.Tensor) + else 0 + ) + actual_int_workspace_bytes += int_workspace_bytes + if int_workspace_bytes >= FLASHINFER_DEFAULT_INT_WORKSPACE_BYTES: + default_int_workspace_wrappers += 1 + if isinstance(int_workspace, torch.Tensor): + unique_int_buffers.add(int_workspace.data_ptr()) + + float_workspace = getattr(wrapper, "_float_workspace_buffer", None) + if isinstance(float_workspace, torch.Tensor): + unique_float_buffers[float_workspace.data_ptr()] = max( + unique_float_buffers.get(float_workspace.data_ptr(), 0), + _buffer_nbytes(float_workspace), + ) + + reserved_sizes = getattr( + self, "_last_reserved_workspace_sizes", WorkspaceSizes(0, 0) + ) + workspace_state = getattr(self, "_workspace_state", None) + workspace_state_live_wrappers = ( + len(workspace_state._live_wrappers()) if workspace_state is not None else 0 + ) + + return { + "workspace_wrapper_count": len(unique_wrappers), + "prefill_wrappers": kind_counts.get("prefill", 0), + "noncausal_prefill_wrappers": kind_counts.get("noncausal_prefill", 0), + "decode_wrappers": kind_counts.get("decode", 0), + "decode_cudagraph_wrappers": kind_counts.get("decode_cudagraph", 0), + "cascade_wrappers": kind_counts.get("cascade", 0), + "workspace_state_live_wrappers": workspace_state_live_wrappers, + "actual_int_workspace_bytes": actual_int_workspace_bytes, + "reserved_int_workspace_bytes": reserved_sizes.int_bytes, + "int_workspace_over_reserved_bytes": max( + actual_int_workspace_bytes - reserved_sizes.int_bytes, 0 + ), + "default_int_workspace_wrappers": default_int_workspace_wrappers, + "unique_int_workspace_buffers": len(unique_int_buffers), + "unique_float_workspace_buffers": len(unique_float_buffers), + "unique_float_workspace_bytes": sum(unique_float_buffers.values()), + } + @staticmethod def _get_flashinfer_trtllm_api_decode_kernel() -> FlashInferDecodeKernel: if current_platform.is_device_capability(90): @@ -983,10 +1353,13 @@ def _get_prefill_wrapper( ) if self._noncausal_prefill_wrapper is None: self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( - self._get_workspace_buffer(), + self._get_workspace_buffer( + self._native_initial_workspace_buffer_size() + ), get_kv_cache_layout(), backend="auto", ) + self._register_workspace_wrapper(self._noncausal_prefill_wrapper) return self._noncausal_prefill_wrapper if self._prefill_wrapper is None: @@ -1000,10 +1373,13 @@ def _get_prefill_wrapper( # the wrapper; fa2/fa3 do not support nvfp4. backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper( - self._get_workspace_buffer(), + self._get_workspace_buffer( + self._native_initial_workspace_buffer_size() + ), get_kv_cache_layout(), backend=backend, ) + self._register_workspace_wrapper(self._prefill_wrapper) assert self._prefill_wrapper is not None return self._prefill_wrapper @@ -1026,7 +1402,9 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): # the wrapper; fa2/fa3 do not support nvfp4. backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto" decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( - self._get_workspace_buffer(), + self._get_workspace_buffer( + self._native_initial_workspace_buffer_size() + ), get_kv_cache_layout(), use_cuda_graph=use_cudagraph, paged_kv_indptr_buffer=paged_kv_indptr, @@ -1038,6 +1416,7 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): use_tensor_cores=True, backend=backend, ) + self._register_workspace_wrapper(decode_wrapper) # save the decode wrapper if use_cudagraph: @@ -1050,10 +1429,428 @@ def _get_decode_wrapper(self, batch_size: int, use_cudagraph: bool = False): def _get_cascade_wrapper(self): if self._cascade_wrapper is None: self._cascade_wrapper = MultiLevelCascadeAttentionWrapper( - 2, self._get_workspace_buffer(), get_kv_cache_layout() + 2, + self._get_workspace_buffer(self._default_workspace_buffer_size()), + get_kv_cache_layout(), ) + self._register_workspace_wrapper(self._cascade_wrapper) return self._cascade_wrapper + def _get_tensor_core_decode_backend( + self, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, + ) -> str: + if _is_float8_dtype(q_data_type) or _is_float8_dtype(kv_data_type): + return determine_attention_backend( + self.device, + PosEncodingMode.NONE.value, + False, # use_fp16_qk_reductions + False, # use_custom_mask + q_data_type, + kv_data_type, + ) + return "fa2" + + def _get_prefill_backend( + self, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, + use_custom_mask: bool, + ) -> str: + return determine_attention_backend( + self.device, + PosEncodingMode.NONE.value, + False, # use_fp16_qk_reductions + use_custom_mask, + q_data_type, + kv_data_type, + ) + + def _get_prefill_workspace_size_func( + self, + *, + q_data_type: torch.dtype, + kv_data_type: torch.dtype, + paged_kv_indptr_dtype: torch.dtype, + use_custom_mask: bool, + window_left: int, + ) -> tuple[str, Any] | None: + backend = self._get_prefill_backend(q_data_type, kv_data_type, use_custom_mask) + try: + # Keep this module lookup aligned with the current plan() calls, + # which pass self.head_dim for both QK and VO dimensions. + module = get_batch_prefill_module( + backend, + q_data_type, + kv_data_type, + self.model_config.dtype, + paged_kv_indptr_dtype, + self.head_dim, + self.head_dim, + PosEncodingMode.NONE.value, + window_left != -1, + (self.logits_soft_cap or 0.0) > 0, + False, # use_fp16_qk_reduction + ) + workspace_size = getattr(module, "workspace_size", None) + if workspace_size is None: + return None + return backend, workspace_size + except Exception: + logger.debug( + "Failed to resolve FlashInfer prefill workspace size helper.", + exc_info=True, + ) + return None + + def _call_prefill_workspace_size( + self, + *, + backend: str, + workspace_size: Any, + qo_indptr_cpu: torch.Tensor, + paged_kv_indptr_cpu: torch.Tensor, + paged_kv_last_page_len_cpu: torch.Tensor, + causal: bool, + window_left: int, + fixed_split_size: int | None, + disable_split_kv: bool, + ) -> WorkspaceSizes | None: + try: + kv_lens_arr_cpu = get_seq_lens( + paged_kv_indptr_cpu, + paged_kv_last_page_len_cpu, + self.page_size, + ) + device_buffer = torch.empty(0, dtype=torch.uint8, device=self.device) + args = [ + device_buffer, + qo_indptr_cpu, + paged_kv_indptr_cpu, + kv_lens_arr_cpu, + int(qo_indptr_cpu[-1].item()), # total_num_rows + len(qo_indptr_cpu) - 1, # batch_size + self.num_qo_heads, + self.num_kv_heads, + self.page_size, + False, # enable_cuda_graph + self.head_dim, + self.head_dim, + causal, + window_left, + ] + if backend == "fa2": + args.extend( + [ + fixed_split_size or -1, + disable_split_kv, + 0, # num_colocated_ctas + ] + ) + return _parse_workspace_sizes(workspace_size(*args)) + except Exception: + logger.debug( + "Failed to calculate FlashInfer prefill workspace size.", + exc_info=True, + ) + return None + + def _get_prefill_workspace_size( + self, + *, + qo_indptr: torch.Tensor, + paged_kv_indptr: torch.Tensor, + paged_kv_last_page_len: torch.Tensor, + causal: bool, + window_left: int, + use_custom_mask: bool = False, + fixed_split_size: int | None = None, + disable_split_kv: bool = False, + ) -> WorkspaceSizes | None: + q_data_type = self.q_data_type_prefill + kv_data_type = self.kv_cache_dtype + helper = self._get_prefill_workspace_size_func( + q_data_type=q_data_type, + kv_data_type=kv_data_type, + paged_kv_indptr_dtype=paged_kv_indptr.dtype, + use_custom_mask=use_custom_mask, + window_left=window_left, + ) + if helper is None: + return None + backend, workspace_size = helper + return self._call_prefill_workspace_size( + backend=backend, + workspace_size=workspace_size, + qo_indptr_cpu=qo_indptr.to("cpu"), + paged_kv_indptr_cpu=paged_kv_indptr.to("cpu"), + paged_kv_last_page_len_cpu=paged_kv_last_page_len.to("cpu"), + causal=causal, + window_left=window_left, + fixed_split_size=fixed_split_size, + disable_split_kv=disable_split_kv, + ) + + @staticmethod + def _get_workspace_query_len_candidates(max_query_len: int) -> list[int]: + # FlashInfer split-K workspace can peak on short cached-prefill tails. + # Probe those densely, then sample larger chunks sparsely. + dense_limit = min(max_query_len, 256) + query_lens = list(range(1, dense_limit + 1)) + if max_query_len > dense_limit: + query_len = 512 + while query_len < max_query_len: + query_lens.append(query_len) + query_len *= 2 + query_lens.append(max_query_len) + return query_lens + + def _make_decode_workspace_inputs( + self, + *, + batch_size: int, + num_pages: int, + last_page_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + batch_arange = torch.arange(batch_size + 1, dtype=torch.int32, device="cpu") + indptr_cpu = batch_arange * num_pages + last_page_len_cpu = torch.full( + (batch_size,), + last_page_len, + dtype=torch.int32, + device="cpu", + ) + return indptr_cpu, last_page_len_cpu + + def _reserve_decode_wrapper_workspace( + self, + *, + batch_size: int, + num_pages: int, + last_page_len: int, + use_cudagraph: bool, + ) -> WorkspaceSizes: + decode_wrapper = self._get_decode_wrapper(batch_size, use_cudagraph) + indptr_cpu, last_page_len_cpu = self._make_decode_workspace_inputs( + batch_size=batch_size, + num_pages=num_pages, + last_page_len=last_page_len, + ) + workspace_sizes = self._get_decode_workspace_size( + decode_wrapper=decode_wrapper, + indptr_cpu=indptr_cpu, + last_page_len_cpu=last_page_len_cpu, + fixed_split_size=self.decode_fixed_split_size, + disable_split_kv=self.disable_split_kv, + ) + if workspace_sizes is None: + return WorkspaceSizes(0, 0) + self._ensure_flashinfer_wrapper_workspace(decode_wrapper, workspace_sizes) + if use_cudagraph: + decode_wrapper._vllm_flashinfer_int_workspace_finalized = True + return workspace_sizes + + def reserve_workspace_for_cudagraph_capture(self) -> int: + self._last_reserved_workspace_sizes = WorkspaceSizes(0, 0) + if self.use_dcp or not self.use_vllm_workspace_manager_for_workspace_buffer(): + return 0 + + max_model_len = self.model_config.max_model_len + scheduler_config = self.vllm_config.scheduler_config + max_num_batched_tokens = min( + scheduler_config.max_num_batched_tokens, + max_model_len, + ) + max_num_seqs = min(scheduler_config.max_num_seqs, max_num_batched_tokens) + if max_num_batched_tokens <= 0 or max_num_seqs <= 0: + return 0 + + helper = self._get_prefill_workspace_size_func( + q_data_type=self.q_data_type_prefill, + kv_data_type=self.kv_cache_dtype, + paged_kv_indptr_dtype=torch.int32, + use_custom_mask=False, + window_left=self.window_left, + ) + if helper is None: + return 0 + + backend, workspace_size = helper + num_pages = cdiv(max_model_len, self.page_size) + last_page_len = max_model_len % self.page_size or self.page_size + max_prefill_workspace_size = WorkspaceSizes(0, 0) + + for batch_size in range(1, max_num_seqs + 1): + max_query_len = max_num_batched_tokens // batch_size + if max_query_len <= 0: + break + batch_arange = torch.arange(batch_size + 1, dtype=torch.int32, device="cpu") + paged_kv_indptr_cpu = batch_arange * num_pages + paged_kv_last_page_len_cpu = torch.full( + (batch_size,), + last_page_len, + dtype=torch.int32, + device="cpu", + ) + query_lens = self._get_workspace_query_len_candidates(max_query_len) + for query_len in query_lens: + qo_indptr_cpu = batch_arange * query_len + workspace_sizes = self._call_prefill_workspace_size( + backend=backend, + workspace_size=workspace_size, + qo_indptr_cpu=qo_indptr_cpu, + paged_kv_indptr_cpu=paged_kv_indptr_cpu, + paged_kv_last_page_len_cpu=paged_kv_last_page_len_cpu, + causal=True, + window_left=self.window_left, + fixed_split_size=self.prefill_fixed_split_size, + disable_split_kv=self.disable_split_kv, + ) + if workspace_sizes is None: + continue + max_prefill_workspace_size = WorkspaceSizes( + max( + max_prefill_workspace_size.float_bytes, + workspace_sizes.float_bytes, + ), + max( + max_prefill_workspace_size.int_bytes, + workspace_sizes.int_bytes, + ), + ( + max_prefill_workspace_size.has_int_bytes + or workspace_sizes.has_int_bytes + ), + ) + + reserved_sizes = max_prefill_workspace_size + if max_prefill_workspace_size.total_bytes > 0: + prefill_wrapper = self._get_prefill_wrapper(causal=True) + self._ensure_flashinfer_wrapper_workspace( + prefill_wrapper, max_prefill_workspace_size + ) + + max_decode_tokens = max_num_seqs + speculative_config = self.vllm_config.speculative_config + if speculative_config is not None: + max_decode_tokens *= 1 + speculative_config.num_speculative_tokens + max_decode_tokens = min(max_decode_tokens, max_num_batched_tokens) + + if max_decode_tokens > 0: + decode_sizes = self._reserve_decode_wrapper_workspace( + batch_size=max_decode_tokens, + num_pages=num_pages, + last_page_len=last_page_len, + use_cudagraph=False, + ) + reserved_sizes = WorkspaceSizes( + max(reserved_sizes.float_bytes, decode_sizes.float_bytes), + reserved_sizes.int_bytes + decode_sizes.int_bytes, + reserved_sizes.has_int_bytes or decode_sizes.has_int_bytes, + ) + + if self.enable_cuda_graph: + for batch_size in self.compilation_config.cudagraph_capture_sizes or []: + if batch_size <= 0 or batch_size > self._decode_cudagraph_max_bs: + continue + decode_sizes = self._reserve_decode_wrapper_workspace( + batch_size=batch_size, + num_pages=num_pages, + last_page_len=last_page_len, + use_cudagraph=True, + ) + reserved_sizes = WorkspaceSizes( + max(reserved_sizes.float_bytes, decode_sizes.float_bytes), + reserved_sizes.int_bytes + decode_sizes.int_bytes, + reserved_sizes.has_int_bytes or decode_sizes.has_int_bytes, + ) + + if reserved_sizes.total_bytes <= 0: + return 0 + + self._last_reserved_workspace_sizes = reserved_sizes + logger.debug( + "Reserved FlashInfer workspace before CUDA graph lock: " + "%.2f MiB float workspace, %.2f MiB dedicated int workspace", + reserved_sizes.float_bytes / (1 << 20), + reserved_sizes.int_bytes / (1 << 20), + ) + return reserved_sizes.total_bytes + + def _get_decode_workspace_size( + self, + *, + decode_wrapper: BatchDecodeWithPagedKVCacheWrapper, + indptr_cpu: torch.Tensor, + last_page_len_cpu: torch.Tensor, + fixed_split_size: int, + disable_split_kv: bool, + ) -> WorkspaceSizes | None: + q_data_type = self.q_data_type_decode + kv_data_type = self.kv_cache_dtype + backend = self._get_tensor_core_decode_backend(q_data_type, kv_data_type) + try: + # Tensor-core decode uses FlashInfer's batch-prefill module template + # for workspace sizing, matching the wrapper planning path below. + module = get_batch_prefill_module( + backend, + q_data_type, + kv_data_type, + self.model_config.dtype, + indptr_cpu.dtype, + self.head_dim, + self.head_dim, + PosEncodingMode.NONE.value, + self.window_left != -1, + (self.logits_soft_cap or 0.0) > 0, + False, # use_fp16_qk_reduction + ) + workspace_size = getattr(module, "workspace_size", None) + if workspace_size is None: + return None + batch_size = len(last_page_len_cpu) + qo_indptr_cpu = torch.arange( + batch_size + 1, dtype=torch.int32, device="cpu" + ) + kv_lens_arr_cpu = get_seq_lens( + indptr_cpu, + last_page_len_cpu, + self.page_size, + ) + device_buffer = torch.empty(0, dtype=torch.uint8, device=self.device) + args = [ + device_buffer, + qo_indptr_cpu, + indptr_cpu, + kv_lens_arr_cpu, + batch_size, # total_num_rows + batch_size, + self.num_qo_heads * self.dcp_world_size, + self.num_kv_heads, + self.page_size, + decode_wrapper.is_cuda_graph_enabled, + self.head_dim, + self.head_dim, + False, # causal + self.window_left, + ] + if backend == "fa2": + args.extend( + [ + fixed_split_size, + disable_split_kv, + 0, # num_colocated_ctas + ] + ) + return _parse_workspace_sizes(workspace_size(*args)) + except Exception: + logger.debug( + "Failed to calculate FlashInfer decode workspace size.", + exc_info=True, + ) + return None + def _compute_flashinfer_kv_metadata( self, num_blocks_np: np.ndarray, @@ -1426,6 +2223,18 @@ def build( o_dtype = ( FP8_DTYPE if self.is_kvcache_nvfp4 else self.model_config.dtype ) + self._ensure_flashinfer_wrapper_workspace( + prefill_wrapper, + self._get_prefill_workspace_size( + qo_indptr=qo_indptr_prefill_cpu, + paged_kv_indptr=paged_kv_indptr_prefill_cpu, + paged_kv_last_page_len=(paged_kv_last_page_len_prefill_cpu), + causal=attn_metadata.causal, + window_left=self.window_left, + fixed_split_size=self.prefill_fixed_split_size, + disable_split_kv=self.disable_split_kv, + ), + ) prefill_wrapper.plan( qo_indptr=qo_indptr_prefill_cpu, paged_kv_indptr=paged_kv_indptr_prefill_cpu, @@ -1480,6 +2289,18 @@ def build( decode_wrapper = self._get_decode_wrapper( num_input_tokens, use_cudagraph ) + indptr_cpu = self.paged_kv_indptr.cpu[: num_input_tokens + 1] + last_page_len_cpu = self.paged_kv_last_page_len.cpu[:num_input_tokens] + self._ensure_flashinfer_wrapper_workspace( + decode_wrapper, + self._get_decode_workspace_size( + decode_wrapper=decode_wrapper, + indptr_cpu=indptr_cpu, + last_page_len_cpu=last_page_len_cpu, + fixed_split_size=self.decode_fixed_split_size, + disable_split_kv=self.disable_split_kv, + ), + ) # Use the persistent buffer with padding length, # instead of the same address but chunked version # in atten_metadata when using cudagraph. @@ -1491,11 +2312,9 @@ def build( ) fast_plan_decode( decode_wrapper, - indptr_cpu=self.paged_kv_indptr.cpu[: num_input_tokens + 1], + indptr_cpu=indptr_cpu, indices=paged_kv_indices, - last_page_len_cpu=self.paged_kv_last_page_len.cpu[ - :num_input_tokens - ], + last_page_len_cpu=last_page_len_cpu, num_qo_heads=self.num_qo_heads * self.dcp_world_size, num_kv_heads=self.num_kv_heads, head_dim=self.head_dim, diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 5c07860b3ba4..6845abb9de0d 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -133,7 +133,7 @@ def init_attn_backend( kernel_block_sizes = prepare_kernel_block_sizes(kv_cache_config, attn_groups) # Phase 3: create metadata builders and determine cudagraph support. - attn_backend_workspace: torch.Tensor | None = None + attn_backend_workspace: object | None = None min_cg_support = AttentionCGSupport.ALWAYS min_cg_attn_backend = None for kv_cache_group_id, groups in enumerate(attn_groups): @@ -149,10 +149,16 @@ def init_attn_backend( ) builder = group.get_metadata_builder(0) if attn_backend_workspace is None: - if hasattr(builder, "_get_workspace_buffer"): + if hasattr(builder, "get_workspace_buffer_state"): + attn_backend_workspace = builder.get_workspace_buffer_state() + elif hasattr(builder, "_get_workspace_buffer"): attn_backend_workspace = builder._get_workspace_buffer() else: - if hasattr(builder, "set_workspace_buffer"): + if hasattr(builder, "set_workspace_buffer_state"): + builder.set_workspace_buffer_state(attn_backend_workspace) + elif isinstance(attn_backend_workspace, torch.Tensor) and hasattr( + builder, "set_workspace_buffer" + ): builder.set_workspace_buffer(attn_backend_workspace) # Check cudagraph support for the attention backend cg_support = builder.get_cudagraph_support( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1dee0d1b9553..75c87820d5af 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -29,7 +29,7 @@ import vllm.envs as envs from vllm.compilation.counter import compilation_counter -from vllm.config import VllmConfig +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import ( get_dcp_group, @@ -50,6 +50,7 @@ from vllm.utils.torch_utils import PIN_MEMORY, STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput @@ -113,6 +114,7 @@ from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin from vllm.v1.worker.utils import KVBlockZeroer, copy_kv_cache_blocks_inplace +from vllm.v1.worker.workspace import lock_workspace logger = init_logger(__name__) @@ -402,7 +404,9 @@ def main_stream(self) -> torch.cuda.Stream: def get_kv_cache_spec(self): return get_kv_cache_spec(self.vllm_config) - def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: + def initialize_kv_cache( + self, kv_cache_config: KVCacheConfig, is_profiling: bool = False + ) -> None: kv_cache_config = deepcopy(kv_cache_config) self.kv_cache_config = kv_cache_config @@ -476,7 +480,7 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.speculator.init_cudagraph_manager(cudagraph_mode) check_attention_cp_compatibility(self.vllm_config) - if isinstance(self.speculator, DraftModelSpeculator): + if isinstance(self.speculator, DraftModelSpeculator) and not is_profiling: # HACK(woosuk) self.speculator.set_attn( self.model_state, @@ -497,7 +501,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.kernel_block_sizes, self.vllm_config, ) - self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) + if not is_profiling: + self.kv_connector = get_kv_connector(self.vllm_config, kv_caches_dict) def _init_kv_zero_meta(self) -> None: """Build KV-block zeroing metadata; invoked from gpu_worker.""" @@ -688,9 +693,290 @@ def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int: # SP is not supported yet. return num_scheduled_tokens + def _init_minimal_kv_cache_for_profiling(self) -> None: + from vllm.v1.core.kv_cache_utils import ( + get_kv_cache_config_from_groups, + get_kv_cache_groups, + ) + + kv_cache_spec = self.get_kv_cache_spec() + KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec) + kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec) + min_blocks = self.compilation_config.max_cudagraph_capture_size or 1 + + # Temporarily allocate just enough KV cache state to instantiate + # attention metadata builders for workspace sizing. + saved_override = self.cache_config.num_gpu_blocks_override + self.cache_config.num_gpu_blocks_override = min_blocks + try: + minimal_config = get_kv_cache_config_from_groups( + self.vllm_config, + kv_cache_groups, + available_memory=0, + ) + finally: + self.cache_config.num_gpu_blocks_override = saved_override + + self.initialize_kv_cache(minimal_config, is_profiling=True) + self.cache_config.num_gpu_blocks = minimal_config.num_blocks + + logger.debug("Initialized minimal KV cache for CUDA graph profiling") + + def _cleanup_profiling_kv_cache(self) -> None: + torch.accelerator.synchronize() + + if hasattr(self, "kv_caches") and self.kv_caches: + for i in range(len(self.kv_caches)): + self.kv_caches[i] = None # type: ignore[assignment] + self.kv_caches.clear() + if hasattr(self, "attn_groups"): + self.attn_groups.clear() + for attr in ( + "attn_groups", + "kv_cache_config", + "block_tables", + "kernel_block_sizes", + "cudagraph_manager", + ): + if hasattr(self, attr): + delattr(self, attr) + self.cache_config.num_gpu_blocks = None + + for layer in self.compilation_config.static_forward_context.values(): + if hasattr(layer, "kv_cache"): + kv_cache = layer.kv_cache + layer.kv_cache = ( + torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else [] + ) + # Clean up quantized KV cache scale views + # (int8_per_token_head, fp8_per_token_head). + if hasattr(layer, "impl"): + if hasattr(layer.impl, "_k_scale_cache"): + layer.impl._k_scale_cache = None + if hasattr(layer.impl, "_v_scale_cache"): + layer.impl._v_scale_cache = None + + gc.collect() + torch.accelerator.empty_cache() + logger.debug("Cleaned up profiling KV cache and CUDA graphs") + + def _reserve_attention_workspace_for_cudagraph_capture(self) -> int: + if not getattr(self, "attn_groups", None): + return 0 + + reserved_before = torch.accelerator.memory_reserved(self.device) + allocated_before = torch.accelerator.memory_allocated(self.device) + requested_total = 0 + builder_records: list[ + tuple[str, int, int, int, int | None, int, dict[str, int] | None] + ] = [] + for groups in self.attn_groups: + for attn_group in groups: + for builder in attn_group.metadata_builders: + builder_reserved_before = torch.accelerator.memory_reserved( + self.device + ) + builder_allocated_before = torch.accelerator.memory_allocated( + self.device + ) + requested_bytes = int( + builder.reserve_workspace_for_cudagraph_capture() or 0 + ) + workspace_debug_info = None + get_workspace_reserve_debug_info = getattr( + builder, "get_workspace_reserve_debug_info", None + ) + if callable(get_workspace_reserve_debug_info): + workspace_debug_info = get_workspace_reserve_debug_info() + builder_reserved_after = torch.accelerator.memory_reserved( + self.device + ) + builder_allocated_after = torch.accelerator.memory_allocated( + self.device + ) + requested_total += max(requested_bytes, 0) + workspace_buffer_ptr = None + workspace_buffer_bytes = 0 + get_workspace_buffer_state = getattr( + builder, "get_workspace_buffer_state", None + ) + if callable(get_workspace_buffer_state): + workspace_state = get_workspace_buffer_state() + workspace_buffer = getattr(workspace_state, "buffer", None) + if isinstance(workspace_buffer, torch.Tensor): + workspace_buffer_ptr = workspace_buffer.data_ptr() + workspace_buffer_bytes = ( + workspace_buffer.numel() + * workspace_buffer.element_size() + ) + builder_records.append( + ( + type(builder).__name__, + requested_bytes, + max( + builder_reserved_after - builder_reserved_before, + 0, + ), + max( + builder_allocated_after - builder_allocated_before, + 0, + ), + workspace_buffer_ptr, + workspace_buffer_bytes, + workspace_debug_info, + ) + ) + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + reserved_after = torch.accelerator.memory_reserved(self.device) + allocated_after = torch.accelerator.memory_allocated(self.device) + reserved_delta = max(reserved_after - reserved_before, 0) + allocated_delta = max(allocated_after - allocated_before, 0) + + if builder_records: + unique_workspace_buffers: dict[int, int] = {} + for ( + _builder_name, + _requested_bytes, + _builder_reserved_delta, + _builder_allocated_delta, + buffer_ptr, + buffer_bytes, + _workspace_debug_info, + ) in builder_records: + if buffer_ptr is not None: + unique_workspace_buffers[buffer_ptr] = max( + unique_workspace_buffers.get(buffer_ptr, 0), buffer_bytes + ) + unique_workspace_buffer_bytes = sum(unique_workspace_buffers.values()) + workspace_debug_infos = [ + workspace_debug_info + for _, _, _, _, _, _, workspace_debug_info in builder_records + if workspace_debug_info is not None + ] + logger.debug( + "Reserved attention workspace before CUDA graph capture: " + "%.2f MiB allocator reserved delta, %.2f MiB allocated delta, " + "%.2f MiB requested by builders, %.2f MiB unexplained " + "(%d builders, %d unique workspace buffers, %.2f MiB unique " + "workspace bytes)", + reserved_delta / (1 << 20), + allocated_delta / (1 << 20), + requested_total / (1 << 20), + max(reserved_delta - requested_total, 0) / (1 << 20), + len(builder_records), + len(unique_workspace_buffers), + unique_workspace_buffer_bytes / (1 << 20), + ) + if workspace_debug_infos: + actual_int_workspace_bytes = sum( + info["actual_int_workspace_bytes"] for info in workspace_debug_infos + ) + reserved_int_workspace_bytes = sum( + info["reserved_int_workspace_bytes"] + for info in workspace_debug_infos + ) + int_workspace_over_reserved_bytes = sum( + info["int_workspace_over_reserved_bytes"] + for info in workspace_debug_infos + ) + unique_wrapper_float_workspace_bytes = max( + ( + info.get("unique_float_workspace_bytes", 0) + for info in workspace_debug_infos + ), + default=0, + ) + logger.debug( + "Reserved attention workspace wrapper breakdown: " + "%d wrappers, %d decode CUDA graph wrappers, " + "%d default-or-larger int workspaces, " + "%.2f MiB actual int workspace, %.2f MiB requested int " + "workspace, %.2f MiB int workspace over request, " + "%.2f MiB max unique float workspace", + sum( + info["workspace_wrapper_count"] + for info in workspace_debug_infos + ), + sum( + info["decode_cudagraph_wrappers"] + for info in workspace_debug_infos + ), + sum( + info["default_int_workspace_wrappers"] + for info in workspace_debug_infos + ), + actual_int_workspace_bytes / (1 << 20), + reserved_int_workspace_bytes / (1 << 20), + int_workspace_over_reserved_bytes / (1 << 20), + unique_wrapper_float_workspace_bytes / (1 << 20), + ) + for ( + builder_name, + requested_bytes, + builder_reserved_delta, + builder_allocated_delta, + workspace_buffer_ptr, + workspace_buffer_bytes, + workspace_debug_info, + ) in builder_records: + logger.debug( + "Reserved attention workspace builder=%s requested=%.2f MiB " + "reserved_delta=%.2f MiB allocated_delta=%.2f MiB " + "workspace_buffer_ptr=%s workspace_buffer=%.2f MiB", + builder_name, + requested_bytes / (1 << 20), + builder_reserved_delta / (1 << 20), + builder_allocated_delta / (1 << 20), + workspace_buffer_ptr, + workspace_buffer_bytes / (1 << 20), + ) + if workspace_debug_info is not None: + logger.debug( + "Reserved attention workspace builder=%s wrappers=%d " + "decode_cudagraph_wrappers=%d " + "default_int_workspaces=%d actual_int=%.2f MiB " + "requested_int=%.2f MiB int_over_request=%.2f MiB " + "unique_int_buffers=%d unique_float_buffers=%d " + "unique_float=%.2f MiB " + "workspace_state_live_wrappers=%d", + builder_name, + workspace_debug_info["workspace_wrapper_count"], + workspace_debug_info["decode_cudagraph_wrappers"], + workspace_debug_info["default_int_workspace_wrappers"], + workspace_debug_info["actual_int_workspace_bytes"] / (1 << 20), + workspace_debug_info["reserved_int_workspace_bytes"] + / (1 << 20), + workspace_debug_info["int_workspace_over_reserved_bytes"] + / (1 << 20), + workspace_debug_info["unique_int_workspace_buffers"], + workspace_debug_info["unique_float_workspace_buffers"], + workspace_debug_info.get("unique_float_workspace_bytes", 0) + / (1 << 20), + workspace_debug_info["workspace_state_live_wrappers"], + ) + + return reserved_delta + def profile_cudagraph_memory(self) -> int: - # NOTE(woosuk): It is TBD whether we keep this API or not. - return 0 + self.cudagraph_memory_persistent_estimate = 0 + self.cudagraph_memory_graph_pool_estimate = 0 + + try: + with set_current_vllm_config(self.vllm_config): + self._init_minimal_kv_cache_for_profiling() + persistent_estimate = ( + self._reserve_attention_workspace_for_cudagraph_capture() + ) + finally: + self._cleanup_profiling_kv_cache() + + self.cudagraph_memory_persistent_estimate = int(persistent_estimate) + logger.info( + "Estimated CUDA graph persistent workspace memory: %.2f GiB", + persistent_estimate / (1 << 30), + ) + return int(persistent_estimate) @torch.inference_mode() def capture_model(self) -> int: @@ -706,6 +992,8 @@ def capture_model(self) -> int: start_time = time.perf_counter() gc.collect() + self._reserve_attention_workspace_for_cudagraph_capture() + torch.accelerator.synchronize() torch.accelerator.empty_cache() start_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -735,6 +1023,7 @@ def capture_model(self) -> int: elapsed_time, cuda_graph_size / (1 << 30), ) + lock_workspace() return cuda_graph_size def _remove_request(self, req_id: str) -> bool: diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 377fb670e395..748cf5424ee3 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6553,16 +6553,111 @@ def _maybe_init_encoder_cudagraph_manager(self) -> None: if self.encoder_cudagraph_manager is not None: logger.info("Initialized EncoderCudaGraphManager for vision encoder") + def _profile_seq_lens_for_cudagraph_memory( + self, + cudagraph_runtime_mode: CUDAGraphMode, + desc_index: int, + desc: BatchDescriptor, + ) -> int | None: + if cudagraph_runtime_mode != CUDAGraphMode.FULL or desc_index != 0: + return None + return min(self.max_model_len, self.max_num_tokens // desc.num_tokens) + + def _requires_separate_cudagraph_memory_profiling(self) -> bool: + for attn_group in self._kv_cache_spec_attn_group_iterator(): + builder_cls = attn_group.backend.get_builder_cls() + if builder_cls.requires_separate_cudagraph_memory_profiling( + self.vllm_config, attn_group.kv_cache_spec + ): + return True + return False + + def _reserve_attention_workspace_for_cudagraph_capture(self) -> int: + if not getattr(self, "attn_groups", None): + return 0 + + reserved_before = torch.accelerator.memory_reserved(self.device) + for attn_group in self._attn_group_iterator(): + for builder in attn_group.metadata_builders: + builder.reserve_workspace_for_cudagraph_capture() + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + reserved_after = torch.accelerator.memory_reserved(self.device) + return max(reserved_after - reserved_before, 0) + + def _profile_cudagraph_memory_separately( + self, + capture_descs: list[tuple[CUDAGraphMode, list[BatchDescriptor]]], + shared_memory_estimate: dict[CUDAGraphMode, int], + per_graph_estimate: dict[CUDAGraphMode, int], + ) -> int: + reserved_before = torch.accelerator.memory_reserved(self.device) + for mode, descs in capture_descs: + profile_descs = descs[:2] + for i, desc in enumerate(profile_descs): + self._warmup_before_cudagraph_capture( + desc, + cudagraph_runtime_mode=mode, + profile_seq_lens=( + self._profile_seq_lens_for_cudagraph_memory(mode, i, desc) + ), + ) + + torch.accelerator.synchronize() + torch.accelerator.empty_cache() + reserved_after = torch.accelerator.memory_reserved(self.device) + persistent_memory_estimate = max(reserved_after - reserved_before, 0) + + for mode, descs in capture_descs: + profile_descs = descs[:2] + mem_samples: list[int] = [] + + for i, desc in enumerate(profile_descs): + mem_before = torch.accelerator.get_memory_info()[0] + self._warmup_and_capture( + desc, + cudagraph_runtime_mode=mode, + profile_seq_lens=( + self._profile_seq_lens_for_cudagraph_memory(mode, i, desc) + ), + num_warmups=0, + ) + torch.accelerator.synchronize() + free_after = torch.accelerator.get_memory_info()[0] + mem_samples.append(mem_before - free_after) + + first_capture = mem_samples[0] + # Use at least 1 MiB per graph for driver overhead. + per_graph = max(mem_samples[1] if len(mem_samples) > 1 else 0, 1 << 20) + + shared_memory_estimate[mode] = first_capture + per_graph_estimate[mode] = per_graph * (len(descs) - 1) + + logger.debug( + "Estimated %s CUDA graph memory with separate profiling: " + "%.2f MiB first-capture + (%d-1) x %.2f MiB per-graph", + mode.name, + first_capture / (1 << 20), + len(descs), + per_graph / (1 << 20), + ) + + return persistent_memory_estimate + @torch.inference_mode() def profile_cudagraph_memory(self) -> int: + self.cudagraph_memory_persistent_estimate = 0 + self.cudagraph_memory_graph_pool_estimate = 0 + with set_current_vllm_config(self.vllm_config): self._init_minimal_kv_cache_for_profiling() saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured + use_separate_profiling = self._requires_separate_cudagraph_memory_profiling() capture_descs = self.cudagraph_dispatcher.get_capture_descs() - # Use a temporary manager for memory profiling. The persistent manager - # is initialized later so it does not keep profiling-only graph state. + # Use a temporary encoder manager for memory profiling so encoder + # profiling graphs do not persist after the estimate is collected. encoder_cudagraph_manager = self._create_encoder_cudagraph_manager() decoder_graphs = sum(len(descs) for _, descs in capture_descs) @@ -6591,6 +6686,11 @@ def profile_cudagraph_memory(self) -> int: ) logger.info("Profiling CUDA graph memory: %s", ", ".join(graph_groups)) + if use_separate_profiling: + logger.info( + "Using separate CUDA graph memory profiling for persistent " + "warmup allocations and graph-pool allocations" + ) # Use a temporary pool for profiling to avoid fragmentation in the main pool. profiling_pool = current_platform.graph_pool_handle() @@ -6603,8 +6703,9 @@ def profile_cudagraph_memory(self) -> int: original_pools[id(instance)] = instance.graph_pool instance.graph_pool = profiling_pool - shared_memory_estimate = {} - per_graph_estimate = {} + shared_memory_estimate: dict[CUDAGraphMode, int] = {} + per_graph_estimate: dict[CUDAGraphMode, int] = {} + persistent_memory_estimate = 0 encoder_memory_estimate = 0 # Cleanup-only guard: CUDA graph capture errors should still propagate @@ -6614,6 +6715,19 @@ def profile_cudagraph_memory(self) -> int: with self._freeze_gc(), graph_capture(device=self.device): torch.accelerator.synchronize() torch.accelerator.empty_cache() + persistent_memory_estimate += ( + self._reserve_attention_workspace_for_cudagraph_capture() + ) + + if use_separate_profiling: + persistent_memory_estimate += ( + self._profile_cudagraph_memory_separately( + capture_descs, + shared_memory_estimate, + per_graph_estimate, + ) + ) + capture_descs = [] for mode, descs in capture_descs: profile_descs = descs[:2] @@ -6648,7 +6762,7 @@ def profile_cudagraph_memory(self) -> int: logger.debug( "Estimated %s CUDA graph memory: " - "%.2f MiB first-capture + (%d-1) × %.2f MiB per-graph", + "%.2f MiB first-capture + (%d-1) x %.2f MiB per-graph", mode.name, first_capture / (1 << 20), len(descs), @@ -6694,7 +6808,10 @@ def profile_cudagraph_memory(self) -> int: ) # Encoder graphs use a manager-local pool at runtime, separate from the # decoder pool, so add their estimate instead of overlaying it. - total_estimate = decoder_estimate + encoder_memory_estimate + graph_pool_estimate = decoder_estimate + encoder_memory_estimate + total_estimate = persistent_memory_estimate + graph_pool_estimate + self.cudagraph_memory_persistent_estimate = int(persistent_memory_estimate) + self.cudagraph_memory_graph_pool_estimate = int(graph_pool_estimate) logger.info( "Estimated CUDA graph memory: %.2f GiB total", total_estimate / (1 << 30), @@ -6723,6 +6840,7 @@ def capture_model(self) -> int: # can reuse the memory pool allocated for the large shapes. set_cudagraph_capturing_enabled(True) with self._freeze_gc(), graph_capture(device=self.device): + self._reserve_attention_workspace_for_cudagraph_capture() torch.accelerator.synchronize() torch.accelerator.empty_cache() start_free_gpu_memory = torch.accelerator.get_memory_info()[0] @@ -6770,7 +6888,7 @@ def capture_model(self) -> int: ) return cuda_graph_size - def _warmup_and_capture( + def _warmup_before_cudagraph_capture( self, desc: BatchDescriptor, cudagraph_runtime_mode: CUDAGraphMode, @@ -6793,6 +6911,22 @@ def _warmup_and_capture( num_active_loras=desc.num_active_loras, profile_seq_lens=profile_seq_lens, ) + + def _warmup_and_capture( + self, + desc: BatchDescriptor, + cudagraph_runtime_mode: CUDAGraphMode, + profile_seq_lens: int | None = None, + allow_microbatching: bool = False, + num_warmups: int | None = None, + ): + self._warmup_before_cudagraph_capture( + desc, + cudagraph_runtime_mode=cudagraph_runtime_mode, + profile_seq_lens=profile_seq_lens, + allow_microbatching=allow_microbatching, + num_warmups=num_warmups, + ) self._dummy_run( desc.num_tokens, cudagraph_runtime_mode=cudagraph_runtime_mode, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 871f2f31c006..d76f8cd17770 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -523,6 +523,12 @@ def determine_available_memory(self) -> int: self.non_torch_memory = profile_result.non_torch_increase self.peak_activation_memory = profile_result.torch_peak_increase self.cudagraph_memory_estimate = cudagraph_memory_estimate + self.cudagraph_memory_persistent_estimate = getattr( + self.model_runner, "cudagraph_memory_persistent_estimate", 0 + ) + self.cudagraph_memory_graph_pool_estimate = getattr( + self.model_runner, "cudagraph_memory_graph_pool_estimate", 0 + ) free_gpu_memory = profile_result.after_profile.free_memory # NOTE(woosuk): Here we assume that the other processes using the same @@ -787,15 +793,23 @@ def compile_or_warm_up_model(self) -> CompilationTimes: and self.cudagraph_memory_estimate > 0 ): GiB = lambda b: round(b / GiB_bytes, 2) - diff = abs(cuda_graph_memory_bytes - self.cudagraph_memory_estimate) + graph_pool_estimate = self.cudagraph_memory_graph_pool_estimate + if graph_pool_estimate == 0: + graph_pool_estimate = self.cudagraph_memory_estimate + diff = abs(cuda_graph_memory_bytes - graph_pool_estimate) logger.info( "CUDA graph pool memory: %s GiB (actual), %s GiB (estimated), " "difference: %s GiB (%.1f%%).", GiB(cuda_graph_memory_bytes), - GiB(self.cudagraph_memory_estimate), + GiB(graph_pool_estimate), GiB(diff), 100 * diff / max(cuda_graph_memory_bytes, 1), ) + if self.cudagraph_memory_persistent_estimate > 0: + logger.info( + "CUDA graph persistent memory: %s GiB (estimated).", + GiB(self.cudagraph_memory_persistent_estimate), + ) if self.cache_config.kv_cache_memory_bytes is None and hasattr( self, "peak_activation_memory" @@ -812,12 +826,18 @@ def compile_or_warm_up_model(self) -> CompilationTimes: # slightly underestimate the memory consumption. # So leave a small buffer (=150MiB) to avoid OOM. redundancy_buffer_memory = 150 * (1 << 20) + cuda_graph_memory_for_sizing = cuda_graph_memory_bytes + if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: + cuda_graph_memory_for_sizing = max( + cuda_graph_memory_for_sizing, + self.cudagraph_memory_estimate, + ) non_kv_cache_memory = ( self.model_runner.model_memory_usage + self.peak_activation_memory + self.non_torch_memory - + cuda_graph_memory_bytes + + cuda_graph_memory_for_sizing ) kv_cache_memory_bytes_to_gpu_limit = ( self.init_snapshot.free_memory @@ -840,7 +860,8 @@ def compile_or_warm_up_model(self) -> CompilationTimes: f"Actual usage is {format_gib(self.model_runner.model_memory_usage)} " f"GiB for weight, {format_gib(self.peak_activation_memory)} GiB " f"for peak activation, {format_gib(self.non_torch_memory)} GiB " - f"for non-torch memory, and {format_gib(cuda_graph_memory_bytes)} " + f"for non-torch memory, and " + f"{format_gib(cuda_graph_memory_for_sizing)} " f"GiB for CUDAGraph memory. Replace gpu_memory_utilization " f"config with `--kv-cache-memory=" f"{kv_cache_memory_bytes_to_requested_limit}` "