ACL 图

ACL 图#

为什么需要 ACL 图？#

在 LLM 推理中，每个 token 需要执行近千次算子。当主机（host）启动算子的速度慢于设备（device）时，会导致主机瓶颈（host bound）。在严重情况下，设备超过一半的时间将处于空闲状态。为了解决这个问题，我们在 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）。当引擎启动时，我们捕获模型前向传播中的所有算子并将其保存为一个图。当请求到达时，我们只需在设备上重放该图并等待结果。

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

填充与分桶#

由于图只能重放之前捕获的算子，而不会进行分片（tiling）或检查图输入，因此我们需要确保图输入的一致性。然而，我们知道模型输入的形状取决于调度器（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 中，prefill_tokens 和 decode_tokens 有不同的计算方法，因此当 MLA 中的一个批次同时包含预填充和解码时，图模式难以处理这种情况。

vLLM 通过分段图模式解决了这个问题。我们使用 eager 模式来启动注意力算子，并使用图来处理其他算子。但这也会带来一些问题：启动算子的开销再次变大。虽然比 eager 模式小得多，但当 CPU 性能较差或 num_tokens 较小时，仍会导致主机瓶颈。

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

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

  2. 当完整图无法工作时，使用分段图作为替代；

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

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

如何实现？#

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

在图模式下，vLLM 会调用 current_platform.get_static_graph_wrapper_cls 来获取当前设备的图模型包装器，因此我们需要做的是在 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 domain）的数量有关。每个通信域都会增加一个图消耗的流。

  • 当在子模块中显式调用多流（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；