AI资讯 / 研究

研究 / 官方

AR+DiT 混合架构的工程实践

LMSYS Blog

混合自回归扩散(AR+DiT)生成框架将 AR 模型的全局语义理解与扩散模型的局部细节精炼相结合,在海报、信息图表等知识密集型视觉任务上表现出色。以 GLM-Image 为代表,该框架先由 9B 视觉语言模型自回归生成语义先验 token,再由 7B DiT 模型经 30-50 步去噪生成高分辨率图像。然而在 SGLang 中高效部署此类混合管线面临严峻挑战:AR 与 DiT 共享同一进程导致架构耦合,无法独立扩展;串行处理并发请求造成算力大量闲置;单体部署迫使 DiT 在非最优配置下运行。为此,Ascend 团队贡献了三个递进式 PR,将系统从单体架构演进为完全解耦的异构分布式架构,在 4 卡 NPU 配置下端到端延迟相比单卡基线降低 77.2%,AR 阶段延迟降低 78.8%。

AR+DiT 混合生成框架的核心思路是「先规划后绘制」:自回归模型负责从文本提示中提取高层语义,扩散模型则在此基础上迭代精炼出高质量图像。这一设计在文字排版、信息图表等端到端扩散模型难以胜任的场景中取得了最优水平的效果。然而,原生部署将 AR 编码器、DiT 去噪器和 VAE 解码器链式耦合在单一工作进程中,带来了架构耦合、资源利用率低和资源分配失衡三大核心痛点。

第一阶段优化(PR #25381)将 AR 后端从 HuggingFace 替换为 SRT,把 AR 阶段从扩散工作进程中解耦为独立的 SRT 服务。AR 与 DiT 分别独立加载权重、拥有各自的调度生命周期,并可独立扩展。AR 服务器现可自行配置张量并行(TP),不再受 DiT 空间并行(SP)策略的约束。这一「独立服务器 + HTTP 调用」的方案将原本复杂的 VLM 重写任务转化为对现有基础设施的复用,大幅降低了耦合度。实测数据显示,仅软件层面的解耦即可将单卡 AR 阶段延迟从 122.8 秒降至 46.6 秒,降幅达 62.1%。

第二阶段优化(PR #30683)针对高并发场景下延迟随请求数线性增长的问题,引入动态批处理与提前返回机制。通过实现 supports_dynamic_batching 和 supports_native_grouped_requests 接口,系统可将并发请求打包为单次前向传播,消除串行执行带来的算力闲置。经过评估,动态批处理仅应用于 AR 阶段——DiT 的每步延迟与批大小成正比,批处理对其吞吐量无净收益。此外,新增的提前返回机制允许每张图像完成后立即返回结果,而非等待整批处理完毕。

第三阶段优化(PR #31320)实现了 AR 与 DiT 工作流的解耦与扇出架构。该方案采用「每设备一个去噪器」的并行 DiT 执行模式,并通过缓冲 AR 结果将 AR 与 DiT 的工作流重叠执行。这一架构使 DiT 能够在 batch=1 的最优配置下运行,同时通过多设备并行最大化吞吐量,从根本上解决了单体部署中资源分配失衡的问题。

从整体性能数据来看,优化效果显著。在 4 卡 NPU 异构配置(AR 采用 TP=4,DiT 采用 SP=4)下,端到端延迟从单卡基线的 154.6 秒降至 35.2 秒,降幅达 77.2%;AR 阶段延迟从 122.8 秒降至 26.1 秒,降幅 78.8%;去噪阶段延迟从 31.6 秒降至 9.0 秒。值得注意的是,旧架构中 AR 阶段在 2 卡 SP 配置下不仅无法受益,反而因通信开销出现 4.1% 的性能回退,而 SRT 路径则真正实现了多卡 TP 加速。这一系列优化为 SGLang 中高效部署混合生成模型提供了可复用的工程范式。

要点

  • 将 AR 后端从 HuggingFace 替换为 SRT 并解耦为独立服务,仅软件层面即可将 AR 阶段延迟降低 62%,配合 4 卡 NPU 并行可实现 78.8% 的降幅
  • AR 与 DiT 对并行策略的需求截然不同:AR 适合张量并行(TP),DiT 适合空间并行(SP)且在 batch=1 时效率最高,单体架构无法同时满足两者
  • 动态批处理应选择性应用——对 AR 阶段有效,对 DiT 阶段因延迟与批大小成正比而无净收益
  • 「每设备一个去噪器」的扇出架构配合 AR 结果缓冲,可将 AR 与 DiT 工作流重叠执行,进一步提升整体吞吐量
  • 三阶段递进式工程优化路径(解耦→动态批处理→异构分布式)为类似混合生成管线的部署优化提供了可参考的方法论
查看原始来源

原始标题:Full-Stack Performance Optimization of AR+DiT in SGL-Diffusion

本文由 DataHub 基于公开来源整理,用于信息发现与摘要阅读;具体事实、数据和后续更新以原始来源为准。