跳到主要内容

数据截至 (上游 commit f775db03aaa8)

05 · 多模态流水线与分布式

本章讲什么: 两个「横向扩展」话题:输入从纯文本扩展到图像/视频/音频时数据怎么流;算力从一张卡扩展到集群时进程怎么摆。都是概述级——每个子话题都够单独一章,本章给地图和主线。

5.1 多模态:图像怎么混进 token 流​

它要解决的小问题​

调度器和注意力内核只认 token id 和 KV。一张图片要进来,得回答两个问题:在哪变成向量(CPU 贵活还是 GPU 贵活)、变成的向量怎么嵌入 token 序列的位置。

思路:CPU 预处理 + 占位 token + prefill 时原地换 embedding​

SGLang 的分工:

  1. tokenizer 进程(CPU): 下载/解码图像,跑 HF processor 得到 pixel values;把图像展开成一串占位 token(如 <image> 重复 N 次),和文本 token 拼成一个序列。
  2. scheduler(GPU): 对调度器来说,多模态请求只是「带 mm_inputs 的普通请求」,照常进 radix 树、照常组批。
  3. 模型前向时(GPU): prefill 批内跑视觉编码器(ViT),把 placeholder token 区间的 embedding 原地替换成视觉特征;decode 步不再重复算。
客户端图像 → TokenizerManager:HF processor → pixel values + 占位 token 序列
      │
      ▼ ZMQ
Scheduler:和纯文本请求一样排队、组批(mm_inputs 随行)
      │
      ▼
ModelRunner forward(prefill):ViT 编码图像 → 替换 placeholder 区间 embedding
      │
      ▼
后续 decode:纯 token 流水线,图像已成 KV cache 里的历史

真实实现​

  • 处理器注册表。 每个支持的 VLM 架构在 PROCESSOR_MAPPING(python/sglang/srt/managers/multimodal_processor.py:13)里登记一个处理器类;get_mm_processor_cls(multimodal_processor.py:44)按 hf_config.architectures 查表,get_mm_processor(multimodal_processor.py:74)实例化。所有处理器的基类是 BaseMultimodalProcessor(python/sglang/srt/multimodal/processors/base_processor.py:205),入口 process_mm_data_async(base_processor.py:806)——异步、带预处理结果缓存(MultimodalPreprocessCache,按 tokenizer worker 数均摊内存)。
  • GPU 侧替换。 模型 forward 里走 general_mm_embed_routine(python/sglang/srt/managers/mm_utils.py:608):只在 prefill 且批内含 mm 输入时,先跑 ViT 得视觉 embedding,再替换 placeholder 区间,之后正常进语言模型。
  • chunked prefill 的坑。 长多模态输入被切块后,后续 chunk 还要用到图像特征——GPU 上的 embedding 是 best-effort 缓存,未命中时回退到 CPU 上的副本(mm_utils.py:689 起的注释解释了为什么要 offload 到 CPU 而不是丢弃)。

诚实边界: 各模型自己的 processor(python/sglang/srt/multimodal/processors/ 下几十个文件)差异很大,本文没有逐一读;ViT 编码与语言模型前向的 overlap / 分离部署(disaggregation/encoder/)是较新的方向,细节本文不展开。

5.2 分布式:一张卡放不下之后​

并行维度一览​

维度拆什么SGLang 里的落点
TP(张量并行)单层权重横切到多卡每 rank 一个 scheduler 进程,NCCL 互联
PP(流水并行)层按段切到多卡scheduler 带 pp_rank,微批在 rank 间接力
DP(数据并行)整模型复制多份,请求分流DataParallelController 路由
EP(专家并行)MoE 专家切到多卡scheduler 带 moe_ep_rank,与 DP attention 组合

进程拓扑在启动时一次定型。 _launch_scheduler_processes(python/sglang/srt/entrypoints/engine.py:820)按 pp_rank × tp_rank 双重循环拉子进程,每个拿到自己的 gpu_id, tp_rank, pp_rank, moe_ep_rank … 去跑 run_scheduler_process(python/sglang/srt/managers/scheduler.py:5378);开了 DP 则改拉一个 DataParallelController 进程,由它再管各 DP rank。

DP 路由:cache-aware 的负载均衡​

DataParallelController(python/sglang/srt/managers/data_parallel_controller.py:138)在 DP > 1 时站在 TokenizerManager 和各 DP rank 之间。路由策略(data_parallel_controller.py:165-170 的 dispatch 表):

  • ROUND_ROBIN(data_parallel_controller.py:769)——轮转,最朴素。
  • TOTAL_REQUESTS / TOTAL_TOKENS——按各 rank 上报的负载快照选最闲的。

「按 token 数选最闲」和第 2 章的 radix 树合起来就是 README 里 "cache-aware load balancer" 的落地:把请求发到最可能命中它前缀、且最闲的那个 rank。

PD 分离:prefill 和 decode 分家​

prefill 是 compute-bound、decode 是 memory-bandwidth-bound,两者的最优并行方式和 GPU 配比不同。PD 分离(prefill-decode disaggregation)把它们拆成两组独立实例,中间用 KV 传输握手:

  • DisaggregationMode(python/sglang/srt/disaggregation/utils.py:101)三态:NULL / PREFILL / DECODE,每个实例启动时定死角色。
  • prefill 侧实例算完 prompt 的 KV,经传输后端(Mooncake / NIXL 等,python/sglang/srt/disaggregation/ 下各有目录)发给 decode 侧;SchedulerDisaggregationPrefillMixin(disaggregation/prefill.py:511)管发送队列,decode 侧对称地收(decode.py:2452 的 mixin)。
  • 交接的纽带是 bootstrap_room:一次请求一个房间号,prefill/decode 两侧靠它对上号(prefill.py:337-381 的 disagg_kv_sender 初始化)。

调度器主线因此有两个变体事件循环;本文第 3 章讲的是合并模式(NULL),分离模式的组批逻辑大量复用同一套 PrefillAdder / radix 树,只是在「请求完成」的语义上多了「KV 已发送」这一站。

TP rank 间的「一个逻辑调度器」​

TP > 1 时,N 个 scheduler 进程跑同样的调度循环、做同样的决定(第 1 章):请求只进 rank 0 的 ZMQ,再经分布式组广播给其余 rank;任何需要数据才能做的决定(grammar 编译完成、DP attention 的 overlap 开关)都必须跨 rank 同步出一致结论——第 4 章的 all_gather_object 就是一例。这是这套架构最脆弱也最有纪律性的地方:所有 rank 的调度器是确定性副本,谁敢发散谁 hang。

5.3 本章代码地图​

主题文件路径符号名
多模态处理器注册python/sglang/srt/managers/multimodal_processor.pyPROCESSOR_MAPPING、get_mm_processor
处理器基类python/sglang/srt/multimodal/processors/base_processor.pyBaseMultimodalProcessor.process_mm_data_async
GPU 侧 embedding 替换python/sglang/srt/managers/mm_utils.pygeneral_mm_embed_routine
并行进程拉起python/sglang/srt/entrypoints/engine.pyEngine._launch_scheduler_processes
scheduler 进程参数python/sglang/srt/managers/scheduler.pyrun_scheduler_process
DP 路由python/sglang/srt/managers/data_parallel_controller.pyDataParallelController、round_robin_scheduler
PD 分离角色python/sglang/srt/disaggregation/utils.pyDisaggregationMode
KV 发送(prefill 侧)python/sglang/srt/disaggregation/prefill.pySchedulerDisaggregationPrefillMixin
KV 接收(decode 侧)python/sglang/srt/disaggregation/decode.pySchedulerDisaggregationDecodeMixin