[BugFix] Revert pre-KV ACL graph profiling (vllm-project#9865) to fix lmhead TP hang (vllm-project#11562)

### What this PR does / why we need it?

This PR reverts vllm-project#9865 / commit
`9099b7f66d123ea704e329d7586333ad8b08db50`, which introduced pre-KV ACL
graph memory profiling in `determine_available_memory()`.

After that change, the combined scenario below can hang during inference
and eventually fail with an HCCL timeout in lmhead TP communication:

- MTP enabled
- `finegrained_tp_config.lmhead_tensor_parallel_size > 0`
- ACL graph enabled
- AIV/HCCL graph communication path enabled

The likely trigger is that the pre-KV graph memory profiling path runs
an additional graph warmup/capture before normal KV cache allocation and
normal graph capture. In the MTP + lmhead TP path, the draft graph can
execute `compute_logits()`, which enters lmhead TP `all_gather` /
`all_to_all`. With AIV enabled, this introduces an extra collective
graph/capture path before the regular runtime path, and can leave the
lmhead TP communication sequence or graph/stream state inconsistent.

### Does this PR introduce _any_ user-facing change?

Yes. This reverts the ACL graph memory estimation added by vllm-project#9865.

KV cache auto-sizing will no longer subtract the estimated ACL graph
pool memory during `determine_available_memory()`. This restores the
previous behavior and may increase the computed KV cache budget compared
with the reverted implementation.

### How was this patch tested?

This PR is a targeted revert.

Validation focus:

- MTP + lmhead TP + ACL graph + AIV inference no longer hangs in lmhead
TP communication.
- Existing initialization and graph capture flow returns to the behavior
before vllm-project#9865.

No new unit test is added because the failure requires multi-node NPU
runtime, HCCL communication, ACL graph capture, MTP, lmhead TP, and AIV
enabled together, which is not covered by local UT.

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@1f486d9

Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

Author: yiz-liu

tests/e2e/pull_request/four_card/test_graph_mode.py

@@ -346,7 +346,8 @@
     "data_parallel_size": 1,
     "enable_expert_parallel": False,
     "golden_answers": {"short": QWEN3_PROMPTS_SHORT_BASELINE, "long": QWEN3_PROMPTS_LONG_BASELINE},
-    "baseline_capture_mem": 0.20,
+    # TODO: it increases after profile graph memory is disabled, invetigate later
+    "baseline_capture_mem": 0.30,
     "capture_mem_tolerance": 1.3,
 }
 

tests/ut/worker/a2/test_worker_multi_instance.py

@@ -18,7 +18,6 @@
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
-from vllm.config import CUDAGraphMode
 from vllm.utils.mem_constants import GiB_bytes
 
 from tests.ut.base import TestBase
