vllm-project/vllm-ascend
[Contribution] [Perf][MRV2] bincount 算子性能优化
Ouverte
#14 158 ouverte le 13 août 2026
help wanted
Métriques du dépôt
- Stars
- (2 637 étoiles)
- Métriques de merge PR
- (Merge moyen 4j 5h) (559 PRs mergées en 30 j)
Description
背景
MRV2 采样路径中，bincount 构建 prompt 的二进制 bitmask（prompt_bin_mask，标记 prompt 中出现过的 token）与输出 token 计数（output_bin_counts），供 apply_penalties 使用。Ascend 侧实现位于 vllm_ascend/worker/v2/sample/penalties.py。
当前实现为 Triton kernel _bincount_kernel，grid (num_tokens, num_blocks)，BLOCK_SIZE=1024。kernel 通过 tl.atomic_or 写 bitmask（token_id // 32 为字索引，token_id % 32 为位），通过 tl.atomic_add 累加 output_bin_counts。两个原子操作在高并发下存在竞争。
存在以下可优化点：
atomic_or（bitmask）与atomic_add（counts）原子操作竞争；BLOCK_SIZE=1024固定；- 每个 token 独立遍历 prompt，重复扫描可合并。
任务
优化 bincount 在 Ascend NPU 上的性能。
- 路径：
vllm_ascend/worker/v2/sample/penalties.py（_bincount_kernel/bincount） - 方向：降低原子竞争（分块规约 / 私有副本合并）、BLOCK_SIZE 自适应、减少重复扫描
验收标准
1. 精度对比（前后）
- 优化前后
prompt_bin_mask与output_bin_counts在相同输入下完全一致（bitmask 位级一致、counts 整数一致）； - 与 vllm-core GPU 基线对齐；
- 端到端（bincount + apply_penalties 联动）输出与优化前一致。
2. 性能对比（前后）
- 单算子 profiling：优化前后 kernel 延迟（多组 num_tokens / prompt_len / prefill_len）；
- NPU profiler 时间线（原子操作开销）；
- 端到端：开启惩罚场景下 decode 延迟对比。
3. 交付件
- PR + 设计说明 + 精度对比表 + 性能数据 + 单测
环境约定
- vllm-ascend：最新 main
- 硬件：Ascend NPU（注明型号 + 卡数 + TP 配置）
- 关联任务池：#9079
- 验收人：@chengduxiaowu
重点关注
- bitmask 为位级语义，优化不得改变任一位；
- 原子规约合并后须保证 counts 正确。
任务周期
- 发布：2026-08-12
- 回收：2026-10-31