背景

在使用昇腾NPU进行大模型推理时，MT FSDP图模式下aclgraph图捕获阶段出现性能劣化现象。具体表现为：在重复捕获多张不同shape的图时，单次图捕获耗时随捕获次数增加持续劣化，严重影响模型加载效率。本文档详细描述了该问题的复现方法、定位过程及解决方案。

问题复现

复现环境配置

为模拟实际场景中重复捕获多张不同shape图的行为，对vLLM图捕获逻辑进行了如下两处代码修改：

1. vllm_code/vllm/vllm/v1/worker/gpu_model_runner.py - 在_capture_cudagraphs的cudagraph_capture_size遍历外层添加循环，模拟重复捕获多张不同shape图的行为。
2. vllm_code/vllm-ascend/vllm_ascend/compilation/acl_graph.py - 在__call__中增加单次图捕获耗时打印，便于观测性能变化。

代码修改说明

capture_cudagraphs增加外层循环

vllm_code/vllm/vllm/v1/worker/gpu_model_runner.py
class GPUModelRunner:

def _capture_cudagraphs(
    self,
    compilation_cases: list[tuple[int, bool]],
    cudagraph_runtime_mode: CUDAGraphMode,
    uniform_decode: bool,
):
    assert (
        cudagraph_runtime_mode != CUDAGraphMode.NONE
        and cudagraph_runtime_mode.valid_runtime_modes()
    ), f"Invalid cudagraph runtime mode: {cudagraph_runtime_mode}"

    # Only rank 0 should print progress bar during capture
    if is_global_first_rank():
        compilation_cases = tqdm(
            compilation_cases,
            disable=not self.load_config.use_tqdm_on_load,
            desc="Capturing CUDA graphs ({}, {})".format(
                "decode" if uniform_decode else "mixed prefill-decode",
                cudagraph_runtime_mode.name,
            ),
        )
    for i in range(100):  # 外层循环，循环100次
        self.model.concrete_aclgraph_entries.clear() # 释放已经捕获到的图
        set_cudagraph_capturing_enabled(True) # 重新使能图捕获
        # We skip EPLB here since we don't want to record dummy metrics
        for num_tokens, activate_lora in compilation_cases: # 内层循环，compilation_cases 入图个数，入图batch_size(num_tokens)
            # We currently only capture ubatched graphs when its a FULL
            # cudagraph, a uniform decode batch, and the number of tokens
            # is above the threshold. Otherwise we just capture a non-ubatched
            # version of the graph
            print(f">>>>>>>>>>>>>>>>>>>>>> i {i} compilation_cases num_tokens {num_tokens} time consuming <<<<<<<<<<<<<<<<<<")
            allow_microbatching = (
                self.parallel_config.use_ubatching
                and cudagraph_runtime_mode == CUDAGraphMode.FULL
                and uniform_decode
                and check_ubatch_thresholds(
                    config=self.vllm_config.parallel_config,
                    num_tokens=num_tokens,
                    uniform_decode=uniform_decode,
                )
            )

            for _ in range(self.compilation_config.cudagraph_num_of_warmups):
                # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.
                # But be careful, warm up with `NONE`is orthogonal to
                # if we want to warm up attention or not. This is
                # different from the case where `FULL` implies capture
                # attention while `PIECEWISE` implies no attention.
                force_attention = cudagraph_runtime_mode == CUDAGraphMode.FULL
                self._dummy_run(
                    num_tokens,
                    cudagraph_runtime_mode=CUDAGraphMode.NONE,
                    force_attention=force_attention,
                    uniform_decode=uniform_decode,
                    allow_microbatching=allow_microbatching,
                    skip_eplb=True,
                    remove_lora=False,
                    activate_lora=activate_lora,
                )
            self._dummy_run(
                num_tokens,
                cudagraph_runtime_mode=cudagraph_runtime_mode,
                uniform_decode=uniform_decode,
                allow_microbatching=allow_microbatching,
                skip_eplb=True,
                remove_lora=False,
                activate_lora=activate_lora,
                is_graph_capturing=True,
            )
    self.maybe_remove_all_loras(self.lora_config)

打印单次图捕获耗时