@@ -51,10 +50,6 @@ def _make_worker(
         worker.model_runner = MagicMock()
         worker.model_runner.model_memory_usage = model_memory_usage
 
-        mock_vllm_config = MagicMock()
-        mock_vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
-        worker.vllm_config = mock_vllm_config
-
         mock_cache_config = MagicMock()
         mock_cache_config.kv_cache_memory_bytes = None
         mock_cache_config.gpu_memory_utilization = requested_memory / init_total_memory
@@ -140,18 +135,14 @@ def test_single_instance_positive_kv_cache(self, mock_logger):
         self.assertGreater(result, 0)
 
     @patch("vllm_ascend.worker.worker.logger")
-    def test_deepseek_v4_compressed_skips_npugraph_memory_profile(self, mock_logger):
-        """DSV4 DSA must not run the pre-KV graph memory profiling path."""
+    def test_determine_available_memory_does_not_profile_npugraph_memory(self, mock_logger):
         total = int(64 * GiB_bytes)
         requested_memory = int(total * 0.9)
         init_free = int(60 * GiB_bytes)
         non_kv_cache = int(1 * GiB_bytes)
 
         worker = self._make_worker(requested_memory, init_free, total)
-        worker.vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
-        worker.model_config.hf_config.model_type = "deepseek_v4"
-        worker.model_runner.use_compress = True
-        worker.model_runner.profile_cudagraph_memory.return_value = int(2 * GiB_bytes)
+        worker.model_runner.profile_cudagraph_memory = MagicMock()
         profile_result = self._make_profile_result(
             free_memory_after=init_free - non_kv_cache,
             non_kv_cache_memory=non_kv_cache,
@@ -162,42 +153,7 @@ def test_deepseek_v4_compressed_skips_npugraph_memory_profile(self, mock_logger)
 
         worker.model_runner.profile_run.assert_called_once()
         worker.model_runner.profile_cudagraph_memory.assert_not_called()
-        self.assertEqual(
-            worker.vllm_config.compilation_config.cudagraph_mode,
-            CUDAGraphMode.FULL_DECODE_ONLY,
-        )
-        self.assertEqual(worker.npugraph_memory_estimate, 0)
-        self.assertEqual(result, requested_memory - non_kv_cache)
-
-    @patch("vllm_ascend.worker.worker.logger")
-    def test_non_deepseek_compressed_still_profiles_npugraph_memory(self, mock_logger):
-        """The DSV4 guard must not disable graph memory profiling globally."""
-        total = int(64 * GiB_bytes)
-        requested_memory = int(total * 0.9)
-        init_free = int(60 * GiB_bytes)
-        non_kv_cache = int(1 * GiB_bytes)
-        npugraph_memory = int(2 * GiB_bytes)
-
-        worker = self._make_worker(requested_memory, init_free, total)
-        worker.vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
-        worker.model_runner.use_compress = True
-        worker.model_runner.profile_cudagraph_memory.return_value = npugraph_memory
-        profile_result = self._make_profile_result(
-            free_memory_after=init_free - non_kv_cache,
-            non_kv_cache_memory=non_kv_cache,
-        )
-
-        with (
-            self._patch_memory_profiling(profile_result),
-            patch(
-                "vllm_ascend.worker.worker.envs_vllm.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS",
-                False,
-            ),
-        ):
-            result = worker.determine_available_memory()
-
-        worker.model_runner.profile_cudagraph_memory.assert_called_once_with()
-        self.assertEqual(worker.npugraph_memory_estimate, npugraph_memory)
+        self.assertFalse(hasattr(worker, "npugraph_memory_estimate"))
         self.assertEqual(result, requested_memory - non_kv_cache)
 
     @patch("vllm_ascend.worker.worker.logger")

tests/ut/worker/a2/test_worker_v1.py

@@ -2,7 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 import torch
-from vllm.config import CacheConfig, CUDAGraphMode, ModelConfig, ParallelConfig, ProfilerConfig, VllmConfig
+from vllm.config import CacheConfig, ModelConfig, ParallelConfig, ProfilerConfig, VllmConfig
 
 from tests.ut.base import TestBase
 
@@ -582,8 +582,6 @@ def test_determine_available_memory_normal_case(
             worker.requested_memory = 10000 * 0.8
             worker.model_runner = MagicMock()
             worker.model_runner.model_memory_usage = 500
-            worker.vllm_config = MagicMock()
-            worker.vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
             worker.cache_config = MagicMock()
             worker.cache_config.gpu_memory_utilization = 0.8
             worker.cache_config.kv_cache_memory_bytes = None
@@ -645,8 +643,6 @@ def test_determine_available_memory_with_non_torch_allocations(
             worker.requested_memory = 10000 * 0.9
             worker.model_runner = MagicMock()
             worker.model_runner.model_memory_usage = 500
-            worker.vllm_config = MagicMock()
-            worker.vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
             worker.cache_config = MagicMock()
             worker.cache_config.gpu_memory_utilization = 0.9
             worker.cache_config.kv_cache_memory_bytes = None
@@ -665,14 +661,8 @@ def test_determine_available_memory_with_non_torch_allocations(
     @patch("torch.npu.mem_get_info")
     @patch("torch.npu.reset_peak_memory_stats")
     @patch("torch.npu.empty_cache")
-    @patch("torch_npu.npu.memory_stats")
     def test_determine_available_memory_memory_profiling_error(
-        self,
-        mock_torch_memory_stats,
-        mock_torch_empty_cache,
-        mock_torch_reset_peak_memory_stats,
-        mock_torch_mem_get_info,
-        mock_memory_profiling,
+        self, mock_torch_empty_cache, mock_torch_reset_peak_memory_stats, mock_torch_mem_get_info, mock_memory_profiling
     ):
         """Test determine_available_memory throws exception on memory profiling error"""
         from vllm_ascend.worker.worker import NPUWorker
@@ -701,16 +691,11 @@ def test_determine_available_memory_memory_profiling_error(
             worker.init_snapshot = mock_init_snapshot
             worker.requested_memory = 10000 * 0.8
             worker.model_runner = MagicMock()
-            worker.model_runner.model_memory_usage = 0
-            worker.vllm_config = MagicMock()
-            worker.vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
             worker.cache_config = MagicMock()
             worker.cache_config.gpu_memory_utilization = 0.8
             worker.cache_config.kv_cache_memory_bytes = None
             worker.device = torch.device("npu:0")
 
-            mock_torch_memory_stats.return_value = {"allocated_bytes.all.peak": 0}
-
             # Test should throw assertion error
             with self.assertRaises(AssertionError) as cm:
                 worker.determine_available_memory()
@@ -760,8 +745,6 @@ def test_determine_available_memory_negative_result(
             worker.requested_memory = 10000 * 0.8
             worker.model_runner = MagicMock()
             worker.model_runner.model_memory_usage = 500
-            worker.vllm_config = MagicMock()
-            worker.vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
             worker.cache_config = MagicMock()
             worker.cache_config.gpu_memory_utilization = 0.8
             worker.cache_config.kv_cache_memory_bytes = None

vllm_ascend/_310p/model_runner_310p.py

@@ -710,7 +710,6 @@ def _check_and_update_cudagraph_mode(
         self,
         attention_backends,
         kv_cache_groups,
-        is_profiling=False,
     ) -> None:
         # 910B does not need this branch because runner/dispatcher query_len are
         # naturally consistent there. 310P ngram needs temporary alignment.

vllm_ascend/compilation/acl_graph.py

@@ -6,7 +6,7 @@
 from collections.abc import Callable
 from contextlib import ExitStack
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any
 from unittest.mock import patch
 
 import torch
@@ -82,13 +82,6 @@ class ACLGraphWrapper:
     guaranteed when VLLM_LOGGING_LEVEL == "DEBUG".
     """
 
-    _all_instances: ClassVar[weakref.WeakSet["ACLGraphWrapper"]] = weakref.WeakSet()
-
-    @classmethod
-    def clear_all_graphs(cls) -> None:
-        for instance in list(cls._all_instances):
-            instance.clear_graphs()
-
     def __init__(
         self,
         runnable: Callable,
@@ -123,8 +116,6 @@ def __init__(
         self.use_eagle = use_eagle
         _acl_graph_wrappers.add(self)
 
-        ACLGraphWrapper._all_instances.add(self)
-
     def __getattr__(self, key: str):
         # allow accessing the attributes of the runnable.
         if hasattr(self.runnable, key):
@@ -139,13 +130,6 @@ def unwrap(self) -> Callable:
         # in case we need to access the original runnable.
         return self.runnable
 
-    @property
-    def cudagraph_wrapper(self) -> "ACLGraphWrapper":
-        return self
-
-    def clear_graphs(self) -> None:
-        self.concrete_aclgraph_entries.clear()
-
     def __call__(self, *args, **kwargs):
         forward_context = get_forward_context()
         batch_descriptor = forward_context.batch_descriptor
@@ -341,13 +325,6 @@ class GraphParams:
 _graph_params: GraphParams | None = None
 
 
-def reset_graph_params():
-    global _graph_params, _draft_graph_params, _draft_graph_prefill_params
-    _graph_params = None
-    _draft_graph_params = None
-    _draft_graph_prefill_params = None
-
-
 def set_graph_params(aclgraph_capture_sizes: list[int]):
     global _graph_params
     if _graph_params is not None:

vllm_ascend/worker/model_runner_v1.py

@@ -17,7 +17,6 @@
 # Adapted from vllm-project/vllm/vllm/worker/gpu_model_runner.py
 #
 
-import gc
 import logging
 import math
 import sys
@@ -114,7 +113,6 @@
 # yapf: disable
 from vllm_ascend.compilation.acl_graph import (
     ACLGraphWrapper,
-    reset_graph_params,
     set_draft_graph_params,
     set_graph_params,
     update_full_graph_params,
@@ -3945,11 +3943,7 @@ def _finalize_dump_data(self, **kwargs) -> None:
 
         self.debugger.step(**kwargs)
 
-    def initialize_kv_cache(
-        self,
-        kv_cache_config: KVCacheConfig,
-        is_profiling: bool = False,
-    ) -> None:
+    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         """
         Initialize KV cache based on `kv_cache_config`.
         Args:
@@ -3963,7 +3957,7 @@ def initialize_kv_cache(
         self.may_add_encoder_only_layers_to_kv_cache_config()
         self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
         # NOTE(cmq): initialize_attn_backend must before using self.attn_groups
-        self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
+        self.initialize_attn_backend(kv_cache_config)
         self.use_hybrid_blocks = len(self.attn_groups) > 1
         # NOTE: Currently, we determine whether we need `num_accepted_tokens` through `MambaSpec`.
         self.need_accepted_tokens = any(
@@ -3989,7 +3983,7 @@ def initialize_kv_cache(
                 self.kernel_block_sizes, list) else self.kernel_block_sizes)
             self.drafter.initialize_attn_backend(kv_cache_config, block_size)
 
-        if has_kv_transfer_group() and not is_profiling:
+        if has_kv_transfer_group():
             get_kv_transfer_group().register_kv_caches(kv_caches)
 
         if self.model_config.enable_return_routed_experts:
@@ -4846,11 +4840,7 @@ def may_reinitialize_input_batch(self, kv_cache_config: KVCacheConfig) -> None:
                 cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
             )
 
-    def initialize_attn_backend(
-        self,
-        kv_cache_config: KVCacheConfig,
-        is_profiling: bool = False,
-    ) -> None:
+    def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
         """
         Initialize the attention backends and attention metadata builders.
         """
@@ -4912,11 +4902,7 @@ def create_attn_groups(
             attention_backend_maps.append(attn_backends[0])
             attention_backend_list.append(attn_backends[1])
 
-        self._check_and_update_cudagraph_mode(
-            attention_backend_list,
-            kv_cache_config.kv_cache_groups,
-            is_profiling=is_profiling,
-        )
+        self._check_and_update_cudagraph_mode(attention_backend_list, kv_cache_config.kv_cache_groups)
 
         for i, attn_backend_map in enumerate(attention_backend_maps):
             self.attn_groups.append(create_attn_groups(attn_backend_map, i))
@@ -5054,7 +5040,6 @@ def _check_and_update_cudagraph_mode(
         self,
         attention_backends: list[set[type[AttentionBackend]]],
         kv_cache_groups: list[KVCacheGroupSpec],
-        is_profiling: bool = False,
     ) -> None:
         min_cg_support = AttentionCGSupport.ALWAYS
         min_cg_attn_backend = None
@@ -5079,7 +5064,6 @@ def _check_and_update_cudagraph_mode(
                 self.parallel_config.tensor_parallel_size,
                 self.kv_cache_config,
                 self.max_num_reqs,
-                is_profiling=is_profiling,
             )
             self.cudagraph_dispatcher.initialize_cudagraph_keys(
                 cudagraph_mode, self.uniform_decode_query_len
@@ -5108,35 +5092,11 @@ def _check_and_update_cudagraph_mode(
 
         # NOTE: Since aclgraph_batch_sizes cannot be determined until here,
         # we set the graph params right before initializing the keys.
-        # Profiling still runs real graph warmup/capture paths, so the NPU-side
-        # graph params must exist there as well.
         if self.use_aclgraph:
             set_graph_params(capture_sizes)
             if self.speculative_config:
                 set_draft_graph_params(capture_sizes)
 
-    def profile_cudagraph_memory(self) -> int:
-        parent_module_name = _get_gpu_model_runner_module_name(self)
-        with _torch_cuda_wrapper(), _replace_gpu_model_runner_function_wrapper(parent_module_name):
-            result = GPUModelRunner.profile_cudagraph_memory(self)
-
-        reset_graph_params()
-
-        # NOTE: This is a serious problem that we maintain two extra copies of the KV cache as the instance
-        # variable of the attention layers, when they are local variables in the upstream vLLM code.
-        # We have to manually clear them here to release memory after profiling. 
-        for layer in self.compilation_config.static_forward_context.values():
-            if hasattr(layer, "impl"):
-                if hasattr(layer.impl, "key_cache"):
-                    layer.impl.key_cache = None
-                if hasattr(layer.impl, "value_cache"):
-                    layer.impl.value_cache = None
-
-        gc.collect()
-        torch.accelerator.empty_cache()
-
-        return result
-
     def capture_model(self) -> int:
         """Capture NPU graphs and return actual graph pool memory bytes consumed."""
         parent_module_name = _get_gpu_model_runner_module_name(self)
@@ -5267,23 +5227,16 @@ def _replace_gpu_model_runner_function_wrapper(target_module_name):
     _encoder_mgr_orig = _vllm_encoder_cudagraph.EncoderCudaGraphManager
     _vllm_encoder_cudagraph.EncoderCudaGraphManager = EncoderAclGraphManager
     target_module = None
-    original_attrs = {}
     try:
         target_module = sys.modules[target_module_name]
-        if hasattr(target_module, "graph_capture"):
-            original_attrs["graph_capture"] = target_module.graph_capture
         setattr(target_module, "graph_capture", graph_capture)  # noqa: B010
-        if hasattr(target_module, "CUDAGraphWrapper"):
-            original_attrs["CUDAGraphWrapper"] = target_module.CUDAGraphWrapper
-            setattr(target_module, "CUDAGraphWrapper", ACLGraphWrapper)  # noqa: B010
         yield
     except Exception as e:
         raise RuntimeError(f"NPUModelRunner failed, error is {e}")
     finally:
         _vllm_encoder_cudagraph.EncoderCudaGraphManager = _encoder_mgr_orig
         if target_module is not None:
-            for attr_name, attr_value in original_attrs.items():
-                setattr(target_module, attr_name, attr_value)  # noqa: B010
+            setattr(target_module, "graph_capture", graph_capture)  # noqa: B010
 
 
 # TODO: remove it when flash_comm1 is removed