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

背景

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