GlmMoeDsaForCausalLM model_type · glm_moe_dsa IndexShare MLA + DSA MTP ×1

GLM-5.2 架构解剖

GLM-5.2 在 ATOM 里没有自己的模型文件——它挂在 atom/models/deepseek_v2.py 上， 子类只覆盖了一张量化排除名映射。真正需要读懂的不是模型定义，而是 config 里那 78 项 indexer_types 如何一路改写构造期的模块树、运行期的 buffer 绑定、PP 的跨 rank 通信，以及 MTP 的草稿循环。

753 B总参数（含 MTP）
~41 B单 token 激活
78 + 1主干层 + MTP 层
21 / 78拥有 indexer 的层
6144hidden_size
1 Mmax_position
2048index_topk

规模

753.3 B 参数，单 token 只碰 41.25 B

总量的 96.2 % 是 MoE 路由专家（75 层 × 256 专家 = 724.8 B）， 画成条形图只会是一根长柱加七条细缝。真正有信息量的是单 token 前向实际碰到的 41.25 B：

MoE 专家8 routed + 1 shared × 75 层 25.48 B61.8 %
注意力 MLA78 层 × 165.0 M 12.87 B31.2 %
Embed + LM head2 × 154880 × 6144 1.90 B4.6 %
稠密 MLP层 0–2 × 226.5 M 0.68 B1.6 %
Indexer（DSA）21 个 full 层 × 9.4 M 0.20 B0.5 %
MoE gate75 层 × 1.57 M 0.12 B0.3 %
激活合计 41.25 B / 总量 753.3 B　→　5.5 % IndexShare 让 indexer 只占激活的 0.5 %（不共享时 1.8 %）
总参数构成参数量占比说明
MoE 路由专家724.78 B96.21 %75 层 × 256 专家 × 3 投影 × 2048 × 6144
注意力 MLA12.87 B1.71 %78 层 × 165.0 M
MTP 块（层 78）9.95 B1.32 %自带完整注意力 + MoE + indexer
MoE 共享专家2.83 B0.38 %75 层 × 37.75 M
Embed + LM head1.90 B0.25 %BF16，未量化
稠密 MLP0.68 B0.09 %层 0–2
Indexer（DSA）0.20 B0.03 %仅 21 个 full 层
MoE gate0.12 B0.02 %75 层

按张量形状累加得 753.3 B，与 checkpoint 的 total_size = 755 617 140 416 字节（FP8 主体 + BF16 embedding/scale）互相印证。

定位

它是 DeepSeek-V3.2 的骨架，配 GLM 自己的尺寸

ATOM 的入口判断只有一句：hasattr(config, "index_topk") 为真 → 这是稀疏 MLA（DSA）模型， 走 DeepSeek-V3.2 的全部代码路径。GLM 的子类薄到只剩一件事：

唯一的子类覆盖

一张量化排除名映射

GLM 的 HF 量化配置写的是 indexers_proj，而 ATOM 的模块路径叫 indexer.weights_proj；不映射的话这条 BF16 投影会被 FP4/MXFP4 回退量化掉。

atom/models/deepseek_v2.py:3444 GlmMoeDsaForCausalLM.quant_exclude_name_mapping

构造期的一次合并

fused_qkv_a_proj

checkpoint 里的注意力权重是拆开的 q_a_proj + kv_a_proj_with_mqa， ATOM 合并成一条 MergedReplicatedLinear，且强制 needs_preshuffled_weight—— 它的前向调用的是 preshuffle blockscale GEMM。

6144 → 2048 + 576 · 16.1 M 参数 / 层

其余全部由 config 数值驱动，没有一处新的算子拓扑。所以这份文档的重心不在"模型怎么写"， 而在那 78 项 indexer_types 在 ATOM 里落到了哪几处代码。

核心机制 · 层调度

78 层里只有 21 层真的算 top-k

这是 GLM-5.2 相对 GLM-5 / 5.1 唯一的结构性改动，也是整份 config 里唯一真正新的东西： indexer_types 把每层标成 full（自己跑 indexer，选出 2048 个 KV） 或 shared（不带 indexer 权重，直接复用上一个 full 层选出来的下标）。 节奏是 前 3 层连续 full，之后每 4 层一个 full。

indexer
mlp
悬停任一格查看该层配置 · 右端虚线格为 MTP 层 78
full — 本层计算 top-k，写入共享 buffer shared — 无 indexer 权重，读上一层结果 dense MLP — intermediate 12288 MoE — 256 专家 / top-8 / 1 共享 MTP 层 78 — 自带完整 indexer

