diff --git a/tests/distributed/test_b12x_fused_all_reduce.py b/tests/distributed/test_b12x_fused_all_reduce.py index 81a0ecd819c1..acb49eb1c204 100644 --- a/tests/distributed/test_b12x_fused_all_reduce.py +++ b/tests/distributed/test_b12x_fused_all_reduce.py @@ -131,6 +131,21 @@ def test_b12x_fused_allreduce_zero_cutoff_disables_support() -> None: assert not custom_allreduce.supports_fused_add_rms_norm() +def test_b12x_channel_checkpoint_delegates_to_runtime() -> None: + custom_allreduce, runtime = make_b12x_custom_allreduce( + allreduce_max_size=64, + fused_max_size=64, + ) + checkpoint = object() + runtime.checkpoint_channels.return_value = checkpoint + + assert custom_allreduce.checkpoint_pcie_channels() is checkpoint + custom_allreduce.rollback_pcie_channels(checkpoint) + + runtime.checkpoint_channels.assert_called_once_with() + runtime.rollback_channels.assert_called_once_with(checkpoint) + + def test_b12x_fused_custom_op_dispatch(monkeypatch) -> None: custom_allreduce = MagicMock() custom_allreduce.try_fused_add_rms_norm.return_value = True diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index c001e212ab5f..6acf33f7dc20 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -506,6 +506,98 @@ def fake_all_reduce(tensor, *, op, group): assert captured["op"] == dist.ReduceOp.MAX +def test_b12x_dcp_channel_rollback_restores_existing_and_closes_new_pools( + monkeypatch, +): + from vllm.v1.attention.ops import dcp_alltoall + + events = [] + + class _FakePool: + def __init__(self, name): + self.name = name + + def checkpoint_channels(self): + events.append(("checkpoint", self.name)) + return f"{self.name}-checkpoint" + + def rollback_channels(self, checkpoint): + events.append(("rollback", self.name, checkpoint)) + + def close(self): + events.append(("close", self.name)) + + device_group = object() + group = _FakeCPGroup(2, device_group) # type: ignore[arg-type] + existing_key = (id(device_group), 0, 64, 512, 576, 64) + new_key = (id(device_group), 0, 64, 576, 576, 64) + foreign_key = (id(object()), 0, 64, 512, 576, 64) + existing = _FakePool("existing") + foreign = _FakePool("foreign") + pools = {existing_key: existing, foreign_key: foreign} + monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_POOLS", pools) + + checkpoint = dcp_alltoall.checkpoint_b12x_dcp_a2a_channels(group) + transient = _FakePool("transient") + pools[new_key] = transient + dcp_alltoall.rollback_b12x_dcp_a2a_channels(checkpoint) + + assert pools == {existing_key: existing, foreign_key: foreign} + assert events == [ + ("checkpoint", "existing"), + ("rollback", "existing", "existing-checkpoint"), + ("close", "transient"), + ] + + +def test_profile_channel_checkpoint_rolls_back_all_b12x_transports(monkeypatch): + from vllm.distributed import parallel_state + from vllm.v1.attention.ops import dcp_alltoall + + events = [] + + class _FakeCommunicator: + def checkpoint_pcie_channels(self): + events.append("checkpoint-tp") + return "tp-checkpoint" + + def rollback_pcie_channels(self, checkpoint): + events.append(("rollback-tp", checkpoint)) + + class _FakeGroup: + def __init__(self, *, world_size, communicator=None): + self.world_size = world_size + self.device_communicator = type( + "DeviceCommunicator", (), {"ca_comm": communicator} + )() + + communicator = _FakeCommunicator() + tp_group = _FakeGroup(world_size=8, communicator=communicator) + pp_group = _FakeGroup(world_size=1, communicator=communicator) + dcp_group = _FakeGroup(world_size=2) + monkeypatch.setattr(parallel_state, "_TP", tp_group) + monkeypatch.setattr(parallel_state, "_PP", pp_group) + monkeypatch.setattr(parallel_state, "_DCP", dcp_group) + monkeypatch.setattr( + dcp_alltoall, + "checkpoint_b12x_dcp_a2a_channels", + lambda group: events.append("checkpoint-dcp") or "dcp-checkpoint", + ) + monkeypatch.setattr( + dcp_alltoall, + "rollback_b12x_dcp_a2a_channels", + lambda checkpoint: events.append(("rollback-dcp", checkpoint)), + ) + + checkpoint = parallel_state.checkpoint_b12x_graph_channels() + parallel_state.rollback_b12x_graph_channels(checkpoint) + + assert events == [ + "checkpoint-tp", + "checkpoint-dcp", + ("rollback-dcp", "dcp-checkpoint"), + ("rollback-tp", "tp-checkpoint"), + ] @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.") def test_b12x_lse_reduce_honors_token_cap(monkeypatch: pytest.MonkeyPatch): from vllm.v1.attention.ops import dcp_alltoall diff --git a/tests/v1/attention/test_b12x_mla_dcp_workspace.py b/tests/v1/attention/test_b12x_mla_dcp_workspace.py index 8fd0a8fd020c..80afe53e095c 100644 --- a/tests/v1/attention/test_b12x_mla_dcp_workspace.py +++ b/tests/v1/attention/test_b12x_mla_dcp_workspace.py @@ -4,8 +4,11 @@ import pytest import torch +from vllm import envs from vllm.model_executor.layers.attention.mla_attention import ( + MLAAttention, _can_use_b12x_dcp_prefill_workspace, + _estimate_dcp_ag_rs_transient_bytes, ) from vllm.v1.attention.backends.mla.b12x_mla_sparse import B12xMLASparseImpl from vllm.v1.attention.ops import common @@ -60,6 +63,110 @@ def test_dcp_workspace_gate_accepts_valid_rows(num_tokens, max_num_tokens): ) +def _make_profile_attention(*, workspace_enabled: bool, pure_a2a: bool = False): + class Backend: + @staticmethod + def get_name(): + return "B12X_MLA_SPARSE" + + class Impl: + dcp_world_size = 6 + dcp_workspace_non_dbo = True + is_sparse = True + _max_batched = 4096 + + attn = object.__new__(MLAAttention) + attn.attn_backend = Backend + attn.impl = Impl() + attn.num_heads = 11 + attn.kv_lora_rank = 512 + attn.qk_rope_head_dim = 64 + attn.v_head_dim = 256 + attn.dcp_project_before_merge = True + attn.dcp_project_before_merge_min_prefill_tokens = 1024 + attn.dcp_a2a = True + attn.dcp_a2a_max_tokens = 0 if pure_a2a else 256 + attn.dcp_a2a_large_backend = "ag_rs" + return attn, workspace_enabled + + +def test_sparse_profile_reserves_largest_non_workspace_ag_rs_batch(monkeypatch): + attn, workspace_enabled = _make_profile_attention(workspace_enabled=True) + monkeypatch.setattr(envs, "VLLM_MEMORY_PROFILE_INCLUDE_ATTN", True) + monkeypatch.setattr( + envs, "VLLM_B12X_MLA_DCP_GATHER_IN_WORKSPACE", workspace_enabled + ) + + expected = _estimate_dcp_ag_rs_transient_bytes( + num_tokens=1024, + local_heads=11, + dcp_world_size=6, + q_head_dim=576, + output_head_dim=512, + kv_lora_rank=512, + v_head_dim=256, + project_before_merge=False, + ) + assert attn._get_sparse_memory_profile_bytes() == expected + + +def test_sparse_profile_accounts_for_projected_fallback_without_workspace(monkeypatch): + attn, workspace_enabled = _make_profile_attention(workspace_enabled=False) + monkeypatch.setattr(envs, "VLLM_MEMORY_PROFILE_INCLUDE_ATTN", True) + monkeypatch.setattr( + envs, "VLLM_B12X_MLA_DCP_GATHER_IN_WORKSPACE", workspace_enabled + ) + + unprojected = _estimate_dcp_ag_rs_transient_bytes( + num_tokens=1024, + local_heads=11, + dcp_world_size=6, + q_head_dim=576, + output_head_dim=512, + kv_lora_rank=512, + v_head_dim=256, + project_before_merge=False, + ) + projected = _estimate_dcp_ag_rs_transient_bytes( + num_tokens=4096, + local_heads=11, + dcp_world_size=6, + q_head_dim=576, + output_head_dim=256, + kv_lora_rank=512, + v_head_dim=256, + project_before_merge=True, + ) + assert attn._get_sparse_memory_profile_bytes() == max(unprojected, projected) + + +def test_sparse_profile_accounts_for_unprojected_full_batch(monkeypatch): + attn, _ = _make_profile_attention(workspace_enabled=False) + attn.dcp_project_before_merge = False + monkeypatch.setattr(envs, "VLLM_MEMORY_PROFILE_INCLUDE_ATTN", True) + monkeypatch.setattr(envs, "VLLM_B12X_MLA_DCP_GATHER_IN_WORKSPACE", False) + + expected = _estimate_dcp_ag_rs_transient_bytes( + num_tokens=4096, + local_heads=11, + dcp_world_size=6, + q_head_dim=576, + output_head_dim=512, + kv_lora_rank=512, + v_head_dim=256, + project_before_merge=False, + ) + assert attn._get_sparse_memory_profile_bytes() == expected + + +def test_sparse_profile_skips_pure_a2a(monkeypatch): + attn, _ = _make_profile_attention(workspace_enabled=True, pure_a2a=True) + monkeypatch.setattr(envs, "VLLM_MEMORY_PROFILE_INCLUDE_ATTN", True) + monkeypatch.setattr(envs, "VLLM_B12X_MLA_DCP_GATHER_IN_WORKSPACE", True) + + assert attn._get_sparse_memory_profile_bytes() == 0 + + @pytest.mark.parametrize("world_size", [2, 3, 4, 6, 8]) def test_cp_lse_ag_out_rs_into_preserves_borrowed_output(monkeypatch, world_size): rank = world_size - 1 diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index 742aafd3890e..af568c5d18d6 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -9,6 +9,7 @@ import os import threading from contextlib import nullcontext +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -17,6 +18,121 @@ os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1" +def test_cudagraph_manager_clear_releases_capture_state(): + from vllm.v1.worker.gpu.cudagraph_utils import ModelCudaGraphManager + + manager = ModelCudaGraphManager.__new__(ModelCudaGraphManager) + manager.graphs = {object(): object()} + manager._graphs_captured = True + manager.breakable_cg_runner = object() + manager.hidden_states = object() + manager.aux_hidden_states = [object()] + manager.intermediate_tensors = object() + + manager.clear() + + assert manager.graphs == {} + assert not manager._graphs_captured + assert manager.breakable_cg_runner is None + assert manager.hidden_states is None + assert manager.aux_hidden_states == [] + assert manager.intermediate_tensors is None + + +def test_memory_profile_destroys_graphs_before_restoring_pools(monkeypatch): + from vllm.v1.worker.gpu import model_runner as model_runner_module + + profile_pool = object() + production_pool = object() + events: list[str] = [] + + class FakeManager: + def __init__(self): + self.pool = production_pool + + def needs_capture(self): + return True + + def capture(self, *args, **kwargs): + assert self.pool is profile_pool + assert wrapper.graph_pool is profile_pool + events.append("capture") + + class FakeWrapper: + def __init__(self): + self.graph_pool = production_pool + + manager = FakeManager() + wrapper = FakeWrapper() + runner = model_runner_module.GPUModelRunner.__new__( + model_runner_module.GPUModelRunner + ) + runner.vllm_config = object() + runner.cudagraph_manager = manager + runner.speculator = None + runner.lora_config = None + runner.model = object() + runner.model_state = object() + runner.input_buffers = object() + runner.intermediate_tensors = object() + runner.block_tables = object() + runner.attn_groups = object() + runner.kv_cache_config = object() + runner.use_aux_hidden_state_outputs = False + runner._init_minimal_kv_cache_for_profiling = lambda: None + runner.maybe_setup_dummy_loras = lambda _: nullcontext() + runner._zero_cudagraph_capture_kv_blocks = lambda: None + + def cleanup(): + events.append("cleanup") + assert manager.pool is profile_pool + assert wrapper.graph_pool is profile_pool + + runner._cleanup_cudagraph_memory_profile = cleanup + + memory_info = iter(((1000, 0), (900, 0), (950, 0))) + monkeypatch.setattr( + model_runner_module, "set_current_vllm_config", lambda _: nullcontext() + ) + monkeypatch.setattr( + model_runner_module, + "current_platform", + SimpleNamespace( + graph_pool_handle=lambda: profile_pool, + get_global_graph_pool=lambda: production_pool, + ), + ) + monkeypatch.setattr( + model_runner_module.CUDAGraphWrapper, "_all_instances", [wrapper] + ) + monkeypatch.setattr( + model_runner_module.BreakableCUDAGraphWrapper, "_all_instances", [] + ) + monkeypatch.setattr(model_runner_module.gc, "collect", lambda: None) + monkeypatch.setattr(torch.accelerator, "empty_cache", lambda: None) + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + monkeypatch.setattr(torch.accelerator, "get_memory_info", lambda: next(memory_info)) + monkeypatch.setattr( + model_runner_module, + "checkpoint_b12x_graph_channels", + lambda: events.append("checkpoint") or ("channel-checkpoint",), + ) + monkeypatch.setattr( + model_runner_module, + "rollback_b12x_graph_channels", + lambda checkpoint: ( + events.append("rollback") + if checkpoint == ("channel-checkpoint",) + else pytest.fail("rollback received the wrong channel checkpoint") + ), + ) + + assert runner.profile_cudagraph_memory() == 50 + assert events == ["checkpoint", "capture", "cleanup", "rollback"] + assert manager.pool is production_pool + assert wrapper.graph_pool is production_pool + + def test_piecewise_capture_builds_fresh_metadata_for_both_passes(): from vllm.config import CUDAGraphMode from vllm.v1.worker.gpu.cudagraph_utils import ( diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index b064de5a6d08..8cedeb8c7ba3 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/vllm/distributed/device_communicators/custom_all_reduce.py @@ -650,6 +650,34 @@ def capture(self, stream: torch.cuda.Stream | None = None): if not self.disabled and self._pcie_runtime is None: self.register_graph_buffers() + def checkpoint_pcie_channels(self) -> Any | None: + """Snapshot B12X channels before a throwaway graph capture. + + Returns: + An opaque runtime checkpoint, or ``None`` when PCIe all-reduce is + unavailable. + """ + runtime = self._pcie_runtime + checkpoint = getattr(runtime, "checkpoint_channels", None) + if checkpoint is None: + return None + return checkpoint() + + def rollback_pcie_channels(self, checkpoint: Any) -> None: + """Release B12X channels created after ``checkpoint``. + + Args: + checkpoint: Opaque state returned by ``checkpoint_pcie_channels``. + + Raises: + RuntimeError: If the B12X PCIe all-reduce runtime is unavailable. + """ + runtime = self._pcie_runtime + rollback = getattr(runtime, "rollback_channels", None) + if rollback is None: + raise RuntimeError("B12X PCIe all-reduce runtime is unavailable") + rollback(checkpoint) + def _pcie_runtime_stream(self) -> torch.cuda.Stream | None: pinned = self._pcie_capture_stream if pinned is None: diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 668e78ea03c8..e9c1727999ad 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1439,6 +1439,51 @@ def get_pp_group() -> GroupCoordinator: return _PP +def checkpoint_b12x_graph_channels() -> tuple[ + tuple[Callable[[Any], None], Any], ... +]: + """Snapshot B12X channels used by disposable CUDA graph captures.""" + checkpoints: list[tuple[Callable[[Any], None], Any]] = [] + seen_communicators: set[int] = set() + for group in (_TP, _DCP, _PP): + device_communicator = ( + None if group is None else group.device_communicator + ) + communicator = getattr(device_communicator, "ca_comm", None) + if communicator is None or id(communicator) in seen_communicators: + continue + seen_communicators.add(id(communicator)) + checkpoint_fn = getattr(communicator, "checkpoint_pcie_channels", None) + rollback_fn = getattr(communicator, "rollback_pcie_channels", None) + if checkpoint_fn is None or rollback_fn is None: + continue + checkpoint = checkpoint_fn() + if checkpoint is not None: + checkpoints.append((rollback_fn, checkpoint)) + + if _DCP is not None and _DCP.world_size > 1: + from vllm.v1.attention.ops.dcp_alltoall import ( + checkpoint_b12x_dcp_a2a_channels, + rollback_b12x_dcp_a2a_channels, + ) + + checkpoints.append( + ( + rollback_b12x_dcp_a2a_channels, + checkpoint_b12x_dcp_a2a_channels(_DCP), + ) + ) + return tuple(checkpoints) + + +def rollback_b12x_graph_channels( + checkpoints: tuple[tuple[Callable[[Any], None], Any], ...], +) -> None: + """Roll back B12X channels after disposable graphs are synchronized.""" + for rollback, checkpoint in reversed(checkpoints): + rollback(checkpoint) + + _DP: GroupCoordinator | None = None diff --git a/vllm/envs.py b/vllm/envs.py index a5f870bf875c..55af9cf151ea 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -322,6 +322,7 @@ VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: bool = False VLLM_ELASTIC_EP_DRAIN_REQUESTS: bool = False VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: bool = True + VLLM_MEMORY_PROFILE_INCLUDE_ATTN: bool = False VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False VLLM_XPU_USE_SAMPLER_KERNEL: bool = True @@ -2184,6 +2185,10 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS": lambda: bool( int(os.getenv("VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS", "1")) ), + # Include backend-declared transient attention buffers in the profile peak. + "VLLM_MEMORY_PROFILE_INCLUDE_ATTN": lambda: bool( + int(os.getenv("VLLM_MEMORY_PROFILE_INCLUDE_ATTN", "0")) + ), # NIXL EP environment variables "VLLM_NIXL_EP_MAX_NUM_RANKS": lambda: int( os.getenv("VLLM_NIXL_EP_MAX_NUM_RANKS", "32") diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 983fe831675a..d87241456b91 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -318,6 +318,47 @@ def _can_use_b12x_dcp_prefill_workspace( ) +def _estimate_dcp_ag_rs_transient_bytes( + *, + num_tokens: int, + local_heads: int, + dcp_world_size: int, + q_head_dim: int, + output_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + project_before_merge: bool, +) -> int: + """Upper-bound simultaneously live eager DCP AG/RS attention tensors.""" + if num_tokens <= 0 or local_heads <= 0 or dcp_world_size <= 1: + return 0 + + bf16_bytes = 2 + fp32_bytes = 4 + global_heads = local_heads * dcp_world_size + + gathered_query = num_tokens * global_heads * q_head_dim * bf16_bytes + attention_output = num_tokens * global_heads * output_head_dim * bf16_bytes + # CUDA communicator materializes a head-major contiguous RS input, then an + # output and its token-major contiguous return while attention_output lives. + reduce_scatter = attention_output + ( + 2 * num_tokens * local_heads * output_head_dim * bf16_bytes + ) + gathered_lse = (dcp_world_size + 1) * num_tokens * global_heads * fp32_bytes + gathered_w_uv = ( + global_heads * kv_lora_rank * v_head_dim * bf16_bytes + if project_before_merge + else 0 + ) + return ( + gathered_query + + attention_output + + reduce_scatter + + gathered_lse + + gathered_w_uv + ) + + def _extract_single_layer_index(layer_name: str) -> int | None: int_vals = [int(part) for part in layer_name.split(".") if part.isdecimal()] return int_vals[0] if len(int_vals) == 1 else None @@ -697,6 +738,78 @@ def __init__( compile_native=True, ) + def _get_sparse_memory_profile_bytes(self) -> int: + if ( + not envs.VLLM_MEMORY_PROFILE_INCLUDE_ATTN + or self.attn_backend.get_name() != "B12X_MLA_SPARSE" + or self.impl.dcp_world_size <= 1 + ): + return 0 + + max_tokens = int(getattr(self.impl, "_max_batched", 0)) + if max_tokens <= 0: + return 0 + + # Pure A2A does not enter the allocating NCCL AG/RS path. Hybrid A2A + # enters it immediately above the configured small-batch cap. + if self.dcp_a2a: + if self.dcp_a2a_max_tokens <= 0 or self.dcp_a2a_large_backend == "a2a": + return 0 + first_ag_rs_row = self.dcp_a2a_max_tokens + 1 + else: + first_ag_rs_row = 1 + + project_threshold = self.dcp_project_before_merge_min_prefill_tokens + workspace_start = max(1025, project_threshold + 1) + workspace_eligible = ( + workspace_start <= max_tokens + and _can_use_b12x_dcp_prefill_workspace( + enabled=envs.VLLM_B12X_MLA_DCP_GATHER_IN_WORKSPACE, + project_before_merge=self.dcp_project_before_merge, + dcp_use_b12x=False, + num_tokens=workspace_start, + max_num_tokens=max_tokens, + non_dbo_workspace=getattr(self.impl, "dcp_workspace_non_dbo", False), + is_sparse_impl=self.impl.is_sparse, + backend_name=self.attn_backend.get_name(), + is_capturing=False, + ) + ) + last_ag_rs_row = ( + min(max_tokens, workspace_start - 1) if workspace_eligible else max_tokens + ) + if first_ag_rs_row > last_ag_rs_row: + return 0 + + candidates: list[tuple[int, bool]] = [] + if not self.dcp_project_before_merge: + candidates.append((last_ag_rs_row, False)) + else: + unprojected_rows = min(last_ag_rs_row, project_threshold) + if unprojected_rows >= first_ag_rs_row: + candidates.append((unprojected_rows, False)) + if not workspace_eligible and last_ag_rs_row > project_threshold: + candidates.append((last_ag_rs_row, True)) + + return max( + ( + _estimate_dcp_ag_rs_transient_bytes( + num_tokens=num_tokens, + local_heads=self.num_heads, + dcp_world_size=self.impl.dcp_world_size, + q_head_dim=self.kv_lora_rank + self.qk_rope_head_dim, + output_head_dim=( + self.v_head_dim if projected else self.kv_lora_rank + ), + kv_lora_rank=self.kv_lora_rank, + v_head_dim=self.v_head_dim, + project_before_merge=projected, + ) + for num_tokens, projected in candidates + ), + default=0, + ) + @property def chunked_prefill_workspace_size(self) -> int: if self._chunked_prefill_workspace_size is None: @@ -828,6 +941,19 @@ def forward_impl( device=k_c_normed.device, dtype=k_c_normed.dtype, ) + else: + profile_workspace_bytes = self._get_sparse_memory_profile_bytes() + if profile_workspace_bytes > 0: + _ = torch.empty( + (profile_workspace_bytes,), + device=k_c_normed.device, + dtype=torch.uint8, + ) + logger.info_once( + "Including %.2f MiB of B12X sparse DCP transient " + "memory in the profile peak", + profile_workspace_bytes / (1 << 20), + ) # The zero fill is required when used with DP + EP # to ensure all ranks within a DP group compute the diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py index f8d986034a57..ce78133bb273 100644 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ b/vllm/v1/attention/ops/dcp_alltoall.py @@ -150,6 +150,37 @@ def _get_b12x_dcp_a2a_pool( return pool +def checkpoint_b12x_dcp_a2a_channels( + cp_group: GroupCoordinator, +) -> tuple[int, dict[Any, tuple[Any, Any]]]: + """Snapshot DCP pools before a throwaway graph capture.""" + group_id = id(cp_group.device_group) + checkpoints = { + key: (pool, pool.checkpoint_channels()) + for key, pool in _B12X_DCP_A2A_POOLS.items() + if key[0] == group_id + } + return group_id, checkpoints + + +def rollback_b12x_dcp_a2a_channels( + checkpoint: tuple[int, dict[Any, tuple[Any, Any]]], +) -> None: + """Restore DCP pools after their profiling graphs have been destroyed.""" + group_id, checkpoints = checkpoint + for key, pool in list(_B12X_DCP_A2A_POOLS.items()): + if key[0] != group_id: + continue + saved = checkpoints.get(key) + if saved is None: + pool.close() + del _B12X_DCP_A2A_POOLS[key] + continue + saved_pool, channel_checkpoint = saved + if pool is not saved_pool: + pool.close() + _B12X_DCP_A2A_POOLS[key] = saved_pool + saved_pool.rollback_channels(channel_checkpoint) def _try_b12x_dcp_lse_reduce( cp_attn_out: torch.Tensor, cp_attn_lse: torch.Tensor, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index ddc4d417dc69..d36d1c5ff1b0 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -436,6 +436,12 @@ def decode_descs( def needs_capture(self) -> bool: return len(self._capture_descs) > 0 + def clear(self) -> None: + """Release captured graphs and reset this manager for a later capture.""" + self.graphs.clear() + self._graphs_captured = False + self.breakable_cg_runner = None + @torch.inference_mode() def capture( self, @@ -743,6 +749,12 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None: super().capture(create_forward_fn, progress_bar_desc) + def clear(self) -> None: + super().clear() + self.hidden_states = None + self.aux_hidden_states.clear() + self.intermediate_tensors = None + def run_fullgraph( self, desc: BatchExecutionDescriptor ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b84cdd9096bd..e830ddc1b968 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -29,13 +29,17 @@ import torch.nn as nn import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper from vllm.compilation.counter import compilation_counter -from vllm.config import VllmConfig +from vllm.compilation.cuda_graph import CUDAGraphWrapper +from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import ( + checkpoint_b12x_graph_channels, get_dcp_group, get_pp_group, prepare_communication_buffer_for_model, + rollback_b12x_graph_channels, ) from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger @@ -44,6 +48,7 @@ ) from vllm.model_executor.model_loader import get_model_loader from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.utils.math_utils import cdiv @@ -51,6 +56,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.utils import record_function_or_nullcontext from vllm.v1.worker.cp_utils import check_attention_cp_compatibility @@ -791,9 +797,163 @@ 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 = ( + min(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size) + or 1 + ) + + 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) + self.cache_config.num_gpu_blocks = minimal_config.num_blocks + + def _cleanup_cudagraph_memory_profile(self) -> None: + torch.accelerator.synchronize() + if self.cudagraph_manager is not None: + self.cudagraph_manager.clear() + if self.speculator is not None: + self.speculator.clear_cudagraphs() + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + + if hasattr(self, "kv_caches"): + self.kv_caches.clear() + if hasattr(self, "attn_groups"): + self.attn_groups.clear() + if hasattr(self, "kv_cache_config"): + del self.kv_cache_config + for attr in ("block_tables", "kernel_block_sizes"): + if hasattr(self, attr): + delattr(self, attr) + + self.kv_connector = NO_OP_KV_CONNECTOR + self.kv_block_zeroer = None + self.cudagraph_manager = None + self.verification_capacity_manager = None + 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 [] + ) + + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() + def profile_cudagraph_memory(self) -> int: - # NOTE(woosuk): It is TBD whether we keep this API or not. - return 0 + with set_current_vllm_config(self.vllm_config): + self._init_minimal_kv_cache_for_profiling() + + assert self.cudagraph_manager is not None + if not self.cudagraph_manager.needs_capture(): + self._cleanup_cudagraph_memory_profile() + return 0 + + saved_num_cudagraph_captured = compilation_counter.num_cudagraph_captured + profiling_pool = current_platform.graph_pool_handle() + managers = [self.cudagraph_manager] + if self.speculator is not None: + managers.extend(self.speculator.get_cudagraph_managers()) + original_manager_pools = {id(manager): manager.pool for manager in managers} + for manager in managers: + manager.pool = profiling_pool + + wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + original_wrapper_pools = { + id(wrapper): wrapper.graph_pool for wrapper in wrappers + } + for wrapper in wrappers: + wrapper.graph_pool = profiling_pool + + gc.collect() + torch.accelerator.empty_cache() + torch.accelerator.synchronize() + start_free_gpu_memory = torch.accelerator.get_memory_info()[0] + graph_channel_checkpoints = () + try: + # Both target and draft captures can allocate graph-owned B12X + # channels. Snapshot before either capture so the disposable + # profiling pass cannot leave stale channels behind. + graph_channel_checkpoints = checkpoint_b12x_graph_channels() + with self.maybe_setup_dummy_loras(self.lora_config): + self.cudagraph_manager.capture( + self.model, + self.model_state, + self.input_buffers, + self.intermediate_tensors, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + has_lora=self.lora_config is not None, + use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, + lora_capture_hook=create_lora_capture_hook(self.lora_config, self), + progress_bar_desc="Profiling CUDA graph memory", + ) + if self.speculator is not None: + self.speculator.capture() + self._zero_cudagraph_capture_kv_blocks() + end_free_gpu_memory = torch.accelerator.get_memory_info()[0] + gross_cuda_graph_size = max(start_free_gpu_memory - end_free_gpu_memory, 0) + finally: + try: + # Destroy disposable graphs while every manager and wrapper still + # points at the private pool that owns their allocations. + try: + self._cleanup_cudagraph_memory_profile() + finally: + rollback_b12x_graph_channels(graph_channel_checkpoints) + finally: + for manager in managers: + manager.pool = original_manager_pools[id(manager)] + wrappers = list(CUDAGraphWrapper._all_instances) + list( + BreakableCUDAGraphWrapper._all_instances + ) + for wrapper in wrappers: + original_pool = original_wrapper_pools.get(id(wrapper)) + if id(wrapper) in original_wrapper_pools: + wrapper.graph_pool = original_pool + else: + wrapper.graph_pool = current_platform.get_global_graph_pool() + del profiling_pool + compilation_counter.num_cudagraph_captured = ( + saved_num_cudagraph_captured + ) + + free_after_cleanup = torch.accelerator.get_memory_info()[0] + retained_pool_size = max(start_free_gpu_memory - free_after_cleanup, 0) + # A CUDA graph private pool can retain physical pages after its graph + # objects are destroyed. memory_profiling observes those pages as + # non-torch memory, so only return the remaining capture cost here. + cuda_graph_size = max(gross_cuda_graph_size - retained_pool_size, 0) + logger.info( + "Estimated MRV2 CUDA graph memory: %.2f GiB additional " + "(%.2f GiB captured, %.2f GiB retained and counted as non-torch)", + cuda_graph_size / (1 << 30), + gross_cuda_graph_size / (1 << 30), + retained_pool_size / (1 << 30), + ) + return int(cuda_graph_size) @torch.inference_mode() def capture_model(self) -> int: diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index b025346fba64..f798ef59db6f 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -132,6 +132,16 @@ def capture(self) -> None: progress_bar_desc="Capturing decode CUDA graphs", ) + def get_cudagraph_managers(self) -> tuple[SpeculatorCudaGraphManager, ...]: + return tuple( + manager + for manager in ( + self.prefill_cudagraph_manager, + self.decode_cudagraph_manager, + ) + if manager is not None + ) + @torch.inference_mode() def propose( self, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index f9e9ca4b04b5..4859d74bcb15 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -203,6 +203,15 @@ def capture(self) -> None: progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", ) + def get_cudagraph_managers(self) -> tuple[DFlashCudaGraphManager, ...]: + if self.query_cudagraph_manager is None: + return () + return (self.query_cudagraph_manager,) + + def clear_cudagraphs(self) -> None: + super().clear_cudagraphs() + self._captured_backbone_outputs.clear() + def _warmup_prepare_inputs_kernel(self) -> None: if self.draft_kv_cache_group_id < 0: return diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 7a27b46de4b8..da42d3e416ff 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -27,6 +27,7 @@ from vllm.v1.worker.utils import AttentionGroup if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager from vllm.v1.worker.gpu.spec_decode.dspark.online_sts import DSparkOnlineSTS logger = init_logger(__name__) @@ -53,6 +54,13 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: def capture(self) -> None: pass + def get_cudagraph_managers(self) -> tuple["CudaGraphManager", ...]: + return () + + def clear_cudagraphs(self) -> None: + for manager in self.get_cudagraph_managers(): + manager.clear() + @abstractmethod def propose( self,