跳到正文
LMSYS Blog·· 1 天前精选AI 评分64

SGLang 为 816B 开源多模态 MoE 模型 Muse Spark Open 提供 Day-0 支持

Blog SGLang and Miles Add Day-0 Support for Muse Spark Open - Model: Muse Spark Open is an open-weight multimodal mixture-of-experts model. It uses an 816B-parameter text transformer backbone, with approximately 42 billion active parameters per token. It accep... SGLang Team October 15, 2026

AI 导读

SGLang 团队宣布对 Meta 开源多模态 MoE 模型 Muse Spark Open(816B 参数、约 42B 激活、1M token 上下文)提供 Day-0 推理支持,权重提供 NVFP4 和 BF16 两种格式。

推荐理由

原文给出了多平台推理与 RL 训练支持的完整细节和实测吞吐数据,读者可以据此评估大规模部署该模型的可行路径。

正文 · AI 翻译

亮点

  • 模型:Muse Spark Open 是一个开放权重的多模态混合专家模型。它采用 816B 参数的文本 Transformer 主干,每个 token 约有 420 亿激活参数。它接受文本、图像、视频和音频输入,并支持 1M token 上下文窗口。检查点提供 NVFP4 和 BF16 格式。
  • 性能:在 OpenHands 多轮智能体编码回放中,使用 NVFP4 检查点、运行于 8 块 NVIDIA B300 GPU 上,SGLang 在并发为 1 时达到每 GPU 每秒 14.8K 总 token,每用户每秒 543 token(p50),在并发为 128 时达到每 GPU 每秒 80.3K 总 token,每用户每秒 61 token(p50)。
  • 硬件:目标部署平台包括 NVIDIA Blackwell GPU(GB300、B300、GB200、B200)、NVIDIA H200、AMD MI350X、MI355X,以及搭载 P-core 的 Intel® Xeon® 6。
  • 特性:序列并行 + NVFP4 MegaMoE 加速预填充,多模态预处理优化,以及 DFlash V2 投机解码提升交互性。

Muse Spark Open 架构

Muse Spark Open 将稀疏语言解码器与视觉和音频编码器相结合,把多模态输入引入共享的 token 生成流水线。

Muse Spark Open maps text, image/video and audio inputs into a shared decoder with GQA and sliding-window/global attention; its MoE block routes full-width inputs while pooling expert features, normalizing each expert output before the weighted sum and restoring the model width with an unpooler.

多模态编码器为共享语言解码器提供输入。展开的 MoE 模块展示了聚合前的特征宽度压缩和逐专家归一化。

GQA 与滑动窗口注意力

针对其 1M token 上下文窗口,Muse Spark Open 将分组查询注意力(GQA)与滑动窗口和全局注意力的混合机制相结合。GQA 在查询头分组之间共享键/值头,从而减少 KV 缓存占用。滑动窗口层关注有界的近期上下文,而周期性的全局层则保留对完整序列的访问。这在无需每一层都维持全上下文注意力的前提下,平衡了局部注意力效率与长距离信息流动。

文本、音频、图像和视频

Muse Spark Open 将文本、音频、图像和视频作为输入处理并送入其语言解码器。图像和视频帧共享一个 NaViT 视觉编码器,而波形音频则通过专用的音频编码器。它们的输出被投影到解码器的隐藏维度,并与文本嵌入并列放置,使语言模型能够生成基于多种模态的响应。

带池化器、反池化器和逐专家归一化的 MoE

MoE 模块在其路由的 SwiGLU 专家周围加入了学习得到的池化器和反池化器投影。池化器将每个 token 的隐藏向量从模型宽度压缩到较小的专家宽度;反池化器在专家聚合后恢复模型宽度。这些投影在改变特征宽度的同时保持 token 数量不变。路由器从全宽度表示中选择专家,池化后的专家输入在执行专家前进行归一化。

一个关键细节是逐专家 RMSNorm:每个被选中专家的输出在乘以路由权重并与其他专家求和之前先进行归一化。随后对合并后的向量进行 unpool 操作。这一顺序很重要,因为对合并后的输出进行归一化会改变模型的计算。

功能支持

运行时优化

序列并行

序列并行(SP)沿 token 维度将激活值划分到各个张量并行(TP)rank 上,因此每个 rank 只对序列中属于自己的部分执行 token 本地操作。我们将 Korthikanti 等人,Reducing Activation Recomputation in Large Transformer Models 中描述的方法适配到 Muse Spark Open prefill。

