它与旗舰 GLM-5.3 同名不同构。旗舰是 78 层的 glm_moe_dsa,在 ATOM 里没有自己的模型文件;
Flash 是 GLM-5 家族里唯一拥有独立模型文件的成员——
atom/models/glm5_next.py,1257 行:45 层里 34 层线性注意力、11 层稀疏 MLA,
残差是 4 路宽的,全模型没有位置编码,稀疏索引按 4 个 token 一池打分。
先分清对象
它新在哪
inputs_embeds 在 embedding 处就被展开成
[T, 4, 4096],一路带到第 45 层,最后用无权重的均值塌回
[T, 4096]。每个子层进出都要收一次、放一次。qk_rope_head_dim == 0,
MLA 不转 rope,indexer 也不转;位置信息来自 KDA 层的因果卷积与递推。与旗舰 GLM-5.3 的关系
只有名字是共享的。旗舰 GLM-5.3(GlmMoeDsaForCausalLM,78 层 / 6144 / 256 专家 /
753.3 B)的 config 与 GLM-5.2 逐字段相同,挂在 deepseek_v2.py 上;
Flash 的 45 层混合结构、四路残差与 k-pool 索引在那边一样都没有。
checkpoint 是原生多模态的(另带 24 层视觉塔),ATOM 只 serve 文本路径,
model.visual.* 在加载时按前缀跳过。
规模
下面每个数字都来自 62 个 safetensors 分片的头部(读 shape 相加,fp8 的
weight_scale_inv 单列,不计入参数量)。官方标称 “320B 总 / 18B 激活”
对应文本侧含 MTP 的 320.78 B 与实测 17.38 B。
激活量是怎么凑出来的
注意力占激活量的 35 %,远高于旗舰的 31 %—— 因为 KDA 层每层 137.7 M 参数全部参与,没有稀疏可言。
量化布局
q_a/q_b/kv_a/o_proj。kv_b_proj、
lm_head、embedding、所有 norm 与 mHC 参数。weight_scale_inv 的形状即分块数,例如专家
gate_proj [2048, 4096] → scale [16, 32]。核心机制 · 层调度
layer_types 把 45 层排成一个严格的 4 拍循环:3 层 KDA,1 层 DSA,
从层 3 起每 4 层一个 DSA,最后一层(44)是 KDA。mlp_layer_types 是另一条独立的
schedule:前 3 层稠密 MLP,之后全是 MoE。
图 1 · 层调度
checkpoint 里的直接证据
layers.0.self_attn(KDA):q/k/v_proj、q/k/v_conv1d、
A_log、dt_bias、f_a/f_b_proj、g_a/g_b_proj、
b_proj、o_norm、o_proj——没有任何 q_a/kv_a。layers.3.self_attn(DSA):q_a_proj、q_b_proj、
kv_a_proj_with_mqa、kv_b_proj、o_proj 与
indexer.*——没有任何 conv1d。hc_attn_* / hc_ffn_* 形状完全一致,
所以 mHC 与注意力种类无关。这条 schedule 省了什么
KV cache 只由 11 层承担而不是 45 层。BF16 下每 token
11 × 576 × 2 B = 12 672 B;若 45 层全是 MLA 则是 51 840 B。
代价是 34 层各自持有一份每请求的递推状态(见“缓存账”一节),
它随并发数而不是随上下文长度增长,也正因如此
prefix caching 必须关掉。
核心机制 · 残差
这是与 ATOM 现有模型差别最大的一处。Glm5NextHyperConnection
与 DeepSeek-V4 的 Block 在数学上完全一致——同样的 sigmoid 门、
同样的 Sinkhorn 日程(含特殊的第一轮)、同样的 HC_POST_MULT = 2.0——
连 checkpoint 的张量名(hc_attn_fn / hc_attn_base /
hc_attn_scale)都是 Block 期待的名字,
hc_attn_fn 的 [24, 16384] 正好是 hc_split_sinkhorn 的
mixes 布局。
图 2 · mHC 四路残差
x + sublayer(x),而是子层输出与旧残差按两组学出来的权重整体替换。comb 是每 token 一个的 4×4 双随机矩阵,所以四条 stream 在每一层都被重新混合一次——它们不是四份副本,而是四条各自演化、持续交换信息的通路。mHC 参数为什么平铺在 layer 上
参数的名字来自它的属性路径。checkpoint 里它们是
layers.N.hc_attn_fn 这样平铺在层上的,
所以 ATOM 也必须把它们声明在 Glm5NextDecoderLayer 上,
而不是塞进 Glm5NextHyperConnection 那个 helper——
埋进子模块会把名字改成 hc.fn,然后静默地停留在初始值。
结构全图
图 3 · 结构全图
[T, 4096] 残差,而是 [T, 4, 4096]。四个编号站点里,① 和 ③ 是同一段代码,只是换了一组 hc_* 参数;收放的机制见图 2。分支 A · 34 层的主力
Glm5NextKDAAttention 直接继承 Kimi-K3 的 KimiKDAAttention。
两者结构上是同一层——同样把 q/k/v_conv1d 分开存、同样的每头 A_log、
每通道 dt_bias、低秩遗忘门——唯一的差别是输出门:
Kimi 一步投影(g_proj),GLM 拆成秩 128 的两步(g_b_proj @ g_a_proj)。
图 4 · KDA 数据通路
q|k|v|g|b|f_a 是 forward 里硬编码的边界,与 process_weights_after_loading 的拼接顺序必须一致。TP8 下 in_proj 变成 [4232, 4096](4×1024 + 8 + 128),ssm_state 变成 [8, 128, 128]。为什么折叠输出门是精确的
g = g_b_proj(g_a_proj(x)),两步之间没有激活函数,
所以 W_g = W_b @ W_a 在 fp32 下乘一次就等价。ATOM 在
process_weights_after_loading 里把结果写进
in_proj.weight[3·lp : 4·lp],再把两个因子的存储清空。
这样父类的状态缓存、TP 切分与 CUDA graph 处理全部原样复用。
加载完整性靠 _loaded_input_shards 记账——
q/k/v 三个 checkpoint 张量指向同一个参数,
参数级的加载报告分辨不出“q 到了”和“q/k/v 都到了”。
transformers 侧的一个真 bug
checkpoint 的 modules_to_not_convert 写的是
model.layers.N.self_attn.f_a_proj,但真实 key 前缀是
model.language_model.layers.,而且 glm5_next 的转换表会在
FP8 量化器跑之前把它改名成 self_attn.forget_gate.f_a_proj。
两处都对不上 → 全部 68 个遗忘门线性层(34 层 × 2)被包进
FP8Linear 却仍持 BF16 权重、配一个新初始化的
weight_scale_inv。用 transformers 直接跑这个 checkpoint 的人,
KDA 的衰减是静默损坏的。ATOM 不走这条路,不受影响。
分支 B · 11 层的 MLA
config 写的是 qk_rope_head_dim = 0、mla_use_nope = true。
最自然的实现是让 rope 那一半是零宽切片——这也是本次移植最初的选择,
而它在两个互相独立的方向上都是错的。
图 5 · NoPE 的两种实现
_ROPE_PAD 是 0 还是 64。看起来更“干净”的那个是错的——而且错得不显眼:prefill 路径对 head-dim 是通用的,所以只做 prefill 的逐层比对会给一个全绿的结果。零宽切片错在哪 · 其一:静默算错
paged MLA entry 按 kv_lora_rank + qk_rope_head_dim 算出来是 512,
而 aiter 的 asm decode kernel 是按 576 宽的查询编译的。
它只在 gfx1250 路径上断言这个 576,cfg_mla_asm 的分发表又从不按 head_size 选择,
于是 gfx950 上这次不匹配被算出来而不是被拒绝。
更麻烦的是它只影响 decode:prefill 走 flash_attn_varlen,
对 head-dim 是通用的,所以任何只覆盖 prefill 的逐层比对都会放行一个坏模型。
零宽切片错在哪 · 其二:直接崩
KV_PeDim == 0 让每一处 tl.arange(0, KV_PeDim) 变成
arange(0, 0),Triton 在编译期就拒绝;上游 aiter 的三个
gather_kv_b_proj kernel 对此没有任何保护。
它只在命中前缀缓存的 prefill 路径上触发,
所以单条 prompt 的 demo 能过,一上并发就死于
NameError('kv_pe_data is not defined')。
两个容易漏掉的配套
其一,填充后的宽度必须用 config.mla_kv_entry_dim 声明给缓存分配器:
KimiMLAGDNBackend 会遮蔽普通 MLA 分配器、直接按原始 config 算 pool 大小,
少了这个声明就会 pool 按 512 建、写入按 576 走,服务在启动时死于
shape '[..., -1, 576]' is invalid。
其二,这个 pad 必须由 _ZeroRopePad 追加,而它刻意不是
nn.Module——把 q_b_proj 包进一个 Module 会在参数路径里
插进一层,权重就再也加载不上了。
核心机制 · 稀疏索引
DSA 层的 indexer 不给每个 token 存一行 key。连续
index_kpool = 4 个 token 压成 一个 cache entry,
top-k 因此在池粒度上跑——index_topk / index_kpool = 512 个候选池——
每个被选中的池再展开回它覆盖的 4 个 token 位置。
还没压满的“尾池”按 index_kpool_always_select_tail 永远入选,
所以最新的 token 绝不会被丢掉。
图 6 · k-pool 索引
k 与 gate 必须活过产生它们的那一步。ATOM 把它们塞进 KDA 已经拥有的每请求状态槽而不是另开一个分页缓存,于是免费继承了状态池的生命周期、fork 与搬迁语义。少了搬迁那一步,一个被重定位的请求就会读到前一个请求留在新槽里的半截池——只有在负载下池边界移动时才会暴露的那种损坏。block size 为什么被改成 64
deepgemm_fp8_paged_mqa_logits 只在 preshuffle 布局下算得对,
而该布局要求每个 block 的行数是 16 的倍数。
一个 B token 的 block 需要 B / 4 个索引行,
所以最小可行的 B 就是 4 × 16 = 64。
Config 在架构判定处直接改写 kv_cache_block_size,
因为 BlockManager 与 slot_mapping 都假定全局只有一个 block size。
index cache 1584 B/token → 396 B/token
off-switch 为什么必须报错而不是降级
ATOM_GLM5_KPOOL=0 同时关掉池化写入、池化打分与池化选择,
让两套实现成为同一个选择的两个真正独立的实现——
这正是短上下文下的 A/B 能成为检验而不是同义反复的原因。
但超过 index_topk 之后两者的选择会真正分叉,
此时返回 token 粒度的回退结果是静默地错,而不只是慢,
所以 _assert_kpool_regime 在那里直接抛 NotImplementedError。
权重实测
checkpoint 里所有张量都在 model.language_model.* / model.visual.* 之下,
只有 lm_head 在顶层。ATOM 的 weights_mapping 因此只有一条规则:
"model.language_model." → "model.";model.visual.* 由
skip_weight_prefixes 丢弃,MTP 层 45 由 loader 的“层号越界”过滤自动丢弃。
| checkpoint 张量(去掉前缀) | dtype | shape | 说明 |
|---|---|---|---|
| embed_tokens.weight | BF16 | [154880, 4096] | 634.4 M,每 rank 全表 |
| layers.N.hc_attn_fn / hc_ffn_fn | BF16 | [24, 16384] | 24 = (2+4)×4,16384 = 4×4096;ATOM 侧声明为 fp32,加载时上转 |
| layers.N.hc_*_base / hc_*_scale | F32 | [24] / [3] | 激活前的加性偏置 / pre·post·comb 三个缩放 |
| KDA 层(0,1,2,4,5,6,8,…,44 共 34 层) | |||
| self_attn.q_proj / k_proj / v_proj | BF16 | [8192, 4096] | 各自映射到融合 in_proj 的分片 0/1/2 |
| self_attn.q_conv1d / k_conv1d / v_conv1d | BF16 | [8192, 1, 4] | 每通道因果卷积,kernel = 4 |
| self_attn.g_a_proj / g_b_proj | BF16 | [128, 4096] / [8192, 128] | 秩 128 输出门,load 后折叠进 in_proj 分片 3 |
| self_attn.f_a_proj / f_b_proj | BF16 | [128, 4096] / [8192, 128] | 低秩遗忘门 |
| self_attn.A_log / dt_bias | F32 | [64] / [8192] | 每头 / 每通道 |
| self_attn.b_proj.weight | BF16 | [64, 4096] | beta,每头一个标量 |
| self_attn.o_norm.weight / o_proj.weight | BF16 | [128] / [4096, 8192] | 门控 RMSNorm + 输出投影 |
| DSA 层(3,7,11,…,43 共 11 层) | |||
| self_attn.q_a_proj / q_b_proj | F8_E4M3 | [1536, 4096] / [16384, 1536] | scale [12,32] / [128,12] |
| self_attn.kv_a_proj_with_mqa | F8_E4M3 | [512, 4096] | 没有 rope 分量,输出宽度就是 kv_lora_rank |
| self_attn.kv_b_proj.weight | BF16 | [32768, 512] | 64 × (256 + 256),不量化,被 MLA 吸收 |
| self_attn.o_proj.weight | F8_E4M3 | [4096, 16384] | 64 头 × v_head_dim 256 |
| indexer.wq_b / wk / weights_proj | BF16 | [4096, 1536] / [128, 4096] / [32, 4096] | 32 头 × 128;整个 indexer 都在量化排除名单里 |
| indexer.index_kpool_compress_gate | BF16 | [128, 4096] | checkpoint 存成裸矩阵,映射规则补上 .weight |
| indexer.index_kpool_compress_ape | BF16 | [4, 128] | 每 slot 的可学习偏置,进 softmax 前相加 |
| FFN | |||
| mlp.gate_proj / up_proj(层 0–2) | F8_E4M3 | [12288, 4096] | scale [96, 32] |
| mlp.gate.weight / e_score_correction_bias | BF16 / F32 | [288, 4096] / [288] | ATOM 用 nn.Linear(dtype=fp32),与参考实现一致 |
| mlp.experts.{0..287}.{gate,up}_proj | F8_E4M3 | [2048, 4096] | 每层 288 份 · 单专家 25.17 M |
| mlp.experts.*.down_proj | F8_E4M3 | [4096, 2048] | 合并成 gate_up_proj / down_proj 两个大张量 |
| mlp.shared_experts.* | F8_E4M3 | 同上 | 1 个,与 routed 输出相加 |
| 不加载 | |||
| layers.45.*(MTP) | — | eh_proj [4096, 8192] | 完整 DSA + MoE + enorm/hnorm/shared_head.norm,7.43 B |
| model.visual.* | — | 24 层 · hidden 1024 | image 448 / patch 14 / merge 2 / 时序 patch 2 → out 4096,0.56 B |
packed_modules_mapping 为什么可以是层无关的
Kimi-K3 必须逐 KDA 层枚举 .q_proj → .in_proj,因为它的全注意力层也有一个
g_proj 不能被折叠。GLM-5.3-Flash 不需要:它的 MLA 层用的是
q_a_proj / q_b_proj / kv_a_proj_with_mqa / kv_b_proj,
没有一个包含 .q_proj、.k_proj、.v_proj 子串(匹配锚定在前导的点上)。
于是一张层无关的表就够了——而且只有层无关的表才能挂在类上,
让 model_runner 在模型构造之前读到它。
这个顺序很要紧:从 __init__ 里二次重映射量化配置会破坏它的 layer pattern,
并静默地把每一个注意力投影都标成已量化。
显存
混合架构把显存开销劈成两块。分页部分(MLA KV + 索引 key)只由 11 个 DSA 层贡献, 按 token 计费;每请求部分(KDA 递推状态 + 尾池)由 34 个 KDA 层贡献, 按并发数计费,与上下文长度无关。
| 缓存 | 形状 / 计费单位 | 大小 | 备注 |
|---|---|---|---|
| MLA KV(分页) | 11 层 × 576 × dtype / token | 12 672 B/token BF16 · FP8 时减半 |
若 45 层全是 MLA 则是 51 840 B |
| 索引 key(分页) | 11 层 × 144 B / 4 token | 396 B/token | 未池化时 1 584 B;144 = (128 + 4) 向上对齐到 16 |
| 分页合计 | 每 token | 13 068 B ≈ 12.8 KiB | 1 M 上下文 ≈ 12.8 GiB(BF16 KV) |
| 一个 block(64 token) | MLA 811 008 B + index 25 344 B | 836 352 B | block size 由 index_kpool × 16 定出 |
| KDA conv 状态 | 34 层 × [3, 24576/tp] BF16 | 0.60 MiB/请求 TP8 |
24576 = 128×64×2 + 128×64 |
| KDA 递推状态 | 34 层 × [64/tp, 128, 128] F32 | 17.0 MiB/请求 TP8;TP4 为 34.0 |
必须 fp32,aiter 内核原样读回 |
| k-pool 尾池 | 11 层 × [slot, 2, 4, 128] BF16 | 22 KiB/请求 | 与 KDA 状态共用同一个槽 id |
| 每请求合计 | TP8 | 17.6 MiB | 并发 32 时约 0.55 GiB |
prefix caching 必须关
KDA 的递推状态是每请求的,不是每 block 的,无法像分页 KV 那样在请求间共享。
带前缀命中的 prefill 会 fork 状态槽:读槽(state_fork_src)与写槽不同,
尾池缓存必须跟着一起搬——relocate_state_slots 就是干这个的。
ATOM 侧
45 层里几乎每一块难的东西,ATOM 都已经因为别的模型存在了。真正新写的只有 k-pool 索引:
model_ops/glm5_next/indexer.py 负责分发,kpool.py 负责内核,
geometry.py 负责那几个 CPU 可测的几何契约。
| GLM-5.3-Flash 的部件 | ATOM 里已有的东西 | 贴合度 |
|---|---|---|
| KDA 线性注意力 | kimi_k3.KimiKDAAttention + aiter kimi_delta_attn | 非常近。同样分开存的 q/k/v_conv1d、每头 A_log、每通道 dt_bias、
低秩遗忘门,连 gate_lower_bound 都已经在读。仅需折叠输出门 |
| mHC 超连接 | sparse_attn_v4.hc_split_sinkhorn + deepseek_v4.Block | 数学上完全一致,连 checkpoint 的张量名都是它期待的名字;
dim = 4096 也满足 aiter 融合 mhc_pre/mhc_post 的
% 512 == 0 约束 |
| k-pool DSA 索引 | 新写 | DeepSeek-V4 的 Compressor 也按 compress_ratio=4 池化并带 ape 项,
但它是重叠的、带 RoPE 的;GLM 这个不重叠、无 RoPE |
| MLA | model_ops/attention_mla.py 经 MLAModules | 需要 _ROPE_PAD = 64 与 rope_is_zero_pad,见图 5 |
| 288 专家 sigmoid / noaux_tc | model_ops/fused_moe、glm4_moe.py、deepseek_v2.py | 直接可用 |
| Block FP8 128×128 | 既有 DeepSeek block-FP8 路径 | 直接可用 |
| clamped SwiGLU | model_ops/swiglu_oai.swiglu_oai_split | alpha=1, beta=0, limit=10.0 等价于参考实现的 silu(clamp(gate))·clamp(up) |
三个调试开关
ATOM_GLM5_KPOOL=0(默认 1)—— 同时关掉池化写入、池化打分与池化选择。
超过 2048 时直接报错,不静默降级。ATOM_GLM5_FORCE_DENSE_MLA=1(默认 0)—— 关掉稀疏。
2048 以内本就等价,任何输出差异就定位在 indexer / top-k 而不是 MLA 本身。ATOM_GLM5_DISABLE_FUSED_MHC=1(默认 0)—— 把 aiter 的
mhc_pre/mhc_post 换成 torch 参考路径,用于逐层对数。明确拒绝的组合
Config 在识别到 Glm5Next* 架构后会直接抛错而不是降级:
PCP > 1、DCP > 1、投机解码、TBO。理由是这些特性的池化索引 metadata / 状态布局
还没有对应的实现。kpool 的 dispatch 里另有两道同样性质的关卡:
DCP/PCP world size > 1、以及 max_seqlen_q > 1 的投机解码。
MTP 草稿层与多模态输入同样尚未接通。
读码路线
| # | 文件 · 位置 | 看什么 |
|---|---|---|
| 1 | atom/models/glm5_next.py · 模块 docstring | 三个决定:池粒度选择、64 lane 零 rope、输出门折叠。先读这段,后面全是它的展开 |
| 2 | _normalize_glm5_next_config() | config 别名怎么补齐;glm5_kda_layers 是 0-based(Kimi-K3 是 1-based,别照抄);
三处 raise 分别守住 NoPE、index 几何、层划分的完整性 |
| 3 | Glm5NextDecoderLayer.forward | 四个站点的顺序,以及 mHC 参数为什么平铺在 layer 上而不是塞进 helper |
| 4 | Glm5NextHyperConnection.pre / post_expand | 融合内核与 torch 参考两条路径并列,逐行对得上(图 2) |
| 5 | Glm5NextMLAAttention.__init__ / forward | _ZeroRopePad、mla_kv_entry_dim、
以及 indexer 为什么要在 mla_attn 之前被显式调用 |
| 6 | model_ops/glm5_next/kpool.py · 前 115 行 | 三个 *_ref 参考实现就是正确性的定义:池化 softmax 的轴、
Hadamard 的 1/√128、ue8m0 的幂次缩放 |
| 7 | model_ops/glm5_next/indexer.py | prefill 与 decode 两条分发;_kpool_write_completed_pools 如何从 tail cache 里
补上跨 chunk 的槽 |
| 8 | model_ops/attentions/kimi_mla_gdn_attn.py · 172–300 | 索引缓存行数、尾池 buffer 与 KDA 状态槽共享生命周期 |
| 9 | atom/config.py · glm5_kpool_block_size / 2035 附近 | block size 为什么被改成 64,以及被明确拒绝的并行特性 |
| 10 | recipes/GLM-5.3-Flash.md | bring-up 期发现的 4 个 bug(3 个在上游)与参考 oracle 的搭法 |