跳到正文
Hacker News· frostbyte7·· 8 天前AI 评分32

Kolibri 架构设计:从第一性原理看架构权衡

Designing Kolibri: Architecture Trade-Offs from First Principles

AI 导读

Aleph Alpha 研究团队拆解自回归语言模型的架构选择对训练与部署成本的影响,聚焦参数与 FLOPs 分配、FLOPs 随上下文长度和序列混合器状态的扩展方式,并分析近期开源权重模型。文章以 Apache 2.0 许可发布的 Kolibri 和 Kolibri Origin 为参照,提供交互工具供配置模型并对比。

正文 · AI 翻译

Warm orange gradient with translucent stepped blocks in the lower left and upper right corners and a soft pale glow in the centre

研究

Pit NeitemeierAlessio Serra

07/10/2026

简而言之

本博文探讨自回归语言模型中的主要架构选择对训练和部署成本的影响,重点关注参数和 FLOPs 的分配、FLOPs 如何随上下文长度和序列混合器1状态2大小而变化。我们推导这些量,并分析它们在近期开放权重架构中的变化。文中还包含一个交互式工具,可配置你自己的模型并与近期开放权重发布版本进行比较。


引言

语言模型的架构决定了其容量和计算如何分配。其中许多选择会同时影响质量、训练成本和推理服务成本。某个选择对质量的影响必须通过实证来衡量,但成本则直接由配置决定:

  • 总参数量决定了权重内存占用,进而决定了部署模型所需的最低硬件。
  • 每 token 的 FLOPs乘以训练 token 范围3,决定了训练计算量。它们也是预填充成本的有用代理指标,而预填充通常是计算受限的。
  • 序列混合器状态2, 例如基于注意力的序列混合器的键值(KV)缓存,决定了解码批次中能容纳多少序列、多少上下文,以及每个解码步骤必须读取多少字节。

其中哪一个成为推理服务瓶颈取决于工作负载。小批量、短上下文解码受限于读取权重。预填充和短查询的大批量解码受限于计算。长上下文解码受限于读取序列混合器状态。

在本博文中,我们推导 MoE 解码器的每个量,并分析近期的开放权重模型:

  • 参数分配:活跃参数 如何在 FFN(在我们分析的所有模型中都是专家混合(MoE)层)、序列混合器(例如注意力)和语言建模(LM)头之间分配。
  • FLOP 分配:每 token 的 FLOPs 如何在各组件之间分配,并如何随序列长度扩展。
  • 序列混合器状态:每个序列持有多少状态,每个解码步骤读取多少状态,取决于序列长度和序列混合器类型。

我们在下方的交互式浏览器中可视化这些量,让你可以更改配置并查看对这三个维度的影响。

架构配置

下面的探索器是我们在设计模型架构时使用的内部工具的简化子集。它帮助我们在确定配置之前量化不同架构选择之间的权衡。受 Sebastian Raschka 的 LLM Architecture Gallery 启发,我们比较了近期开放权重的架构。我们用两个 Aleph Alpha 模型初始化探索器,即 Kolibri Origin 和 Kolibri,后者最近以 Apache 2.0 许可证发布。4 它们为本文讨论的设计选择提供了两个具体的参考点。你可以使用滑块修改任一配置,或用 + 添加自己的配置。

总计
78.1B
每 token 激活
3.5B
FLOPs / token · 16K
28.3G
状态 · 16K
189M

参数分配

在 MoE 架构中,总参数和激活参数可能相差一个数量级。总参数主要影响内存占用,因此决定模型的最低部署要求。激活参数通常占训练和预填充期间的大部分 FLOPs。在小批量下,它们也在很大程度上决定了每个解码步骤从内存中读取多少权重字节。5

如果你调整滑块,会注意到除了全注意力/SWA 划分和窗口大小之外,每个滑块都会影响总参数或激活参数。下面我们推导每个组件对这两个总数的贡献。

