/ AI Infra
[Infra-14] vLLM-Omni: From Token Scheduling to a Heterogeneous Generative Runtime
A systems-oriented guide to how vLLM-Omni unifies autoregressive, diffusion, encoder, and audio-decoder workloads through stages, orchestration, transport, hierarchical scheduling, and disaggregated deployment.
过去讨论 vLLM，核心问题通常是：怎样更高效地生成下一个 token？
围绕这个问题，vLLM 建立了成熟的推理系统：Paged KV Cache、continuous batching、prefix caching、chunked prefill、speculative decoding，以及 TP、PP、DP、EP 等并行策略。它们虽然分布在不同模块中，却共享同一个基本假设：模型执行最终可以抽象为 prefill 与 decode 构成的自回归循环。
Prompt
|
v
Prefill
|
v
Decode token 0 -> token 1 -> token 2 -> ...但 Omni 模型正在打破这个假设。一个请求可能先经过多模态 Encoder 和自回归 Thinker，再由 Talker 生成声学 token，最后交给非自回归的音频 Decoder；一个文生图请求则可能依次经过 Text Encoder、DiT 的多轮去噪和 VAE Decoder。这些阶段不仅模型不同，计算范式、状态、batch 维度和并行策略也可能完全不同。
text / image / audio
|
v
Thinker -- hidden states --> Talker
|
acoustic codes
|
v
Code2Wav
|
v
waveform因此，vLLM-Omni 不是简单地给 vLLM 增加图片或音频输入，而是在回答一个更大的系统问题：
当一个生成模型内部同时包含 AR、Diffusion、Encoder 和 Decoder 等异构计算阶段时，runtime 应该如何描述、调度、通信、部署和服务化整张推理图？
理解这个变化，可以先记住一句话：
vLLM 是以 token 为中心的高性能 LLM runtime；vLLM-Omni 在复用 vLLM 的同时，把系统抽象提升到了 stage-centric 的异构生成 pipeline。
1. 从 Token-Centric 到 Stage-Centric#
传统 vLLM 的主干可以简化为：
Request
|
v
AsyncLLMEngine
|
v
Scheduler
|
v
Executor / Worker
|
v
ModelRunner
|
v
ModelScheduler 看到的是 request 的 prefill 和 decode 状态，资源管理围绕 token budget、KV Cache block 和 sequence 展开。即使模型能够接收图片，其视觉特征通常也会在进入语言模型后参与同一条 AR 执行路径；这仍然属于多模态理解，不等于多模态生成 runtime。
Omni workload 则可能长这样：
Text-to-image:
Prompt -> Text Encoder -> DiT step T -> ... -> DiT step 0 -> VAE -> Image
Omni speech generation:
Multimodal input -> Thinker -> Talker -> Code2Wav -> Audio如果把这些逻辑全部塞进一个巨大的 Model.forward()，系统会立刻失去几个重要能力：
- 不同子模型无法独立设置并行度、量化方式和设备；
- AR 与 Diffusion 被迫共享并不合适的 scheduler 抽象；
- 中间结果只能在一次单体调用内部传递，难以跨进程或跨机器部署；
- 上下游难以形成流式流水，慢阶段也无法独立扩容；
- 故障、取消、背压和生命周期只能在整条 pipeline 上粗粒度处理。
vLLM-Omni 为此引入 Stage：把一条 Omni 推理链拆成若干具有独立执行语义的阶段。
Request
|
v
AsyncOmniEngine
|
v
Orchestrator
|
+---------------+---------------+
| | |
v v v
Stage 0 Stage 1 Stage 2
Thinker Talker Code2Wav
| | |
AR runtime AR runtime non-AR runtime这里的 Stage 不是一个普通 Python module，也不只是神经网络中的若干 layer。更准确地说，它是一个 执行隔离、局部调度和部署的边界。
同一条 pipeline 中，各 Stage 可以分别拥有：
- 独立的 engine、scheduler、worker 和 batch；
- 不同数量的 replica，以及不同的 TP、DP 或其他并行配置；
- 不同设备、进程乃至机器 placement；
- 不同精度和量化策略；
- AR、Diffusion 或其他非自回归执行范式。
例如：
Stage 0: Thinker
AR, TP=4, KV Cache, continuous batching
|
| hidden states
v
Stage 1: Talker
AR, TP=2, KV Cache
|
| acoustic tokens
v
Stage 2: Code2Wav
non-AR, TP=1一旦 Stage 成为一等公民，系统优化的对象便不再只是单个请求的 token 序列，而是由多个异构服务节点组成的执行图。
2. 四个核心组件如何分工#
Stage 只是逻辑边界。要让它真正运行，还需要解决 pipeline 生命周期、跨阶段路由、每个阶段的部署，以及中间 tensor 的传输。vLLM-Omni 用四个核心组件划分这些职责：
AsyncOmniEngine -> 整个 pipeline 的入口与生命周期
Orchestrator -> request 的跨 Stage 状态与路由
StageRuntime -> 单个 Stage 的实例、placement 与生命周期
OmniConnector -> Stage 之间的数据传输这四个名字容易混在一起，最简单的区分方式是：
AsyncOmniEngine 管整个系统如何启动和对外服务
Orchestrator 管一个请求接下来去哪里
StageRuntime 管某个 Stage 在哪里、以什么实例活着
OmniConnector 管中间数据怎样从这里到那里2.1 AsyncOmniEngine：Pipeline 的 Composition Root#
AsyncOmniEngine 可以类比 vLLM 中的异步引擎入口，但它面对的不是单个 AR model，而是完整的 multi-stage pipeline。
它负责把几个原本分散的问题组合起来：
Load pipeline and deployment configuration
|
v
Initialize Stage runtimes and orchestration
|
v
Start background event loops
|
v
Accept request and submit it to the pipeline
|
v
Collect final or streaming output
|
v
Return to API layer因此，它更像整个 Omni graph 的 composition root：向上承接 Python API 或 OpenAI-compatible server，向下建立 Orchestrator、StageRuntime 和各类 engine 的关系。
2.2 Orchestrator：跨 Stage 的控制面#
如果一个请求要按 Thinker -> Talker -> Code2Wav 执行，系统必须持续知道：
request_id = 123
current_stage = Talker
upstream output = ready
downstream route = Code2Wav replica 0
stream chunk = 7
cancellation = false这些不是任一 Stage 内部 scheduler 能单独回答的问题。Orchestrator 因而负责维护 cross-stage request state，并根据 pipeline 拓扑将输出路由到下一阶段。
可以把它理解成一种层次化调度关系：
Orchestrator
inter-stage scheduling
|
+--------------+--------------+
| | |
v v v
Scheduler Scheduler Scheduler
Thinker Talker DiT
| | |
tokens tokens denoise jobs这里最重要的职责边界是：
Orchestrator 决定请求在 Stage 之间如何推进；Stage 内部的 scheduler 决定本阶段这一轮执行哪些工作。
Orchestrator 也不应承担设备 placement。路由和部署是相关但不同的问题：前者处理 request state machine，后者处理进程与资源生命周期。把两者分开，才能独立调整部署而不改写模型逻辑。
2.3 StageRuntime：让一个 Stage 真正“活起来”#
StageRuntime 管理的是 Stage 的运行实体，包括 replica 创建、进程生命周期、readiness、设备亲和性、placement、关闭和故障处理等。
Logical Stage
|
+-- replica 0 -> process / devices / engine
+-- replica 1 -> process / devices / engine
`-- replica 2 -> process / devices / engine这使 vLLM-Omni 可以区分两类配置：
PipelineConfig
描述模型逻辑上有哪些 Stage，以及它们如何连接
DeployConfig
描述每个 Stage 有多少实例、放到什么资源上运行也就是：
Model topology != Deployment topology同一条逻辑 pipeline 可以在开发机上全部 colocate，也可以在生产环境中拆到不同机器；Talker 吞吐不足时，可以只增加 Talker replica，而无需复制整套 Thinker 和 Code2Wav。
2.4 OmniConnector：跨 Stage 的数据面#
Stage 分开以后，中间数据可能是 token、hidden states、latent、音频 code 或其他 tensor。它们既可能在同机设备间传递，也可能跨进程、跨节点搬运。
OmniConnector 为此提供 Stage-to-Stage transport：
Stage A Stage B
GPU 0-3 GPU 4-5
| ^
| hidden states |
+-------- OmniConnector -------+从系统设计角度，可以用 control plane / data plane 建立直觉：
Orchestrator = control plane: 去哪里、何时推进、状态是什么
Connector = data plane: 数据如何 put/get 和完成传输这不是说二者在实现中毫无交互，而是强调 Connector 不应该自行决定 pipeline routing。保持 transport-only，才能让传输后端和调度策略分别演进。
3. 一次 Omni Request 如何穿过系统#
以包含文本、图片和音频输入，并最终生成语音的请求为例，可以把完整路径压缩为：
HTTP / Python request
|
v
API Server
|
v
AsyncOmniEngine
|
v
Orchestrator: current_stage = Thinker
|
v
Thinker StageRuntime
-> AR scheduler
-> vLLM worker
-> model runner
|
| hidden representation
v
OmniConnector
|
v
Orchestrator: current_stage = Talker
|
v
Talker StageRuntime
-> AR scheduler
-> acoustic token generation
|
| acoustic codes
v
OmniConnector
|
v
Orchestrator: current_stage = Code2Wav
|
v
Code2Wav StageRuntime
-> non-AR decoder
|
v
OmniRequestOutput -> Client这条路径揭示了三种不同粒度的状态：
- Pipeline state：一个请求当前在哪个 Stage、接下来去哪里、是否已经结束；
- Stage state：请求在当前 engine 中排队、运行、暂停或输出到什么位置；
- Transfer state：哪些 chunk 已产生、哪些已消费、buffer 是否满、传输是否完成。
普通 vLLM 主要处理第二类状态中的 AR 部分。vLLM-Omni 的新增复杂度，恰恰来自第一类和第三类状态，以及三者之间的一致性。
4. AR Runtime 复用 vLLM，Diffusion Runtime 重新建模#
vLLM-Omni 并没有重新实现成熟的 LLM serving 栈。对于 AR Stage，它继续复用或扩展 vLLM 已有的：
- Scheduler 与 continuous batching；
- KV Cache、block manager 与 prefix caching；
- Executor、Worker 和 ModelRunner；
- attention backend、量化与分布式执行；
- chunked prefill 和 speculative decoding 等优化。
这意味着对 vLLM scheduler、KV Cache、GDN/Mamba state、TP/EP 或 kernel 的理解，可以直接迁移到 Omni pipeline 的 AR Stage 中。
vLLM-Omni
|
+-- AR Stage --------> reuse vLLM runtime
|
`-- Diffusion Stage -> dedicated diffusion runtime但 Diffusion 无法自然套进 token scheduler。两类 workload 的迭代轴完全不同：
Autoregressive:
token 0 -> token 1 -> token 2 -> ...
Diffusion:
noise x_T -> x_(T-1) -> x_(T-2) -> ... -> imageAR scheduler 关心：
prefill or decode
number of scheduled tokens
KV Cache blocks
sequence lengthDiffusion scheduler 更关心：
denoising timestep
latent shape and resolution
batch composition
sequence/context parallelism
reuse across adjacent timesteps所以 vLLM-Omni 在 vllm_omni/diffusion 中建设独立的 engine、scheduler、worker、model runner、pipeline、attention backend 和 cache backend。这里的 cache 也不等同于 LLM KV Cache：它可能复用相邻 denoise step 中变化较小的 block feature，以减少重复计算。
这反映了 vLLM-Omni 更深的一层目标：统一的不是某一种 scheduler，而是让不同 runtime 能在同一个 Stage/Pipeline 抽象下协作。
5. Async Chunk：从串行链路变成异构流水线#
仅仅把模型拆成多个 Stage 并不会自动提升性能。如果每个 Stage 都等待上游产生完整结果，端到端延迟仍是各阶段延迟之和：
time ---------------------------------------------------->
Thinker [================]
Talker [==============]
Code2Wav [========]若上游能够逐段产生数据，下游也能逐段消费，就可以把中间结果切成 chunk，在 Stage 之间异步传输：
time ---------------------------------------------------->
Thinker [C0][C1][C2][C3]
Talker [C0][C1][C2][C3]
Code2Wav [C0][C1][C2][C3]这形成的是 heterogeneous pipeline parallelism。它不同于传统 Transformer PP：后者通常把一组连续 layer 切到不同设备，而这里串联的是不同子模型、不同计算范式，甚至不同服务实例。
异步 chunk 同时带来一整组系统问题：
(request_id, chunk_id)
producer / consumer progress
buffer ownership and lifetime
out-of-order arrival
backpressure
cancellation and cleanup
end-of-stream semantics如果 Thinker 生产速度长期快于 Talker，无限制 buffer 会耗尽 host 或 device memory；如果简单阻塞上游，又可能让 Thinker worker 的 batch 和设备利用率下降。因此，chunk transfer 不是一个 send(tensor) 就能解决的功能，而是一套生产者—消费者协议。
当前设计进一步通过 OmniChunkTransferAdapter 管理 chunk 生命周期和异步 put/get，把 Connector 收敛为更纯粹的 transport。这种分层很重要：
OmniChunkTransferAdapter
-> request/chunk lifecycle, buffering, async semantics
OmniConnector
-> data transport官方 async chunk 设计展示了 Thinker、Talker 与 Code2Wav 由串行执行转为流水执行后，端到端 latency 和首个输出时间都能明显提前。不过具体收益取决于各 Stage 时长、chunk 粒度、传输成本和能否真正流式消费，不能简单等同于 Stage 数量带来的线性加速。
6. Omni 实际上拥有两级调度#
vLLM-Omni 的一个关键系统特征，是同时存在 inter-stage 与 intra-stage 两个调度层级。
Global level
Orchestrator
-> request routing, stage transition, cross-stage state
Local level
Stage Scheduler
-> tokens, denoise steps, local batching, cache and worker execution这种分层避免了一个全知 scheduler 同时理解所有模型细节，但也产生了新的协调问题。
6.1 跨 Stage 重新组 Batch#
Thinker 中一起执行的请求，未必会同时到达 Talker：
Thinker batch: [A, B, C]
A -> ready -----+
B -> slow +--> Talker batch: [A, C, D, ...]
C -> ready -----+因此，batch 不是沿 pipeline 原样传递，而会在每个 Stage 的输入边界被重新形成。局部 scheduler 希望积累更大的 batch 以提高吞吐，Orchestrator 则还要考虑端到端 latency、输出顺序和下游拥塞。
6.2 Backpressure#
假设各阶段稳态吞吐为：
Thinker 100 requests/s
Talker 40 requests/s
Code2Wav 200 requests/s如果入口持续按 100 requests/s 注入，Talker 前的队列必然增长。局部看，Thinker 的利用率很好；全局看，系统正在积累无法及时交付的中间结果。
这要求全局层感知 pipeline bottleneck，并通过 admission control、buffer watermark、上游限速、更多 Talker replica 或动态 placement 等机制控制压力。
6.3 公平性与优先级#
一个请求可能已经在 Stage 0 消耗大量计算，却在 Stage 1 后面等待；另一个短请求刚刚进入系统。全局策略需要决定：优先完成在途请求以释放状态，还是优先短任务改善尾延迟。仅优化每个 Stage 的局部吞吐，不一定能优化端到端 SLO。
因此，未来真正有研究价值的问题不是简单地“给每个 Stage 加一个 scheduler”，而是：
Global orchestration 如何向局部 scheduler 传递拥塞、优先级和成本信息，又不破坏各 runtime 的独立性？
7. Stage 解耦如何改变部署拓扑#
当模型拓扑和部署拓扑分离后，一条 pipeline 可以从单机 colocated 部署平滑扩展到多机 disaggregated deployment：
Machine A
GPU 0 GPU 1 GPU 2 GPU 3
Thinker TP=4
|
| OmniConnector
v
Machine B
GPU 0 GPU 1
Talker TP=2, replicas=2
|
v
Machine C
GPU 0
Code2Wav这类解耦的价值不仅是“模型放得下”，还包括：
- 按 Stage 的实际吞吐独立扩缩容；
- 让 AR、DiT 和 Decoder 使用更匹配的硬件；
- 为不同 Stage 配置独立量化和并行策略；
- 隔离故障与发布，减少整条 pipeline 同步升级的成本；
- 通过 replica 和 placement 调整流水线平衡。
代价也同样明确：中间 tensor 可能远大于 token id，跨节点传输会增加带宽和延迟成本；每个 Stage 的独立队列会产生额外排队；故障恢复还必须处理已经产生但尚未消费的中间状态。
所以，是否 disaggregate 不能只看单个 Stage 的算力需求，而应比较：
saved compute / better utilization
versus
transfer + queueing + coordination overhead8. 它与 TP、EP、PP 和 PD 分离是什么关系#
Stage 不是对既有并行策略的替代，而是位于更高层的组合维度。
传统 Pipeline Parallelism 常把同一模型的连续 layer 切开：
layers 0-20 -> layers 21-40Omni Stage 则切开不同子模型或执行范式：
Thinker -> Talker -> Code2Wav
Encoder -> DiT -> VAE每个 Stage 内部仍可继续使用 TP、DP、EP 或 PP：
Omni pipeline
|
+-- Thinker: TP=4, EP=8
+-- Talker: TP=2, replicas=2
`-- DiT: sequence parallel=4Prefill/Decode disaggregation 也属于 Stage 内部的 AR 优化。二者可以正交组合：
Omni pipeline
Thinker Talker
| |
+------+------+ +------+------+
| | | |
Prefill Decode Prefill Decode因此，一套完整 deployment topology 未来可能同时包含：
Stage x TP x DP x EP x PP x P/D split x replicas这里也可以理解文档中的 E/P/D/G fully disaggregation：Encoding、Processing、Decoding、Generation 等阶段不再被要求在同一个 engine 或同一组设备中从头运行到尾，而可拆到独立资源池，通过中间传输连接。
Encoder Pool
|
v
Processing Pool
|
v
Generator Pool
|
v
Decoder Pool它延续了现代 LLM Infra 从 monolithic inference 走向 disaggregated inference 的趋势，只是把解耦边界从一个 AR 模型内部扩展到了完整的异构生成图。
9. 真正困难的系统问题#
Stage-based runtime 给了系统清晰的抽象，但它不会自动解决所有问题。相反，它让一批过去隐藏在单体 forward() 中的问题变成显式的 Infra 课题。
9.1 异构资源分配#
给定 8 张卡，Thinker=4, Talker=2, Code2Wav=2 很直观，却未必最优。资源比例应该由各阶段 service time、batch efficiency、通信成本和目标 workload 共同决定。
pipeline throughput
~= min(stage capacity)瓶颈 Stage 多一张卡可能显著提高端到端吞吐，非瓶颈 Stage 多一张卡则可能几乎没有收益。静态配置之外，更难的是怎样在 workload 变化时安全调整 replica 和并行度。
9.2 跨阶段状态与容错#
如果 Talker 崩溃，Thinker 已生成的 hidden states 如何处理？系统需要明确：
- 中间结果是否可重放或需要重新计算；
- retry 应从当前 Stage 还是 pipeline 起点开始；
- Connector 中未消费的 chunk 如何清理；
- 取消信号怎样传播到所有上下游 Stage；
- streaming response 已经交付一部分时采用什么失败语义。
这比普通无状态 RPC 更接近一个带长生命周期状态的分布式 dataflow。
9.3 端到端可观测性#
只看每个 worker 的 kernel time 或 tokens/s 已经不够。系统还需要串起同一个 request 的：
admission wait
-> Stage 0 queue / compute / transfer
-> Stage 1 queue / compute / transfer
-> Stage 2 queue / compute
-> first output / completion否则，一个 Stage 看似利用率不足，可能是上游没有及时供给；某次延迟尖峰看似来自计算，实际可能是 Connector buffer 的 backpressure。
9.4 全局与局部目标不一致#
局部 scheduler 的最佳 batch 可能导致下游 burst；全局为降低 TTFP 选择更小 chunk，又可能让通信和 launch overhead 过高。系统最终需要优化的是端到端 throughput、TTFP、latency 与资源成本，而不是任一 Stage 的孤立指标。
10. 如何阅读 vLLM-Omni 源码#
如果目标是理解系统，而不是只增加一个模型定义，不建议一开始扎进具体 modeling_*.py。更有效的路线是沿着 request 和数据各走一遍。
第一层：先建立架构词汇#
先读 Architecture Overview，回答四个问题：
Stage 是什么边界？
Orchestrator 保存什么状态？
StageRuntime 管什么生命周期？
Connector 传什么、不决定什么？如果这四个职责还混在一起，后面的进程和异步代码会很难定位。
第二层：从 AsyncOmniEngine 追请求入口#
从公开 API 或 server 进入 AsyncOmniEngine，跟踪：
request creation
-> request id and input preprocessing
-> submit to orchestrator
-> background loop
-> output aggregation / streaming这一层的目标不是看懂每一行，而是找到 pipeline 的 composition root 与主 event loop。
第三层：跟踪 Orchestrator 的状态机#
选择一个真实模型 pipeline，记录每次 Stage 完成后：
谁接收 output？
谁决定 next stage？
request state 存在哪里？
错误与 cancellation 怎样传播？
最终 output 在哪里组装？这一步会把抽象的“跨 Stage 调度”落到具体队列、协程和对象上。
第四层：沿 Connector 追 tensor#
对一个中间 tensor 追踪：
producer output
-> serialization / metadata
-> put
-> transport backend
-> get
-> consumer input同时区分 routing metadata、transfer metadata 和模型 tensor，观察同机与跨节点路径是否采用不同后端。
第五层：研究 Async Chunk 与背压#
在理解完整 output 的 stage transfer 后，再进入 chunked streaming：
chunk 如何编号？
何时认为 stream 结束？
buffer 在哪里？
producer 快于 consumer 时发生什么？
request 取消后谁负责清理？这是最适合学习异步流水、生命周期和并发 correctness 的部分。
第六层：对比 AR 与 Diffusion Scheduler#
最后选择一条 AR pipeline 和一条 Diffusion pipeline，对照它们的 request state、batch 维度、cache 和 step loop。重点不在于记类名，而是回答：为什么 Stage 抽象能够统一二者，而 scheduler 抽象不能？
可以先建立下面的源码地图，再根据实际版本核对目录和类名：
vllm_omni/
|-- entrypoints/ API and service entrypoints
|-- engine/ AsyncOmniEngine and engine composition
|-- core/ orchestration and stage execution
|-- distributed/ connector and communication
|-- worker/ worker-side execution
|-- model_executor/ model runners and model integration
|-- diffusion/ non-AR diffusion runtime
|-- attention/ attention backends
|-- quantization/ quantization support
`-- outputs/ Omni output structures项目仍在快速演进，源码路径和类职责可能随版本调整，因此阅读时应以所用 commit 的 design docs 与代码为准，而不是把目录结构当成稳定 API。
11. 总结#
vLLM-Omni 最重要的价值，不是支持列表里又多了多少多模态模型，而是重新定义了推理 runtime 的边界：
vLLM
How to efficiently generate tokens?
vLLM-Omni
How to efficiently execute a heterogeneous generative pipeline?它通过 Stage 把 AR、Diffusion、Encoder 和 Decoder 变成可独立执行与部署的单元，再通过四类组件让整条 pipeline 运转起来：
AsyncOmniEngine -> 系统入口与生命周期
Orchestrator -> 跨 Stage 控制与请求状态
StageRuntime -> Stage 实例与部署生命周期
OmniConnector -> 中间数据传输在 Stage 内，成熟的 AR workload 继续复用 vLLM；无法套进 token loop 的 Diffusion 则拥有自己的 runtime。异步 chunk 进一步让异构子模型从串行调用变成 producer-consumer pipeline，而两级调度又把优化问题推向 global orchestration 与 local scheduling 的协同。
最终，可以用这张图保持心智模型：
Request
|
v
AsyncOmniEngine
|
v
Orchestrator
global scheduling
|
+-------------------+-------------------+
| | |
v v v
Stage A Stage B Stage C
Thinker Talker DiT
| | |
Scheduler Scheduler Scheduler
| | |
Worker Worker Worker
| | |
ModelRunner ModelRunner Diffusion Pipeline
| | |
+---------- OmniConnector --------------+
Stage-level parallelism
x
TP / DP / EP / PP
x
P/D split
x
async streaming对于 Infra 工程师，最值得继续追踪的也正是这些新边界：Orchestrator 如何维护全局状态，StageRuntime 如何管理异构资源，Connector 与 chunk adapter 如何建立流式语义，以及 cross-stage scheduler 如何处理 batching、backpressure、公平性和容错。
当这些问题成为主角，vLLM-Omni 就不再只是“支持多模态的 vLLM”，而开始接近一套面向异构生成模型的通用 inference runtime。