来自 checkpoint 的直接证据

model.layers.0.self_attn. 下有 indexer.wq_b / wk / weights_proj / k_norm； model.layers.5.self_attn. 下这四个张量完全不存在。 shared 层不是"算了但丢弃"，是磁盘上就没有权重。

省了多少

indexer 前向从 78 次降到 21 次（÷3.7）。index cache 也同步减少： 每 token 每 full 层 index_head_dim + 4 = 132 B（uint8 装 fp8 + scale）， 21 层共 ≈ 2.7 KB/token，而不是 10.3 KB/token。

结构全图

从 input_ids 到 logits，再放大其中一层

左边是完整模型栈，最左侧那条窄色带按层高比例画出 78 层的 indexer 调度（琥珀=full，灰=shared）。 右边把层 3–77 中的任意一层完全展开，每个模块标了权重形状与参数量。 注意 fused_qkv_a_proj 那一格——它是一次 [2624, 6144] 的 GEMM， 输出 按列切成三份（2624 = 2048 + 512 + 64），不是投影到某个三维形状； 三条分叉箭头上的胶囊标签就是各自的宽度。

完整模型栈 GlmMoeDsaForCausalLM · 753.3 B input_ids [T] embed_tokens 154880 × 6144 · BF16 951.5 M · 每 rank 全表 layers 0–2 · dense MLP 6144 ↔ 12288 0.68 B layers 3–77 · MoE × 75 层 256 routed (top-8) + 1 shared · d_ff 2048 742.6 B 参数 激活 25.5 B 每层 self_attn 165.0 M 每层 mlp 9.70 B 21 层带 indexer 57 层复用（IndexShare） KV cache 576 / token / 层 index cache 132 B（仅 full） ↓ 右侧展开其中一层 结构逐层相同 model.norm RMSNorm(6144) lm_head 154880 × 6144 · BF16 951.5 M · TP 分片 logits [T, 154880] layers.78 · MTP enorm · hnorm · shared_head.norm eh_proj [6144 × 12288] 完整注意力 + MoE + indexer · 9.95 B idx 77 末层 hidden + 下一 token embed DecoderLayer 展开（层 3–77 之一） 单层 9.88 B（full 索引层） / 9.87 B（shared 索引层） residual stream · 6144 self_attn · MLA + DSA input_layernorm · RMSNorm(6144) ATOM 可与 all-reduce + quant 融成一个 kernel fused_qkv_a_proj [2624, 6144] 一次 GEMM，输出按列切三份 2624 = 2048 + 512 + 64 · 16.1 M q_a_layernorm RMSNorm(2048) kv_a_layernorm RMSNorm(512) RoPE 交错 · θ=8e6 q_b_proj 2048 → 64×256 33.6 M kv_b_proj 512 → 64×(192+256) 14.7 M KV cache latent 512 + 64 576 B/token/层 sparse MLA 64 heads · q [64, 256] × latent KV [576] → out [64, 256] 只在 top-k 选中的 2048 个 KV 上计算 o_proj [6144, 16384] 100.7 M — 单层最大的注意力权重 Indexer · DSA 仅 21 个 full 层拥有此模块 wq_b 2048 → 32×128 8.4 M wk 6144 → 128 0.79 M weights_proj 6144→32 0.20 M k_norm = LayerNorm fp32（含 bias） rope 交错 · fp8 per_1x128 (ue8m0) 写 index cache · 132 B/token/层 9.4 M → top-k = 2048 _sparse_kv_indices_gpu int32 [tokens × 2048] · 全部层共用 aiter_mla.py:384 分配 · :522 重绑 57 个 shared 层：indexer = None 0 参数 · 0 计算 · skip_topk = True 跳过上面那个框，直接用此 buffer mlp · MoE post_attention_layernorm · RMSNorm(6144) gate [256, 6144] · 1.57 M sigmoid + e_score_correction_bias noaux_tc → top-8 · scale 2.5 · fp32 路由 shared_experts × 1 6144 ↔ 2048 · 37.7 M 与路由专家双流并行 routed_experts × 256 每专家 gate_proj / up_proj [2048, 6144] + down_proj [6144, 2048] 9.66 B 参数 每 token 只激活 8 个 → 302 M combine = routed + shared （scale 2.5 已在 topk 权重里乘过） q_c 2048 kv_c 512 k_pe 64 hidden 6144 bf16 q_c 2048 full 层 writes reads · 决定这 64 个头读哪 2048 个 KV + + residual add residual add → 下一层 [T, 6144]
右侧琥珀色那条 lane 是全图的重心：Indexer 只存在于 21 个 full 层（9.4 M 参数， 占单层 9.88 B 的 0.095 %），它把 top-k 写进 _sparse_kv_indices_gpu； 57 个 shared 层的 self.indexer 是 None，跳过整个框直接读同一块 buffer。 两类层的其余部分——注意力、MoE、残差——逐字节相同。 稀疏性由模型级标志 mla_modules.is_sparse 决定，而不是"本层是否有 indexer"， 这正是 shared 层能在没有 indexer 的情况下仍然跑稀疏 MLA 的原因 （atom/model_ops/attention_mla.py:486）。

