Understanding Muse Spark Open

Source: https://memory-bound.vercel.app/blog/muse-spark-open-architecture/

一周前，这个出现在了我的动态流中：

我们发起一个 $100,000 的挑战赛，看看谁能将 Muse Spark 跑得最快：openkernelchallenge.dev (https://openkernelchallenge.dev/) — Daniel Gross (@danielgross) 2026年10月22日 (https://twitter.com/danielgross/status/2113351902402613558)

它所指的模型在仅仅一周前才发布。10月15日，Meta 发布了 Muse Spark Open——Muse 系列中的全尺寸成员——并开放了权重。

挑战赛

这样一个大小的模型无法在单卡上运行。它运行在一个由 8 块 H200 组成的节点上，而挑战赛提出的问题是：你能把它跑多快？ 赛道有两个。

Kernel 赛道将模型拆成五个部分，要求你让每一部分都比 SGLang (https://github.com/sgl-project/sglang) 更快。SGLang 是 Meta 发布该模型时所附带的推理引擎。每个部分是一个独立的问题，每个问题在单块 H200 上运行：

1. 稠密投影（即普通矩阵乘法）。
2. 解码阶段的注意力。
3. 路由专家——模型 97% 参数所在的部分。
4. 整个解码器层在解码阶段的性能。
5. 整个解码器层在预填充阶段的性能——即模型读取 prompt 的阶段。

你提交一个 Python 文件。评测系统在一块 H200 上运行它，检查其输出是否与参考实现一致，然后与 SGLang 计时对比。得分 1.00× 意味着你追平了；高于 1.00× 说明你节省了一些时间。Kernel 赛道的提交截止日期为 11 月 4 日。

Serving 赛道在第二天，即 11 月 5 日开放，这是真正与奖金相关的赛道：你直接修改 SGLang 本身，评测系统在完整的节点上运行整个模型。

深入 Muse

预填充层的题目说明中写着“大部分时间花在路由专家上“。我复制了 SGLang 的基准代码，运行了评测系统的分析工具，结论一致：大约 65% 的时间花在专家内部。

然后我意识到，我无法在脑海中勾勒出专家对一个 token 做了什么。我能读懂规格说明中的公式。我能看到权重矩阵的形状。但我没有一个清晰的意象：一个 token 进入层、被路由、被一个 4-bit 矩阵相乘、然后从另一端输出。在不理解这些的情况下优化它，感觉就像蒙着眼睛跑步。

所以，这似乎是逐块了解这个模型的一个绝佳借口。

总览

在深入细节之前，先看一下模型架构的整体概览。

一个 prompt 从三个入口之一进入。文本经过 tokenizer；图像和视频经过感知编码器（perception encoder）；音频经过音频编码器（audio encoder）。无论走哪个入口，输出都是一样的：一个 token 列表，每个 token 是一个维度为 8,192 的向量。这个列表，[T, 8192]，就是贯穿下面所有部分的形状。

然后是 62 层，每层执行两个操作：注意力（Attention）——token 查看它之前的 token 并从中提取信息；以及混合专家（Mixture of Experts）——token 独立处理。一个路由器对 256 个小网络进行评分，选出最好的 8 个，token 仅在这 8 个中运行。这就是为什么模型总参数为 816B，但只有 42B 激活参数。这是本次挑战中模型最有趣的部分——它耗时最长，也有最大的优化空间。

你可能还会注意到，在两者之间有一个残差流（residual stream）：图中左侧的橙色线。每个块将结果加到它上面，而不是替换它。这是来自 ResNet (https://arxiv.org/abs/1512.03385) 的思想，允许某一层在“什么都不做“是正确操作时什么都不做。

textimagevideoaudiotokenizervocabulary 202,048perception encoder1.8B · 50 layersaudio encodertoken sequence[T, 8192]up to 1M tokenslanguage modelgated attention128 q heads · 16 kv headslearned residual gate mixture of experts256 experts · 8 activelearned residual gate × 62output layervocabulary 202,048next tokeninside one layerRMSNormattentionlooks at earlier tokens + learned residual gateRMSNormmixture of expertsa router picks 8 of 256 + learned residual gateto the next layer

左图：Meta 绘制的模型全貌。右图：62 层中的一层。侧面的箭头是残差：每个块的输出通过一个学习门控混合回残差流中，而非简单相加。

experts (×62)799Battention (×62)13.5BMoE projections (×62)4.2Bembeddings3.3Bperception encoder1.8Brouters (×62)0.13B

参数分布。面积与参数量成正比；专家持有 97% 的参数。数据来自模型卡片中的维度信息。

另外一个有趣的点是，并非所有 62 层都做相同的事情。其中 46 层只对最后 2,048 个 token 的窗口进行注意力计算，另外 16 层对整个上下文进行注意力计算。它们按照固定的四步模式交替排列：滑动窗口（sliding）、全局（full）、滑动窗口（sliding）、滑动窗口（sliding）。

012345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061sliding · 46 layers · last 2,048 tokensfull · 16 layers · whole context

按顺序排列的 62 层。模式每四层重复一次：sliding、full、sliding、sliding。

单层概览

让我们看看 62 层中一层的内部。如果将其视为一个黑盒，一层就是一个函数：接收一组向量列表，返回相同形状的向量列表。输入是残差流 s：每个 token 一个 8,192 维的数字向量，输出也是相同形状。这一层对残差流进行两次编辑：一次通过注意力，一次通过混合专家。

s [T, 8192] fp32the residual stream, from the previous layerα₁ · sRMSNorm · (1 + w)h bf16attentionsliding or fulla [T, 8192] +learned residual gate s ← α₁ · s + β₁ · rms(a)α₂ · sRMSNorm · (1 + w)h bf16mixture of experts8 of 256 expertsf [T, 8192] +learned residual gate s ← α₂ · s + β₂ · rms(f)s [T, 8192] fp32the same stream, edited twice, to the next layers

62 层中的一层。橙色路径是残差流；它从不被替换，只会与每个分支的输出通过学习门控混合。

残差流进入

s \in \mathbb{R}^{T \times 8192}

来自上一层的残差流，每个 token 一个 8,192 维的向量。我们跟随其中一个 token。

使用 ‹ › 或方向键，或点击图中的方框。

在每个分支看到残差流之前，它会先被归一化：

\text{rms}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \varepsilon}}, \qquad h = \text{rms}(s) \cdot (1 + w)

RMSNorm 将向量除以其均方根，因此无论 s 的规模漂移到什么程度，分支接收到的向量始终约为 \sqrt{d} 大小。学习到的 w（每个维度一个）允许模型在输入时对某些维度进行放大或缩小。将其写成 1 + w 意味着未训练的归一化是恒等变换，因为 w 初始值为 0。\varepsilon 仅为数值稳定性而存在，d 是向量维度，即 8,192。

这被称为 pre-norm：归一化在分支上，而不是在主干上。残差流本身永远不会被就地归一化。

注意力，然后是门控

归一化后的 h 经过注意力，返回 a，同样为 [T, 8192]。在标准 transformer 中，下一步是 s \leftarrow s + a。Muse 做了略有不同的操作：

s \leftarrow \alpha_1 \cdot s + \beta_1 \cdot \text{rms}(a)

有两个区别。第一，a 在相加之前先被归一化，因此无论注意力输出了什么，结果的大小始终一致。第二，我们不是简单相加，而是加权求和：\alpha_1 决定保留多少旧的 s，\beta_1 决定加入多少新信息。两者都是学习得到的，且每个维度各有一个值，因此 8,192 个维度中每个维度都有自己的混合比例。Meta 将此称为学习残差门控（learned residual gate）。原因是：在普通加法中，s 会在每一层不断增大，经过 62 层后可能膨胀得非常厉害。有了门控，它就能保持在固定范围内。

专家，然后是另一个门控

后半部分的工作方式完全一样。我们再次对 s 进行归一化（这次用它自己的 w），将结果 h 送入混合专家。返回的是 f，同样为 [T, 8192]，然后通过第二个门控混合回残差流：

s \leftarrow \alpha_2 \cdot s + \beta_2 \cdot \text{rms}(f)

\alpha_2 和 \beta_2 是另一对学习到的值，每个维度一个，与之前类似。经过这一步，s 已被编辑两次——一次由注意力完成，一次由专家完成——然后它进入下一层。

62 层之间的差异

所有 62 层都有相同的结构，变化的是权重：每一层有自己的 w、\alpha、\beta，有自己的注意力矩阵和自己的 256 个专家。结构上唯一的区别就是我们之前看到的：其中 46 层的注意力是滑动窗口（sliding），16 层是全局（full）。其余完全相同，且每一层都有混合专家，没有稠密层。

注意力块内部

sthe residual streamRMSNormhattentionaRMSNormmixture of expertsto the next layer

注意力在层中的位置：它读取归一化后的残差流 h，返回 a，由门控将其混合回残差流中。

Muse 中的注意力就是你已经熟知的注意力，只是在上面加了几处修改。完整写出，这个块是这样的：

attentionhWqkvgQK-normRoPEKV cacheq·kT / √64masksoftmax× v ⊙ σ(g)Wo

• qkvg 1 · 第四个投影，g (https://memory-bound.vercel.app/blog/muse-spark-open-architecture/#one-matrix-four-pieces)
• 2 · 128 个查询头，16 个键/值头 (https://memory-bound.vercel.app/blog/muse-spark-open-architecture/#128-query-heads-16-key-and-value-heads)
• 3 · QK-norm (https://memory-bound.vercel.app/blog/muse-spark-open-architecture/#qk-norm)
• 4 · 滑动窗口层和全局层 (https://memory-bound.vercel.app/blog/muse-spark-open-architecture/#sliding-and-full-layers)
• 5 · 输出门控 (https://memory-bound.vercel.app/blog/muse-spark-open-architecture/#the-output-gate)
• 6 · KV 缓存 (https://memory-bound.vercel.app/blog/muse-spark-open-architecture/#the-kv-cache-at-a-million-tokens)

注意力块。带编号的部分是 Muse 与教科书注意力不同的地方；每个编号链接到下文对应的章节。

如果你对注意力机制本身还不熟悉，Jay Alammar 的注意力可视化指南 (https://jalammar.github.io/visualizing-neural-machine-translation-mechanics-of-seq2seq-models-with-attention/) 是一个很好的起点。

一个矩阵，四个部分

首先，h 乘以一个矩阵 W_{qkvg}，结果被切分成四个部分：查询 q、键 k、值 v 和第四个 g——我们最后再讲它。q 有 128 个头，k 和 v 各有 16 个头，每个头都是 64 个数字。128 \times 64 = 8192，所以查询保持了模型的完整宽度。

attentionhWqkvg8,192 in 8,432 outQK-normRoPEKV cacheq·kT / √64masksoftmax× v ⊙ σ(g)Wo8,192 × 8,192

带尺寸标注的注意力块（单个 token）。两个矩阵画成网格；KV 缓存每个 token 增长一行。

Inside × W_qkvg，一个小示例：6 个 token，每个 16 个数字，乘以一个 16 × 40 矩阵。token 在左侧面，矩阵在顶面，结果在正面。每个结果数字对应左侧面的一行和顶面的一列的交汇处：16 个乘积沿该线求和。

128 个查询头，16 个键值头

拥有 128 个查询头但只有 16 个键和值头，这被称为分组查询注意力（grouped-query attention, GQA）：每 8 个查询头共享同一组 k 和 v。原因在于内存。所有过去的 token 的键和值都必须存储（这就是 KV 缓存），而该开销取决于 k 和 v 头的数量，而非 q 头的数量。从 128 降到 16，缓存缩小了 8 倍。Muse 还以 FP8 格式存储，每个数字一个字节。更多关于 GQA 的信息请参见引入它的论文 (https://arxiv.org/pdf/2305.13245)。

attentionhWqkvg8,192 in 8,432 outQK-normRoPEKV cacheq·kT / √64masksoftmax× v ⊙ σ(g)Wo8,192 × 8,192

一个头组的注意力块：8 个 q 头（8 条细线）、1 个 k 头和 1 个 v 头，它们全部共享。Muse 共有 16 个这样的组。

QK-norm

然后 q 和 k 在点积之前先被归一化。这被称为 QK-norm。它将 q k^T 保持在有界范围内，使 softmax 不会因为出现过大的分数而坍缩到单一 token 上。它几乎没有开销，而且大多数近期模型都在使用。

q1q2q3q4kbefore: every head has its own size sizes from 0.55 to 2.3 ÷ its size q1q2q3q4kafter QK-norm: all the same size every arrow ends on the circle

QK-norm。q 和 k 的每个头都是一个向量，这里画成一条箭头，长度代表大小。将每个头除以自身的大小后，它们全部落在同一个圆上：方向保持不变，大小变得相同，因此没有任何一个分数能够超出范围。

滑动窗口层和全局层

每一层的掩码都是因果掩码（causal）——一个 token 只能看到它之前的 token。在滑动窗口层中，它还会截断超过 2,048 个 token 的更早部分。这就是滑动窗口层和全局层之间唯一的区别。

fulla token sees every token before it keys, tokens 0 … 11 queries sliding, window 4only the last 4, including itself keys, tokens 0 … 11 queries

每个查询可以看到哪些键。第 i 行代表 token i 的提问，填充的单元格表示它可以查看那个 token。在真实模型中窗口为 2,048，而非 4。

RoPE——告诉注意力两个 token 之间距离多远的旋转操作——仅在 46 个滑动窗口层中应用。在 16 个全局层中完全没有应用，因此这些层不知道位置，只知道内容。这种混合有无位置编码层的方式由 Cohere (https://arxiv.org/abs/2501.18795) 提出，Llama 4 (https://ai.meta.com/blog/llama-4-multimodal-intelligence/) 也采用了它：滑动窗口层已经了解了位置信息，而没有位置编码的层 (https://arxiv.org/abs/2305.19466) 在处理超出训练长度的序列时表现更好。

q at position 2 k at position 5

q 的一个头和 k 的一个头上的 RoPE，每个 8 个数字（在 Muse 中为 64）。每对数字画成一条箭头，RoPE 根据 token 位置乘以该对的“速度“来旋转它：数字变了，但箭头长度不变。下图展示了每对数字对 q · k 的贡献。同时平移两个 token，下方的数值不变：只有它们之间的相对距离起作用。

输出门控