vllm-project/vllm-ascend

[Contribution] [Perf][MRV2] bincount 算子性能优化

開放

#14,158 建立於 2026年8月13日

 (3 則留言) (0 個反應) (0 位負責人)C++ (2,048 個分叉)github user discovery
help wanted

倉庫指標

星標
 (2,637 顆星)
PR 合併指標
 (平均合併 4天 5小時) (30 天內合併 559 個 PR)

描述

背景

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_maskoutput_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

貢獻者指南