Prhub

#22361 [Whisper] Batch encoder forward for concurrent prefill requests

sgl-project/sglang · 作者 JustinTong0323 · 合并时间 2026-04-12 14:15

分析状态 已生成
文件变更 1提交数 2 · 评论 12
代码增减 +21 / -16
performance diffusion run-ci deepseek npu

执行摘要

Whisper 模型编码器从串行改为批量执行,提升高并发预填充吞吐量。

根据PR body中的描述,当预填充批次包含N个新的Whisper请求时,编码器原本在for循环中串行运行N次。每个编码器调用在B200上约7ms,在GB300上约18ms,导致5个请求的批次会阻塞调度器35-90ms,期间无法执行解码批次。性能分析显示编码器在高并发时占调度器时间的35%,串行执行直接限制了吞吐量扩展。

该PR值得精读,尤其是对于关注模型推理性能优化的工程师。关键设计决策包括:1) 识别编码器为瓶颈并量化其开销;2) 利用编码器天然的批次兼容性实现无交互的批量执行;3) 输出重塑策略以适配下游KV缓存。建议结合基准测试数据理解不同硬件平台上的收益差异。

讨论亮点

review中仅有一次实质性讨论:gemini-code-assist[bot]建议将encoder_position_ids张量直接创建在目标设备上以避免不必要的CPU到GPU传输。该建议被采纳,体现在最终代码中使用了device=features_batch.device。yhyang201批准了PR,无其他争议。

实现拆解

主要修改位于python/sglang/srt/models/whisper.py的forward方法。关键改动包括:1) 添加设备放置修复(从#22293 cherry-pick),确保输入特征和位置ID移动到正确设备;2) 重构编码器逻辑,收集所有未缓存请求的特征到features_to_encode列表;3) 使用torch.cat将特征拼接为批次张量features_batch;4) 单次调用self.encoder处理整个批次;5) 将输出从[N, seq_len, dim]重塑为[N*seq_len, dim]以适应下游交叉注意力KV缓存。

文件 模块 状态 重要度
python/sglang/srt/models/whisper.py models modified 9.0

分析完成后,这里会展示 LLM 生成的相对完整源码片段和详细注释。

关键符号

forward

评论区精华

编码器位置 ID 张量的设备创建优化 性能

gemini-code-assist[bot] 建议将 encoder_position_ids 直接创建在目标设备上,避免不必要的 CPU 到 GPU 传输。

结论:建议被采纳,代码中已使用 device=features_batch.device。 · 已解决

风险与影响

风险较低:1) 编码器本身完全支持批次处理(Conv1d、embed_positions、32个transformer层、LayerNorm均沿批次维度操作),无跨请求交互,逻辑正确性有保障;2) 包含#22293的设备放置修复,避免设备不匹配问题;3) 基准测试显示WER(12.77-12.78%)未变化,准确性无影响;4) 修改范围局限在单个文件的forward方法,影响面可控。潜在风险:批次张量拼接可能增加内存峰值,但在音频特征维度固定且批次大小有限的情况下影响不大。

对系统:显著提升Whisper模型在高并发预填充场景下的吞吐量,尤其在GB300等编码器开销较大的平台上效果明显(64并发时+28.9%)。减少调度器阻塞时间,改善整体资源利用率。对用户:提升语音识别服务的并发处理能力,降低延迟。对团队:展示了针对特定模型组件的性能优化模式,可作为类似优化的参考。

低风险变更 已通过基准测试验证

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

执行摘要

该PR优化了SGLang中Whisper模型的前向传播,将编码器从串行执行改为批量处理,显著提升了高并发预填充场景下的吞吐量。在GB300平台上64并发时吞吐量提升达28.9%,而B200平台由于编码器开销较小提升约3%。修改范围集中,风险低,已通过准确性(WER未变)和性能基准测试验证。

功能与动机

问题背景:当预填充批次包含多个新的Whisper请求时,编码器原本在for循环中串行运行,每个调用在B200上约7ms、GB300上约18ms。这导致调度器被长时间阻塞(例如5个请求阻塞35-90ms),期间无法执行解码批次,成为吞吐量扩展的瓶颈。性能分析显示,编码器在高并发时占调度器时间的35%。

优化目标:通过批量执行编码器调用,减少调度器阻塞时间,提升高并发下的吞吐量。

实现拆解

主要修改位于 python/sglang/srt/models/whisper.pyforward 方法:

  1. 设备放置修复:cherry-pick了PR #22293的修复,确保输入特征和位置ID移动到正确设备:
    python device = self.conv1.weight.device input_features = input_features.to(device=device) position_ids = position_ids.to(device=device)

  2. 特征收集逻辑重构

    • 将原本的 encoder_list 循环改为收集未缓存特征到 features_to_encode 列表。
    • 跳过已缓存或无效的请求。
  3. 批量编码器执行
    python if features_to_encode: features_batch = torch.cat(features_to_encode, dim=0) encoder_len = features_batch.shape[-1] // 2 encoder_position_ids = torch.arange(encoder_len, device=features_batch.device) batched_output = self.encoder(features_batch, encoder_position_ids, forward_batch) encoder_hidden_states = batched_output.reshape(-1, batched_output.shape[-1])

  4. 输出适配:将批量输出从 [N, seq_len, dim] 重塑为 [N*seq_len, dim],以适配下游交叉注意力KV缓存。

关键设计决策:利用编码器天然的批次兼容性(Conv1d、位置嵌入、32个transformer层、LayerNorm均支持批次维度),实现无跨请求交互的批量处理,确保逻辑正确性。

评论区精华

review讨论较为简洁,仅有一次实质性建议:

gemini-code-assist[bot]:"For better performance and to avoid an unnecessary CPU-to-GPU transfer, you can create the encoder_position_ids tensor directly on the target device."

该建议被采纳,最终代码中使用了 device=features_batch.device 直接创建张量,避免了额外的设备间传输开销。yhyang201批准了PR,无其他争议。

风险与影响

技术风险

  • 低风险:编码器本身完全支持批次处理,无跨请求交互,逻辑正确性有保障。
  • 设备兼容性:包含#22293的设备放置修复,避免设备不匹配问题。
  • 准确性:基准测试显示WER(12.77-12.78%)在所有配置下未变化。
  • 内存:批次张量拼接可能轻微增加内存峰值,但音频特征维度固定且批次大小有限,影响可控。

影响范围

  • 性能提升:在GB300平台上效果显著(64并发时吞吐量+28.9%),B200平台提升较小(+3%),反映了不同硬件上编码器开销的差异。
  • 调度优化:减少调度器阻塞时间,改善整体资源利用率。
  • 用户价值:提升语音识别服务的并发处理能力和响应速度。

关联脉络

该PR与历史PR #22293 直接相关,后者提供了设备放置修复,被cherry-pick到当前优化中作为基础。从近期历史PR看,SGLang仓库持续进行多平台(AMD、NPU、Intel XPU)性能优化和CI改进,本PR延续了针对特定模型组件(如Whisper编码器)的精细化性能调优趋势。同时,标签中的 diffusiondeepseeknpu 反映了该优化可能跨多个模型和平台受益,体现了代码复用和平台扩展的协同演进。

参与讨论