78.1B3.5B4.4%
30.6B3.3B10.7%
31.6B3.2B10.2%
120.7B12.2B10.1%
549.3B55.0B10.0%
125.1B6.0B4.8%
1.57T48.8B3.1%
551.5B15.8B2.9%
2.78T104.2B3.7%
308.8B14.8B4.8%
116.8B5.1B4.4%
20.9B3.6B17.3%
743.4B40.3B5.4%
426.2B24.7B5.8%

总计统计文本解码器存储的每个权重。每 token 激活统计每个 token 经过的权重:序列混合器、top-k 路由专家和共享专家、稠密 FFN 层、路由器、归一化和 LM head;输入嵌入是查表,不计入。

图 1. 每个模型的总参数和激活参数,激活参数不包括输入嵌入。可按任意列排序,点击模型可在所有图表中高亮显示。

FFN:稠密层、路由专家和共享专家

路由专家数量和专家宽度决定 FFN 的容量。Top-k 和共享专家决定一个 token 激活多少该容量。

一个 SwiGLU 专家使用三个线性投影:gate、up 和 down。gate 和 up 投影将隐藏维度 d 映射到中间维度 de,而 down 投影将 de 映射回 d。这给出每个专家 3×d×de 个参数。将此数乘以专家数量和 MoE 层数,得到路由总参数。对于激活路由参数,我们将 E 替换为 k,因为只有 top-k 专家会被使用。

