动态分块流水线并行 (CPP)#

摘要:CPP 使用基于性能分析的动态分块技术,均衡每块延迟,消除 PP 场景中的流水线气泡。

背景#

问题陈述#

在流水线并行 (PP) + 分块预填充场景中,长序列被分割成固定大小的块,依次通过流水线。由于自注意力的 O(n²) 计算复杂度,随着前缀序列增长,相同大小的块处理时间会越来越长

Chunk 1 (history=0):     ██████         → Time T1
Chunk 2 (history=4K):    ████████       → Time T2 > T1
Chunk 3 (history=8K):    ██████████     → Time T3 > T2
Chunk 4 (history=12K):   ████████████   → Time T4 > T3

这种时间差异会在流水线阶段间传播,导致空闲等待(流水线气泡)增加,并显著降低 GPU 利用率。

解决方案概述#

动态分块流水线并行采用先性能分析,再预测的策略:

Fixed Chunking (equal chunk size, unequal time):

          Stage 0  |■■■■|■■■■■■|■■■■■■■■|■■■■■■■■■■|
          Stage 1  |    |■■■■  |■■■■■■  |■■■■■■■■  |■■■■■■■■■■|
                        ↑ bubble  ↑ bubble   ↑ bubble

Dynamic Chunking (unequal chunk size, equal time):

          Stage 0  |■■■■■■|■■■■■■|■■■■■■|■■■■■■|
          Stage 1  |      |■■■■■■|■■■■■■|■■■■■■|■■■■■■|
                          ↑ no bubble — stages stay in sync

核心思想借鉴自 SGLang 的动态分块机制,并增加了在线校准等增强功能。

设计#

二次延迟模型#

由于 O(n²) 的自注意力机制,Transformer 预填充延迟随序列长度呈二次增长:

\[f(l) = a \cdot l^2 + b \cdot l + c\]

其中:

  • \(a \cdot l^2\):注意力开销(二次)

  • \(b \cdot l\):线性操作(FFN,投影)

  • \(c\):固定开销(内核启动)

启动阶段:性能分析#

在引擎初始化期间,系统对实际模型性能进行分析:

  1. 采样:从 base_chunk_size 到接近 0 均匀采样 64 种不同的块大小

  2. 执行:对每种块大小执行真实模型前向传播,并精确测量延迟(毫秒)

  3. 拟合:使用最小二乘法拟合二次模型

  4. 目标设定:基于 base_chunk_size 计算每块目标延迟

在 PP 模式下,所有工作节点执行前向传播以保持同步,但仅使用第一个 PP 等级的计时结果进行调度决策。

运行时阶段:动态预测#

给定当前前缀长度 \(L\) 和目标延迟 \(T = f(\text{base\_chunk\_size}) - f(0)\),系统求解下一个块大小 \(x\)

\[f(L + x) - f(L) = T\]

展开为:

\[a \cdot x^2 + (2aL + b) \cdot x - T = 0\]

使用求根公式求解:

\[x = \frac{-(2aL + b) + \sqrt{(2aL + b)^2 + 4aT}}{2a}\]

结果经过后处理:

  1. 平滑:使用 smooth_factor 将预测块大小与 base_chunk_size 混合

  2. 对齐:向下取整到 page_size 的倍数(最小 64)

  3. 约束:不超过 max_model_len - history_lenmax_num_scheduled_tokens

在线校准#

由于性能分析仅覆盖到 max_num_batched_tokens 的序列(通常比实际工作负载短),系统在运行时持续优化模型。

扩展模型(双变量):

\[f(C, H) = a \cdot C(C+H) + b \cdot (C+H) + c\]

其中 \(C\) 是块大小,\(H\) 是前缀历史长度。

每批次后,记录特征向量 [Σ(C+H)·C, Σ(C+H), N] 和实际执行时间。一旦累积足够的数据点(5-30),使用最小二乘法更新模型参数。

架构#

关键组件#

组件

位置

职责

ChunkSizePredictor

vllm_ascend/core/profiling_chunk_predictor.py

二次模型拟合与预测

ProfilingChunkManager

vllm_ascend/core/profiling_chunk_predictor.py

管理性能分析工作流和预测器

Scheduler

vllm_ascend/core/scheduler_profiling_chunk.py

集成 CPP 调度

EngineCore

vllm_ascend/patch/platform/patch_profiling_chunk.py

启动性能分析,记录执行时间

NPUWorker

vllm_ascend/worker/worker.py

执行真实前向传播性能分析

NPUModelRunner

vllm_ascend/worker/model_runner_v1.py

profile_cpp=True 模式

工作流#

┌─────────────────────────────────────────────────────────────┐
│                    Startup Phase                            │
├─────────────────────────────────────────────────────────────┤
│  1. EngineCore.init() triggers profiling                    │
│  2. ProfilingChunkManager samples 64 chunk sizes            │
│  3. NPUWorker executes forward passes                       │
│  4. ChunkSizePredictor fits quadratic model                 │
│  5. Target latency = f(base_chunk_size) - f(0)              │
└─────────────────────────────────────────────────────────────┘
                           ↓
┌─────────────────────────────────────────────────────────────┐
│                    Runtime Phase                            │
├─────────────────────────────────────────────────────────────┤
│  For each prefill chunk:                                    │
│    1. Scheduler queries ChunkSizePredictor                  │
│    2. Given history length L, solve for optimal chunk size  │
│    3. Apply smoothing and alignment                         │
│    4. Execute chunk                                         │
│    5. Record actual timing for online calibration           │
│    6. Update model if enough samples collected              │
└─────────────────────────────────────────────────────────────┘

与 SGLang 对比#

特性

SGLang 动态分块

动态分块流水线并行

性能分析方法

预设二次函数

启动时真实前向传播性能分析

模型拟合

\(f(l) = a \cdot l^2 + b \cdot l + c\)

相同 + 在线校准 \(f(C,H)\)

在线更新

基于历史的拟合

准确性

在不同硬件上可能产生偏差

适配实际硬件性能

启动开销

约64次前向传播(数十秒)

约束条件#

  • 需要流水线并行:必须设置--pipeline-parallel-size > 1

  • 需要分块预填充:必须启用--enable-chunked-prefill

  • 与均衡调度不兼容:不能启用VLLM_ASCEND_BALANCE_SCHEDULING

  • 启动开销:性能分析阶段会增加数十秒的初始化时间

  • 内存:无额外运行时内存开销;性能分析复用现有的dummy_run机制

参考资料#