动态分块流水线并行 (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 预填充延迟随序列长度呈二次增长:
其中:
\(a \cdot l^2\):注意力开销(二次)
\(b \cdot l\):线性操作(FFN,投影)
\(c\):固定开销(内核启动)
启动阶段:性能分析#
在引擎初始化期间,系统对实际模型性能进行分析:
采样:从
base_chunk_size到接近 0 均匀采样 64 种不同的块大小执行:对每种块大小执行真实模型前向传播,并精确测量延迟(毫秒)
拟合:使用最小二乘法拟合二次模型
目标设定:基于
base_chunk_size计算每块目标延迟
在 PP 模式下,所有工作节点执行前向传播以保持同步,但仅使用第一个 PP 等级的计时结果进行调度决策。
运行时阶段:动态预测#
给定当前前缀长度 \(L\) 和目标延迟 \(T = f(\text{base\_chunk\_size}) - f(0)\),系统求解下一个块大小 \(x\):
展开为:
使用求根公式求解:
结果经过后处理:
平滑:使用
smooth_factor将预测块大小与base_chunk_size混合对齐:向下取整到
page_size的倍数(最小 64)约束:不超过
max_model_len - history_len和max_num_scheduled_tokens
在线校准#
由于性能分析仅覆盖到 max_num_batched_tokens 的序列(通常比实际工作负载短),系统在运行时持续优化模型。
扩展模型(双变量):
其中 \(C\) 是块大小,\(H\) 是前缀历史长度。
每批次后,记录特征向量 [Σ(C+H)·C, Σ(C+H), N] 和实际执行时间。一旦累积足够的数据点(5-30),使用最小二乘法更新模型参数。
架构#
关键组件#
组件 |
位置 |
职责 |
|---|---|---|
ChunkSizePredictor |
|
二次模型拟合与预测 |
ProfilingChunkManager |
|
管理性能分析工作流和预测器 |
Scheduler |
|
集成 CPP 调度 |
EngineCore |
|
启动性能分析,记录执行时间 |
NPUWorker |
|
执行真实前向传播性能分析 |
NPUModelRunner |
|
|
工作流#
┌─────────────────────────────────────────────────────────────┐
│ 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机制