这一点很重要,因为仅靠张量并行会让 Muse Spark Open 的 pooler、unpooler 和边界层归一化在每个 rank 上处理相同的完整 token 批次。使用 SP 后,这 P 个 rank 中的每一个在这些区域处理约 S/P 个 token,同时保留每个 token 的完整隐藏向量。这将投影和归一化工作分散到各个 rank 上,并减少长提示 prefill 期间冗余的激活存储。

在注意力输出处,reduce-scatter 生成 token 分片。归一化、路由和池化在这些分片上本地执行。NVFP4 MegaMoE 随后通过其集成的 all-to-all dispatch 和 combine 将序列并行与专家并行连接起来:dispatch 将每个 token 发送到持有其被选中专家的 rank;专家计算其输出,并在路由加权之前应用逐专家 RMSNorm;combine 将完整的加权结果返回到该 token 的原始 rank。Unpool 及随后的归一化/残差操作保持序列分片状态,直到为下一个 QKV 投影提供输入的 all-gather。

Four TP/EP ranks keep token shards, dispatch them to NVFP4 MegaMoE expert owners, and combine normalized expert outputs back to the original ranks before local unpooling.

我们在四块 GB300 GPU 上测量了三种配置,使用 NVFP4 目标模型、DFlash V2 草稿模型、TP4/EP4 拓扑、FP8 KV 缓存和 16K prefill 分块。在每个测试点,我们发送八个请求组,每组包含八个请求,名义输入 token 数为 80K,输出 token 数为 20。同一组内的请求共享一个 73,360 token 的生成前缀;不同组使用独立生成的前缀。最大客户端并发数从 1 变化到 64。

p90 TTFT for TP4, TP4 with prefill SP, and TP4 with prefill SP and MegaMoE across maximum concurrency 1 to 64 on four GB300 GPUs with DFlash V2.

在此次扫描中,SP + MegaMoE 在所有七个并发级别下将 p90 TTFT 相比 TP4 降低了 28.3–41.9%。在并发数为 64 时,p90 TTFT 从 29.03 秒降至 19.47 秒,降低了 32.9%。

多模态优化

Muse Spark Open 的多模态流水线在 CPU 上使用 Rust 或 Pillow 缩放对图像和视频帧进行预处理,随后使用 SGLang 优化的注意力后端进行视觉编码。当相同图像在对话轮次中重复出现时,会复用缓存结果,避免重复计算。

SGLang 为 Muse Spark Open 支持独立的编码器、prefill 和解码 worker(EPD Disaggregation)。编码器处理图像和视频,然后将生成的特征传递给 prefill worker。解码 worker 接收 KV 缓存并生成响应。分离这些阶段使部署能够根据工作负载将资源分配给媒体处理、提示处理和生成。

使用 DFlash V2 的投机解码

SGLang 为 Muse Spark Open 支持使用 DFlash V2 的推测解码,构建在 DFlash 的并行草稿架构之上。独立的草稿模型使用 Muse Spark Open 的隐藏状态来填充自己的 KV 缓存,然后在单次前向传播中提出一个未来 token 块。它共享目标的 token 嵌入和 LM head,Muse Spark Open 在一次前向传播中验证所提出的块,允许多个被接受的 token 一起推进。

DFlash V2 增加了动态卷积,在草稿块内跨相邻位置混合信息,并加入了一个轻量级候选选择器,考虑所提出 token 之间的依赖关系。SGLang 将草稿和候选选择集成到同一个 CUDA graph 中,并支持重叠调度以减少 GPU 操作之间的 CPU 调度开销。此路径可通过 --speculative-algorithm DFLASH 使用兼容的 Muse Spark Open DFlash V2 检查点。

对于采样的 DFlash V2 提议,SGLang 还支持 块验证,通过 --speculative-use-block-verification 启用。它使用目标和草稿概率联合评估草稿前缀,可以接受比逐 token 验证更长的前缀。验证器保留最长被接受的前缀,并在部分接受后采样一个修正 token,或在完全接受后采样一个目标模型奖励 token,从而保持目标采样分布。

兼容性

为了提升真实服务场景中的性能,我们还支持 Muse Spark Open 配合以下 SGLang 服务特性:

  • 预填充/解码分离(PD 分离): 将预填充和解码工作进程分离,并在它们之间传输 KV 状态,可减少两个阶段之间的干扰,有助于满足首 token 时间(TTFT)和每输出 token 时间(TPOT)的服务级别目标(SLO)。
  • HiCache: GPU/主机 KV 缓存层次结构通过更高的前缀缓存命中率提升服务性能,复用缓存的全注意力和滑动窗口状态,避免重复的预填充计算。
  • 可中断 CUDA graph: 捕获的预填充计算在 eager 注意力调用周围重放,减少 CPU 启动开销,同时保持注意力边界的灵活性。