构造期

_indexer_weights_shared()

按 indexer_types[layer_id] == "shared" 判定，直接把 self.indexer 置为 None，不建任何参数。

atom/models/deepseek_v2.py:406 · 2612

运行期

metadata builder 统一重绑

遍历 static_forward_context，把每个带 sparse_kv_indices_buffer 属性的模块和 impl 都指向同一块 GPU buffer（DCP 的两个补偿 buffer 一起绑）。

atom/model_ops/attentions/aiter_mla.py:522-533

流水并行

跨 rank 传递 top-k

若某 rank 的首层是 shared，它需要上一 rank 的选择结果。 ModelRunner 预先算好 _pp_recv_needs_sparse / _pp_send_needs_sparse， 只在必要时收发。

atom/model_engine/model_runner.py:2956-3010 · 3133

不变量

第一层不能是 shared

否则它会读到一块从未写入的 buffer。SGLang 插件在 setup 时显式断言这一点； GLM-5.2 的 config 用连续三个 full 开头正好满足。

atom/plugin/sglang/models/glm52_dsa.py:60-70

尺寸对照

GLM-5 → 5.1 → 5.2：骨架没动，只加了 IndexShare 和长上下文

把三代 GLM-5 和 DeepSeek-V3.2 并排放，就能看清 ATOM 为什么能一个文件全吃下： 所有差异都是 config 数值，没有一处是新的算子拓扑。

config 字段GLM-5 / 5.1GLM-5.2DeepSeek-V3.2说明
num_hidden_layers787861+1 层 MTP（index 78）
hidden_size614461447168—
num_attention_heads6464128头少一半，头维更宽
q_lora_rank204820481536indexer 的 wq_b 也吃这个秩
kv_lora_rank512512512latent KV 宽度 512+64=576
qk_nope / qk_rope192 / 64192 / 64128 / 64qk_head_dim = 256
v_head_dim256256128V 比 K 的 nope 部分还宽，DeepSeek 系没有这种配置
n_group / topk_group1 / 11 / 18 / 4分组路由退化为全局 top-8
index_n_heads323264index_head_dim 三者都是 128
index_topk_freq—4—GLM-5.2 引入
indexer_types—78 项—逐层 full/shared，权威于 freq
index_share_for_mtp_iteration—true—草稿步之间复用 top-k
max_position_embeddings202 7521 048 576163 840GLM-5.2 迈进 1M
rope_theta1e68e61e4 + YaRN×40GLM 不用 YaRN，纯 default rope
rope_interleavetruetrue未设→ is_neox_style = false
indexer_rope_interleavetruetrue未设（默认 neox）GLM 把 indexer rope 也改成交错
vocab_size154 880154 880129 280embed 与 lm_head 各 951 M 参数

rope 的两个开关容易搞反

_is_neox_rope_style() 的语义是 interleave 为真 ⇒ neox 为假，两者互斥。 主 MLA rope 在 DeepSeek 缺省下就是交错的（default_is_neox=False）， 而 indexer rope 在 DeepSeek 缺省下是 neox（default_is_neox=True）—— GLM-5.x 用 indexer_rope_interleave=true 把后者也翻成交错。 两个默认值方向相反，是这段代码最容易读错的地方。 atom/models/deepseek_v2.py:308-322

权重实测

从 safetensors 头部读到的真实形状

下面每一行都是从 GLM-5.2-FP8 的 141 个分片头里直接读出来的，不是从 config 推的。 整个 checkpoint total_size = 755 617 140 416 字节，与 753 B 参数量互相印证。

