SGLang 赋能扩散大模型：即日支持 LLaDA 2.0

5 约9分钟 LMSYS
我们兴奋地推出 SGLang 中的 Diffusion Large Language Model (dLLM) 框架设计与实现。通过利用现有的 ChunkedPrefill 机制，该系统实现了无缝集成、无需核心架构变更、继承现有推理优化，并提供最大灵活性让用户自定义扩散解码算法。LLaDA 作为首款 dLLM，展现出优异的数据理解能力和更快推理速度，尤其在低延迟小批量场景。面对大规模 dLLM 如 100B LLaDA2.0-flash 的评估与 RL 后训练挑战，现有的推理引擎不足以支撑。我们在 SGLang 中引入 Block Diffusion 支持，利用 Chunked-Prefill 管道，仅微调关键组件，实现高效批处理与流式输出，显著提升吞吐量（如 LLaDA2.0-flash-CAP 达 935 tokens/s）。（128字）

TL;DR

我们兴奋地介绍 SGLang 内 Diffusion Large Language Model (dLLM) 框架的设计与实现。通过利用现有的 ChunkedPrefill 机制，该系统实现了：

  • 无缝集成：内置于 SGLang 生态，无需核心架构变更。
  • 继承性能：受益于现有推理优化技术。
  • 最大灵活性：用户可完全自定义扩散解码算法。

背景

动机

今年年初，LLaDA 作为首款 Diffusion Large Language Model (dLLM) 亮相，迅速吸引学术界与产业界的关注。这项由中国人民大学与蚂蚁集团合作成果，展示了 dLLM 独特执行范式的优异数据理解能力，且在低延迟小批量场景下，推理速度超越 Auto-Regressive (AR) 模型。

随着 dLLM 参数规模扩大，我们观察到类似 AR LLM 的 scaling-law 效应。为追求更强 dLLM，我们训练了 100B 参数的 LLaDA2.0-flash 模型。

然而，在训练 LLaDA2.0-flash 过程中，我们面临一系列 AI 基础设施工程挑战，尤其是模型评估与 RL 后训练的效率与稳定性。

挑战

现有 dLLM 推理引擎无法满足大规模模型的评估与 RL 后训练需求。例如，Fast-dLLM 虽优秀，但更适合算法研究，缺乏生产级服务能力，如批处理、调度、RL 生态集成与并行。

相比之下，SGLang 是当下最受欢迎的 LLM 推理引擎，具有多重优势：

  1. 生产就绪：已在数千家公司部署，提供成熟可靠工程能力。
  2. 技术领先：集成众多先进推理优化，社区持续贡献新优化。
  3. 完整生态：与 RL 后训练生态高度集成，尤其分布式权重 GPU P2P 更新。

但 SGLang 目前仅支持 Auto-Regressive 计算范式，未适配扩散计算。因此，挑战在于：在不破坏现有架构前提下，为 dLLM 引入支持，让其受益 SGLang 所有优化。

设计

关键洞察

基于 dLLM 当前发展，我们提炼几点关键洞察：

  1. 由于 Bidirectional Attention Diffusion 计算成本高且 KV Cache 利用率低，主流 dLLM 正转向 Block Diffusion 架构。
  2. Block Diffusion 计算模式与 SGLang 现有 Chunked-Prefill 过程高度相似。
  3. 不同于 AR 模型，扩散语言模型需多种解码策略，支持专用接口灵活自定义。

架构

我们利用 SGLang 现有 Chunked-Prefill 管道，实现 Block Diffusion LLM 计算支持。该方法无缝集成 dLLM，无需改动核心框架，直接受益 SGLang 积累的推理优化。

主要执行流程图

如图所示，我们对 SGLang 的修改极为克制，仅触及核心边缘。原有 generate request 执行流不变，主要利用并修改 Chunked Prefill，聚焦 prefill adder 与 chunked reqs 两组件。

SGLang 中 Chunked Prefill 旨在最大化 GPU 利用率，单 chunk 大小通常 2K-16K tokens。但 dLLM 解码以 block 级分割（如 LLaDA2.0 的 32 tokens block）。若沿用单大请求逻辑，将浪费 GPU 性能，故需解决高效批处理。我们修改 chunked reqs 与 prefill adder，支持单计算周期内多 Diffusion Block 处理。

此外，在解码执行层，我们在 TP Worker 与 Model Runner 间插入扩散算法抽象层：

  • Worker 识别 Diffusion 模型时，进入专用分支。
  • 调用 Diffusion 算法的 run 函数。
  • 内部通过前向迭代循环驱动 Model Runner，直至完整 Block 解码。

Attention Mask

因果掩码对比图

Block Diffusion 与 Chunk Prefill 单次前向传播的最大差异在于 attention mask 处理：

  • Block Diffusion 使用 block-wise causal mask。
  • AR 模型 Chunk Prefill 使用传统 token-wise causal mask。

Block Diffusion 可视为 Chunk Prefill 的功能扩展。单次前向涉及两部分计算，最终输出拼接：

  1. Context Query：当前 Q_curr 与现有 KV Cache 双向注意力，与 Chunk Prefill 完全相同，确保当前 block 关注历史信息。
  2. Intra-Block Query：当前 Q_curr 与自身 KV 计算。
    • Block Diffusion 使用双向注意力。
    • Chunk Prefill 使用 causal Mask。

可视化 Q_curr attention mask：

  • Chunk Prefill (causal) 为梯形/三角掩码。
  • Block Diffusion (双向) 为矩形掩码。

流式输出演示

以下动画对比 LLaDA2.0-flash-CAP (100B / BF16) 与 gpt-oss-120B (117B / MXFP4) 的流式输出。LLaDA2.0-flash-CAP 使用 SGLang dLLM TP8 于 8×H20，gpt-oss-120B 使用标准 AR 流程同硬件。

任务：在 10 种编程语言实现快速排序——扩散 LLM 擅长的场景。如图，LLaDA2.0-flash-CAP 吞吐量达 935 tokens/s，远超 gpt-oss-120B 的 263 tokens/s。

LLaDA2.0-flash-CAP 与 gpt-oss-120B 输出对比动画

SGLang dLLM 支持如 AR 模型般的流式输出，但以 block (如 32 tokens) 为单位输出。

dLLM 流式输出动画

使用方法

启动命令示例

python3 -m sglang.launch_server \
  --model-path inclusionAI/LLaDA2.0-mini \ # 示例 HF/本地路径
  --dllm-algorithm LowConfidence \
  --dllm-algorithm-config ./config.yaml \ # 可选，使用算法默认配置
  --host 0.0.0.0 \
  --port 30000
注意：使用 --dllm-algorithm-config 高级配置选定 --dllm-algorithm。此特性将配置与代码解耦，便于用户自定义算法统一入口传参。

客户端代码示例

如其他支持模型，dLLM 可通过 REST API 或离线引擎 API 使用。

Curl 生成请求示例：

curl -X POST "http://127.0.0.1:30000/generate" \
     -H "Content-Type: application/json" \
     -d '{
        "text": [
            "<role>SYSTEM</role>detailed thinking off<|role_end|><role>HUMAN</role>Write the number from 1 to 128<|role_end|><role>ASSISTANT</role>",
            "<role>SYSTEM<"
KTransformers加速SGLang的混合推理 02-04 SGLang-Diffusion 两个月的进展 02-04 SGLang Pipeline Parallelism：百万Token上下文扩展与性能突破 02-04