内核优化

我们减少了内核启动和内存流量,替换了 BF16 专家循环,将复制的投影分布到多个 GPU 上,并重叠独立工作。该图显示了这些更改在 TP4/EP4、8K 输入 token、1K 输出 token 且无推测解码时的累积效果。

  • 融合逐元素链。 我们将 FFN 后归一化、门控残差加法和注意力前归一化合并为一个内核。Q/K 归一化通过行步长直接读取融合后的 QKV 张量,避免布局复制,同时查询缩放被折叠到注意力 softmax 缩放中。

  • 分两步替换 BF16 专家循环。 第 0 层和第 61 层使用 BF16 专家权重。我们首先将 64 个本地专家批处理为分组 GEMM,将主要中间结果限制在 64 MiB。然后我们切换到 FlashInfer 的 trtllm_bf16_routed_moe,仅计算被路由的专家工作,并在加载时将专家后归一化绑定到 finalize 操作。

  • 融合归一化并移除路由转换。 CUDA 边界内核在两次 RMS 归约过程中将每一行保留在寄存器中。专家前归一化被折叠进 NVFP4 量化,复用量化器的行访问,避免单独的读/写遍历。路由还以专家内核所消费的格式输出打包的 ID 和权重。

  • 融合 RoPE 和 KV 缓存更新。 Q/K 归一化同时应用 RoPE 并写入解码 KV 缓存,消除独立内核和重复的内存遍历。

  • 将路由和 norm/gate 移至 CUDA。 CUDA 内核处理 top-k 路由,并将逐头归一化与门控结合,包括快速倒数实现。此步骤还启用了共享图流。Verify-KV 融合有利于推测验证路径;在此无推测基准测试中它处于非活动状态。

  • 在行内并行化 RMS 归一化。 一组 CTA 通过分布式共享内存协作,比每行单个 CTA 暴露更多并行性。

  • 先分片 unpooler,再分片 pooler。 这些投影最初在每个张量并行 rank 上重复完整的矩阵计算。现在每个都使用 BF16 调度器计算输出的一部分,随后通过 all-gather 组装完整结果。我们对 pooler 分片使用 TGV,通过双流重叠改进了最终配置。

  • 重叠独立的 MoE 工作。 路由器 GEMM 和 top-k 选择可以独立于 pooler 和激活准备运行。我们将路由放在侧流上,并在专家执行之前立即汇合两个分支。

  • 调优 BF16 投影调度。 我们为解码期间使用的小 token 数量和矩阵形状选择调优过的 GEMM 实现。Warp split-K 内核通过将归约维度拆分到多个 warp 来加速投影。

  • 更早启动依赖内核。 我们通过将启动触发点移到小内核最后一次输入加载之后,为它们调优 PDL。

Cumulative kernel optimizations at BS1, TP4/EP4, 8K input and 1K output tokens, without speculative decoding

MegaMoE 功能与优化

DeepGEMM 的 MegaMoE 将调度、两个专家 GEMM、激活和合并融合为一个持久内核。跨 rank 传输使用对称内存,使通信能够与计算重叠,而无需单独的 all-to-all 启动。

支持 Muse Spark Open 需要 60 层中的 NVFP4 专家权重和激活、两层中的 BF16 专家,以及在路由加权求和之前进行逐专家 RMS 归一化。MegaMoE 还必须读取 TensorRT-LLM 的常驻权重布局,以便两个后端可以共享权重。

NVFP4 作为第三种 MMA 类型

NVFP4 路径在 fp8xfp4 和 mxfp4xmxfp4 之外增加了 nvfp4xnvfp4。它保留稠密 FP4 操作数,但使用每 16 个元素的 UE4M3 块缩放,而不是 MXFP4 的每 32 个元素的 UE8M0 缩放。逐 token 的 FP32 输入缩放应用于第一个 GEMM 的累加器,并且预调度支持 BF16 到 NVFP4 的量化。

此外,作为此扩展的一部分,我们重构了该类。现在 MMA 类型在对称缓冲区分配时显式指定,并在整个执行路径中传递,使格式选择更易于理解。

融合的逐专家后归一化

每个专家输出必须在应用其路由权重之前进行归一化。由于第二个 GEMM 将隐藏行分布到多个 CTA,其尾声无法独立计算整行 RMS。因此归一化在合并阶段运行,此时对等写入在 NVLink Barrier 之后可见。