张量dtypeshape拆解
model.embed_tokens.weightBF16[154880, 6144]未量化；ATOM 对 GLM 默认整表复制到每个 rank
…self_attn.q_a_proj.weightF8_E4M3[2048, 6144]与 kv_a 合并成 fused_qkv_a_proj
…self_attn.kv_a_proj_with_mqa.weightF8_E4M3[576, 6144]512 latent + 64 rope
…self_attn.q_b_proj.weightF8_E4M3[16384, 2048]64 × (192+64)
…self_attn.kv_b_proj.weightF8_E4M3[28672, 512]64 × (192 nope + 256 v)
…self_attn.o_proj.weightF8_E4M3[6144, 16384]64 × 256 → 6144；单层最大的注意力权重
…indexer.wq_b.weightF8_E4M3[4096, 2048]32 index 头 × 128
…indexer.wk.weightF8_E4M3[128, 6144]只有一份 key，32 个 index 头共享（MQA 式）
…indexer.weights_proj.weightBF16[32, 6144]保持 BF16，量化排除名单里
…indexer.k_norm.{weight,bias}BF16[128]带 bias，是 LayerNorm(fp32) 而非 RMSNorm
…mlp.experts.N.gate_proj.weightF8_E4M3[2048, 6144]×256 专家 ×3 投影 = 9.66 B / 层
…layers.0.mlp.gate_proj.weightF8_E4M3[12288, 6144]仅 layer 0–2 的稠密 MLP
…layers.78.eh_proj.weightBF16[6144, 12288]MTP：concat(enorm(e), hnorm(h)) → 6144

参数账

注意力 165.0 M/层 · indexer 9.4 M/full 层 · MoE 块 9.703 B/层 · 稠密 MLP 226.5 M/层
→ 主干 743.4 B + MTP 9.95 B ≈ 753.3 B，单 token 激活 ≈ 41 B。

KV 账

MLA latent KV：(512+64) × 2 B = 1152 B 每层每 token，78 层 ≈ 89.9 KB/token（BF16）， FP8 KV cache 时减半。index cache 另计 ≈ 2.7 KB/token。

ATOM 侧

为 GLM-5.2 加的七处优化

这些都不是模型定义的一部分——它们是 ATOM 为了让这套结构在 MI355X 上跑得快而加的， 每一处都有独立的环境变量可以回退。

优化内容门控
融合 indexer 内核 q-rope + fp8 量化 + k-cache 写入合成一个 kernel。准入条件是 index_head_dim == 128 且 qk_rope_head_dim == 64， GLM 与 DeepSeek-V3.2 同时满足，所以数学上与逐算子路径等价 ATOM_ENABLE_GLM_FUSED_INDEXER
wk + weights_proj GEMM 合并 IndexerWkWeightsProjLinear 把 [128,6144] 与 [32,6144] 拼成一次 GEMM 再 split。GLM checkpoint 用的是标准 indexer.wk / indexer.weights_proj 名字，合并加载路径直接可用 ATOM_ENABLE_DS_INDEXER_
QK_ROPE_CACHE_FUSION
复制词表 embedding ReplicatedEmbedding 整表存每个 rank，省掉 embedding 后的 all-reduce。 因为 tie_word_embeddings=false，embedding 与仍然分片的 lm_head 解耦，查表结果与"分片 + masked + all-reduce"逐位相同。代价是每 rank 多 951 M 参数 ATOM_REPLICATE_VOCAB_EMBED
stable_topk TP > 1 且是 GLM-5.2 时启用稳定排序，保证各 rank 选出完全相同的 KV 下标。 分数打平时不稳定排序会让各 rank 分歧，而 IndexShare 会把这个分歧传播到后面三层 自动（model_type + tp>1）
量化名映射 HF 量化配置写 indexers_proj，ATOM 模块路径叫 indexer.weights_proj； 不映射的话这条 BF16 投影会被 FP4/MXFP4 回退量化掉 子类常量
MoE 双流 共享专家与路由专家在独立 CUDA stream 上并行，注册成对 Dynamo 不透明的 custom op， 避免 torch.compile 把 dummy-run 的分支烤进图里 ATOM_DUAL_STREAM_MOE_
TOKEN_THRESHOLD
fused_qkv_a_proj checkpoint 里拆开的 q_a_proj + kv_a_proj_with_mqa 合成一条 MergedReplicatedLinear，并强制 needs_preshuffled_weight （其前向调的是 preshuffle blockscale GEMM） 默认开

注意点

两个容易踩的地方

其一 · freq 回退公式与 config 的 offset 约定不一致

