vllm-project/vllm-ascend

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

Offen

#14.160 geöffnet am 13.08.2026

 (2 Kommentare) (0 Reaktionen) (0 zugewiesene Personen)C++ (2.048 Forks)github user discovery
help wanted

Repository-Metriken

Stars
 (2.637 Sterne)
PR-Merge-Metriken
 (Durchschn. Merge 4T 5h) (559 gemergte PRs in 30 T)

Beschreibung

背景

MRV2 结构化输出(structured output / grammar)路径中,_apply_grammar_bitmask_kernel 将 grammar bitmask 应用到 logits 上,把被屏蔽位置置为 -inf。Ascend 侧实现位于 vllm_ascend/worker/v2/structured_outputs.py,经 patch/worker/patch_v2/patch_triton.py 替换 vllm-core GPU 版本。

当前实现注释明确指出:上游 GPU 版 BLOCK_SIZE=8192 在 Ascend NPU 上会溢出 UB(Unified Buffer),而缩小 BLOCK_SIZE 又使 grid 不稳定,故保留 BLOCK_SIZE=8192 并以 BLOCK_SIZE_SUB=1024 子分块循环处理。kernel 内已有 # TODO: Optimize the kernel performance with NPU profiling data.

存在以下可优化点:

  • for sub_offset in tl.range(0, BLOCK_SIZE, BLOCK_SIZE_SUB) 子分块循环是 UB 溢出绕过方案,可结合 NPU UB 容量重新设计分块;
  • bitmask 解包(packed >> arange & 1)与 apply(置 -inf)可融合 / 向量化优化;
  • grid 与 BLOCK_SIZE 选择需 profiling 标定。

任务

优化 _apply_grammar_bitmask_kernel 在 Ascend NPU 上的性能(对应代码内已有的 TODO)。

  • 路径:vllm_ascend/worker/v2/structured_outputs.py
  • 方向:结合 NPU profiling 重新标定 BLOCK_SIZE / 子分块、融合 bitmask 解包与 apply、提升向量核利用率

验收标准

1. 精度对比(前后)

  • 优化前后应用 bitmask 后的 logits 在相同输入下完全一致(被屏蔽位置均为 -inf,其余不变);
  • 与 vllm-core GPU 基线对齐;
  • 结构化输出(JSON / regex / grammar)端到端生成结果与优化前一致(greedy)。

2. 性能对比(前后)

  • 单算子 profiling:优化前后 kernel 延迟(多组 num_bitmasks / vocab_size);
  • NPU profiler 时间线(UB 占用、子分块循环开销);
  • 端到端:结构化输出场景下 decode 延迟对比。

3. 交付件

  • PR + 设计说明 + 精度对比表 + 性能数据 + 单测

环境约定

  • vllm-ascend:最新 main
  • 硬件:Ascend NPU(注明型号 + 卡数 + TP 配置)
  • 关联任务池:#9079
  • 验收人:@chengduxiaowu

重点关注

  • 屏蔽语义为位级(bitmask 第 i 位 = 0 -> token i 置 -inf),优化不得改变任一位置的屏蔽状态;
  • UB 溢出是原绕过方案的根因,新方案须在最大词表下不溢出。

任务周期

  • 发布:2026-08-12
  • 回收:2026-10-31

Contributor Guide