vllm-project/vllm-ascend

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

Aberta

#14.160 aberto em 13 de ago. de 2026

 (2 comentários) (0 reação) (0 responsável)C++ (2.048 forks)github user discovery
help wanted

Métricas do repositório

Stars
 (2.637 estrelas)
Métricas de merge de PR
 (Mesclagem média 4d 5h) (559 fundiu PRs em 30d)

Description

背景

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

Guia do colaborador