后归一化部分和:为了避免仅仅为了在 combine 中计算其 RMS 而重新读取完整的专家输出行,每个第二个 GEMM 分块会从它已保存在寄存器中的 BF16 值累加平方和的部分和。这些部分和在 token 的原始 rank 上原子累加;随后 combine 为每个选中的专家读取一个标量。

Post Expert Norm

读取 TensorRT-LLM 的权重

共享 TensorRT-LLM 的常驻缓冲区可避免重复的权重和缩放因子,同时保留在解码时使用其 MoE 内核的选项。每种布局需要不同的处理方式:

  • NVFP4 权重:第一个 GEMM 的 row ^ 8 shuffle 会在每个 16 行组内交换 gate 和 up 两半。MegaMoE 按存储原样加载分块,并在 SwiGLU 尾声中选择对应的寄存器槽位,用一次加载取代了每个 K 原子十六次小型 TMA 加载的朴素方式。第二个 GEMM 的行置换由其 tensor-map 描述符处理。

  • NVFP4 块缩放因子:一个 3D tensor map 直接读取第一个 GEMM 的缩放因子;一个 5D map 还会撤销第二个 GEMM 的组内置换。调整共享内存源偏移量以保留现有的 UTCCP 4×32 转置。

  • BF16 权重:TensorRT-LLM 使用 BlockMajorK 存储,物理布局为 [E, K/64, N, 64]。一个 5D tensor map 通过每个 32 行块一次 TMA 加载来处理第一个 GEMM 的组反转,而 SwiGLU 吸收剩余的 gate/up 交换。第二个 GEMM 使用一个 4D map。

MegaMoE reads TRT-LLM layouts

TRT-LLM 与 MegaMoE 的权重和缩放因子布局。TMA 寻址和 SwiGLU 寄存器选择在不复制常驻缓冲区的情况下协调了这些布局。

NVFP4 重新量化倒数表

在两个 GEMM 之间,SwiGLU 输出被重新量化为 NVFP4。此前,每个 16 元素组使用 __frcp_rn 计算块缩放因子的倒数。__frcp_rn 并不便宜:它是 IEEE 就近舍入,因此包含一个 MUFU 种子、一个牛顿精化链以及特殊情况处理。由于缩放因子字节只有 256 种编码,共享内存中的一张 256 项 FP32 表用查表取代了重复的转换和倒数计算。在孤立的 B300 微基准测试中,共享内存表将缩放因子到倒数的计算在分散访问下加速约 2.9 倍,在 warp 一致访问下加速 4.8 倍。

性能结果

OpenHands 智能体编码扫描

我们在八块 B300 GPU 上以 TP8/EP8 配置,使用 OpenHands 多轮智能体编码工作负载,收集了 Muse Spark Open NVFP4 的性能结果。每段对话以一个约 74K token 的提示开始,每轮输出约 220 个 token,共 13 轮。连续轮次复用前缀,总体缓存命中率约 92%。我们将并发从 1 扫描到 128 段对话,并绘制每 GPU token 吞吐量(tok/s/GPU)与 p50 交互性(tok/s/user)的关系。

Muse Spark NVFP4 throughput per GPU versus p50 interactivity on eight B300 GPUs with DFlash v2.

AgentX 智能体编码扫描

我们在八块 B300 GPU 上以 TP8/EP8 配置,使用 AgentX——SemiAnalysis 用于回放记录的多轮智能体编码轨迹的工具——收集了 Muse Spark Open NVFP4 的性能结果。AgentX 复现了真实编码智能体会话的提示大小、前缀共享、子智能体扇出和轮次间时序,因此结果反映了“真实的智能体负载”。会话中的请求共享一个长前缀,因此 96.9% 到 98.0% 的输入 token 命中前缀缓存。我们将并发从 1 扫描到 48 个会话,并绘制每 GPU token 吞吐量(tok/s/GPU)与 p90 交互性(tok/s/user)的关系。在并发 48 时,吞吐量达到 34,509 tok/s/GPU,交互性为 73.0 tok/s/user;单个会话运行速度为 457.8 tok/s/user。

Muse Spark NVFP4 total throughput per GPU versus p90 interactivity on eight B300 GPUs with AgentX.

多平台支持

AMD 支持