vllm_code/vllm-ascend/vllm_ascend/compilation/acl_graph.py
class ACLGraphWrapper:
def __call__(self, *args, **kwargs):
    forward_context = get_forward_context()
    batch_descriptor = forward_context.batch_descriptor
    aclgraph_runtime_mode = forward_context.cudagraph_runtime_mode

    if aclgraph_runtime_mode == CUDAGraphMode.NONE or aclgraph_runtime_mode != self.runtime_mode:
        # CUDAGraphMode.NONE could mean the profile run, a warmup run, or
        # running without aclgraphs.
        # We do not trigger capture/replay if the runtime mode is not
        # matches. This enables properly dispatching to the correct
        # CUDAGraphWrapper when nesting multiple instances with different
        # runtime modes.
        return self.runnable(*args, **kwargs)

    if batch_descriptor not in self.concrete_aclgraph_entries:
        # create a new entry for this batch descriptor
        self.concrete_aclgraph_entries[batch_descriptor] = ACLGraphEntry(batch_descriptor=batch_descriptor)

    entry = self.concrete_aclgraph_entries[batch_descriptor]

    if entry.aclgraph is None:
        if self.aclgraph_options.debug_log_enable:
            # Since we capture aclgraph for many different shapes and
            # capturing is fast, we don't need to log it for every
            # shape. E.g. we only log it for the first subgraph in
            # piecewise mode.
            logger.debug("Capturing a aclgraph on (%s,%s)", self.runtime_mode.name, entry.batch_descriptor)
        # validate that aclgraph capturing is legal at this point.
        validate_cudagraph_capturing_enabled()

        input_addresses = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
        entry.input_addresses = input_addresses
        time_start = time.perf_counter() # 统计图捕获耗时 start
        aclgraph = torch.npu.NPUGraph()

        with ExitStack() as stack:
            if self.aclgraph_options.gc_disable:
                # during every model forward for piecewise aclgraph
                # mode, we will capture many pieces of aclgraphs
                # (roughly one per layer). running gc again and again
                # across layers will make the aclgraph capture very slow.
                # therefore, we only run gc for the first graph,
                # and disable gc for the rest of the graphs.
                stack.enter_context(patch("gc.collect", lambda: None))
                stack.enter_context(patch("torch.npu.empty_cache", lambda: None))

            # mind-exploding: carefully manage the reference and memory.
            forward_context.capturing = True
            with torch.npu.graph(aclgraph, pool=self.graph_pool):
                output = self.runnable(*args, **kwargs)
                if self.aclgraph_options.weak_ref_output:
                    # by converting it to weak ref,
                    # the original `output` will immediately be released
                    # to save memory. It is only safe to do this for
                    # the last graph in piecewise aclgraph mode, because
                    # the output of the last graph will not be used by
                    # any other acl graph.
                    output = weak_ref_tensors(output)
        time_end = time.perf_counter() # 统计图捕获耗时 end
        print(f"time consuming : {time_end - time_start}") # 打印耗时
        # here we always use weak ref for the workspaces
        # to save memory
        global _graph_params
        global _draft_graph_params
        weak_ref_workspaces(_graph_params)
        weak_ref_workspaces(_draft_graph_params)

        # here we always use weak ref for the output
        # to save memory
        entry.output = weak_ref_tensors(output)
        entry.aclgraph = aclgraph

        compilation_counter.num_cudagraph_captured += 1

        # important: we need to return the output, rather than
        # the weak ref of the output, so that pytorch can correctly
        # manage the memory during acl graph capture
        return output

    if self.is_debugging_mode:
        # check if the input addresses are the same
        new_input_addresses = [x.data_ptr() for x in args if isinstance(x, torch.Tensor)]
        assert new_input_addresses == entry.input_addresses, (
            f"Input addresses for aclgraphs are different "
            f"during replay. Expected {entry.input_addresses}, "
            f"got {new_input_addresses}"
        )

    logger.info_once("Replaying aclgraph")
    # In async scheduling or multi-threaded (MT) scenarios, it is possible that
    # the CPU's record event (from update_attn_params) for the iteration i completes
    # before the grph replay of iteration i-1.
    # To ensure proper ordering, we must call synchronize here before replaying,
    # so that update_attn_params only executes after the previous graph replay has fully completed.
    # If we do not in main model and in full-graph mode when using merge-eagle-graph,
    # we do not need to synchronize.
    use_eagle = (
        self.vllm_config.speculative_config.method in ("eagle", "eagle3")
        if self.vllm_config.speculative_config
        else False
    )
    if self.runtime_mode != CUDAGraphMode.FULL or not forward_context.is_draft_model or not use_eagle:
        torch.npu.synchronize()
    entry.aclgraph.replay()
    return entry.output

整网测试脚本

模型启动时需增加整图配置"cudagraph_mode": "FULL"：

