[GG] Fix MRV2 CUDA graph and sparse DCP memory profiling

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

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

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

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 (