ACL图

ACL图#

为什么需要ACL图？#

在大语言模型推理时，每个令牌需要近千次算子执行。当主机侧（CPU）的算子启动速度慢于设备侧（NPU）的计算速度时，会造成主机瓶颈（Host Bound）。在严重情况下，设备侧会有一半以上的时间处于空闲状态。为了解决这个问题，我们在LLM推理中使用了图模式。

eager mode:

host:   |  launch op1  |  launch op2  |  launch op3  |  launch op4  |  launch op5  |

device:                | run op1 |free| run op2 |free| run op3 |free| run op4 |free| run op5 |

        | <-----                           total time                                 -----> |

graph mode:

host:   |  launch graph  |

device:                  | run op1 | run op2 | run op3 | run op4 | run op5 |

        | <-----                    total time                      -----> |

如何使用ACL图？#

在V1 Engine中，ACL图模式默认启用，只需确保enforce_eager没有被设置为True。更多详细信息请参阅：图模式指南

它是如何工作的？#

简而言之，图模式分两步工作：捕获和重放。当引擎启动时，我们会捕获模型前向传播中的所有算子并将其保存为一个图。当请求到来时，我们只需在设备上重放该图，并等待结果。

但实际上，图模式并非如此简单。

填充与分桶#

由于图只能重放之前捕获的算子，无法进行动态分块和检查图输入，我们需要确保图输入的一致性。但我们知道，模型输入的形状取决于调度器（Scheduler）调度的请求，我们无法保证这种一致性。

显然，我们可以通过捕获最大形状并将所有模型输入填充到该形状来解决这个问题。但这会带来大量冗余计算并导致性能下降。因此，我们可以捕获多个不同形状的图，并将模型输入填充到最接近的图，这将大大减少冗余计算。但是，当max_num_batched_tokens非常大时，需要捕获的图数量也会变得非常大。但我们知道，当输入张量形状很大时，计算时间会很长，此时图模式并非必需。因此，我们需要做的是：

  1. 设置一个阈值；

  2. 当num_scheduled_tokens大于阈值时，使用eager_mode；

  3. 在阈值以下的范围内捕获多个图；

|    graph1    |
|           graph2           |
|                    graph3                    |
|                              graph4                              |    # the threshold

| input1 | pad |    # use graph1
|           input2           |  # don't need pad
|                      input3                      |      pad      |    # use graph4
|                                    input4                                    |    # use eager mode

分段图与完整图#

由于当前LLM中注意力层的复杂性不断增加，我们无法确保所有类型的注意力都可以在图中运行。在MLA（Multi-Head Latent Attention）中，预填充令牌和解码令牌具有不同的计算方法，因此当一个批次中同时包含预填充和解码请求时，图模式难以处理这种情况。

vLLM通过分段图模式解决了这个问题。我们使用Eager模式启动注意力相关的算子，而用图处理其他部分。但这也会带来一些问题：启动算子的开销再次变大（尽管比完整的Eager模式小很多），当CPU性能较差或num_tokens较小时，仍可能导致主机瓶颈。

总而言之，我们需要同时支持分段图和完整图模式。

  1. 当注意力可以在图中运行时，我们倾向于选择完整图模式以获得最佳性能；

  2. 当完整图不可用时，使用分段图作为替代方案；

  3. 当分段图性能不佳且完整图模式被阻塞时，将预填充和解码请求分离，并在仅解码的情况下使用完整图模式。因为当一个批次包含预填充请求时，通常num_tokens会相当大，不会造成主机瓶颈。

目前，由于流资源限制，我们在分段图模式下只能支持少量分桶，这会导致冗余计算，并可能使得性能低于Eager模式。

它是如何实现的？#

vLLM已经在图模式中实现了大部分模块。更多详细信息请参阅：CUDA图

在图模式下，vLLM会调用current_platform.get_static_graph_wrapper_cls来获取当前设备的图模型包装器。因此，我们需要在昇腾平台上实现图模式包装器：ACLGraphWrapper。

vLLM已为所有模型添加了support_torch_compile装饰器。该装饰器会替换模型类的__init__和forward接口。当调用forward时，将执行ACLGraphWrapper内部的代码，并执行如上所述的捕获或重放操作。

使用分段图时，我们只需遵循上述流程。但在完整图模式下，由于注意力的复杂性，有时我们需要在执行前更新注意力算子的参数。因此，我们为完整图模式实现了update_attn_params和update_mla_attn_params函数。在前向传播时，不同算子之间会复用内存，因此我们无法在前向传播之前更新注意力算子的参数。在ACL图中，我们使用torch.npu.graph_task_update_begin和torch.npu.graph_task_update_end来实现这一点，并使用torch.npu.ExternalEvent来确保参数更新和算子执行之间的顺序。

可调试性/可观测性（DFX）#

流资源限制#

目前，由于ACL图的限制（每个图至少需要一个独立的流），我们最多只能捕获1800个图。这个数字受限于流的总数（2048个），我们保留了248个流作为缓冲。此外，有许多因素会影响分桶数量：

  • 分段图会根据注意力层将模型划分为num_hidden_layers + 1个子模块。每个子模块都是一个单独的图，需要消耗流资源。因此，与完整图模式相比，分段图模式下的分桶数量非常紧张。

  • 图所需的流数量与通信域（comm domains）的数量有关。每个通信域都会增加一个图所消耗的流。

  • 当在子模块中显式调用多流（multi-stream）时，会消耗额外的流。

关于ACL图和流还有一些其他规则。目前，我们使用update_aclgraph_sizes函数来计算最大分桶数，并更新graph_batch_sizes以确保流资源充足。

我们将在未来扩展流资源限制。

当前限制#

  1. 目前不支持FULL和FULL_AND_PIECEWISE模式；

  2. 当同时使用ACL图、MTP（多任务并行）且num_speculative_tokens > 1时，由于vLLM v0.11.0不支持此情况，我们需要显式设置cudagraph_capture_sizes。

  3. 目前不支持use_inductor；

目录