系列 《扩散模型推理基础设施:从一次去噪到一个生成服务》 第 8 / 10 篇 深入
系列总览 — 为什么这样组织、读它需要什么、读完能做什么 · 幻灯片 — 整个系列的精华,一份可分享的 deck
学习地图 — AI-Infra 工程师学习地图 第 11 / 14 个系列
- 负载画像——一次生成在 GPU 上发生什么
- 单卡执行——attention 后端、编译、FP8 / INT4 与 offload
- 跨步冗余——TeaCache、First-Block Cache 一族的缓存与跳步
- 视频——长序列 attention 的账与稀疏化
- 多卡并行——序列并行、CFG 并行与 PipeFusion,为什么不是张量并行
- 少步与自回归——把步数变成系统参数
- serving 形态——请求、批、三段分离、附件、异步任务与成本
- 三个引擎的对照导读——同一张图的请求在 SGLang Diffusion、vLLM-Omni 与 xDiT 里各走过什么
- 配置、评测与排障——从一张卡的推导到一条伪影的排查
- 系列总结与通关自测
上一篇:serving 形态——请求、批、三段分离、附件、异步任务与成本下一篇:配置、评测与排障——从一张卡的推导到一条伪影的排查
前七篇每篇末尾有一张对照表,指向三个引擎里实现同一机制的文件与类。这一篇把这些点连成线:一个”生成一张 FLUX 1024² 的图”的请求,在 SGLang Diffusion、vLLM-Omni 与 xDiT 里各自从哪里进来、经过哪些进程、哪些类、哪些函数,到哪里出去。三条线走完,三个引擎的取向就清楚了:SGLang Diffusion 把扩散塞进了 LLM serving 的结构(scheduler、worker、kernel 栈、warmup、CUDA graph 都复用),vLLM-Omni 为全模态模型设计了 stage 流水线(一个请求可以经过 LLM → DiT → TTS,扩散是其中一种 stage),xDiT 只做并行(把 diffusers 的 pipeline 包一层、换掉 attention、建好并行组,没有服务层)。而 diffusers 是三者共同的底座——模型定义、调度器、pipeline 的接口都从它来,三个引擎要么直接复用它的 pipeline、要么按它的接口重写”原生”版本。
本文只讨论结构——每个组件在哪、谁调谁;每个机制的原理在前七篇。版本:diffusers v0.40.0、SGLang v0.5.19(sglang.multimodal_gen)、vLLM-Omni v0.28.0(vllm_omni.diffusion)、xDiT 2026-09-02 的主线 commit 07572e7(xfuser)。引用只到目录与类 / 函数名,不引行号。
本篇要回答的核心问题是:
一个
/v1/images/generations请求从进 HTTP 到返回 base64,在三个引擎里各经过哪些函数?1 它们在进程模型、pipeline 抽象、并行组、调度上各做了什么不同的选择?2 为什么会这样选?3
一、总览
1. 一个生图请求在三个引擎里各走哪条路
%% 图:SGLang Diffusion 的路径:http_server → AsyncSchedulerClient → rank 0 的 Scheduler → 每个 worker 进程的 GPUWorker → ComposedPipelineBase
flowchart TB
subgraph SG["SGLang Diffusion(sglang.multimodal_gen)"]
direction TB
SG1["http_server.py
openai/image_api.py:/v1/images/generations"] --> SG3["AsyncSchedulerClient
(scheduler_client.py,ZMQ)"] --> SG4["Scheduler(rank 0 的 worker 进程)
(managers/scheduler.py)
Req · 动态批处理准入 · warmup · 广播给其它 rank"] --> SG5["GPUWorker × num_gpus(每个 worker 进程一个)
(managers/gpu_worker.py)"] --> SG6["ComposedPipelineBase
(pipelines_core/composed_pipeline_base.py)
stages:text_encoding → latent_preparation → denoising → decoding"] --> SG7["runtime/models/ 原生 DiT
runtime/layers/:USPAttention · fused ops · cache"]
end
classDef sg fill:#eef6ff,stroke:#5b8fd6,color:#222
class SG1 sg
%% 图:vLLM-Omni 的路径:vllm serve --omni → AsyncOmni → Orchestrator → stage 0 的 DiffusionEngine → MultiprocDiffusionExecutor → DiffusionWorker
flowchart TB
subgraph VO["vLLM-Omni(vllm_omni.diffusion)"]
direction TB
VO1["vllm serve MODEL --omni
entrypoints/openai/:/v1/images/generations"] --> VO2["AsyncOmni(entrypoints/async_omni.py)
API server 进程 + Orchestrator 线程"] --> VO3["stage 0:InlineStageDiffusionClient(单 stage 单副本)
或 StageDiffusionClient → StageDiffusionProc(stage_diffusion_client.py / stage_diffusion_proc.py)"] --> VO4["DiffusionEngine
(diffusion_engine.py)
RequestScheduler / StepScheduler(sched/)"] --> VO5["MultiprocDiffusionExecutor → WorkerProc → DiffusionWorker
(executor/ · worker/diffusion_worker.py)"] --> VO6["DiffusionModelRunner
(worker/diffusion_model_runner.py)"] --> VO7["models/flux/pipeline_flux.py + flux_transformer.py
attention/ · distributed/ · cache/ · offloader/"]
end
classDef vo fill:#eefaf0,stroke:#4d9a5c,color:#222
class VO1 vo
%% 图:xDiT 的路径:torchrun 用户脚本 → EngineConfig → xDiTParallel 包装 diffusers pipeline → 初始化并行组 → 替换 attention processor
flowchart TB
subgraph XD["xDiT(xfuser)"]
direction TB
XD1["torchrun --nproc_per_node=N
用户脚本"] --> XD2["xFuserArgs.create_config() → EngineConfig
(config/args.py,含 init_distributed_environment)"] --> XD3["xDiTParallel(pipe, engine_config)
(parallel.py)"] --> XD4["xFuserPipelineWrapperRegister
→ xFuserFluxPipeline(model_executor/pipelines/pipeline_flux.py)
包装 diffusers 的 FluxPipeline"] --> XD5["initialize_runtime_state → initialize_model_parallel
(core/distributed/runtime_state.py / parallel_state.py)"] --> XD6["xFuserFluxAttentionWrapper 替换 attention processor
(model_executor/models/transformers/transformer_flux.py)
USP:core/long_ctx_attention/"] --> XD7["diffusers 的 transformer 原样执行
每个 rank 一份"]
end
classDef xd fill:#fff7e0,stroke:#c98a00,color:#222
class XD1 xd
| 维度 | SGLang Diffusion | vLLM-Omni | xDiT |
|---|---|---|---|
| 定位 | 图像 / 视频生成的 serving 框架,SGLang 的一个子包 | 全模态(LLM + 扩散 + 语音 + 动作)serving,vLLM 的扩展 | 扩散 Transformer 的并行推理引擎(库),无服务层 |
| 进程模型 | HTTP server 进程 → 每卡一个 worker 进程(Scheduler + GPUWorker),rank 0 收 ZMQ 再广播给其它 rank,与 SGLang LLM 同构 |
|
torchrun SPMD:每个 rank 跑同一段脚本、同一个 pipeline |
| pipeline 抽象 | 原生 ComposedPipelineBase = stage 列表(文本编码、latent 准备、去噪、解码…);--backend diffusers 回退到 diffusers pipeline |
按模型的原生 pipeline(models/<model>/pipeline_*.py)+ diffusers_adapter;--diffusion-load-format diffusers 回退 |
xFuserPipelineBaseWrapper 包装 diffusers pipeline,替换 attention 与部分层 |
| 模型定义 | runtime/models/ 原生重写(用自己的 layers/) |
models/<model>/*_transformer.py 原生重写 |
直接用 diffusers 的模型 |
| 并行 | USP(UlyssesAttention / USPAttention)、CFG 并行、TP、FSDP 推理、跨节点 SP |
USP / Ring、CFG 并行、TP、HSDP、PP、VAE patch 并行 | USP、CFG 并行、PipeFusion、DP、Parallel VAE——源头 |
| 调度 / 批处理 | Scheduler + dynamic_batch_admission(形状兼容的同构批) |
RequestScheduler(请求级同构批)/ StepScheduler(--step-execution:step 级批,兼容请求可在步边界进出) |
无(一个进程组一个请求) |
| 缓存 / 加速 | Cache-DiT、TeaCache、Spectrum、CFG gating、breakable CUDA graph、fused ops、--quality |
TeaCache / MagCache / StepCache / Cache-DiT 后端(hook)、prompt embedding 缓存、regional compile | TeaCache / FBCache adapter、DiTFastAttn、compile |
| API | OpenAI 兼容 /v1/images、/v1/videos、/v1/meshes、realtime |
OpenAI 兼容 /v1/images/generations、/v1/chat/completions 扩展 |
无(Ray 示例) |
| 分离 | disaggregation:encoder / denoiser / decoder / server 角色 | stage-based:天然分离 | 无 |
| 自回归 / 实时 | causal DMD pipelines、KV cache、realtime/ 会话、KV 量化 |
diffusion_kv/(分页 KV) |
pipeline_causal_wan.py |
| 硬件 | NVIDIA / AMD / Intel XPU / Ascend / MPS / MUSA | NVIDIA / AMD(ROCm)/ Intel XPU / Ascend(NPU)/ MUSA | NVIDIA、AMD(ROCm,含 FlyDSL kernel),另有 MUSA / Ascend / MPS 的设备分支 |
2. 本文的章节安排
| 章 | 主题 | 内容 |
|---|---|---|
| 二 | diffusers:共同的底座 | pipeline / 模型 / 调度器 / hook 的接口;三个引擎各拿了什么 |
| 三 | SGLang Diffusion 的一条路 | 从 image_api.py 到 GPUWorker 到 stage 到 USPAttention |
| 四 | vLLM-Omni 的一条路 | 从 --omni 到 stage 进程到 DiffusionEngine 到 DiffusionModelRunner |
| 五 | xDiT 的一条路 | 从 torchrun 到 xDiTParallel 到 attention processor 替换 |
| 六 | 逐机制对照 | 并行组、attention 后端、缓存 hook、offload、量化、LoRA、warmup 在三者里的位置 |
| 七 | 分歧点 | 三种取向各来自什么、各付了什么代价;该选哪个 |
| 八 | 本文小结 | |
| 九 | 自测 | 5 道题 |
二、diffusers:共同的底座
1. 四个接口
diffusers 给扩散推理定义了四层可替换的接口,三个引擎都在这四层上工作:
| 层 | diffusers 里 | 内容 | 引擎怎么用 |
|---|---|---|---|
| pipeline | pipelines/flux/pipeline_flux.py 的 FluxPipeline.__call__ |
编码 prompt → 准备 latent 与时间步 → 去噪循环 → VAE 解码 → 后处理 | xDiT 包装它;SGLang / vLLM-Omni 按同样的段落重写成 stage 或原生 pipeline,并保留 --backend diffusers 回退 |
| 模型 | models/transformers/transformer_flux.py(FluxTransformer2DModel)、models/autoencoders/、文本编码器来自 transformers |
block 结构、attention、adaLN、RoPE | xDiT 原样用、只换 attention processor;SGLang / vLLM-Omni 用自己的 layers/ 重写(为了 fused ops、USP、量化、CUDA graph) |
| attention 分派 | models/attention_dispatch.py 的 AttentionBackendName(flash / _flash_3_hub / sage / flex / native / …);模型上用 set_attention_backend(...) 选 |
后端枚举与分派 | 各引擎有自己的 selector,语义相同 |
| hook | hooks/:HookRegistry、ModelHook、StateManager;group_offloading.py、layerwise_casting.py、first_block_cache.py、taylorseer_cache.py、mag_cache.py、context_parallel.py、tensor_parallel.py |
不改模型代码的 forward 拦截 | “不改单个模型代码”的缓存与 offload 形态被三者沿用;vLLM-Omni 的缓存后端就是 hook,xDiT 的 TeaCache / FBCache 则是把 block 列表换成带缓存的包装模块 |
2. 权重格式
Hugging Face Hub 上的 diffusers 格式 checkpoint(transformer/、vae/、text_encoder/ 子目录 + model_index.json)是三个引擎的输入。SGLang 的 runtime/loader/(weight_load_plan.py、component_loaders/、transformer_load_utils.py)与 vLLM-Omni 的 model_loader/(diffusers_loader.py、checkpoint_adapters/、host_weight_loader.py)都是”读 diffusers 格式、映射到自己的模块名”——与 RL 后训练基础设施系列第四篇的”HF 名 ↔ Megatron 名”是同一类工作。原生重写模型的代价之一就是要维护这层映射(SGLang 的 add-model 指南里大半篇幅在讲它)。
三、SGLang Diffusion 的一条路
1. 进程
sglang serve --model-path black-forest-labs/FLUX.1-dev --port 30010 起的进程:
- HTTP server 进程(主进程跑
launch_http_server_only):runtime/entrypoints/http_server.py(FastAPI)openai/image_api.py的@router.post("/generations"):build_sampling_params→prepare_request(entrypoints/utils.py)构造Req(pipelines_core/schedule_batch.py:prompt、size、steps、seed、cache 参数…;LoRA 不在请求里,是服务级的POST /v1/set_lora)process_generation_batch→AsyncSchedulerClient.forward(scheduler_client.py):ZMQ 发到 rank 0 的 Scheduler,等回OutputBatchserver_warmup.py:HTTP 起来后run_async_client_warmup按--warmup-resolutions发 dummy 请求(编译 / CUDA graph 捕获),Scheduler 一侧由SchedulerWarmupMixin处理
- worker 进程 × num_gpus:
launch_server(runtime/launch_server.py)给每张卡起一个mp.Process(target=run_scheduler_process)(run_scheduler_process在runtime/managers/gpu_worker.py),进程里是一个Scheduler(runtime/managers/scheduler.py,class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisaggMixin))加一个GPUWorker(self.worker)- rank 0 的 Scheduler 绑 ZMQ ROUTER 收 Req,
_can_dynamic_batch判断队列里的 Req 与当前请求的采样参数签名是否相同,BatchAdmissionController(dynamic_batch_admission.py)再按--batching-max-size(默认 1,即不合批)/--batching-config的上限决定收几个 recv_reqs用broadcast_pyobj在 SP / CFG / TP 组内把请求广播给其它 rank,每个 rank 的 Scheduler 让自己的 GPUWorker 跑同一批;rank 0 的return_result把OutputBatch经 ZMQ 回给 HTTP 进程GPUWorker:持有一个ComposedPipelineBase实例 + 并行组的一个 rankforward_context.py:ForwardContext传每次前向的状态(current_timestep、attn_metadata、当前的Req)memory_managers/:component_residency(三段的 offload)、layerwise_offload、memory_occupation_controller
- rank 0 的 Scheduler 绑 ZMQ ROUTER 收 Req,
DiffGenerator(entrypoints/diffusion_generator.py)是离线的 Python 入口:launch_server(launch_http_server=False)只起 worker 进程,generate()经sync_scheduler_client发同样的请求;HTTP 路由不经过它
与 SGLang 的 LLM 路径(tokenizer manager → scheduler → TP worker)同构:HTTP 与调度分进程,每个 rank 一个进程、各跑一份 Scheduler,rank 0 收 ZMQ 再 broadcast_pyobj 给同组 rank。扩散没有 tokenizer manager(文本编码在 GPU 上作为 stage),多了 warmup 与 batching 准入。
2. pipeline = stage 列表
ComposedPipelineBase(pipelines_core/composed_pipeline_base.py)把一次生成拆成 pipelines_core/stages/ 里的 stage 对象顺序执行:
- InputValidationStage → TextEncodingStage(+ 可选 ImageEncodingStage / ConditionEncodingStage)
- LatentPreparationStage → TimestepPreparationStage(
add_standard_t2i_stages里就是这个顺序) - DenoisingStage(或 DmdDenoisingStage / CausalDMDDenoisingStage)
- DecodingStage → 后处理
runtime/pipelines/flux.py 就是”FLUX 用哪些 stage、每个 stage 用哪个模块”的声明;comfyui_*_pipeline.py 是接 ComfyUI 格式 checkpoint 的变体;diffusers_pipeline.py 是 --backend diffusers 的回退——把整个 diffusers pipeline 当一个 stage。pipelines_core/executors/(sync_executor、parallel_executor、pipeline_executor)决定 stage 怎样在 worker 上执行;disaggregation 模式下 stage 按 RoleType(encoder / denoiser / decoder)落到不同的 worker 组,disaggregation/transport/ 传中间张量。
DenoisingStage 里的循环:每步 → 取 CFG 分支(distributed/cfg_policy.py)→ 调 runtime/models/ 里的原生 DiT forward(TeaCache 的跳步判断在 DiT 基类 CachableDiT 混入的 TeaCacheMixin 里,cache/teacache.py;Cache-DiT 走 cache_dit_integration.py)→ 采样器 scheduler.step()(调度器实现在 runtime/models/schedulers/)。
3. 模型与层
runtime/models/ 是原生重写的 DiT(FLUX、Wan、Qwen-Image…),用 runtime/layers/ 的部件:attention/layer.py 的 UlyssesAttention / USPAttention / LocalAttention(后端由 attention/selector.py 按 --attention-backend 选,attention/backends/ 是 FA / SDPA / Sage / STA / VSA 的实现)、fused_scale_shift_gate.py(adaLN 融合)、rotary_embedding/、layernorm.py、linear.py(含量化)、lora/、kvcache/(因果模型)。并行组在 distributed/parallel_state.py / group_coordinator.py / parallel_groups.py(SP 组、CFG 组、TP 组),usp.py 与 sp_shard_utils.py 负责按 token 切分与 all-to-all 重排。breakable_cuda_graph/runner.py 把 DiT 前向分段捕获。
4. 出口
DecodingStage 调 VAE(tiling 配置来自 --vae-config.* 一组参数),OutputBatch 经 scheduler 回 HTTP 进程;process_generation_batch 用 save_outputs(entrypoints/utils.py)编码成 PNG / MP4 文件,image_api.py 再按 response_format 读成 b64 返回,或经 openai/storage.py 的 CloudStorage 上传后返回 URL。postprocess/ 是可选的超分(Real-ESRGAN)与插帧(RIFE)。/v1/videos 的 job 状态在 openai/stores.py。
四、vLLM-Omni 的一条路
1. 进程:stage
vllm serve black-forest-labs/FLUX.1-dev --omni --port 8091 起的进程:
- API server 进程:
entrypoints/openai/(FastAPI,/v1/images/generations、/v1/chat/completions扩展)AsyncOmni(entrypoints/async_omni.py,实现 vLLM 的EngineClient接口)→AsyncOmniEngine起 Orchestrator 线程(engine/orchestrator.py)- 按
vllm_omni/deploy/<model>.yaml的 stage 配置编排;FLUX 这类纯扩散模型没有 yaml,走_create_default_diffusion_stage_cfg,只有一个 diffusion stage,就是 stage 0 - 只有一个 stage、一个副本时用
InlineStageDiffusionClient(diffusion/inline_stage_diffusion_client.py):DiffusionEngine在本进程的线程池里跑,省一跳 ZMQ - 否则
StageDiffusionClient(diffusion/stage_diffusion_client.py)经 ZMQ PUSH / PULL 把请求发给 diffusion stage 进程
- diffusion stage 进程:
StageDiffusionProc(diffusion/stage_diffusion_proc.py),多 stage、多副本,或 stage-based CLI 时才有- stage-based CLI:每个 stage 用
--stage-id单独启动,经--omni-master-address/--omni-master-port汇合,所以 stage 可以放到别的主机上 --stage-overrides按 stage 覆盖配置(JSON),stage 用哪几张卡写在它的runtime.devices里
- stage-based CLI:每个 stage 用
DiffusionEngine(diffusion/diffusion_engine.py),在上面哪个进程里都一样:RequestScheduler/StepScheduler(sched/):请求级或 step 级的批,兼容键分别是RequestBatchSamplingParamsKey/StepBatchSamplingParamsKey- 执行器:单卡是
UniProcDiffusionExecutor(worker 就在本进程);多卡是MultiprocDiffusionExecutor(executor/multiproc_executor.py),起WorkerProc× N 进程(worker/diffusion_worker.py),经共享内存MessageQueue派任务
WorkerProc进程:DiffusionWorker→DiffusionModelRunner(worker/diffusion_model_runner.py)- 持有原生 pipeline(
models/flux/pipeline_flux.py)与 transformer(flux_transformer.py) - 并行组:
distributed/parallel_state.py(USP / Ring / CFG / TP / HSDP / PP;VAE patch 并行在vae_patch_parallel.py) - 缓存:
cache/selector.py的get_cache_backend选 cache_dit / tea_cache / mag_cache / step_cache 后端(CacheBackend),以 hook 挂到 transformer 上 - offload:
offloader/(sequential / layerwise / distributed layerwise) - 因果模型:
diffusion_kv/(KV manager、PagedAttention 适配)
- 持有原生 pipeline(
stage 是 vLLM-Omni 的核心抽象:一个全模态请求(比如 Qwen3-Omni:文本 + 图 → LLM → 语音)经过多个 stage,每个 stage 是一个独立的引擎(AR 的用 vLLM 本体,扩散的用 DiffusionEngine),stage 间用 OmniConnector 传张量、可以放在不同的 GPU 上独立扩缩。对纯扩散模型这套结构看起来”过重”(只有一个 stage,默认还 inline 在 API 进程里),但它让第七篇的三段分离几乎免费:把文本编码器或 VAE 做成另一个 stage 就是加一段 yaml。
2. 引擎
DiffusionEngine 与 vLLM 的 EngineCore 同构:一个调度循环,每轮从 RequestScheduler(请求级)或 StepScheduler(step 级——每轮把 batch 里每个请求各推进一步,_StepProgress 跟踪每个请求走到第几步,兼容的新请求可以在步边界加入)取一批,交给 executor 执行一步或整个请求,收回 RunnerOutput / BatchRunnerOutput(worker/utils.py)。DiffusionExecutionMode(REQUEST_BATCH / STEP_BATCH,由 --step-execution 决定,默认请求级)决定用哪种调度器。sched/sigma_schedule.py 的 DMD2SigmaSchedule 是少步蒸馏模型的时间步表。
3. 模型
models/<model>/ 每个模型一个目录:pipeline_*.py(生成流程)、*_transformer.py(原生 DiT,用 diffusion/layers/ 与 attention/)、*_mixin.py(并行与缓存的模型特定钩子)。models/diffusers_adapter/ 是通用回退(--diffusion-load-format diffusers)。models/dmd2/ 是少步蒸馏模型的共享逻辑。quantization/、lora/(manager.py、loader.py、layers/)、host_weight_runtime/(权重留在主机内存、按需上卡)挂在模型加载路径上。
4. 出口
DiffusionModelRunner 的输出经 output_formatter.py(format_diffusion_outputs)整理成 OmniRequestOutput 回到 API server 进程,OpenAI 层(image_api_utils.py 的 encode_image_base64)按 response_format 返回 b64 或文件。
五、xDiT 的一条路
1. 进程:torchrun
torchrun --nproc_per_node=8 my_flux.py --model FLUX.1-dev --ulysses_degree 4 --pipefusion_parallel_degree 2 --height 1024 --width 1024 --num_inference_steps 28
│ 8 个 rank 各跑一遍脚本:
├─ xFuserArgs.from_cli_args(args).create_config() → EngineConfig, InputConfig(config/args.py、config/config.py)
│ · torch.distributed 还没初始化时先 init_distributed_environment()
│ · ParallelConfig:dp / cfg / ulysses / ring / pipefusion 度数 + num_pipeline_patch;RuntimeConfig:compile、parallel_vae、cache
├─ pipe = FluxPipeline.from_pretrained(...) ← diffusers 原生 pipeline
├─ xDiTParallel(pipe, engine_config, input_config) (parallel.py)
│ · xFuserPipelineWrapperRegister.get_class(pipe) → xFuserFluxPipeline(model_executor/pipelines/pipeline_flux.py)
│ · 包装:xFuserPipelineBaseWrapper.__init__(model_executor/pipelines/base_pipeline.py)
│ · _init_runtime_state → RuntimeState._check_distributed_env(core/distributed/runtime_state.py)
│ → initialize_model_parallel(core/distributed/parallel_state.py)建 dp / cfg / sp / tp / pp 组
│ · _convert_transformer_backbone → xFuserFluxTransformer2DWrapper(model_executor/models/transformers/transformer_flux.py)
│ 把每个 block 的 attn 包成 xFuserFluxAttentionWrapper,processor 换成 xFuserFluxAttnProcessor
│ (通用的 xFuserAttentionWrapper / xFuserJointAttnProcessor2_0 在 model_executor/layers/attention_processor.py)
│ · USP 的 attention 实现在 core/long_ctx_attention/(ring / hybrid);PipeFusion 的 patch 循环在 pipeline wrapper 里
│ · use_parallel_vae → _convert_vae:distvae 的 parallelize_decoder
│ · self.pipe.prepare_run(input_config):按形状先跑 3 步 warmup(含 compile)
├─ xDiTParallel(prompt=...) → 每个 rank 跑 pipeline,attention 处发生 all-to-all / P2P
└─ save():is_dp_last_group() 的 rank 存图
my_flux.py 指一个按 xDiTParallel 接口写的用户脚本:官方的 examples/flux_example.py 在 2026-02(#641)删掉了,仓库现在自带的入口是 xdit 命令(xfuser/cli.py → xfuser/runner.py 的 xFuserModelRunner),按 --model 在注册表里找到 FLUX 的模型类,里面同样是 create_config 建并行环境、xFuserFluxPipeline 包装 diffusers pipeline。
没有服务层、没有调度器:一个进程组就是一个请求(或 DP 度数个请求)。要做服务,用 xfuser/ray/ 的 Ray 示例或自己包一层。
2. 包装而不是重写
xDiT 的核心设计是包装 diffusers:模型定义、pipeline 流程、调度器都原样用,只在两处介入——(1)替换 attention processor:diffusers 的 Attention 模块允许注入 processor,xDiT 的 processor 在调 attention 之前做 USP 的 all-to-all(或 Ring)、在 PipeFusion 模式下用 stale K / V;(2)pipeline wrapper 重写去噪循环:CFG 并行的分支分派、PipeFusion 的 patch 流水线、Parallel VAE 的切分。model_executor/layers/ 里另有 usp.py、fp8_linear.py、mxfp4_linear.py、fused_qk_rope_* 等可选替换层,core/cache_manager/ 是 PipeFusion 的 stale 激活缓存,model_executor/cache/ 接 TeaCache / FBCache。
好处:新模型只要 diffusers 有 pipeline,写一个 wrapper 就能并行;代价:diffusers 的 eager 实现本身不快(第二篇:MFU 0.31),xDiT 的绝对性能依赖 torch.compile,且 diffusers 的接口变化会打断 wrapper(compat.py 的存在就是为此)。
六、逐机制对照
前七篇末尾各表的汇总,按请求路径的顺序:
| 环节 | SGLang Diffusion | vLLM-Omni | xDiT | 备注 |
|---|---|---|---|---|
| 入口 | openai/image_api.py、video_api.py、realtime/ |
entrypoints/openai/;AsyncOmni |
用户脚本 / xfuser/ray/ |
|
| 请求对象 | pipelines_core/schedule_batch.py 的 Req / OutputBatch |
diffusion/request.py;sched/interface.py 的 NewRequestData / CachedRequestData |
InputConfig |
|
| 调度与合批 | managers/scheduler.py + dynamic_batch_admission.py |
sched/request_scheduler.py / step_scheduler.py;RequestBatchSamplingParamsKey / StepBatchSamplingParamsKey |
— | 都按兼容键同构批;vLLM-Omni 的 step 模式允许步边界进出 |
| 进程 / 执行器 | GPUWorker 进程组,ZMQ |
MultiprocDiffusionExecutor → WorkerProc → DiffusionWorker |
torchrun rank | |
| 并行组 | distributed/parallel_state.py、group_coordinator.py、parallel_groups.py |
distributed/parallel_state.py、group_coordinator.py、sp_plan.py |
core/distributed/parallel_state.py、group_coordinator.py、runtime_state.py |
三者的 GroupCoordinator 都从 vLLM 的 GroupCoordinator 演化 |
| 序列并行 | layers/attention/layer.py:UlyssesAttention、USPAttention;layers/usp.py、distributed/sp_shard_utils.py |
|
core/long_ctx_attention/{ring,hybrid}/;model_executor/layers/usp.py |
|
| CFG 并行 | distributed/cfg_parallel_utils.py、cfg_policy.py |
distributed/cfg_parallel.py |
pipeline wrapper 内 | |
| PipeFusion | — | distributed/pipeline_parallel.py(PP) |
core/cache_manager/ + pipeline wrapper |
xDiT 独有 |
| TP / FSDP | --tp-size;loader/fsdp_load.py |
--tensor-parallel-size;distributed/hsdp.py |
— | |
| attention 后端 | layers/attention/selector.py、backends/ |
attention/selector.py、backends/、ops/ |
core/distributed/attention_backend.py;core/fast_attention/、sparge_attention/、vsa_attention.py |
|
| 跨步缓存 | runtime/cache/teacache.py、cache_dit_integration.py、spectrum.py |
cache/base.py(CacheBackend、CachedTransformer)、teacache/、magcache/、stepcache/、cachedit/、selector.py |
model_executor/cache/(utils.py 的 TeaCachedTransformerBlocks / FBCachedTransformerBlocks,adapters/ 按模型接入) |
|
| offload | managers/memory_managers/:component_residency*.py、layerwise_offload*.py、host_memory_budget.py |
offloader/:sequential_backend.py、layerwise_backend.py、distributed_layerwise_backend.py、module_residency.py |
沿用 diffusers | |
| 编译 / CUDA graph | --enable-torch-compile;breakable_cuda_graph/runner.py |
compile.py(regionally_compile,按重复 block 编译) |
xfuser/compile/ |
|
| 量化 |
|
quantization/(含 hsdp_fp8.py);--quantization |
layers/fp8_linear.py、mxfp4_linear.py;core/distributed/fp4_quantize.py |
|
| LoRA | pipelines_core/lora/、layers/lora/ |
lora/manager.py、loader.py、layers/ |
diffusers 的 load_lora_weights |
|
| VAE |
|
distributed/vae_patch_parallel.py、distributed/autoencoders/;--vae-use-tiling |
distvae 的 parallelize_decoder(Parallel VAE);Ray 路径用 xFuserVAEWrapper |
|
| warmup | server_warmup.py、warmup_request_builder.py;--warmup-mode、--warmup-resolutions |
DiffusionEngine.run_startup_warmup(启动时一个 512² 的 dummy 请求) |
prepare_run |
|
| 分离 | disaggregation/:roles.py、orchestrator.py、dispatch_policy.py、transport/ |
stage 本身 | — | |
| 因果 / KV | layers/kvcache/、stages/causal_denoising.py、realtime/ |
diffusion_kv/(manager.py、paged_attention_adapter.py);调度器的 KVPrefetchJob |
pipeline_causal_wan.py |
|
| 权重加载 | loader/:weight_load_plan.py、component_loaders/、weight_readers/、rank_local_checkpoint.py |
model_loader/:diffusers_loader.py、checkpoint_adapters/、host_weight_loader.py、hub_prefetch.py |
diffusers from_pretrained |
|
| profiling | profiler.py、benchmarks/(--perf-dump-path) |
profiler/;--log-stats、pipeline profiler |
benchmark 脚本 |
七、分歧点
1. 三种取向来自三种出发点
| SGLang Diffusion | vLLM-Omni | xDiT | |
|---|---|---|---|
| 出发点 | 把 SGLang 成熟的 serving 结构(scheduler、worker、kernel、warmup、CUDA graph、多平台)复用到扩散上 | 服务全模态模型:一个请求经过多种模型,需要 stage 流水线;扩散是必须支持的一种 stage | 研究并行方法(USP、PipeFusion)并给出参考实现 |
| 模型来源 | 原生重写(为了 fused ops、量化、CUDA graph、多平台) | 原生重写(为了与 vLLM 的 worker / KV / 量化栈一致) | 直接用 diffusers |
| 支持新模型的成本 | 高:重写模型 + 权重映射 + stage 声明(有 add-model 指南与 diffusers 回退) | 高:同上(有 add-diffusion-model 指南与 adapter 回退) | 低:写一个 pipeline wrapper |
| 单卡性能 | 高(fused ops、--quality high、BCG) |
高 | 依赖 diffusers + compile |
| 多卡 | USP / CFG / TP / FSDP,跨节点 | USP / Ring / CFG / TP / HSDP / PP / VAE patch | 最全:+ PipeFusion、Parallel VAE,弱互联 |
| serving | 完整:API、批处理、warmup、分离、job、realtime | 完整:API、stage 分离、批处理 | 无 |
| 分离 | 角色制 disaggregation | stage 制,天然 | — |
| 付出的代价 | 与 SGLang 主线耦合(版本、依赖);模型覆盖靠社区重写 | stage 结构对纯扩散模型偏重;模型覆盖同样靠重写 |
|
2. 三个共同点
- 并行组的写法趋同:三者的
parallel_state.py/GroupCoordinator几乎可以互换——都是 vLLM 那套”按维度建 NCCL 组、rank 映射、collective 封装”的写法。序列并行的实现也趋同(Ulysses 的 all-to-all + FlashAttention、Ring 的 P2P 环),xDiT 的 USP 论文与实现是共同来源。 - 缓存与 offload 都不改单个模型的代码:vLLM-Omni 的缓存后端是 hook(
cache/teacache/hook.py),xDiT 把 FLUX 的 block 列表换成带缓存的包装模块,SGLang 把 TeaCache 放进所有 DiT 共用的基类CachableDiT——都是 diffusershooks/那种”在模型外面包一层”的思路。 - 同构批:有调度器的两个都是”形状兼容才合批”。默认都是整批开始整批结束;只有 vLLM-Omni 打开
--step-execution时,StepScheduler让兼容键相同的请求在步边界加入、各自走完就离开——这是受形状约束的 step 级连续批,目前只有实现了denoise_step的少数 pipeline(Qwen-Image、HunyuanImage-3 等,不含 FLUX)支持。形状必须相同是负载性质决定的(第七篇)。
3. 该选哪个
| 场景 | 选择 |
|---|---|
| 对外的图像 / 视频生成服务,NVIDIA 单节点或多节点 | SGLang Diffusion 或 vLLM-Omni;两者的模型覆盖与性能互有胜负,按目标模型是否有原生实现选 |
| 全模态服务(LLM + 生成 + 语音在一个请求里) | vLLM-Omni(stage 流水线为此设计) |
| 已有 vLLM 基础设施、想复用运维 | vLLM-Omni |
| 已有 SGLang 基础设施;Apple(MPS)上跑 | SGLang Diffusion(AMD / Ascend / Intel XPU 两者都有平台层,按目标模型的支持情况选) |
| 新模型 diffusers 刚出、两个引擎还没原生支持 | 两者的 diffusers 回退(性能打折)或 xDiT 包装 |
| PCIe / 以太网集群上的多卡 | xDiT(PipeFusion)或等 SGLang / vLLM-Omni 的 PP |
| 研究并行方法 / 做 benchmark | xDiT |
| 只是单卡跑图 | diffusers 本身 + compile + hook 缓存 |
八、本文小结
| 项 | SGLang Diffusion | vLLM-Omni | xDiT |
|---|---|---|---|
| 一句话 | 扩散进 LLM serving 的结构 | 扩散作为全模态流水线的一个 stage | diffusers 之上的并行包装 |
| 路径 | image_api → AsyncSchedulerClient → Scheduler → GPUWorker → ComposedPipelineBase 的 stage → runtime/models + layers/attention/USPAttention |
AsyncOmni → Orchestrator → InlineStageDiffusionClient(或 StageDiffusionClient → StageDiffusionProc)→ DiffusionEngine + Step/RequestScheduler → MultiprocDiffusionExecutor → DiffusionWorker → DiffusionModelRunner → models/flux |
xFuserArgs → xDiTParallel → xFuserFluxPipeline 包装 → initialize_model_parallel → attention processor 替换 → diffusers transformer |
| 进程 | HTTP 进程 + 每卡一个 worker 进程 | API 进程(+ 可选的 stage 进程)+ 执行器 | torchrun SPMD |
| 模型 | 原生重写 | 原生重写 | diffusers 原样 |
| 独有 | 动态批处理准入、breakable CUDA graph、--quality、realtime 会话、KV 量化、MPS 支持 |
stage 分离、PP、VAE patch 并行、step 级批、分页 diffusion KV | PipeFusion、Parallel VAE、弱互联 |
| 共有 | USP / CFG 并行、模型外挂式缓存与 offload、同构批、vLLM 式并行组、diffusers 权重格式 |
下一篇
最后一篇把前八篇变成决策与运维:给定模型、分辨率、步数、GPU 与 SLO,按什么顺序推出配置;性能与质量各怎么评;确定性从哪里丢;常见故障(VAE OOM、FP8 NaN、缓存伪影、重编译风暴、SP hang)各先查什么;以及系列总结。
九、自测
-
一个新的 DiT 模型刚在 diffusers 里合入 pipeline,三个引擎里哪个能最快跑起多卡并行?为什么?代价是什么?
-
SGLang Diffusion 与 vLLM-Omni 都有调度器,为什么没有一个做 LLM 式的连续批处理(请求随时进出 batch)?它们的批在什么条件下形成?
答案
扩散每步是整段序列的完整前向,batch 里的请求必须每步形状与执行策略相同才能拼成一个张量;且 compute-bound 下 batch 不提吞吐,动态进出的收益为零。所以两者都是同构批:SGLang 的
Scheduler._can_dynamic_batch要求采样参数签名(除 prompt、seed 等少数字段外的全部SamplingParams,含 height / width / frames / steps / CFG)全同、dynamic_batch_admission再管批大小与代价上限,vLLM-Omni 的StepBatchSamplingParamsKey同样要求形状、CFG、LoRA 等全同,默认整批开始整批结束;vLLM-Omni 打开--step-execution后,StepScheduler让兼容请求在步边界加入、走完就离开(各请求可处在不同步数),但仍要求兼容键相同,且只有实现了denoise_step的 pipeline 支持。详见第六章、第七章。 -
vLLM-Omni 对一个纯扩散模型也用 stage 结构(orchestrator + 一个 diffusion stage),看起来多了一层。这层在什么时候变成优势?
-
三个引擎的序列并行实现为什么高度相似?相似的部分具体是哪几块?
答案
共同来源:xDiT 的 USP 论文与实现(Ulysses all-to-all 换 head + Ring P2P 传 K/V 的二维组合),以及 vLLM 的
GroupCoordinator/parallel_state写法(按维度建 NCCL 组、rank 映射、collective 封装)。相似的块:parallel_state.py+group_coordinator.py(三者几乎同名同构)、按 token 切分与 all-to-all 重排(SGLangusp.py/sp_shard_utils.py,vLLM-Omnia2a_permute.py/sp_sharding.py,xDiTlong_ctx_attention/)、attention 层的 USP 包装。详见第六章、第七章。 -
你要在两台以太网互联的 8×L40 上部署 FLUX 服务。三个引擎各能提供什么、缺什么?
答案
以太网跨节点只有 PipeFusion 可行(第五篇)。xDiT 有 PipeFusion + Ulysses 的混合并行,但没有服务层(要自己包 API、队列、job)。SGLang Diffusion 有完整服务层与跨节点 SP(Ring),但没有 PipeFusion——跨节点 Ring 在以太网上每步 1.6 GB 不可行,只能节点内 8 卡 USP、两节点各自独立成实例(DP)。vLLM-Omni 有 PP(
pipeline_parallel.py)与服务层,PP 是否达到 PipeFusion 的 patch 流水线效果要实测。实际最常见的选择:放弃跨节点切请求,两节点各 8 卡 USP 独立服务,吞吐反而更高(第五篇:DP 是吞吐最优)。详见第七章。
-
SGLang Diffusion:
openai/image_api.py的/generations路由构造Req→AsyncSchedulerClient(ZMQ)→ rank 0 的Scheduler(合批准入、warmup,再广播给其它 rank)→ 各 rank 的GPUWorker→ComposedPipelineBase顺序执行TextEncodingStage/LatentPreparationStage/DenoisingStage(每步经 CFG policy、缓存决策、runtime/models的原生 DiT、USPAttention)/DecodingStage→OutputBatch回 HTTP →save_outputs编码成文件 → b64 或经storage.py上传返回 URL。vLLM-Omni:OpenAI 层 →AsyncOmni→ Orchestrator →InlineStageDiffusionClient(或StageDiffusionClient→StageDiffusionProc)→DiffusionEngine(RequestScheduler/StepScheduler)→MultiprocDiffusionExecutor→WorkerProc/DiffusionWorker→DiffusionModelRunner→models/flux/pipeline_flux.py+flux_transformer.py(并行组、缓存 hook、offloader)→output_formatter→ OpenAI 层返回。xDiT:无 HTTP;torchrun每个 rank 跑脚本 →xFuserArgs→xDiTParallel→xFuserFluxPipeline包装 diffusers 的FluxPipeline(initialize_model_parallel、把每个attn包成xFuserFluxAttentionWrapper、processor 换成xFuserFluxAttnProcessor、prepare_runwarmup)→ 调用 →save。详见第三章、第四章、第五章。 ↩ -
进程模型:SGLang 是 HTTP 进程 + 每卡一个 worker 进程(各含 Scheduler 与 GPUWorker,与 SGLang LLM 同构),vLLM-Omni 是 API 进程里的 orchestrator 编排 stage(单 stage 单副本时扩散引擎 inline 在 API 进程,否则是独立的
StageDiffusionProc,可跨 GPU / 主机),引擎内再有执行器与 worker 进程,xDiT 是 torchrun 的 SPMD rank。pipeline 抽象:SGLang 的ComposedPipelineBase= stage 列表 + 原生模型(diffusers 回退),vLLM-Omni 的models/<model>/pipeline_*.py原生 pipeline(diffusers adapter 回退),xDiT 包装 diffusers pipeline 只换 attention processor 与去噪循环。并行组:三者的parallel_state/GroupCoordinator同源,xDiT 多 PipeFusion 与 Parallel VAE,vLLM-Omni 多 PP / VAE patch,SGLang 多跨节点 SP;FSDP / HSDP 式的权重分片两者都有。调度:SGLangScheduler+dynamic_batch_admission,vLLM-OmniRequestScheduler/StepScheduler+ 兼容键,xDiT 无——两个有调度器的都是同构批,默认整批开始整批结束,只有 vLLM-Omni 的 step 模式允许兼容请求在步边界进出。详见第一章、第六章。 ↩ -
出发点不同:SGLang 想把已有的 serving 结构(scheduler、kernel 栈、warmup、CUDA graph、多平台)复用到新负载上,所以原生重写模型以接入自己的 layers;vLLM-Omni 要服务全模态模型(一个请求经过 LLM、DiT、TTS 多种模型),stage 流水线是必需的抽象,扩散只是一种 stage,原生重写是为了与 vLLM 的 worker / 量化 / KV 栈一致;xDiT 的目标是研究并验证并行方法(USP、PipeFusion),包装 diffusers 让它以最低成本覆盖最多模型、专注并行本身。代价对应:SGLang / vLLM-Omni 支持新模型成本高(重写 + 权重映射),xDiT 没有服务层且性能受 diffusers 限制。详见第七章。 ↩
系列 《扩散模型推理基础设施:从一次去噪到一个生成服务》 第 8 / 10 篇 深入
系列总览 — 为什么这样组织、读它需要什么、读完能做什么 · 幻灯片 — 整个系列的精华,一份可分享的 deck
学习地图 — AI-Infra 工程师学习地图 第 11 / 14 个系列
- 负载画像——一次生成在 GPU 上发生什么
- 单卡执行——attention 后端、编译、FP8 / INT4 与 offload
- 跨步冗余——TeaCache、First-Block Cache 一族的缓存与跳步
- 视频——长序列 attention 的账与稀疏化
- 多卡并行——序列并行、CFG 并行与 PipeFusion,为什么不是张量并行
- 少步与自回归——把步数变成系统参数
- serving 形态——请求、批、三段分离、附件、异步任务与成本
- 三个引擎的对照导读——同一张图的请求在 SGLang Diffusion、vLLM-Omni 与 xDiT 里各走过什么
- 配置、评测与排障——从一张卡的推导到一条伪影的排查
- 系列总结与通关自测
上一篇:serving 形态——请求、批、三段分离、附件、异步任务与成本下一篇:配置、评测与排障——从一张卡的推导到一条伪影的排查
下一个系列:AI 平台工程:资源层与交付层 进阶 · 9 篇,约 5.8 小时
本文由 arganzheng 创作,采用 CC BY 4.0 许可协议。在保留原文作者、署名以及完整原文链接(https://arganzheng.life/diffusion-engines-compared-sglang-diffusion-vllm-omni-xdit.html)的前提下,欢迎各种形式的转载、翻译或商业引用。
COMMENTS
评论存放在 GitHub Discussions, 用 GitHub 账号登录即可发表,支持 Markdown。 想针对正文某句话说?选中那段文字,点浮出的「评论」即可划线评论;觉得哪里写错了,发表时勾上「同时提交 Issue」。 有人回复你时 GitHub 会按你的通知设置发邮件,不用守在这里。