compilation_config={"cudagraph_capture_sizes":[1,2,4,8,12,16,24,32,40,48,96],"cudagraph_mode": "FULL"},

Python测试脚本

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates how to achieve reproducibility in vLLM.

Main article: https://docs.vllm.ai/en/latest/usage/reproducibility.html
"""

import os
import random
import time
import torch

os.environ['OMP_NUM_THREADS'] = '1'  # 限制OpenMP线程数为1
os.environ['MKL_NUM_THREADS'] = '1'

os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

os.environ["VLLM_ASCEND_ENABLE_NZ"] = "0"

os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:False"

from vllm import LLM, SamplingParams

# os.environ["VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY"] = "0"

SEED = 10

def main():

    # os.environ["VLLM_TORCH_PROFILER_DIR"] = "./vllm_profile_eager_qwen3-0.6b-npu-no-batch-inv"

    llm = LLM(model="/path/to/model",
              tensor_parallel_size=8,
              #   enforce_eager=True,
              dtype="bfloat16",  # not everything is supported
              max_num_seqs=24,
              distributed_executor_backend="mp",
              #   expandable_segments:True,
              #   async_scheduling=True,
              #   scheduler_config={"async_scheduling":True},
              compilation_config={"cudagraph_capture_sizes":[1,2,4,8,12,16,24,32,40,48,96],"cudagraph_mode": "FULL"},
              #   max_num_batched_tokens=100, 
              #   logprobs_mode="processed_logprobs", 
              #   profiler_config = {
              #     "profiler": "torch",
              #     "torch_profiler_dir": "./vllm_profile_eager_qwen3-0.6b-npu_0.14.0",
              #     "torch_profiler_with_stack": True},
              load_format="dummy",
              seed=SEED)

    prompt = "How are you?"

    free, total = torch.npu.mem_get_info()
    print(f"Free memory before sleep: {free / 1024**3:.2f} GiB")
    # record npu memory use baseline in case other process is running
    used_bytes_baseline = total - free
    # llm = LLM("Qwen/Qwen2.5-0.5B-Instruct", enable_sleep_mode=True)
    sampling_params = SamplingParams(temperature=0, max_tokens=10)
    output = llm.generate(prompt, sampling_params)

    print(f"output[0].outputs[0].text {output[0].outputs[0].text}")

    print(random.randint(0, 100))

if __name__ == "__main__":
    main()

问题定位

日志分析

通过循环处理100次后的日志数据，可以观察到单次图捕获耗时呈线性增长趋势：

• 第1次：单次图捕获耗时较短
• 第45次：单次图捕获耗时明显增加
• 第100次：单次图捕获耗时增长显著
  
  

Profiling分析

通过Profiling工具进一步定位，发现耗时增长主要集中在rtMemcpy操作上，且该操作的调用次数随图捕获次数增加而持续增长。

在这里插入图片描述

根因分析

通过Profiling数据深入分析，定位到以下根因：

1. 资源未释放：aclgraph模式下，每个通信算子都有独立的tag，触发aicpu资源刷新流程时会执行h2d memcpy操作。在多图重复捕获场景下，资源释放机制存在适配问题：HCCL难以感知已捕获aclgraph图的资源释放时机，导致前几张图申请的资源未能及时销毁和释放。
2. 资源链表增长：由于资源未及时释放，资源链表持续变长，导致rtMemcpy调用次数不断增加，最终造成图捕获耗时持续劣化。

规避方案

针对上述问题，推荐采用AIV模式进行规避。通过设置环境变量启用AIV模式，可有效避免资源链表增长问题：

export HCCL_OP_EXPANSION_MODE='AIV'

该方案通过将通信算子下沉到AIV执行，避免了aicpu资源刷新流程中的h2d memcpy操作，从而解决了资源链表持续增长导致的性能劣化问题。

详细说明请参考：HCCL_OP_EXPANSION_MODE环境变量说明

解决方案

CANN 在26Q2的发布版本 9.1.0中已经修复了该问题，可以升级到9.1.0版本进行升级使用：

https://www.hiascend.com/cann/download?versionId=796&ids=d802%2Ch0501%2Ch0601%2Ch0702&currentTab=0

总结

本文档针对MT FSDP图模式下aclgraph图捕获性能劣化问题，完整描述了问题复现方法、定位过程及解决方案。可通过临时设置HCCL_OP_EXPANSION_MODE='AIV'环境变量规避性能问题，同时建议升级CANN 9.1.0版本解决该问题