GLM5-744B 模型结构拆解和昇腾profilling分析

作者昇腾实战派知识地图https://blog.csdn.net/Lumos_Lovegood/article/details/161455142背景概述GLM-5 是智谱AI推出的第五代大语言模型采用混合专家MoE架构约 7450 亿总参数256 个专家、每 token 激活 8 个稀疏率 5.9%。GLM总共有744B参数40B的激活参数。维度参数值说明模型类型glm_moe_dsaGLM 混合专家模型带动态稀疏注意力隐藏层维度6,144模型主干宽度层数78Transformer 层总数注意力头数64多头注意力机制的头数前馈层维度12,288稠密层的前馈网络大小词汇表大小154,880支持约 15.5 万个 token最大序列长度202,752约 20 万 token超长上下文能力精度bfloat16训练 / 推理使用的数值精度本文对GLM5-744B 模型结构进行了拆解并以Atlas 800I A3混部的GLM5为例分析其profiling特征。硬件设备信息Atlas 800I A3推理服务器2. 智谱模型结构介绍本文以GLM5模型结构为例讲GLM5进行可视化处理如图1所示图中模型共有78层其中dense层三层Moe层75层。已知Hiddensize6144num_attention_heads64head_dim256q_lora_rank 2048index_topk 2048。2.1 DSA层输入shape如图1所示输入在经过一个RMSnorm后的shape为[BS6144]进入attention前QueryQuery经过Linear后会经过下投影降维到q_lora_rank此时的shape为[BS2048]。在经过一个上投影和RMSNorm后升维成[BS64*256]。此时的Query会被split成两部分。用作内容的部分shape为[BS64*192]其中64代表head_dim192代表qk_nope_head_dim。用作位置的部分shape为[BS64*64]其中前一个64代表head_dim后一个64代表qk_rope_head_dim。用作内容的部分的再经过Q absorb残差吸收线性变幻后shape为[BS64*512]其中64代表head_dim512代表 kv_lora_rank它会和经过位置编码的部分进行拼接最终shape为[BS64*576]其中576为qk_rope_head_dimkv_lora_rank。key valuekey value经过Linear后会经过下投影降维到kv_lora_rankqk_rope_head_dim此时的shape为[BS576]。此时kv会被分离成用作位置的shape和用作内容的shape。在和历史KV cache拼接后和位置编码后的位置信息进行拼接得到了[BS576]。在这里我们会利用indexer对得到的所有KV做筛选只保留top k 2048个KV值丢弃剩下的KV值最后我们会得到k值的shape为[B, 20481576]其中2048为index_topk576为qk_rope_head_dimkv_lora_rank。由于使用的是MQA结构所以K和V的头数为1。进入attention后在完成attention计算后的shape为[BS642048]其中64是num_attention_heads 2048是index_topk。在经过两个线性化层后最终shape为[B,S,hidden_size]2.2MOE层路由打分输入来自注意力层的 hidden states门控网络计算每个专家的分数logits→ Softmax 得到概率选出Top-K 个专家GLM5 的 K 为8专家计算稀疏激活只把当前 token 送给被选中的 K 个专家做前向同时有一个共享专家不参与路由选择必须经过其他专家不参与计算加权合并用门控输出的概率做权重把 K1 个专家的结果加权相加作为 MoE 层输出传给下一层Moe层的所有shape包括输入到输出都是 [B,S,hidden_size] 只有这样才可以做加权相加。2.3 FFN层FFN(x)down(Swish(gate(x))⊙up(x))FFN层输入[B,S,6144]6144hidden_size是会升维成[B,S,12288]其中12288是2*hidden_size同时在另一边gate 对每一个维度输出一个 0~1 左右的权重 然后逐元素相乘Swish(gate)⊙up(x)。得到结果后最终降维到[B,S,6144]输出3. DSAMLA流程与源码解析首先一句话简介 DSA MLAMQA lora Lightning Indexer Top‑kDSA计算逻辑可以分为8个部分3.1 阶段1DSA初始化SFA的QKV初始化主要在AscendSFAImpl类SFA会获取MLA给定的参数self.num_heads num_heads # 注意力头的总数Q总头数如64 self.head_size head_size # 每个头的维度大小如128 self.scale float(scale) # 注意力缩放因子 1/sqrt(head_dim) self.num_kv_heads num_kv_heads # KV的头数GQA模式比Q头少如8 self.kv_cache_dtype kv_cache_dtype # KV Cache存储的数据类型fp16/int8等 self.q_proj kwargs[q_proj] if self.q_lora_rank is None else kwargs[q_b_proj] # Q投影层MLA低秩分支使用q_b_proj self.fused_qkv_a_proj kwargs.get(fused_qkv_a_proj) # MLA低秩A投影QKV共享 self.kv_b_proj kwargs[kv_b_proj] # KV的升维B投影层 self.o_proj kwargs[o_proj] # 注意力输出投影层 self.indexer kwargs[indexer] # DSA稀疏索引器LightningIndexer self.kv_a_proj_with_mqa kwargs.get(kv_a_proj_with_mqa) # 支持GQA/MQA的KV低秩A投影 self.kv_a_layernorm kwargs.get(kv_a_layernorm) # KV低秩特征的LayerNorm归一化 self.q_a_layernorm kwargs.get(q_a_layernorm) # Q低秩特征的LayerNorm归一化 self.num_queries_per_kv self.num_heads // self.num_kv_heads # GQA1个KV对应几个Q头 self.tp_size get_tensor_model_parallel_world_size() # 张量并行总卡数 self.tp_rank get_tp_group().rank_in_group # 当前卡的TP并行组内编号3.2 阶段2MLA 低秩 QKV 投影q_c 代表Q 低秩特征Query Low-Rank Featurekv_no_split代表未拆分的 KV 低秩特征KV Low-Rank Feature这一阶段主要为了获取这两个值。qkv_lora self.fused_qkv_a_proj(hidden_states)[0] q_c, kv_no_split qkv_lora.split([self.q_lora_rank, self.kv_lora_rank self.qk_rope_head_dim], dim-1) q_c self.q_a_layernorm(q_c)3.3 阶段3生成 k_liIndexer 用的轻量 Key相关代码k_li, k_li_scale self.indexer_select_pre_process(xhidden_states, coscos, sinsin) k_li, _ self.wk(x) # [b,s,7168] [7168,128] [b,s,128] k_li self.k_norm(k_li).unsqueeze(1) k_li k_li.view(-1, 1, self.head_dim)3.4 阶段4生成完整 KV 并写入 KV Cache把低秩 KV → 升维成存入kvcache的 KV并最终用于SFA计算存入kv_cache [0] 和 kv_cache [1]K_nopeK_pe带 RoPE相关代码k_pe, k_nope self.exec_kv(kv_no_split, cos, sin, kv_cache, slot_mapping, attn_metadata) def exec_kv( self, kv_no_split: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, kv_cache: tuple, slots: torch.Tensor, attn_metadata: M, ): B kv_no_split.shape[0] N self.num_kv_heads S 1 # npu_kv_rmsnorm_rope_cache needs [B, N, S, D] kv_no_split kv_no_split.view(B, N, S, self.kv_lora_rank self.qk_rope_head_dim) cache_mode PA if self.enable_dsa_cp: ... else: torch_npu.npu_kv_rmsnorm_rope_cache( kv_no_split, self.kv_a_layernorm.weight, # type: ignore[union-attr] cos, sin, slots.to(torch.int64), kv_cache[1], kv_cache[0], epsilonself.kv_a_layernorm.variance_epsilon, # type: ignore[union-attr] cache_modecache_mode, ) return None, None )3.5 阶段5多卡并行TP/CP处理fused_kv_no_split, kv_ag_handle all_gather_async(...)3.6 阶段6Q 升维 RoPE 位置编码ql_nope, q_pe self._q_proj_and_k_up_proj(q_c) q_pe self.rope_single(q_pe, cos, sin)3.7 阶段7 LightningIndexer 核心计算在这一步骤中我们会调用在阶段获取之前存在kv cache中的k_li并且生成新的q_li计算他们的相关性。3.7.1 LightningIndexer核心计算理论LightningIndexer基于一系列操作得到每一个 token 对应的 Top-k kk个位置。对于某个 token 对应的 Index QueryQ i n d e x ∈ R g × d Q_{index}\in\R^{g\times d}Qindex​∈Rg×d给定上下文 Index KeyK i n d e x ∈ R S k × d , W ∈ R g × 1 K_{index}\in\R^{S_{k}\times d},W\in\R^{g\times 1}Kindex​∈RSk​×d,W∈Rg×1其中g gg为 GQA 对应的 group size此处为64d dd为每一个头的维度此处为128S k S_{k}Sk​是上下文的长度LightningIndexer的具体计算公式如下Top- k { [ 1 ] 1 × g [ ( W [ 1 ] 1 × S k ) ⊙ ReLU ( Q i n d e x K i n d e x T ) ] } \text{Top-}k\left\{[1]_{1\times g}\left[(W[1]_{1\times S_{k}})\odot\text{ReLU}\left(Q_{index}K_{index}^T\right)\right]\right\}Top-k{[1]1×g​[(W[1]1×Sk​​)⊙ReLU(Qindex​KindexT​)]}可拆分为如下计算流程计算矩阵乘法S Q i n d e x K i n d e x T S Q_{index}K_{index}^TSQindex​KindexT​【*Q_index K_index.T算相关性我和谁关系好】计算激活函数S ′ ReLU ( S ) S\text{ReLU}(S)S′ReLU(S)【*ReLU去掉负分只看正面关系】计算广播乘法S W ( W [ 1 ] 1 × S k ) ⊙ S ′ S_W(W[1]_{1\times S_{k}})\odot SSW​(W[1]1×Sk​​)⊙S′【*W 加权重要的人分数翻倍VIP 加成】沿G轴进行Reduce操作S c o r e [ 1 ] 1 × g S W Score[1]_{1\times g} S_WScore[1]1×g​SW​【*求和多头合并成最终排名所有头数值加起来】对S c o r e ScoreScore进行Top- k \text{Top-}kTop-k计算即获取数值排序前k kk个的结果并返回其对应的 Index【*Top-k选出最重要的 2048 个位置选分最高的】3.7.2 indexer打分整体梳理q_li和q_c本质上上一样的q_c 是 “低秩压缩特征”不能直接打分必须投影成 q_li才能和 k_li 维度对齐、多头对齐、空间对齐完成相似度计算。获取方法输入q_c→ 变成 q_li从kv_cache[2]取 k_li 来自步骤三indexer_select_post_process关键代码如下def indexer_select_post_process(self, x, q_c, kv_cache, attn_metadata, ...): # 准备query q_li self.wq_b(q_c) # [b,s,1536] [1536,64*128] [b,s,64*128] q_li rope_forward_triton_siso(q_li, cos, sin) # 调用lightning_indexer进行稀疏选择 topk_indices torch.ops._C_ascend.npu_lightning_indexer( queryq_li, keykv_cache[2], # 使用预处理结果 weightsweights, ... ) return topk_indices3.8 阶段8 执行稀疏注意力 SFA只使用top2048 个 KV计算注意力核心的稀疏注意力计算在_execute_sparse_flash_attention_process()方法中实现sfa_v1.py:1042-1065def _execute_sparse_flash_attention_process( self, ql_nope, q_pe, kv_cache, topk_indices, attn_metadata, actual_seq_lengths_query, actual_seq_lengths_key ): block_table attn_metadata.block_table kv kv_cache[0] key_rope kv_cache[1] attn_output torch.ops._C_ascend.npu_sparse_flash_attention( queryql_nope, keykv, valuekv, sparse_indicestopk_indices, scale_valueself.scale, sparse_block_size1, block_tableblock_table, actual_seq_lengths_queryactual_seq_lengths_query, actual_seq_lengths_kvactual_seq_lengths_key, query_ropeq_pe, key_ropekey_rope, layout_queryTND, layout_kvPA_BSND, sparse_mode3, ) return attn_output4.Moe结构GLM5的moe结构使用的是标准MOE结构不分组不添加额外分数加权这部分可以参考之前Qwen的moe结构不再额外赘述。(Qwen3.5 MoE模型结构拆解 - WIKI)4.1 Moe结构transformers实现代码位置transformers仓实现源码class DeepseekV3MoE(nn.Module): A mixed expert module containing shared experts. def __init__(self, config): super().__init__() self.config config self.experts DeepseekV3NaiveMoe(config) self.gate DeepseekV3TopkRouter(config) self.shared_experts DeepseekV3MLP( configconfig, intermediate_sizeconfig.moe_intermediate_size * config.n_shared_experts ) self.n_routed_experts config.n_routed_experts self.n_group config.n_group self.topk_group config.topk_group self.norm_topk_prob config.norm_topk_prob self.routed_scaling_factor config.routed_scaling_factor self.top_k config.num_experts_per_tok def forward(self, hidden_states): residuals hidden_states orig_shape hidden_states.shape router_logits self.gate(hidden_states) topk_indices, topk_weights self.route_tokens_to_experts(router_logits) hidden_states hidden_states.view(-1, hidden_states.shape[-1]) hidden_states self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape) hidden_states hidden_states self.shared_experts(residuals) return hidden_states┌──────────────────────────────┐ │ 输入 hidden_states │ │ [b, s, d] │ └──────────────┬───────────────┘ │ ▼ ┌──────────────────────────────┐ │ Reshape 压扁维度 │ │ [b*s, hidden_dim] │ └──────────────┬───────────────┘ │ ┌───────┴───────┐ │ │ ▼ ▼ ┌──────────────┐ ┌───────────────────┐ │ 路由Gate模块 │ │ 共享专家SharedExpert │ │ DeepseekV3 │ │ 全局所有Token都走 │ │ TopkRouter │ │ 独立MLP计算 │ └──────┬───────┘ └──────────┬──────────┘ │ │ ▼ │ ┌────────────────────────┐ │ │ 分组路由逻辑 │ │ │ 1.专家分 n_group 组 │ │ │ 2.选 topk_group 个组 │ │ │ 3.组内选 top_k 个专家 │ │ │ 4.输出专家索引路由权重 │ │ └──────┬─────────────────┘ │ │ │ ▼ │ ┌──────────────────────────────┐ │ 稀疏专家 Experts 计算 │ │ DeepseekV3NaiveMoe │ │ 按选中专家分流、MLP推理 │ │ 路由权重加权聚合输出 │ └──────────────┬───────────────┘ │ ▼ ┌──────────────────────────────┐ │ 稀疏输出 共享输出 相加融合 │ └──────────────┬───────────────┘ │ ▼ ┌──────────────────────────────┐ │ Reshape 还原维度 │ │ [b, s, d] │ └──────────────┬───────────────┘ │ ▼ ┌──────────────────────────────┐ │ MoE 模块输出 │ └──────────────────────────────┘forward 流程gate(x) → weights8个专家的权重、indices8个专家的编号遍历本卡负责的专家对路由到该专家的 token 调用 expert(x[idx]) * weights[idx]累加到 yshared_experts(x) → z对所有 token 都计算多卡时 all_reduce(y) 汇总路由专家输出返回 y z5. 关键算子分析5.1. 算子总览下面给出单个decode过程算子分析图可以参考5.2. 关键算子归档所有算子通过torch_binding.cpp注册到PyTorch框架中profiling算子名称算子功能算子代码实现算子注册调用MoeGatingTopK在MOE架构中从所有计算结果分组里挑选前Topk个专家点我查看点我查看aclnnMoeInitRoutingCustom_MoeInitRoutingCustom_MoeInitRoutingCustom将输入token按专家索引展开为后续分发做准备点我查看点我查看MoeDistributedDispatchV2(已合入MOE融合算子)将token分发到对应的专家rank同时处理量化和通信优化点我查看点我查看MoeDistributedCombineV2(已合入MOE融合算子)将分散在各rank的结果按照原始的topk权重进行加权合并还原回token的原始顺序点我查看点我查看LightningIndexerVllm稀疏索引计算接受QKV并通过权重加权的相似度计算找出最重要的稀疏块索引点我查看点我查看SparseFlashAttention稀疏注意力计算利用第一阶段生成的稀疏索引计算被选中的重要块点我查看点我查看MoeTokenUnpermuteMoetoken重排点我查看[点我查看]( vllm-ascend/vllm_ascend/ops/fused_moe/token_dispatcher.py at main · vllm-project/vllm-ascend (github.com))5.3 lightning_indexer和SFA算子tilling和pipline分析参考cann-recipes-infer/docs/design/mtp_design.md-代码预览-cann-recipes-infer:基于 CANN 平台的 LLM 与多模态模型推理优化样例项目 - AtomGit | GitCode6. profilling分析以下是一个Atlas 800I A3混部的GLM5profilling分析场景为长序列, 128k输入1k输出mtp3chunkedprefill6. 1总体概览从最大维度看当前profilling采集到了两轮prefill和多轮decode。可以看到模型主要的耗时都在prefill阶段长序列瓶颈主要在prefill侧6. 2decode侧分析我们可以尝试把单个decode打开看6. 3 prefill侧分析我们也可以尝试把单个prefill打开看