跳到正文
SemiAnalysis· Kimbo Chen·· 19 天前AI 评分52

GLM5.3 稀疏注意力如何影响 HBM 内存占用

How GLM5.3 Sparse Attention Affects HBM Memory Usage

AI 导读

SemiAnalysis 分析 GLM-5.3 的稀疏注意力对内存的影响,指出稀疏注意力只降低 SDPA 阶段的 KV cache 内存与带宽需求,并不减少整体内存容量占用,top-k 选择仍需完整上下文驻留 HBM。

正文 · AI 翻译

稀疏注意力如何影响 DRAM/NAND 内存

稀疏注意力如何影响内存的 TAM,包括 HBM 和 NAND?稀疏注意力选择最相关的 top-k 个 token 进行关注,从而减少核心缩放点积注意力(SDPA)操作期间的内存消耗和带宽需求。然而,效率提升在实践中并不会直接转化为整体内存节省。具体而言,top-k 选择操作通常需要完整上下文位于 HBM 中,因此稀疏注意力并不能消除内存容量瓶颈。

稀疏注意力的吞吐量受内存容量瓶颈限制。来源:HiSparse

为克服这一限制,SGLang 团队设计了 HiSparse,一种分层内存系统,主动将 KV 缓存条目从设备 HBM 卸载到主机 DRAM。HiSparse 的行为类似于 LRU(最近最少使用)缓存,当 top-k 选择缓存未命中时,它将 token 从 DRAM 加载到 HBM,并根据 LRU 淘汰策略将 token 从 HBM 卸载到 DRAM。为减少缓存未命中延迟,HiSparse 将第 N 层的 KV 缓存加载与第 N-1 层的执行重叠(逐层重叠),这在其先前工作 HiCache 中引入。

逐层重叠。来源:CachedAttention

借助 HiSparse,SGLang 在高并发和长上下文场景下大幅提升吞吐量,代价是 top-k 缓存未命中的 I/O 开销。

稀疏注意力减少了 SDPA 操作中的 KV 缓存内存和带宽需求,但并未减少整体内存容量使用。此外,HiSparse 表明系统优化可以克服稀疏注意力的内存容量限制,因此仅凭稀疏注意力的内存特征不足以描绘系统 KV 缓存效率的全貌。为更全面地展示稀疏注意力模型的服务情况,这里我们解释 Z.ai 的 GLM-5 模型系列的设计,以及它们在实践中如何被服务。

GLM5.3 Agentic 推理服务

在我们的实时基准测试板 InferenceX 中,GB300 在每秒 150 token 的响应速度下,实现了我们对比中最低的建模服务成本。GB300 每 GPU 处理更多 token,其更高的每小时成本仍使其在成本效率上领先于 GB200。GB300 结果使用 Dynamo-TRT-LLM,而 GB200 结果使用 Dynamo-SGLang。

我们使用9 月 28 日 AgentX 快照结果来帮助说明 GLM-5.3 的服务影响。我们使用 InferenceX 的超大规模拥有和运营基础设施模型来估算成本。总 token 包括输入和生成的输出,包括跨轮次复用的缓存输入历史。流式速度使用 p90 交互性衡量;首 token 时间单独评估。

在每秒 150 token 的速度下,GB200 每百万总 token 成本约为 $0.044,而运行 ATOM 的 MI355X 为 $0.049。在相同的流式速度目标下,成本约低 12%。这些值是从测量曲线估算的,而非在恰好每秒 150 token 下单独测试得出。

在本次对比中,我们采用 InferenceX 估算的超大规模基础设施拥有和运营成本。总 token 数包括输入和生成的输出,其中包含从缓存中复用的输入历史。

来源:InferenceX | SemiAnalysis

在 Dynamo-SGLang 的结果中,GB300 在此处每 GPU 每秒可处理约 13950 个总 token,相比 GB200 的 11873 个总 token 具有 17.5% 的吞吐量优势。然而,我们假设 GB300 每 GPU 小时成本为 2.31 美元,而 GB200 为 1.86 美元。在这一目标下,更高的每小时成本足以抵消 GB300 的吞吐量领先优势。

