ACL 图

ACL Graph#

为什么我们需要 ACL Graph？#

在 LLM 推理时，每个 token 都需要近千次算子执行，当主机启动算子的速度慢于设备时，会导致主机瓶颈。在严重的情况下，设备会空闲超过一半的时间。为了解决这个问题，我们在 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 Graph？#

ACL Graph 在 V1 Engine 中默认启用，只需确保 enforce_eager 未设置为 True。更多详情请参见：图模式指南

它是如何工作的？#

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

但实际上，图模式并没有那么简单。

填充和分桶#

由于图只能重放之前捕获的算子，而无法进行切片和检查图输入，我们需要确保图输入的统一性。但我们知道模型的输入形状取决于调度器（Scheduler）调度的请求，我们无法确保其统一性。

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

  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 会很大，不会导致主机瓶颈。

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

它是如何实现的？#

vLLM 已经在图模式下实现了大部分模块。您可以在以下网址找到更多详细信息：CUDA Graphs

在图模式下，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 Graph 中，我们使用 torch.npu.graph_task_update_begin 和 torch.npu.graph_task_update_end 来实现这一点，并使用 torch.npu.ExternalEvent 来确保参数更新和算子执行之间的顺序。

DFX#

流资源限制#

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

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

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

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

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

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

限制#

  1. FULL 和 FULL_AND_PIECEWISE 目前不支持；

  2. 当使用 ACL Graph 和 MTP 且 num_speculative_tokens > 1 时，由于 vLLM 在 v0.11.0 中不支持此场景，我们需要显式设置 cudagraph_capture_sizes。

  3. use_inductor 目前不支持；

目录