fix(memory): profile GLM DCP attention before KV cache sizing

tests/v1/worker/test_cp_utils.py

@@ -10,6 +10,27 @@
 from vllm.v1.worker.cp_utils import should_skip_dcp_context_attention
 
 
+@pytest.mark.parametrize(
+    "kwargs,expected",
+    [
+        ({}, 0),
+        ({"single_request_prefill": True}, 16),
+        ({"create_mixed_batch": True}, 16),
+    ],
+)
+def test_dcp_dummy_context_covers_single_request_prefill(kwargs, expected):
+    arguments = {
+        "dcp_world_size": 4,
+        "cp_kv_cache_interleave_size": 4,
+        "has_kv_cache_config": True,
+        "create_mixed_batch": False,
+        "is_graph_capturing": False,
+        "uniform_decode": False,
+    }
+    arguments.update(kwargs)
+    assert cp_utils.get_dcp_dummy_context_len(**arguments) == expected
+
+
 def test_skip_gate_only_for_zero_context():
     assert should_skip_dcp_context_attention(torch.zeros(3, dtype=torch.int32))
     assert not should_skip_dcp_context_attention(

tests/v1/worker/test_gpu_model_runner.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import gc
+from contextlib import nullcontext
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -1754,6 +1755,74 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
             runner.initialize_kv_cache(kv_cache_config)
 
 
+@pytest.mark.parametrize("dummy_run_fails", [False, True])
+def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up(
+    monkeypatch: pytest.MonkeyPatch,
+    dummy_run_fails: bool,
+):
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.model_config = SimpleNamespace(
+        architecture="Glm5NextForConditionalGeneration"
+    )
+    runner.dcp_world_size = 4
+    runner.vllm_config = object()
+    runner.max_num_tokens = 4096
+    events: list[object] = []
+
+    runner._init_minimal_kv_cache_for_profiling = lambda: events.append("init-kv")
+
+    def dummy_run(*args, **kwargs):
+        events.append(("dummy-run", args, kwargs))
+        if dummy_run_fails:
+            raise RuntimeError("expected DCP profile failure")
+        return torch.empty(1), torch.empty(1)
+
+    runner._dummy_run = dummy_run
+    runner._sync_device = lambda: events.append("sync")
+    runner._cleanup_profiling_kv_cache = lambda: events.append("cleanup")
+
+    monkeypatch.setattr(
+        gpu_model_runner_module,
+        "set_current_vllm_config",
+        lambda _: nullcontext(),
+    )
+
+    if dummy_run_fails:
+        with pytest.raises(RuntimeError, match="expected DCP profile failure"):
+            runner.profile_glm_dcp_attention()
+    else:
+        runner.profile_glm_dcp_attention()
+
+    assert events[0] == "init-kv"
+    assert events[1] == (
+        "dummy-run",
+        (4096,),
+        {
+            "force_attention": True,
+            "skip_eplb": True,
+            "is_profile": True,
+            "single_request_prefill": True,
+        },
+    )
+    assert events[-1] == "cleanup"
+    if not dummy_run_fails:
+        assert events[-2] == "sync"
+
+
+def test_glm_dcp_attention_profile_skips_non_glm_and_dcp1():
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.model_config = SimpleNamespace(architecture="OtherArchitecture")
+    runner.dcp_world_size = 4
+    runner._init_minimal_kv_cache_for_profiling = Mock()
+    runner.profile_glm_dcp_attention()
+    runner._init_minimal_kv_cache_for_profiling.assert_not_called()
+
+    runner.model_config.architecture = "Glm5NextForConditionalGeneration"
+    runner.dcp_world_size = 1
+    runner.profile_glm_dcp_attention()
+    runner._init_minimal_kv_cache_for_profiling.assert_not_called()
+
+
 class TestInitFp8KvScalesHybridModels:
     """Verify init_fp8_kv_scales handles heterogeneous kv_caches entries.
 

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -116,3 +116,89 @@ def stage_write(self, *_args):
         )
 
     assert block_tables.num_blocks.np[0, 1] == 3
+
+
+@pytest.mark.parametrize("dummy_run_fails", [False, True])
+def test_glm_dcp_attention_profile_uses_single_request_and_cleans_up(
+    monkeypatch: pytest.MonkeyPatch,
+    dummy_run_fails: bool,
+):
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.model_config = SimpleNamespace(
+        architecture="Glm5NextForConditionalGeneration"
+    )
+    runner.dcp_size = 4
+    runner.cp_interleave = 4
+    runner.max_num_tokens = 4096
+    events: list[object] = []
+
+    monkeypatch.setattr(
+        model_runner_module,
+        "_init_minimal_kv_cache_for_profiling",
+        lambda _: events.append("init-kv"),
+    )
+    monkeypatch.setattr(
+        model_runner_module,
+        "_teardown_profiling_state",
+        lambda _: events.append("cleanup"),
+    )
+
+    def dummy_run(*args, **kwargs):
+        events.append(("dummy-run", args, kwargs))
+        if dummy_run_fails:
+            raise RuntimeError("expected DCP profile failure")
+        return torch.empty(1), torch.empty(1)
+
+    runner._dummy_run = dummy_run
+    monkeypatch.setattr(torch.accelerator, "synchronize", lambda: events.append("sync"))
+
+    if dummy_run_fails:
+        with pytest.raises(RuntimeError, match="expected DCP profile failure"):
+            runner.profile_glm_dcp_attention()
+    else:
+        runner.profile_glm_dcp_attention()
+
+    assert events[0] == "init-kv"
+    assert events[1] == (
+        "dummy-run",
+        (4096,),
+        {
+            "context_len": 16,
+            "skip_eplb": True,
+            "is_profile": True,
+            "single_request_prefill": True,
+            "profile_all_kv_cache_groups": True,
+        },
+    )
+    assert events[-1] == "cleanup"
+    if not dummy_run_fails:
+        assert events[-2] == "sync"
+
+
+@pytest.mark.parametrize(
+    ("architecture", "dcp_size"),
+    [("OtherArchitecture", 4), ("Glm5NextForConditionalGeneration", 1)],
+)
+def test_glm_dcp_attention_profile_skips_irrelevant_configurations(
+    monkeypatch: pytest.MonkeyPatch,
+    architecture: str,
+    dcp_size: int,
+):
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.model_config = SimpleNamespace(architecture=architecture)
+    runner.dcp_size = dcp_size
+    initialized = False
+
+    def record_initialization(_):
+        nonlocal initialized
+        initialized = True
+
+    monkeypatch.setattr(
+        model_runner_module,
+        "_init_minimal_kv_cache_for_profiling",
+        record_initialization,
+    )
+
+    runner.profile_glm_dcp_attention()
+
+    assert not initialized

tests/v1/worker/test_gpu_worker.py

@@ -80,14 +80,27 @@ def test_startup_plan_apply_gate(plan_env):
     assert explicit.cache_config.kv_cache_memory_bytes == 7 * GiB_bytes
 
 
-def test_b12x_warmup_precedes_cudagraph_memory_profile(monkeypatch):
+@pytest.mark.parametrize(
+    "final_free_memory,expected_available_memory",
+    [(80, 80), (75, 75)],
+)
+def test_b12x_warmup_precedes_cudagraph_memory_profile(
+    monkeypatch,
+    final_free_memory,
+    expected_available_memory,
+):
     """B12X launch modules must be resolved before descriptor capture begins."""
-    events = []
+    events: list[object] = []
+
+    def profile_cudagraph_memory():
+        events.append("profile_cudagraph_memory")
+        return 0
 
     model_runner = SimpleNamespace(
         model_memory_usage=0,
         profile_run=lambda: events.append("profile_run"),
-        profile_cudagraph_memory=lambda: events.append("profile_cudagraph_memory") or 0,
+        profile_glm_dcp_attention=lambda: events.append("profile_glm_dcp_attention"),
+        profile_cudagraph_memory=profile_cudagraph_memory,
     )
     profile_result = SimpleNamespace(
         total_consumed=10,
@@ -108,6 +121,7 @@ def fake_memory_profiling(*args, **kwargs):
         model_runner=model_runner,
         init_snapshot=SimpleNamespace(free_memory=100, total_memory=100),
         requested_memory=90,
+        device="cuda:0",
         model_config=SimpleNamespace(multimodal_config=None),
         parallel_config=SimpleNamespace(),
         vllm_config=SimpleNamespace(
@@ -120,6 +134,11 @@ def fake_memory_profiling(*args, **kwargs):
 
     monkeypatch.setattr(gpu_worker, "maybe_apply_startup_plan", lambda worker: None)
     monkeypatch.setattr(gpu_worker, "memory_profiling", fake_memory_profiling)
+    monkeypatch.setattr(
+        gpu_worker,
+        "MemorySnapshot",
+        lambda **_kwargs: SimpleNamespace(free_memory=final_free_memory),
+    )
     monkeypatch.setattr(
         gpu_worker,
         "current_platform",
@@ -140,7 +159,8 @@ def fake_memory_profiling(*args, **kwargs):
 
     assert events == [
         "profile_run",
+        "profile_glm_dcp_attention",
         ("b12x_warmup", (8, 4)),
         "profile_cudagraph_memory",
     ]
-    assert available == 80
+    assert available == expected_available_memory

vllm/v1/worker/cp_utils.py

@@ -88,11 +88,16 @@ def get_dcp_dummy_context_len(
     create_mixed_batch: bool,
     is_graph_capturing: bool,
     uniform_decode: bool,
+    single_request_prefill: bool = False,
 ) -> int:
     if (
         dcp_world_size <= 1
         or not has_kv_cache_config
-        or not (create_mixed_batch or (is_graph_capturing and uniform_decode))
+        or not (
+            create_mixed_batch
+            or (is_graph_capturing and uniform_decode)
+            or single_request_prefill
+        )
     ):
         return 0
     return dcp_world_size * cp_kv_cache_interleave_size

vllm/v1/worker/gpu/model_runner.py

@@ -92,6 +92,8 @@
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
     ModelCudaGraphManager,
+    _init_minimal_kv_cache_for_profiling,
+    _teardown_profiling_state,
     normalize_model_token_inputs,
 )
 from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -697,6 +699,8 @@ def _dummy_run(
         context_len: int = 0,
         skip_eplb: bool = False,
         is_profile: bool = False,
+        single_request_prefill: bool = False,
+        profile_all_kv_cache_groups: bool = False,
         **kwargs,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         if skip_attn and not is_profile:
@@ -705,8 +709,9 @@ def _dummy_run(
             )
 
         # Create a dummy scheduler output.
-        num_reqs = min(num_tokens, self.max_num_reqs)
+        num_reqs = 1 if single_request_prefill else min(num_tokens, self.max_num_reqs)
         if uniform_decode:
+            assert not single_request_prefill
             # HACK(lucas): for now since the worker is shared between MRV1 and MRV2,
             # and for spec-decode with MTP we want to make sure the dummy runs use
             # 1+num_speculative_tokens we use max here, this will likely be eventually
@@ -754,6 +759,7 @@ def _dummy_run(
                 skip_attn_for_dummy_run=skip_attn,
                 is_profile=is_profile,
                 context_len=context_len,
+                profile_all_kv_cache_groups=profile_all_kv_cache_groups,
             )
         self.kv_connector.set_disabled(False)
 
@@ -876,6 +882,37 @@ def profile_run(self) -> None:
         self.reset_encoder_cache()
         gc.collect()
 
+    @torch.inference_mode()
+    def profile_glm_dcp_attention(self) -> None:
+        """Measure the GLM DCP query-gather peak before KV cache sizing.
+
+        The general activation profile deliberately skips attention and splits
+        its token budget across many requests. A GLM sparse-MLA prefill can put
+        the complete scheduler budget in one request, causing the DCP query
+        all-gather to require substantially more temporary memory. Bind a
+        minimal split cache, execute that shape, then release all temporary
+        cache and backend state before production cache allocation.
+        """
+        if (
+            self.model_config.architecture != "Glm5NextForConditionalGeneration"
+            or self.dcp_size <= 1
+        ):
+            return
+
+        _init_minimal_kv_cache_for_profiling(self)
+        try:
+            self._dummy_run(
+                self.max_num_tokens,
+                context_len=self.dcp_size * self.cp_interleave,
+                skip_eplb=True,
+                is_profile=True,
+                single_request_prefill=True,
+                profile_all_kv_cache_groups=True,
+            )
+            torch.accelerator.synchronize()
+        finally:
+            _teardown_profiling_state(self)
+
     def post_kv_cache_wake_up(self) -> None:
         self.block_tables.init_block_table_layout_tensors()
 
@@ -1507,6 +1544,7 @@ def execute_model(
         skip_attn_for_dummy_run: bool = False,
         is_profile: bool = False,
         context_len: int = 0,
+        profile_all_kv_cache_groups: bool = False,
     ) -> ModelRunnerOutput | IntermediateTensors | None:
         if not dummy_run:
             # Update the request states.
@@ -1627,7 +1665,7 @@ def execute_model(
             )
             assert block_tables is not None
             attn_groups = self.attn_groups
-            if dummy_run and is_profile:
+            if dummy_run and is_profile and not profile_all_kv_cache_groups:
                 # Mamba layers take a cheap warmup path with no metadata;
                 # attention metadata is still built so those kernels tune.
                 attn_groups = [