与 AMD 相比,在每秒 100 token 时,ATOM 的成本比 GB200 低约 13%。在 125 token 时,GB200 成本低约 5%,在 150 token 时低约 12%。在这三个目标上,两个系统都没有统一的成本优势。

仅计算生成的输出,在每秒 150 token 的响应速度下,GB200 每百万输出 token 的成本为 5.92 美元,比 MI355X ATOM 的 6.68 美元低约 11%。这一指标对于分析代理工作负载下的成本很有用,因为代理会反复复用较长的输入历史,而生成的新文本相对较少。

来源:InferenceX | SemiAnalysis

在每秒 150 token 目标两侧的 GB200 测量中,p90 TTFT(首 token 时间)约为 14–19 秒,而 MI355X ATOM 约为 1.1–1.2 秒。一些 GB200 最低成本运行的首次 token 等待时间要长得多。

因此,我们进行第二次比较,仅使用实际测试过且达到每秒至少 150 token、p90 TTFT 至多两秒的配置。下图显示,搭载 ATOM 的 MI355X 每百万总 token 成本为 0.0607 美元,而搭载 Dynamo-SGLang 的 B200 为 0.0666 美元。在两秒首 token 限制下,ATOM 成本低约 9%。如果将限制放宽到十秒,搭载 Dynamo-TRT-LLM 的 GB300 便符合条件,成本为 0.0451 美元,比 ATOM 的 0.0607 美元低约 26%。

来源:InferenceX | SemiAnalysis

优化

在 GLM 的 5.2 / 5.3 模型中,其注意力架构降低了内存需求,而长时间运行的代理仍需要保留较早的对话历史。GLM 通过两种方式降低处理这些历史的成本:KV 压缩减少了每个 token 存储的缓存状态,稀疏注意力减少了每次注意力操作读取该状态的数量。然后,服务引擎将决定在哪里存储缓存以及如何高效检索它。

这些 B200 结果显示了在 GPU 内存之外可以发生多少复用。当并发从 8 个请求增加到 16 个请求时,从 GPU 内存复用的提示 token 占比从 90.3% 下降到 54.8%。其中大部分复用转移到主机内存,从 6.0% 上升到 40.3%。GPU 缓存下降的大部分被主机内存的复用所抵消,使整体缓存命中率在所有并发级别下都保持在 95% 以上。

来源:SemiAnalysis InferenceX;InferenceX B200 行 440962、440961 和 440959;工作流运行 33683520699

除了缓存管理之外,改变服务引擎执行预填充和解码的方式也可以改变性能。两项独立的 GLM-5.2 研究说明了这一点:

  • vLLM 将第一个解码步骤保持在一致的 CUDA graph 执行路径上,在一项 NVFP4 研究中将平均 TPOT(每输出 token 时间)从约 40 ms 降至 22 ms。

  • 在另一个 8 卡 MI355X 长上下文工作负载上,ATOM 引擎跨流水线阶段分块处理 prefill,总吞吐量提升了 98%,并将中位 TTFT 从 28.6 秒降至 8.7 秒。

TileRT

对于 TileRT 上的 GLM5.3,AMD MI355X 是首个获得支持的平台。TileRT 是一款低延迟 LLM 推理引擎,它将解码编译为单个持久内核,从而减少启动开销,并重叠计算、内存访问和通信。

它优先考虑提升每个用户的 token 生成速度,而非最大批量吞吐量,同时由 vLLM 处理 prefill。

在 AgentX 上,FP8 TileRT MI355X 的 P90 交互性达到最佳 FP4 MI335X 配置的 2 倍。与 GB300 NVL72 相比,其交互性提升了 40%。这些服务指标过去需要专用硬件才能实现。

然而,TTFT 表现欠佳,仍有改进空间,例如优化 KV 传输、支持 FP4 以及增大批量大小。

DeepSeek 稀疏注意力

GLM-5(以及 GLM-5.x 模型)是一个总参数量 744B、激活参数量 40B 的混合专家模型。对于每个 token,它有 1 个共享专家,并路由到 256 个专家中的 8 个,稀疏度为 32。它采用了 DeepSeek 稀疏注意力,我们将在本节中讨论。

DeepSeek 稀疏注意力(DSA)在 DeepSeek V3.2 中引入,由两个组件构成:一个选择 top K token 的 闪电索引器,以及一个 稀疏多头潜在注意力(MLA)。