ATOM 的回退公式是 max(layer_id − offset, 0) % index_topk_freq != 0， offset 默认取 1。而 GLM-5.2 的 config 写的是 index_skip_topk_offset = 3——代入这个公式，max(3−3,0) % 4 == 0 会把 layer 3 判成 full，与 indexer_types 里的 shared 相反，整张表整体偏一层。

当前无害，因为 _should_skip_index_topk() 优先读 indexer_types，公式只是兜底。但如果将来出现只给 index_topk_freq / index_skip_topk_offset、不给 indexer_types 的 checkpoint，回退路径会选错层——而且不会报错， 只会让某些层读到错误的 KV 选择。遇到这类权重时先确认 indexer_types 是否存在。 atom/models/deepseek_v2.py:361-404

其二 · 上下文长度不到 2048 时 indexer 是 no-op

index_topk = 2048，低于这个长度时 top-k 会选中全部 KV，稀疏注意力退化成稠密。 真正的收益要到长上下文才出现——短序列的 benchmark 看不出 IndexShare 的任何价值， 调优时别拿 2K 以内的 ISL 去衡量它。

投机解码

MTP 层自带 indexer，但草稿步之间也共享

层 78 是一个完整的 nextn 块：enorm / hnorm / eh_proj / shared_head.norm，加上一整套注意力、MoE，以及自己的一套 indexer 权重。 这一点两边代码都写了注释强调：index_share_for_mtp_iteration 说的不是"MTP 复用主干的 top-k"， 而是"多个草稿步之间复用第 0 步的 top-k"。

draft step 0

set_skip_topk(False)

MTP 的 indexer 正常跑，为被草稿的位置选出 2048 个 KV。

draft step 1…N

set_skip_topk(True)

跳过 indexer，直接复用第 0 步写进 buffer 的下标。

实现位置：atom/spec_decode/eagle_proposer.py:343-357，配合 atom/models/deepseek_mtp.py:245 的 set_skip_topk()。 层判定还额外挡了一道：layer_id ≥ num_hidden_layers 时一律不 skip， 避免 MTP 层被主干的 full/shared 表误判成 shared 而读到空 buffer。

已验证的部署形态

仓库里跑通过的两套配置

单机 · Agentic

4×MI355X TP4 + MTP

  • checkpoint GLM-5.2-MXFP4，FP8 KV cache
  • MTP 3 个草稿 token，接受率 0.6633 → 每次前向约 2.99 token
  • GPU 前缀缓存 + 200 GiB LMCache CPU 层
recipes/Agentic-GLM-5.2.md

PD 分离 · atomesh

PP4×TP1 prefill + TP4 decode

  • Mooncake RDMA 传 KV，atomesh 路由
  • VLLM_PP_LAYER_PARTITION=18,20,20,20
  • 前轻分层是因为 rank 0 还扛着 951 M 参数的 embedding，均分会 OOM
recipes/mesh/GLM-5.2.md

另有一条 SGLang 插件路径：atom/plugin/sglang/models/glm52_dsa.py 刻意不安装 SGLang 自己的 MLA 前端，而是保留 ATOM 原生 MLAAttention，这样 full 层写共享物理下标 buffer、 shared 层读同一块 buffer 的机制才能成立。

读码路线

按这个顺序看，一小时能过完

要理解的东西文件 · 位置
模型注册与类体atom/models/deepseek_v2.py:3444 GlmMoeDsaForCausalLM
full/shared 判定atom/models/deepseek_v2.py:361 _should_skip_index_topk · :406 _indexer_weights_shared
rope 交错/neox 解析atom/models/deepseek_v2.py:308 _is_neox_rope_style
Indexer 模块与前向atom/models/deepseek_v2.py:2154 class Indexer · :2288 forward_impl
top-k 内核（prefill / decode / DCP）atom/models/deepseek_v2.py:1536 sparse_attn_indexer
MLA 构造：rope、indexer、MLAModulesatom/models/deepseek_v2.py:2404-2700
稀疏标志与 buffer 读取atom/model_ops/attention_mla.py:486-500 · :1493
共享 buffer 分配与重绑atom/model_ops/attentions/aiter_mla.py:384 · :522
PP 跨 rank 传 top-katom/model_engine/model_runner.py:2956 · 3133
MTP 的 skip_topk 切换atom/spec_decode/eagle_proposer.py:343 · atom/models/deepseek_mtp.py:245
MoE：sigmoid + noaux_tc + 双流atom/models/deepseek_v2.py:1077-1310
ATOM 专属优化的门控开关atom/utils/envs.py ATOM_ENABLE_GLM_FUSED_INDEXER · ATOM_REPLICATE_VOCAB_EMBED