ACL 图

ACL 图#

为什么我们需要 ACL 图？#

在大语言模型推理中，每个词元（token）都需要近千次算子执行。当主机（host）启动算子的速度慢于设备（device）计算速度时，会导致主机端瓶颈。在严重情况下，设备有超过一半的时间处于空闲状态。为了解决这个问题，我们在LLM推理中使用了图（graph）技术。

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 图？#

ACL 图在 V1 引擎中默认启用，只需确保 enforce_eager 未设置为 True 即可。更多详情请参阅：图模式使用指南

工作原理#

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

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

填充（Padding）与分桶（Bucketing）#

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

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

设置一个阈值；
当 num_scheduled_tokens 大于该阈值时，使用 eager_mode（即时执行模式）；
在阈值以下的范围内捕获多个图；

|    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-Query Latent Attention）中，预填充词元（prefill_tokens）和解码词元（decode_tokens）的计算方式不同，因此当一个批次中同时包含MLA的预填充和解码任务时，图模式难以处理这种情况。

vLLM 通过分段图（piecewise graph）模式解决了这个问题。我们使用即时执行模式来启动注意力相关的算子，而对于其他算子则使用图模式。但这也会带来一些问题：启动算子的开销再次变大（虽然仍比纯即时执行模式小很多），当CPU性能较差或 num_tokens 较小时，仍可能导致主机端瓶颈。

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

当注意力计算可以在图模式下运行时，我们倾向于选择完整图模式以获得最佳性能；
当完整图模式不适用时，使用分段图作为替代方案；
当分段图性能不佳且完整图模式受阻时，分离预填充和解码任务，并在 仅解码（decode_only） 的场景下使用完整图模式。这是因为当一个批次包含预填充请求时，通常 num_tokens 会很大，不会引起主机端瓶颈。

目前，由于流（stream）资源限制，我们在分段图模式下只能支持少数几个桶（bucket），这会导致冗余计算，并且相比即时执行模式可能导致性能下降。

如何实现？#

vLLM 已经实现了图模式下的大部分模块。您可以在以下链接查看更多细节：CUDA 图

在图模式下，vLLM 会调用 current_platform.get_static_graph_wrapper_cls 来获取当前设备的图模型封装器（wrapper）。因此，我们需要做的是在 Ascend 平台上实现这个图模式封装器：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 以确保流资源充足。

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

限制#

目前不支持 FULL 和 FULL_AND_PIECEWISE 模式；
当同时使用 ACL 图和 MTP，且 num_speculative_tokens > 1 时，由于 vLLM v0.11.0 不支持此情况，我们需要显式设置 cudagraph_capture_sizes。
目前不支持 use_inductor；