闪电索引器

闪电索引器在功能上类似于一种轻量级注意力机制。查询和键被降维投影到较低维度,其中索引器查询为多头,索引器键为单头。由于我们只需要查询和键的关系分数(logits)来选择 top K token,索引器不需要值嵌入,也不需要用于归一化 logits 的 softmax 操作。相反,它计算查询-键点积,然后经过 ReLU 非线性。接着,索引器将每个 token 的索引器分数计算为所有查询头上 logits 的加权和,我们用它来选择 top K token。这意味着尽管索引器有多个头,token 选择在多个注意力头之间是共享的。在选择 top K token 时,我们将关注少于 K 个 token 的查询保留为稠密注意力,否则选择 top K。

闪电索引器的概念示例。来源:SemiAnalysis
索引器分数是所有索引器查询头的加权和。来源:SemiAnalysis

稀疏 MLA

正如我们在 Kimi K3 文章中所解释的,MLA 以两种模式运行:多头注意力(MHA)模式和多查询注意力(MQA)模式。两者都有取舍:MHA 模式 FLOPs 更低,但内存开销高 42 倍;MQA 模式内存开销更低,但 FLOPs 最高高出 3.4 倍。DSA 使用 MQA 模式,因为稀疏注意力关注的 token 更少,从而缓解了较高的 FLOPs 使用。然而在实践中,存在一个序列长度阈值,低于该阈值时 MHA 模式实际上比 MQA 模式更高效。其直觉是,在较短的序列长度下,内存加载时间不再占主导地位,而 MHA 模式可能因 FLOPs 更低而具有优势。vLLM 实现了这一特性,并为数据并行配置了 2K 到约 5K 序列长度的 MHA 模式,为张量并行配置了 2K 到约 77K 序列长度的 MHA 模式。2K 的下限是因为 top K 为 2048,且如上所述,低于 2K 的注意力是稠密的。

MHA(绿色)可以比 MQA(红色)更快。来源:vLLM PR #48770

MLA 修改

比较 GLM-5 和 DeepSeek V3.2 的 MLA 维度配置,我们看到最显著的差异是查询头数量和查询键头维度。

GLM-5 技术报告提到 DeepSeek 根据 H800 的 roofline 选择了查询头数量。这很可能指的是MLA 在解码期间的 SDPA 算术强度大致处于 H800 的脊点。推导如下。假设

  • L:序列长度

  • d:头维度

  • r:RoPE 维度

  • H:头数量

  • b:每个参数的字节数

在 MQA 模式下,FLOPs 为

2 * H * L * (d+r)  # P = Q [H, d+r] @ K.T [d+r, L]
2 * H * L * d      # O = P [H, L] @ V [L * d]

而加载的内存为

b * H * (d+r)  # Q
b * L * (d+r)  # K and V

因此算术强度为

(2 * H * L * (2 * d + r)) / (b * (d+r) * (L + H))
= (2 * H * L * (2 * d + r)) / (b * L * (d + r))  # L >> H
= (2 * H * (2 * d + r)) / (b * (d + r))

如果我们代入 DeepSeek V3.2 配置 d = 512,r = 64,b = 2),我们得到

