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

2026年02月04日 3 约9分钟 LMSYS

LMSYS SGLang dLLM LLaDA 2.0 扩散大模型推理优化

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 推理引擎，具有多重优势：

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

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

设计

关键洞察

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

由于 Bidirectional Attention Diffusion 计算成本高且 KV Cache 利用率低，主流 dLLM 正转向 Block Diffusion 架构。
Block Diffusion 计算模式与 SGLang 现有 Chunked-Prefill 过程高度相似。
不同于 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 的功能扩展。单次前向涉及两部分计算，最终输出拼接：

Context Query：当前 Q_curr 与现有 KV Cache 双向注意力，与 Chunk Prefill 完全相同，确保当前 block 关注历史信息。
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<"