vllm-project/vllm-ascend
[Contribution] [Perf][MRV2] apply_penalties 算子性能优化
Aperta
#14.157 aperta il 13 ago 2026
help wanted
Metriche repository
- Star
- (2637 stelle)
- Metriche merge PR
- (Merge medio 4g 5h) (559 PR mergiate in 30 g)
Descrizione
背景
MRV2 采样路径中，apply_penalties 对 logits 施加 repetition / frequency / presence 惩罚。Ascend 侧实现位于 vllm_ascend/worker/v2/sample/penalties.py，经 patch/worker/patch_v2/patch_triton.py 替换 vllm-core GPU 版本。
当前实现为 Triton kernel _penalties_kernel，grid (num_tokens, num_blocks)，BLOCK_SIZE=4096 分块遍历词表。kernel 内对每个 token 有一段顺序循环 for prev_pos in tl.range(pos) 累加前序 draft token 的 bin count（O(pos) 每程序），再施加三类惩罚。受限于 "NPU 不支持链式 or" 做了 use_penalty = use_penalty or ... 拆分。
存在以下可优化点：
for prev_pos in tl.range(pos)顺序累加 draft token 计数，长序列 / 多 draft token 下串行开销大；BLOCK_SIZE=4096固定；- 早期返回（无惩罚）逻辑可进一步减少无效 launch。
任务
优化 apply_penalties 在 Ascend NPU 上的性能。
- 路径：
vllm_ascend/worker/v2/sample/penalties.py（_penalties_kernel/apply_penalties） - 方向：消除 / 并行化 draft token 计数顺序循环、BLOCK_SIZE 自适应、减少无效 launch
验收标准
1. 精度对比（前后）
- 优化前后施加惩罚后的 logits 在相同输入下一致（数值误差 ≤ 1e-5）；
- 与 vllm-core GPU 基线对齐；
- 开启各惩罚组合（rep / freq / pres 及混合）的端到端输出与优化前一致（greedy + non-greedy）。
2. 性能对比（前后）
- 单算子 profiling：优化前后 kernel 延迟（多组 num_tokens / vocab_size / pos 长度）；
- NPU profiler 时间线对比；
- 端到端：开启惩罚场景下 decode 延迟对比。
3. 交付件
- PR + 设计说明 + 精度对比表 + 性能数据 + 单测
环境约定
- vllm-ascend：最新 main
- 硬件：Ascend NPU（注明型号 + 卡数 + TP 配置）
- 关联任务池：#9079
- 验收人：@chengduxiaowu
重点关注
- rep/freq/pres 各组合须与 GPU 基线逐位对齐；
- 顺序循环消除后须保证 draft token bin count 结果一致。
任务周期
- 发布：2026-08-12
- 回收：2026-10-31