Day-0 启用

  • 平台:MI350X、MI355X Instinct GPU
  • 两个 Meta 模型 ckpt 均受支持:meta-models/Muse-Spark-Open-816B-A42B(BF16)和 meta-models/Muse-Spark-Open-816B-A42B-NVFP4。
  • 支持 AMD/Quark aot 量化 MXFP4 模型:amd/Muse-Spark-Open-816B-A42B-Quark-MXFP4——一个与 Meta NVFP4 模型逐项对应的 MXFP4 量化模型。
  • 注意力内核:用于 decode 的 aiter unified-attention 3D,配合 page-16 CK batch-prefill。
  • FlyDSL MoE:expert GEMM 加上一个 finalize,在单个内核中完成 unpermute、per-expert RMSNorm 和加权求和。
  • 内核融合在 decode 路径上:跨层 FFN 边界、融合 Q/K RMSNorm、split-K router GEMV、用于 pooler/unpooler 的 aiter 调优 GEMM。
  • Prefill 优化:将 atomic-add expert sort 替换为 LDS 内核。
  • 需要 aiter 补丁:page-scatter VM fault 修复、unified-attention 3D 以及 FlyDSL Muse MoE 内核。docker/muse_spark/muse-spark-rocm.Dockerfile 会在 aiter 内核预构建之前应用它们,因此由它构建的镜像无需运行时打补丁。仅在基于纯 ROCm 镜像构建的容器中使用 3rdparty/amd/muse_spark/apply_aiter_patches.sh。

Day-0 性能(MI355X | MI350X)

amd/Muse-Spark-Open-816B-A42B-Quark-MXFP4,8x MI355X TP8/EP8,IL8192/OL1024,随机范围比例 0.8,每点 4x 实测请求,radix cache 开启(带预热运行)。

triton attn + triton moe(基线):

并发数平均 ITL(ms)平均 TPOT(ms)总吞吐量(tok/s)
414.5514.552,384.71
815.4315.424,317.23
1617.8417.847,460.37
3222.1422.1312,396.21

aiter attn + flydsl moe:

并发数平均 ITL(ms)平均 TPOT(ms)总吞吐量(tok/s)
412.2912.292,812.75
813.8813.874,777.36
1616.0416.038,311.38
3219.7619.7513,880.93

aiter attn + flydsl moe + DFlash V2(真实接受率;ITL:每验证步):

并发数平均 ITL(ms)平均 TPOT(ms)总吞吐量(tok/s)
422.167.194,544.71
827.118.467,542.14
1632.4210.6212,056.10
3245.4215.1117,081.46

准确率

基准测试MXFP4(TP4)MXFP4(TP8)NVFP4 Requant(TP4)BF16(TP8)NVFP4 参考(GB300)
GSM8K97.42%97.57%97.95%97.73%97.80%
AIME26(30×8)pass@1 95.83%,pass@8 100%pass@1 95.83%,pass@8 100%pass@1 95.42%,pass@8 100%pass@1 95.83%,pass@8 100%pass@1 95.83%,pass@8 100%
GPQA Diamond89.90%89.90%88.38%88.89%90.04%
MindCube tinybench81.33%80.19%81.24%82.86%81.62%
Video-MME short89.67%88.89%89.11%88.78%89.78%

参考来源为 GPQA Diamond 和 Video-MME 的模型卡、MindCube 已发布的 SGLang 测试结果,以及 GSM8K 的接受门槛。MindCube 使用了 SGLang 测试配置:ff_rsn、reasoning_effort=high、300k token 预算。一段 24 kHz 波形通过 input_audio 和 audio_url 均被正确转录,且无音频对照组正确拒绝。

所有模态均已确认可用:文本、图像、视频、音频和长上下文。

服务器脚本

export SGLANG_USE_AITER=1
export SGLANG_USE_AITER_UNIFIED_ATTN=1
export MUSE_MOE_BACKEND=flydsl

# BF16 model:  meta-models/Muse-Spark-Open-816B-A42B
# MXFP4 model: amd/Muse-Spark-Open-816B-A42B-Quark-MXFP4
python3 -m sglang.launch_server \
  --model-path {MODEL_PATH} \
  --tp-size 8 --ep-size 8 \
  --attention-backend aiter --page-size 16 \
  --kv-cache-dtype fp8_e4m3 \
  --mem-fraction-static 0.85 \
  --max-running-requests 128 \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path RadixArk/Muse-Spark-Open-DFlash2 \
  --speculative-draft-attention-backend triton \
  --speculative-num-draft-tokens 8 \
  --reasoning-parser muse --tool-call-parser muse \
  --host 0.0.0.0 --port {PORT}