Prouted=Lmoe×3×d×de×{Etotalkactive

同样地,我们可以推导出共享专家的参数。注意,多个共享专家等价于一个宽度更大的共享专家 ds,因为所有共享专家始终处于激活状态,且它们的输出会被合并。 因此这里我们只考虑单个共享专家。

Pshared=Lmoe×3×d×ds

许多 MoE 解码器将其前几层保持为稠密层6。它们可以有自己的宽度 dff。在上面的滑块中,dff 与激活专家宽度绑定,因此 dff=k×de,这是此处分析的模型(例如 MiMo-V2.5、MiniMax M3、Kolibri Origin)中常见的选择。

Pdense=Ldense×3×d×dff

路由器将完整的残差流 d 映射为 E 个 logits,因此其规模随专家数量和隐藏维度增长。然而,它只占总参数中很小的一部分,我们将其计入 FFN。

Prouter=Lmoe×E×d

序列混合器:注意力投影

一层的 Query、Key、Value 和 Output 投影取决于隐藏宽度、查询宽度和 KV 宽度。它们的总和随后随注意力层数缩放。

在 GQA 层中,我们将宽度为 d 的隐藏状态投影为宽度为 dq 的查询和两个宽度为 dkv 的 KV 张量,然后将注意力结果从 dq 投影回 d。

PQKVO=Lattn×[ddq+2ddkv+dqd]=2Lattnd(dq+dkv)

在组大小 g=hq/hkv 下,MHA (g=1)和 MQA(hkv=1,因此 g=hq)是 GQA 的两个极端情况。

归一化

每个解码器层通常带有两个(仅 pre-norm)到四个(sandwich norm)RMSNorm,其缩放向量宽度为 d,外加 LM 头之前的一个最终归一化。下面的系数四假设使用 sandwich norm。

Pnorms=(4L+1)×d

嵌入与 LM 头

词表大小 V 与隐藏宽度 d 的乘积同时决定了 LM 头和嵌入。后者只影响总参数,因为它仅决定一个轻量的查找操作,而不是矩阵乘法。

未绑定的输入嵌入和 LM head 各自包含 d×V 个权重。绑定它们会存储一个共享矩阵而非两个,但并不会从训练计算中移除稠密的词表投影。你可以使用嵌入绑定开关来验证仅总参数量如何变化。

Plookup=Phead=d×V

序列混合器FFN · 活跃专家 + 稠密层LM head归一化

图 2. 按组件划分的每 token 活跃参数量。选择哪个组件先堆叠以及模型如何排序。

各架构间的参数分配

如 图 2 所示,分配给序列混合器的活跃参数比例在所选架构间差异显著,从大约 18% 到 49% 不等。

按 总参数量、活跃参数量或 活跃与总参数之比对模型排序,均未显示出与序列混合器占比有明显关系。序列混合器占比的选择主要取决于经验测量的质量。此外,在保持状态大小不变的情况下增加序列混合器计算量会提高算术强度,并使主要受内存限制的内核更高效。

序列混合器与 FFN 参数之间的比例主要由注意力投影的相对宽度和 FFN 的中间活跃宽度控制。使用上面推导的参数量,

Pseq=2Lattn d (dq+dkv)Pffn=3d[Lmoe(k de+ds)+Ldense dff]

因此它们的相对分配为

PseqPffn=2Lattn(dq+dkv)3[Lmoe(k de+ds)+Ldense dff]

残差流宽度 d 被约去。因此仅改变 d 不会改变这一划分,而改变注意力或专家宽度会直接改变分配。

最后,LM head 的占比随着模型增大而缩小。大多数权重矩阵随宽度呈二次增长,例如 d×dq 且 dq∝d,而 LM head(d×V)仅呈线性增长,因为词表大小 V 通常保持不变。你可以通过 将 LM head 先堆叠并按活跃参数量对模型排序清楚地看到这一点。

FLOP 分配

每 token 的训练计算有两个主要组成部分:参数 FLOPs 和注意力 FLOPs。

参数 FLOPs 来自激活值与权重矩阵的相乘,与上下文长度无关。注意力 FLOPs 来自 QK⊤ 和 AV 乘积,它们不涉及学习到的权重矩阵,但取决于注意力宽度和上下文长度。

本节中的公式沿用上文定义的参数分配符号。

Ftrain/ token ≈ 6×Pmat 参数 FLOPs + 12×Lfull×dq×n/2 全注意力 + 12×Lswa×dq×min⁡(n,w) 滑动窗口

参数 FLOPs 项仅取决于活跃权重:6 来自三次同等规模的矩阵 乘法,每个权重每个 token 消耗两个 FLOPs,一次乘法加 一次加法。前向传播(fwd)计算 Y=XW。 反向传播随后需要两个相同形状的乘积。关于 激活的梯度(bwdA)将误差信号传播回 前一层,而关于权重的梯度(bwdW)衡量损失如何依赖于每个权重以及应如何更新以减少 损失。

Y=XW ⁣:2Tdindout∂X=∂YW⊤ ⁣:2Tdindout∂W=X⊤∂Y ⁣:2Tdindout

前面的方程未考虑嵌入查找、归一化和激活,因为在 这个规模下,它们的贡献可以忽略不计。7

注意力 FLOPs 不涉及权重矩阵:QK⊤ 为每个 token 计算一次点积, AV 为每个 token 应用一个值向量。这导致每个查询-键对产生 4×dq FLOPs, 同样的 ×3 乘数对应 fwd、bwdA 和 bwdW。我们统计查询在因果 掩码下关注的键,在全注意力层中平均为 n/2。FLOPs 随序列长度增长, 因为查询关注每个保留的键,即全注意力 层中的所有先前 token,以及在 SWA 层中至多窗口 min⁡(n,w)。注意 GQA 使 dq 和注意力 FLOPs 保持不变:它仅收窄 KV 宽度,从而减小序列混合器状态大小,这将在下一节讨论。

超越全注意力和滑动窗口注意力

下面的闭式推导聚焦于标准和滑动窗口注意力,但 我们比较中的若干架构使用了其他序列混合器。我们在此仅给出这些替代方案的 直觉,并请感兴趣的读者参阅原始论文以了解 其详细公式。

从宏观上看,这些替代方案遵循两种主要思路。稀疏注意力,例如 Native Sparse Attention (NSA),减少了从较长上下文中访问的 token 数量。循环或线性序列 混合器,例如 Mamba-2 和 Gated DeltaNet,则将前缀中的信息压缩为固定大小的状态。近期的模型 常将这些机制与注意力结合在混合架构中。在我们的对比中,DeepSeek V4、GLM-5.3、MiniMax M3 和 Qwen3.8-Flash-Next 使用稀疏注意力,而 Qwen3.8-Flash-Next(Gated DeltaNet)、Kimi K3(Kimi Delta Attention)和 Nemotron 3(Mamba-2) 则将注意力与循环层结合。

FFN

FFN 独立地变换每个 token,并且正如我们之前所见,大部分容量都花在这里。在固定的 top-k 下,未激活的路由专家矩阵会增加 容量,这消耗内存占用但不消耗 FLOPs。

Fffn/ token =6×3×d×[Lmoe×(k×de+ds)+Ldense×dff]

每个 MoE 层在此基础上增加一个小型路由矩阵。

Frouter/ token =6×Lmoe×E×d

序列混合器投影

序列混合器对 token 之间的交互进行建模。其 Query (Q)、Key (K)、Value (V) 和 Output (O) 线性层取决于隐藏宽度、注意力宽度、KV 宽度和深度,但 每个 token 与序列长度无关。

Fproj/ token =6×2×Lattn×d×(dq+dkv)

纯注意力

它包括注意力分数计算 QK⊤ 和值聚合 AV,这取决于查询宽度和保留的键。全注意力层保留 所有先前的 token,因此一个查询平均关注 n/2 个键。

Fattn/ token =12×Lfull×dq×n/2

滑动窗口层最多保留其中 w 个。

Fattn/ token =12×Lswa×dq×min⁡(n,w)

LM Head

输出投影随隐藏宽度乘以词表大小而缩放。

Fhead/ token =6×d×V

纯注意力 + 状态更新序列混合器线性层FFNLM head归一化

图 3。 在所选序列长度下按组件划分的每 token 训练 FLOPs。移动滑块以查看向纯注意力转移的构成变化。

FLOP 分配与上下文长度

调整 图 3 中的序列长度滑块,你可以看到 FLOP 构成随上下文长度发生显著变化。在 16K 上下文下,纯注意力在所选模型中约占训练 FLOPs 的 3–51%。 当上下文接近 1M token 时,这一贡献可能占据主导地位,达到 98%。

参数相关的 FLOPs 随 token 数量线性增长,而全注意力 FLOPs 随序列长度呈二次增长。因此,随着上下文增长,FLOP 构成越来越 偏向纯注意力,在包含许多全注意力层的架构中尤其如此。

在各模型中,16K 序列长度下序列混合器中消耗的 FLOPs 总占比更为受限,大约在 33% 到 65% 之间。随着纯注意力 FLOPs 通过例如稀疏注意力被减少,投影线性层在总 FLOPs 中占据更大比例。 通过减少全注意力层数量并用 滑动窗口注意力替代它们,遵循了这一模式,在 16K 序列长度下将纯注意力占比从 51% 大幅降至 27%,与 相比。

候选参考文献

图 4. 每 token 的模型训练 FLOPs。更低的数值精度不会改变 FLOP 计数,但可以通过减少表示每个值所用的位数来提高硬件吞吐量。

FLOP 扩展与上下文长度

如图 4所示,随着上下文长度增长,计算成本的差异变得越来越 明显。 在整个范围内具有最高的 FLOPs,它结合了较大的激活模型规模和 93 层中的 24 个全注意力 层。 在长上下文下处于相反的一端,其序列混合器中只有一小部分随上下文长度呈二次 增长。

和 在 16K 上下文下需要 相似的 FLOPs,但在 1M token 时 MiMo-V2.5 需要约 5 倍:其全注意力 层关注每个先前的 token,而 DeepSeek V4.1-Flash 只关注固定数量的 选定、压缩的 token,并仅用轻量索引器对完整前缀进行评分。

遵循同样的 趋势。与 相比,它使用滑动窗口注意力,全注意力层数量约少 5 倍,因此随着上下文长度增加, 计算优势不断扩大。

在此比较中,在长上下文下达到最低的 FLOPs/token。此外,与 一起,它是随上下文长度变化最平缓的曲线,从 4K 到 1M 分别仅增加 13% 和 33%。

将 FLOPs 乘以 token 范围 D 得到总训练 FLOPs,Ctraining≈D×Ftrain/ token 8,再除以你实现的训练 FLOP/s 即可估算训练时长。请考虑真实训练成本远比 FLOPs 复杂得多。此处未涉及的其他主要因素包括计算与通信的 重叠、权重量化、内核效率和内核 启动开销、激活重计算以及专家不均衡。

FLOPs 是预填充成本的有用代理,但它们并不能直接转化为更快的 解码,在解码中内存移动变得更为重要。这引出了下一节关于 序列混合器状态大小的内容。

序列混合器状态大小

在每一个解码步骤中,新的查询必须与先前 token 的表示进行交互。那些键和值表示在之前的步骤中已经计算完成,并且永远不会改变,因此我们可以将它们缓存起来,而不是重新投影,从而用 FLOPs 换取 HBM 容量和更多的 I/O。注意力计算本身不受影响,因为在每一步我们只是追加当前 token 的 K 和 V 向量。

限制解码速度的因素取决于具体情形:一个解码步骤将每个权重矩阵与每个序列仅一个 token 相乘,因此在短上下文和小批量下,它主要受限于将模型权重从 HBM 流式传输到计算单元(内存受限);对于 MoE,读取的权重是至少被一个序列选中的每个专家的权重,随着批量增大,这接近于所有专家。在大批量和短上下文下,权重复用提高了算术强度,计算可能成为限制(计算受限)。一旦上下文变长,瓶颈就变成读取每个生成 token 状态的 I/O 成本,与权重不同,该状态不会在序列之间共享(内存受限)。这些只是一般情形:限制因素取决于操作的算术强度以及硬件规格,特别是带宽(字节/秒)和吞吐量(FLOP/s),这两者还取决于数值精度。要更深入地了解这些主题,请阅读我们的文章 从第一性原理看推理性能。

Mstate=c×B×b×dkv×Lkv×nret

缓存的表示

标准注意力存储单独的键和值向量(c=2)。多头潜在注意力(MLA,在 DeepSeek-V2 中引入)则每个 token 只缓存一个向量(c=1):一个低秩内容潜在向量加上其解耦 RoPE 路径所需的位置分量。

批量

在朴素实现中,批量的每个序列 B 都有自己独立的状态。前缀缓存则跨序列复用共享前缀(例如系统提示),因此该状态的一份副本可服务多个请求。

每元素字节数

缓存中的值通常以 1 字节(FP8)或 2 字节(BF16)存储。诸如 KIVI 和 KVQuant 之类的方法可降至 2–4 位,但在生产服务中尚不常见。循环状态则通常保留为 FP32。

KV 宽度

KV 头数乘以头维度决定了 dkv。GQA 在一组查询头之间共享一个 KV 头,而 MQA 对所有查询使用单个 KV 头。

注意力层

默认情况下,每个注意力层保留自己的状态。混合架构用滑动窗口层或循环层替换大多数全注意力层,而诸如跨层注意力(CLA)之类的研究方向让成组的层共享某一层的 K 和 V。例如,DeepSeek V4.1-Flash 在四层中计算其压缩缓存,并让中间的层复用它。

保留的 Token

全注意力保留整个前缀。一种常见且稳健的变体是滑动窗口,它只保留最新的 w 个 token。选择性注意力,例如 NSA,在每一步只读取前缀的一个稀疏子集。这减少了每步读取的字节数,但没有减少存储的字节数,因为任何块之后都可能被选中。

循环层

Mamba-2 和 Gated DeltaNet 层将整个前缀折叠为每个头的固定状态,GDN 为 dk×dv 个元素,Mamba-2 为 dh×dstate 个元素,再加上短卷积的最后几个输入。一个 Qwen3.8-Flash-Next GDN 层保留 48 个 128×128 的头,即 0.79M 个元素,相当于其一个注意力层为约 680 个 token 缓存的内容(K、V 和索引器键),但在任何上下文长度下都相同。

候选参考文献

图 5. 单个序列的状态元素(Mstate,其中 b = 1):缓存中的总量(实线)和一次解码步骤读取的量(虚线,仅在两者不同处显示)。每个全注意力层都随上下文增长;有界层则在其窗口处封顶或保持固定的循环状态。

跨架构的序列混合器状态

图 5 展示了序列混合器状态的总量以及每个解码步骤读取的状态量,这反映了两种不同的约束。9 总状态大小决定了解码批次中能容纳多少序列和多少上下文,而每个 token 读取的状态量在长上下文、内存受限的场景下与解码速度更直接相关。

这一区别对稀疏注意力尤为重要。许多稀疏机制减少了每一步访问的 token 数量,降低了内存流量和 FLOPs,但没有减少必须存储的总状态。在 1M token 时, 每个序列持有 72B 个元素,但每个解码步骤读取 11B。滑动窗口和循环层同时减少两者,而 KV 压缩或跨层共享等技术也可以缩小存储的状态本身。

这产生了与 FLOPs 截然不同的排名。例如, 的长上下文 FLOPs 低于 ,而后者通过 KV 压缩和复用实现了更小的序列混合器状态。 和 DeepSeek V4.1-Flash 同样处于可比的参数规模,但在状态大小方面差异极大:在 1M token 时,MiniMax M3 每个序列为 72B 个元素,而 DeepSeek V4.1-Flash 为 1.7B。

通过将大多数全注意力层替换为滑动窗口注意力, 也比 有了显著改进,减少了随上下文长度增长的状态:在 1M token 时,它每个序列持有 10.8B 个元素,比 Kolibri Origin 的 53.7B 少 5 倍。

结论

在这篇博文中,我们介绍用于设计 Kolibri 架构的 Model Explorer。我们分析了流行开源模型中参数和 FLOPs 如何在主要模型组件之间分配,以及在不同架构下序列混合器状态如何随序列长度增长。此外,我们推导了 Kolibri 模型的 FLOPs 和参数公式。我们很高兴与社区分享这些内容,以便其他人可以探索架构权衡,并从第一性原理出发推理架构。

致谢

我们感谢 Ahmed Hammam、Fabien Benureau、Jonas Knupp、Simon Thel、Thomas Burns 和 Yasser Jadidi 的深思熟虑的审阅。我们感谢 Alexander Wortmeier 和 Noé Beckerle Vallejo 支持将其发布到我们的网站,以及 Helena Treeck 和 Svenja Fahlisch 确保一切顺利进行。

参考文献

  1. Ainslie, J. et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.
  2. Aleph Alpha Research. Inference performance from first principles. Aleph Alpha blog.
  3. Anthony, Q., Biderman, S., Schoelkopf, H. (2023). Transformer Math 101. EleutherAI blog.
  4. Beltagy, I., Peters, M. E., Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150.
  5. Brandon, W. et al. (2024). Reducing Transformer Key-Value Cache Size with Cross-Layer Attention. arXiv:2405.12981.
  6. Dai, D. et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066.
  7. Dao, T., Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv:2405.21060.
  8. DeepSeek-AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.
  9. Ding, M. et al. (2021). CogView: Mastering Text-to-Image Generation via Transformers. arXiv:2105.13290.
  10. Gu, A., Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.
  11. Hooper, C. et al. (2024). KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization. arXiv:2401.18079.
  12. Liu, Z. et al. (2024). KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. arXiv:2402.02750.
  13. Raschka, S. (2026). LLM Architecture Gallery. sebastianraschka.com.
  14. Shazeer, N. et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv:1701.06538.
  15. Shazeer, N. (2019). Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150.
  16. Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv:2002.05202.
  17. Su, J. et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
  18. Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
  19. Yang, S., Kautz, J., Hatamizadeh, A. (2024). Gated Delta Networks: Improving Mamba2 with Delta Rule. arXiv:2412.06464.
  20. Yuan, J. et al. (2025). Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. arXiv:2502.11089.
  21. Zhang, B., Sennrich, R. (2019). Root Mean Square Layer Normalization. arXiv:1910.07467.
  22. Zheng, L. et al. (2023). SGLang: Efficient Execution of Structured Language Model Programs. arXiv:2312.07104.

来源:Hacker News · aleph-alpha.com