(2 * H * (2 * 512 + 64) / (2 * (512 + 64)) ~= 2 * H

H800 的实际算术强度为 258.2 FLOP/B(实际峰值 865 TFLOP/s,3.35 TB/s 带宽,根据 DeepSeek),代入公式得到 H ~= 128。

GLM-5 的 H = 64,即查询头数量减半,这一事实意味着 GLM-5 可能是为不同的硬件设计的。如果我们代入 GLM-5 的配置(d = 512,r = 64,b = 2,H = 64),我们得到算术强度为 120.8 FLOP/B。因此,我们怀疑 GLM-5 是针对摩尔线程 MTT S4000 优化的,其约为 128 FLOP/B。摩尔线程与 Z.ai 的合作,例如GLM-5.3-Flash 首日支持,印证了我们的推测。

另外,GLM-5 将 QK NoPE 维度从 128 增加到 192,因为消融实验表明该配置在等 FLOP 和等参数约束下更优。这使 SDPA 头维度从 192 增加到 256,在考虑头数量减少后,FLOP 仍降低了 33%。

IndexShare

随着上下文增长,索引器的延迟成本变得不可忽略。为缓解这一成本,Z.ai 在 GLM-5.2 中提出了 IndexShare(又名 IndexCache),其中每 4 个 DSA 层共享一个索引器。

30B DSA 模型在不同上下文长度下的延迟成本。来源:IndexCache

IndexShare 给训练和推理带来了复杂性。标准的 DSA 训练包含两个阶段:密集预热阶段,此时除索引器外的所有模型权重都被冻结,注意力计算为密集计算;随后是稀疏适配阶段,此时整个模型以稀疏注意力进行训练。

在训练目标方面,索引器被训练以对齐注意力头分数之和,使用 KL 散度目标。IndexShare 将训练目标调整为使索引器 logit 分布与共享层间注意力分数的平均分布对齐。

在推理方面,通常索引器会缓存键,就像注意力缓存 KV 一样。由于多个层共享一个索引器,我们需要额外在完整层缓存 top K 索引,以便在共享层复用 top K 选择信息。

这一设计将索引器缓存减少了 75%,将索引器 FLOPs 减少了 75%,并在不同上下文长度以及 prefill/decode 下将吞吐量提升了 1.5 倍到 1.8 倍。

30B DSA 模型的基线配置与 GLM-5.2 配置(红色)对比。来源:IndexCache

后训练设计

流水线

在预训练和中期训练之后,GLM-5 的后训练流水线从监督微调(SFT)开始,接着进行三个 RL 训练阶段(推理 RL、智能体 RL 和通用 RL),最后以同策略跨阶段蒸馏结束,将推理和通用 RL 蒸馏回 SFT 检查点。

GLM-5 后训练流水线。来源:SemiAnalysis

这里我们强调两个细节。第一,技术报告指出推理 RL 模型完全在同策略下训练,这意味着它是同步训练的。同步 RL 训练系统效率较低,但 RL 训练稳定性更高。第二,同策略跨阶段蒸馏类似于多教师同策略蒸馏(MOPD),即学生模型通过多个教师模型的同策略蒸馏进行 RL 训练。然而,与 MOPD 的目标是将专家教师模型合并为一个模型不同,GLM-5 在这里似乎更侧重于在模型学习新能力时恢复旧能力的损失。按照这一假设,我们不确定为什么 Z.ai 只使用推理和通用 RL 检查点作为教师,而不是使用全部三个 RL 检查点。

根据公告博客文章,GLM-5.2 可能已将同策略跨阶段蒸馏替换为标准 MOPD 流水线,Z.ai 将其命名为并行 OPD 训练。Z.ai 宣扬了其规模和训练效率:他们通过约两天的训练,将十多个专家模型合并为一个最终模型。

RL 训练算法

Z.ai 在同步和异步 RL 阶段采用不同的 RL 算法。这里我们比较 RL 目标中的主要差异。

在同步 RL 阶段,Z.ai 采用基于 Group Relative Policy Optimization (GRPO) 的 RL 算法,并进行了现代适配,包括:

  • 移除 KL 正则项,以提高系统效率并允许更激进的梯度更新

  • IcePop:基于训练-推理不匹配比例应用 token 级掩码

  • Clip-Higher:对重要性采样比例使用更高的最大阈值以鼓励探索

相比之下,Z.ai 在异步 RL 阶段采用类似 REINFORCE 的 RL 算法,并加入 Direct Double-Sided Importance Sampling (DIS),一种类似 IcePop 的裁剪机制。下面我们比较两个 RL 阶段的 token 级目标:

两个 RL 阶段的 token 级目标。来源:SemiAnalysis

重要性采样 (IS) 比例比较了当前策略和行为策略下每个采样动作的概率,缩放每个动作对训练目标的贡献。对于 IS 比例,同步 RL 阶段使用上一迭代的训练器策略作为行为策略,而异步 RL 阶段则使用推理策略。在异步训练设置中,策略更新在 rollout 期间发生多次,因此我们需要其间的每个检查点版本来计算 IS 比例。然而,由于高内存和延迟开销,这是不可行的。因此,Z.ai 选择使用推理策略以提高效率和简洁性,代价是数值稳定性挑战。

两个 RL 阶段的 IS 比例。来源:SemiAnalysis

最后,我们重点介绍 RL 训练、交叉蒸馏和长时程 RL 训练中优势计算的主要差异。RL 训练采用 GRPO 标准组归一化优势,交叉蒸馏使用反向 KL 散度。

优势计算比较。来源:SemiAnalysis

单次 Rollout 异步优化

长时程任务会产生极长的轨迹,在 GRPO 背景下,轨迹长度的高方差会产生不稳定的训练信号,可能偏向更长的轨迹。更糟糕的是,由于RL rollout 工作负载对端到端延迟敏感,长时程任务中落后 rollout 的尾部延迟会恶化。为缓解此问题,Z.ai 在 GLM-5.2 中为长时程 RL 训练引入了单次 Rollout 异步优化 (SAO)。SAO 用改进的广义优势估计 (GAE) 替代组归一化优势,灵感来自近端策略优化 (PPO)。GAE 只需单次 rollout 即可计算优势。这消除了每个提示进行多次 rollout 的需要,缓解了 GRPO 的缺点。

生成器节点要求训练器节点等待 rollout 完成。来源:SAO

GAE 通过 values(在强化学习术语中,即当前状态的期望回报)计算,由 value model 估计。具体而言,value model 以 rollout 为输入,为每个 token 生成一个价值估计。SAO 通过从优化中移除观察 token(即工具调用结果),将 GAE 适配到智能体强化学习训练中。此前在 GLM-5 中,跳过这些 token 就足够了;在 SAO 中,我们修改了优势估计公式以绕过观察 token。关于贝尔曼目标和消融实验的更多直觉,读者可参阅 SAO 论文 及其附录 A.1。

在 SAO 中,value model 是一个独立的 LLM,与策略模型并发训练。这意味着训练早期阶段容易因 value model 较弱而导致优势计算的高方差。为应对这一问题,SAO 提出了多种训练 value model 的技术:

  • 双时间尺度更新规则:value model 的更新频率是策略模型的两倍。这意味着在每个策略训练步骤中,对每个数据批次训练两次

  • 注意力参数冻结:冻结 value model 的全部注意力层以提高训练稳定性

  • 扩展预训练:扩展 value model 的预训练语料以初始化 value model,使 value model 能在训练早期阶段稳健地估计价值。遗憾的是,Z.ai 未披露预训练语料的扩展信息

来源:SemiAnalysis

以计算开销和双倍内存占用为代价,SAO 用 value model 取代了组采样,缓解了 GRPO 的延迟问题。

后训练基础设施

Z.ai 在所有 GLM-5 模型中使用 slime——最好的开源强化学习训练框架之一—— 作为后训练基础设施。读者可参阅 slime GitHub 仓库 以了解其设计与特性,这里我们重点介绍两个特性。

强化学习 rollout 的性能特征高度依赖任务,且每个任务都带有不同的工具和奖励函数。为确保系统效率,Z.ai 设计了一个 基于服务器的多任务 rollout 编排器。每个任务作为一个独立的微服务服务器,控制 rollout 和奖励逻辑。该编排器 平衡整体性能,控制每个任务的 rollout 比例和生成速度。集中化还允许跨任务动态调整比例、细粒度进度监控等。Z.ai 声称,在训练 GLM-5 时,该编排器支持 1000 个并发 rollout。

多任务 rollout 编排器的架构图。来源:SemiAnalysis

在 GLM-5.3 的公告中,Z.ai 声称除主机内存外还使用本地存储作为缓存层,从而实现面向 MOPD 的动态教师模型切换与预取。slime PR#1538 通过引入一个 TensorBackuper 类来实现教师切换,该类将教师权重以 PyTorch 张量的形式固定在 CPU 内存中。当一次 rollout 被路由到某个教师时,一个 TensorBackuper 对象会将教师模型复制回 GPU 内存以执行前向传播(实现见此处)。然而,这是由 CPU 支撑的权重交换,我们尚未看到由存储支撑的实现。

在我们的付费墙部分,我们呈现了对 GLM-5.3 网络安全能力的分析。

阅读更多

来源:SemiAnalysis · newsletter.semianalysis.com