# NVFP4 model: meta-models/Muse-Spark-Open-816B-A42B-NVFP4
python3 -m sglang.launch_server \
  --model-path {MODEL_PATH} \
  --quantization quark_mxfp4 \
  --tp-size 4 --ep-size 4 \
  --attention-backend aiter --page-size 16 \
  --kv-cache-dtype fp8_e4m3 \
  --mem-fraction-static 0.85 \
  --max-running-requests 128 \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path RadixArk/Muse-Spark-Open-DFlash2 \
  --speculative-draft-attention-backend triton \
  --speculative-num-draft-tokens 8 \
  --reasoning-parser muse --tool-call-parser muse \
  --host 0.0.0.0 --port {PORT}

Intel Xeon 平台支持

SGLang 在 Intel® Xeon® CPU 平台上支持高效的 INT4 Weight-Only Quantization(WOQ)推理,显著降低模型内存需求,同时利用 Intel® Advanced Matrix Extensions(Intel® AMX)加速 AI 计算。为在不影响模型质量的前提下最大化效率,Intel 使用 Intel AutoRound 量化算法直接从 Muse Spark Open 的原始 BF16 检查点生成 INT4 变体,该算法已在 SGLang 中得到良好支持。

为充分实现 Muse Spark Open 的 INT4 推理优势,SGLang CPU 利用 Intel AMX 加速 Attention(滑动窗口)、BF16 GEMM 和 INT4 WOQ FusedMoE 计算。此外,针对其他一些组件的融合优化,包括所有归一化、RoPE、MoE 路由和 finalizing 操作等,进一步提高了吞吐量并降低了推理开销。

鉴于 Muse Spark Open 模型的规模,SGLang CPU 为在 Intel Xeon 平台上进行离线推理提供了一条实用的部署路径,尤其是在配备 P-core 的 Intel® Xeon® 6 上。借助 Intel Xeon 平台可用的大容量内存,组织可以高效地运行私有离线推理,并在现有数据中心基础设施上扩展部署。

TPU 支持

我们首次通过 SGLang-TorchTPU 为开放权重模型带来第 0 天 TPU 支持,这是一个原生 PyTorch 后端,用于在 TPU 上进行 LLM、VLM 和扩散模型推理。它将 TPU 加速与 SGLang 的 RadixAttention 和 HiCache 等特性相结合,为智能体工作负载实现高性价比、生产就绪的服务。

Muse Spark Open 的 TPU 配方以及 SGLang-TorchTPU 的访问权限即将推出。通过此 Google Form 注册以获取最新更新。

强化学习

可训练的 Muse Spark Open 后端

Miles 通过 Megatron-Core 为 Muse Spark Open 提供分布式训练支持,涵盖该模型的压缩 MoE、归一化和残差门控。每个专家输出独立归一化,然后以 FP32 按路由加权求和进行组合。训练后端保留此计算顺序和这些精度边界,并支持全参数更新的反向传播。

Muse Spark Open 复用 Miles 现有的 Muon 和 Adam 支持进行全参数 RL。Miles 的分布式 Muon 和 Adam 实现支持分块 GPU 更新以及 CPU 或 NVMe 优化器状态卸载,以管理大模型的 CPU 内存使用。

Token 一致的多轮 Rollout

多轮 RL 必须保留模型实际采样的 token 序列。当追加工具结果或后续消息时,重新渲染聊天模板可能会改变现有前缀,导致训练序列偏离原始轨迹。我们将 Muse Spark Open 与 TITO 集成,并修复了一个会追溯更改较早消息结束标记的模板行为。

训练结果

我们在 64 块 H200 GPU 和 32 块 NVIDIA B300 GPU 上,使用 GRPO 在 DAPO-Math-17K 上训练了完整的 Muse Spark Open 语言模型,采用 CPU 卸载的 Muon 或 Adam,并将优化器状态流式写入磁盘。每一步为 32 个提示各采样 16 个响应,共 512 个响应的批次,响应上限为 4,096 个 token,推理强度为 xhigh。Rollout 在 H200 和 B300 上均使用 BF16 检查点。Rollout 路由重放(R3)在训练中复用 rollout 的专家路由,大幅降低了训练与 rollout 之间的对数概率差异:使用 R3 的运行中为 0.0044–0.0071,而未使用时为 0.015–0.020。每次权重同步后,所有 rollout 引擎都持有逐位相同的权重。

Two Muon runs and two Adam runs, one of them on B300: unsmoothed training reward, train–rollout KL, and absolute log-probability difference. The legends identify each GPU type, training topology and whether rollout routing replay is enabled.

在 DAPO-Math-17K 上进行全语言模型 RL:训练奖励、训练–rollout KL 以及平均绝对 token 对数概率差异。所有曲线均显示每步原始值。

