vllm-project/vllm-ascend
[Contribution] [Perf][MRV2] compute_token_logprobs 算子性能优化
Offen
#14.156 geöffnet am 13.08.2026
help wanted
Repository-Metriken
- Stars
- (2.637 Sterne)
- PR-Merge-Metriken
- (Durchschn. Merge 4T 5h) (559 gemergte PRs in 30 T)
Beschreibung
背景
MRV2（Model Runner V2）采样路径中，compute_token_logprobs 负责对采样 token 计算 log-softmax 概率。Ascend 侧实现位于 vllm_ascend/worker/v2/sample/logprob.py，通过 vllm_ascend/patch/worker/patch_v2/patch_triton.py 替换 vllm-core 的 GPU 版本（vllm/v1/worker/gpu/sample/logprob.py）。
当前实现为 Triton kernel _topk_log_softmax_kernel：每个 request 启动一个 program，以固定 BLOCK_SIZE=12944 分块遍历整个词表，两趟计算（max + sum exp）得到 log-softmax，再 gather 指定 token_ids 的 logprob。配套的 compute_topk_logprobs 中 _ranks_kernel 计算 token rank（遍历词表统计 logits > 选中值的个数）。
存在以下可优化点：
BLOCK_SIZE=12944硬编码，对不同词表大小 / batch shape 非最优；- log-softmax 两趟遍历词表，未使用 online softmax 融合；
_ranks_kernel用 int32 累加(logits > x)的逐块比较，长词表下开销大；multibuffer=False、grid(batch_size,)未充分利用 NPU 向量核。
任务
优化 compute_token_logprobs / compute_topk_logprobs 在 Ascend NPU 上的性能。
- 路径：
vllm_ascend/worker/v2/sample/logprob.py - 方向（由开发者结合 NPU profiling 自定）：online softmax 单趟融合、BLOCK_SIZE 自适应、rank 计算并行化、向量核利用率提升
验收标准
1. 精度对比（前后）
- 优化前后
compute_token_logprobs/compute_topk_logprobs输出在相同输入下一致（logprob 数值误差 ≤ 1e-5，token_ranks 完全一致）； - 与 vllm-core GPU 基线（
vllm/v1/worker/gpu/sample/logprob.py）输出对齐（误差 ≤ 1e-5）； - 开启 logprob 的端到端推理输出 token 序列与优化前一致（greedy + non-greedy）。
2. 性能对比（前后）
- 单算子 profiling：优化前后 kernel 延迟对比（多组 batch_size / num_logprobs / vocab_size）；
- NPU profiler 时间线对比（kernel 数量、向量核利用率）；
- 端到端：开启 logprob 场景下 decode 吞吐 / 延迟对比。
3. 交付件
- PR + 设计说明 + 精度对比表 + 性能数据曲线 + 单测
环境约定
- vllm-ascend：最新 main
- 硬件：Ascend NPU（注明型号 + 卡数 + TP 配置）
- 关联任务池：#9079
- 验收人：@chengduxiaowu
重点关注
- logprob 为采样/排序关键路径，精度必须严格对齐，不得改变采样结果；
- 大词表（如 15w+）场景是主要收益点。
任务周期
- 发布：2026-08-12
- 回收：2026-10-31