我们还训练了一个视觉语言任务 Geometry3K,冻结 ViT,使用相同的 Adam 设置和 medium 推理强度。测试准确率在 30 次更新中从 88.2% 上升到 94.9%。

Geometry3K RL with the ViT frozen: test accuracy after 0, 5, …, 30 updates, and per-step train–rollout KL and absolute log-probability difference.

冻结 ViT 的 Geometry3K RL:每五次更新的测试准确率(601 个问题,每个四个样本),以及每步原始训练–rollout KL 和平均绝对 token 对数概率差异。

致谢

这项工作由 SGLang 社区、Meta、AMD、Nvidia 和 Intel 合作完成。

SGLang 社区:Baizhou Zhang、Chunan Zeng、Jimmy Shong、Mao Cheng、Jiajun Li、Yuhao Yang、Alex Nails、Zijie Xia、Liangsheng Yin、Ziang Li

Meta:Sidhartha Reddy Potu、Clement Anthonioz Blanc、Ethan Che、Beto de Paola、Varun Vontimitta

AMD:Thomas Wang、Zijie Chen、Jacky Cheng、Wun-guo Huang、Hai Xiao

Nvidia:Wookje Han、Andrew Briand、Anik Afibuzzaman

Intel:Guobing Chen、Jianan Gu、Fan Zhao

附录

复现

OpenHands 扫描背后的脚本位于 bench-openhands-b300-sweep 分支的 benchmark/muse_spark_openhands 中。一个脚本使用 docker/muse_spark 补丁从 dev 分支构建服务环境,安装客户端,构建数据集,并在一个 8 x B300 节点上运行从 128 降至 1 的并发阶梯。

服务器是该分支的 SGLang。客户端是提交 acd09b44 处的 evalscope,使用其 swe_smith 多轮插件。

工作负载是基于 nebius/SWE-rebench-openhands-trajectories 构建并用 nvidia/OpenScienceReasoning-2 文本填充的 OpenHands 多轮智能体回放:每段对话 13 轮,首轮 74,160 个 token,之后每轮 753 个 token(每个请求约 80k 提示 token),每轮输出 220 个 token,聚合前缀缓存命中率约 92%。DFlash 接受长度通过 SGLANG_SIMULATE_ACC_LEN 固定为 4.5,因此解码成本反映的是该接受长度,而非草稿模型实测质量。

服务器启动命令如下,其中 <cc> 设置为该档位的并发数:

export SGLANG_SIMULATE_ACC_LEN=4.5 SGLANG_SIMULATE_ACC_METHOD=match-expected SGLANG_SIMULATE_ACC_TOKEN_MODE=real-draft-token
export SGLANG_TIMEOUT_KEEP_ALIVE=600 SPT_NOENV=1
export OMP_NUM_THREADS=8 MKL_NUM_THREADS=8 PYTHONNOUSERSITE=1 TZ=UTC
export SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=8192

# concurrency 16 and above: HiCache 160 GB per rank
python3 -m sglang.launch_server \
    --model-path meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --served-model-name meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --quantization modelopt_fp4 \
    --moe-runner-backend flashinfer_trtllm \
    --tp-size 8 \
    --ep-size 8 \
    --attention-backend trtllm_mha \
    --kv-cache-dtype fp8_e4m3 \
    --page-size 128 \
    --mem-fraction-static 0.9 \
    --reasoning-parser muse \
    --tool-call-parser muse \
    --enable-layernorm-sp \
    --enable-megamoe-for-layernorm-sp \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path RadixArk/Muse-Spark-Open-DFlash2 \
    --speculative-draft-attention-backend fa4 \
    --speculative-num-draft-tokens 8 \
    --enable-hierarchical-cache \
    --hicache-size 160 \
    --max-running-requests <cc> \
    --cuda-graph-max-bs-decode <cc> \
    --context-length 131072 \
    --enable-cache-report \
    --tokenizer-worker-num 4 \
    --enable-http2 \
    --host 0.0.0.0 \
    --port 30100

# concurrency 8 and below: no HiCache
python3 -m sglang.launch_server \
    --model-path meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --served-model-name meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --quantization modelopt_fp4 \
    --moe-runner-backend flashinfer_trtllm \
    --tp-size 8 \
    --ep-size 8 \
    --attention-backend trtllm_mha \
    --kv-cache-dtype fp8_e4m3 \
    --page-size 128 \
    --mem-fraction-static 0.9 \
    --reasoning-parser muse \
    --tool-call-parser muse \
    --enable-layernorm-sp \
    --enable-megamoe-for-layernorm-sp \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path RadixArk/Muse-Spark-Open-DFlash2 \
    --speculative-draft-attention-backend fa4 \
    --speculative-num-draft-tokens 8 \
    --max-running-requests <cc> \
    --cuda-graph-max-bs-decode <cc> \
    --context-length 131072 \
    --enable-cache-report \
    --tokenizer-worker-num 4 \
    --enable-http2 \
    --host 0.0.0.0 \
    --port 30100

AgentX 复现

AgentX 扫描使用提交 89b21867 处的 AgentX,采用固定 token 模式下的 semianalysis_cc_traces_weka_062126 轨迹,温度 1、top-p 0.95。每个并发级别发送请求一小时,且为单次运行。总吞吐量统计输入 token(包括前缀缓存命中)加上输出 token,p90 交互性为 1000 除以请求级 p90 的毫秒级 token 间延迟。AgentX 保留所记录轨迹的 token 数量和前缀结构,但生成合成提示文本,因此草稿模型在回放上的接受情况无法反映真实文本。因此我们通过 SGLANG_SIMULATE_ACC_LEN 将 DFlash 接受长度固定为 4.16。该值是实测而非选定的:它是 AgentX 的 token 加权接受长度,1 + 接受的草稿 token 数 / 验证步数 = 1 + 96,198 / 30,415,来自 DFlash v2 草稿模型在 80 个真实编码首轮提示(输入 token 数为 158 至 1,581)上的结果。

服务器启动命令如下。并发 16 及以上启用序列并行、MegaMoE 和 HiCache;在并发 16 时将 <hicache_gb> 设为 120,在并发 32 和 48 时设为 216。

export SGLANG_SIMULATE_ACC_LEN=4.16 SGLANG_SIMULATE_ACC_METHOD=match-expected SGLANG_SIMULATE_ACC_TOKEN_MODE=fixed
export SGLANG_TIMEOUT_KEEP_ALIVE=1800 OMP_NUM_THREADS=8 TVM_FFI_CUDA_ARCH_LIST=10.3a
export SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND=python

# concurrency 16 and above: sequence parallelism, MegaMoE and HiCache
export SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=8192
python3 -m sglang.launch_server \
    --model-path meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --served-model-name meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --quantization modelopt_fp4 \
    --kv-cache-dtype fp8_e4m3 \
    --moe-runner-backend flashinfer_trtllm \
    --tp-size 8 \
    --ep-size 8 \
    --attention-backend trtllm_mha \
    --page-size 128 \
    --mem-fraction-static 0.9 \
    --reasoning-parser muse \
    --tool-call-parser muse \
    --enable-layernorm-sp \
    --enable-megamoe-for-layernorm-sp \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path RadixArk/Muse-Spark-Open-DFlash2 \
    --speculative-draft-attention-backend fa4 \
    --speculative-num-draft-tokens 8 \
    --speculative-use-block-verification \
    --enable-hierarchical-cache \
    --hicache-size <hicache_gb> \
    --max-running-requests 192 \
    --cuda-graph-max-bs-decode 128 \
    --cuda-graph-backend-prefill breakable \
    --swa-full-tokens-ratio 0.4 \
    --enable-cache-report \
    --enable-metrics \
    --tokenizer-worker-num 4 \
    --enable-http2 \
    --host 0.0.0.0 \
    --port 30100

# concurrency 8 and below: no sequence parallelism, MegaMoE or HiCache
python3 -m sglang.launch_server \
    --model-path meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --served-model-name meta-models/Muse-Spark-Open-816B-A42B-NVFP4 \
    --quantization modelopt_fp4 \
    --kv-cache-dtype fp8_e4m3 \
    --moe-runner-backend flashinfer_trtllm \
    --tp-size 8 \
    --ep-size 8 \
    --attention-backend trtllm_mha \
    --page-size 128 \
    --mem-fraction-static 0.9 \
    --reasoning-parser muse \
    --tool-call-parser muse \
    --speculative-algorithm DFLASH \
    --speculative-draft-model-path RadixArk/Muse-Spark-Open-DFlash2 \
    --speculative-draft-attention-backend fa4 \
    --speculative-num-draft-tokens 8 \
    --speculative-use-block-verification \
    --max-running-requests 192 \
    --cuda-graph-max-bs-decode 128 \
    --cuda-graph-backend-prefill breakable \
    --swa-full-tokens-ratio 0.4 \
    --enable-cache-report \
    --enable-metrics \
    --tokenizer-worker-num 4 \
    --enable-http2 \
    --host 0.0.0.0 \
    --port 30100

来源:LMSYS Blog · lmsys.org