HyScale: Scaling Long-Context Training for Learned Sparse Attention

发表时间: 2026-10 · arXiv:2610.32057 (Tencent Hunyuan & University of Edinburgh)

原文: https://arxiv.org/abs/2610.32057

作者/机构：Yangsheng Deng1, Xinjing Huang2, Yulong Wang2, Yunzhuo Liu2, Jun Hu2, Shuyong Tang2, Junchi Zeng2, Mingji Han2, Ray Ying2, Key Zhang2, Luo Mai1 (1 University of Edinburgh, 2 Tencent Hunyuan)

速读

一句话结论
本文提出了一个面向可学习稀疏注意力（Learned Sparse Attention）的长上下文训练系统 HyScale，通过消除全量分数表和引入细粒度重计算机制，在 512 张 H20 GPU 上将 1M 上下文的 GLM-5.1 训练速度相比基线提升了 2.07 倍。

要解决什么问题
可学习稀疏注意力（如 DeepSeek 和 GLM 系列所用）通过让每个 Query 只关注由轻量级索引器选出的 top-$k$ 个 Token，大幅降低了长上下文的计算量。然而，“选择”过程本身却成了新的训练卡点，在 1M 上下文下占用了近一半的训练步时。具体机制上的卡点有两处：第一，显存与访存浪费。现有系统在选择时，会先计算 Query 与所有 $L$ 个历史 Token 的分数，并将其作为一个 $n \times L$ 的完整分数表物化到 GPU 显存中，然后再从中挑出 $k$ 个。对于 1M 上下文，保留的 $k$ 个分数仅占总数的 0.2%，但庞大的完整分数表会造成极大的显存开销和读写延迟。第二，重计算代价高昂。在长上下文训练必须的激活值重计算机制下，现有框架只能做到算子级别的重计算。这意味着在反向传播前，系统必须把前向传播中计算所有 $L$ 个 Token 分数并排序的过程从头到尾再跑一遍。此外，现有的上下文并行（CP）在选择前同步收集所有数据会导致计算阻塞，且稀疏算子按单个 Query 分配线程块，造成了严重的 L2 缓存冗余读取。

怎么做的
核心思路是转变对“选择”操作输出的认知：选择操作输出的不是连续的激活值，而是一组离散的索引。因为 top-$k$ 操作不可导，反向传播根本不需要完整的中间分数表，只需要被选中的那 $k$ 个索引及其分数。基于此，HyScale 引入了两个核心部件来绕开上述卡点：
第一是即时选择算子 HyTopK，用于彻底消除分数表。它在计算分数的同时直接进行过滤，不把全量分数写进显存。为了在训练中实现这一点，它采用“乐观步幅采样”——从当前 Query 自己的上下文中均匀采样一小部分 Token，估算出一个略低的阈值分数。绝大多数情况下，这个阈值能准确放过目标 Token；极少数候选过多或过少导致失败的 Query，则通过无主机同步的轻量级 GPU 内核原地重新计算。最后通过基数排序保证输出的绝对确定性。
第二是选择重放（Selection Replay），一种细粒度的重计算策略。前向传播时，系统只保存选出的 $k$ 个索引 $\mathcal{T}_t$。在重计算时，系统不再重新计算所有 $L$ 个 Token 的分数，而是直接读取这 $k$ 个索引，仅对被选中的条目进行重新打分。其重计算的计算复杂度从 $O(L)$ 骤降到 $O(k)$。索引器损失的梯度更新仅依赖这 $k$ 个分数：

$$\begin{aligned} g_{t,s} = \frac{\partial \mathcal{L}^I}{\partial I_{t,s}} = \begin{cases} \lambda (\hat{p}_{t,s} - p_{t,s}) & s \in \mathcal{T}_t \\ 0 & \text{otherwise} \end{cases} \end{aligned}$$

此外，HyScale 还包含两项系统级设计：一是解耦的 CP 重叠，利用选择操作只读索引器键、稀疏注意力只读 KV 状态的特性，将两者的通信拆分并分别掩盖在不依赖它们的计算中；二是位置感知稀疏算子，利用相邻 Query 往往选择相似 Token 的局部性，让一个线程块同时处理两个相邻的 Query，使其共享加载的键值，大幅降低 L2 缓存访存。

效果如何
实验在 512 张 H20 GPU 和 32 张 H800 GPU 的集群上进行。对比的基线方法是 Megatron-LM，它代表了当前最先进的系统架构路线，使用 cuDNN DSA 后端、完全激活重计算和全收集 CP 通信。在端到端训练实验中，模型采用完整的 GLM-5.1 权重，全局批大小为 4M Tokens，单卡 8K Tokens。在 1M 上下文长度下，HyScale 将 GLM-5.1 的训练速度相比 Megatron-LM 提升了 2.07 倍，且随着上下文变长，加速比持续扩大（从 256K 的 1.23 倍升至 1M 的 2.07 倍），同时保持了几乎相同的显存峰值。该方法不仅适用于 Token 级别的稀疏注意力，也能无缝泛化到块级别架构。在 32 张 H800 GPU 上针对 1M 上下文进行测试，HyScale 将 DeepSeek-V4-Flash 的 CSA 层加速了 2.01 倍，将 Qwen3.8-Flash-Next 的 QSA 层加速了 1.57 倍。关于代价与局限性：选择重放机制需要额外保存索引，会带来单层 $4nk$ 字节的微小显存开销，但作者通过将其异步卸载到主机内存中掩盖了这一成本；HyTopK 为了保证训练的确定性引入了排序操作，但这仅增加了 0.8% 的选择时间；在极少数采样阈值失效触发恢复机制时，耗时也仅增加不到 5%。

一、 主要贡献

长上下文（Long context）已成为大型语言模型的核心能力，推动训练序列长度从 128K 扩展到 1M tokens。在这些长度下，由于每个查询（query）都需要关注所有先前的 token，全注意力的计算量呈二次方增长，主导了训练成本。学习型稀疏注意力（Learned sparse attention）通过将注意力分为“选择（selection）”和“稀疏注意力（sparse attention）”两步来降低成本。然而，“选择”阶段成为了新的训练瓶颈。当前的系统在选择阶段会实体化一个完整的分数表，并在激活重计算（activation recomputation）期间重复执行选择操作。

本文的核心问题在于，现有系统将选择的输出视为标准的连续激活，而实际上它只是一组离散的索引。基于这一洞察，本文提出了 HyScale，一个用于学习型稀疏注意力的长上下文训练系统。其主要贡献包括：

  1. HyTopK：一种基于乐观阈值（optimistic thresholding）的即时（on-the-fly）top-$k$ 选择操作符。它在打分的同时选出 top-$k$ 条目，完全消除了庞大的分数表。
  2. 选择重放（Selection replay）：一种针对选择阶段的细粒度重计算策略。它记录前向传播的选择结果，并在重计算时仅对选中的 $k$ 个条目进行回放，消除了重复选择的开销。
  3. 整体系统优化：HyScale 结合了端到端的系统优化，包括解耦的上下文并行（Context Parallelism, CP）重叠以隐藏通信开销，以及局部性感知的稀疏内核以提升计算效率。

实验表明，在 512 个 H20 GPU 上，HyScale 将 GLM-5.1 在 1M token 序列上的端到端训练速度相比 Megatron-LM 提升了 $2.07\times$。此外，它在 32 个 H800 GPU 上将 DeepSeek-V4-Flash 的 CSA 层和 Qwen3.8-Flash-Next 的 QSA 层在 1M tokens 下分别加速了 $2.01\times$ 和 $1.57\times$。

二、 背景知识与研究动机

学习型稀疏注意力机制。学习型稀疏注意力允许每个查询仅关注由学习型索引器（indexer）选出的少数条目。它分为两步运行。考虑一个长度为 $L$ 的序列，对于位置 $t$ 的查询，设 $\mathcal{U}_t$ 为其可选择的候选条目集。在深度求索稀疏注意力（DSA）中，一个条目是一个较早的 token；在块级设计（如 CSA 和 QSA）中，一个条目是一个较早 token 的块。选择阶段对每个条目打分并保留最高的 $k$ 个：

$$ \mathbf{I}_t = \left( \mathrm{Score}(\mathbf{q}_t^I, \mathbf{k}_u^I) \right)_{u\in\mathcal{U}_t}, \quad \mathcal{T}_t = \mathrm{Top\text{-}}k(\mathbf{I}_t), $$


其中 $\mathbf{q}_t^I$ 是索引器的查询，$\mathbf{k}_u^I$ 是条目 $u$ 的索引器键，$\mathbf{I}_t \in \mathbb{R}^{|\mathcal{U}_t|}$ 是分数向量，$\mathcal{T}_t \subseteq \mathcal{U}_t$ 是大小为 $k$ 的索引集。随后，稀疏注意力仅让查询 $\mathbf{q}_t$ 关注选中的条目：

$$ \mathbf{o}_t = \mathrm{Attn}\big(\mathbf{q}_t, \{\mathbf{c}_u\}_{u\in\mathcal{T}_t}\big), $$
其中 $\mathbf{c}_u$ 是稀疏注意力为条目 $u$ 读取的内容（例如 KV 潜在变量），$\mathbf{o}_t$ 是输出。每个查询的选择成本为 $O(|U_t|)$（最高达 $O(L)$），而稀疏注意力的成本为 $O(k)$，其中 $k \ll L$。

索引器的损失函数。由于 top-$k$ 操作不可微，索引器通过自身的损失函数进行学习：

$$ \mathcal{L}^I = \lambda \sum_t \mathrm{KL}\bigl(\mathbf{p}_t \| \mathrm{Softmax}(\mathbf{I}_{t, \mathcal{T}_t})\bigr), $$


其中 $\lambda$ 是标量权重，$\mathbf{p}_t \in \mathbb{R}^k$ 是查询 $t$ 在 $\mathcal{T}_t$ 上跨头平均的注意力分布，$\mathbf{I}_{t, \mathcal{T}_t} \in \mathbb{R}^k$ 保存了索引器在这些条目上的分数。由于 $\mathbf{p}_t$ 和索引器的输入被分离（detached），该损失仅更新索引器，且仅通过选中的 $k$ 个分数进行反向传播。稀疏注意力仅反向传播到选中的 $\mathbf{c}_u$（详见附录 A）。

深度求索稀疏注意力（DSA）。DSA（参考 DeepSeek-V3.2 论文）将单个 token 作为条目，并使用具有 $H_I$ 个维度为 $d_I$ 的头的轻量级索引器对其进行打分：

$$ \mathrm{Score}(\mathbf{q}_t^I, \mathbf{k}_u^I) = \sum_{h=1}^{H_I} w_{t,h} \mathrm{ReLU}\big(\mathbf{q}_{t,h}^{I\top} \mathbf{k}_u^I\big), $$


其中 $\mathbf{q}_{t,h}^I \in \mathbb{R}^{d_I}$ 是头 $h$ 的查询，所有头共享每个 token 的一个键 $\mathbf{k}_u^I \in \mathbb{R}^{d_I}$，$w_{t,h}$ 对每个头进行加权。其稀疏注意力使用多头潜在注意力（MLA），所有头共享每个 token 的一个压缩 KV 潜在变量，作为 $\mathbf{c}_u$。

块级变体。其他生产级设计选择 token 块，它们与 DSA 的主要区别在于条目 $\mathcal{U}_t$ 和打分函数。DeepSeek-V4 的压缩稀疏注意力（CSA）对每个压缩了四个 token 的条目打分。Qwen3.8-Flash-Next 的 Qwen 稀疏注意力（QSA）对平均池化的四个 token 块打分。MiniMax-M3 的 MiniMax 稀疏注意力（MSA）对每个 token 打分，并在每组查询头的 128 个 token 块内取最大值。

上下文并行（CP）。为了训练单 GPU 无法容纳的超长序列，上下文并行（CP）将长度为 $L$ 的序列分割到一组 $P$ 个 GPU 上，每个 GPU 仅持有 $n = L / P$ 个查询。为全注意力设计的 Ring Attention 和 DeepSpeed-Ulysses 不适用于学习型稀疏注意力。现有系统复用了 Llama 3 的 all-gather KV 方法。在选择之前，前向传播在 CP 组内同步 all-gather 索引器键 $\{\mathbf{k}_u^I\}$ 和稀疏注意力读取的状态 $\{\mathbf{c}_u\}$。反向传播则对它们的梯度执行 reduce-scatter。

激活重计算。反向传播需要前向传播的激活值。在长上下文中，存储这些激活会超出显存，因此训练通常使用激活重计算（参考 Chen 等人，2016）。全激活重计算下，每一层仅保留其输入，并在反向传播前重新运行前向传播。选择性重计算（参考 Korthikanti 等人，2023）则保留更多激活，并仅重新运行层的部分内容（如核心注意力），其操作粒度为整个操作符。

选择成为瓶颈。尽管稀疏注意力本身计算量小，注意力层依然是长上下文训练的瓶颈，其中选择操作占了绝大部分时间。如图 1 所示，在 512 个 H20 GPU 上训练 GLM-5.1 时，排除流水线等待时间，注意力层在 1M tokens 时占据了训练步骤的 $61\%$。其中，选择阶段单独占了步骤的 $47\%$，耗时是稀疏注意力的 $6.6\times$。

图1：在 512 个 H20 GPU 上 Megatron-LM 中 GLM-5.1 的一个训练步骤
图1：在 512 个 H20 GPU 上 Megatron-LM 中 GLM-5.1 的一个训练步骤

低效点 1：实体化分数。现有系统在选出任何内容前会实体化每一个分数（图 2 中的 ①）。打分内核将 $n$ 个本地查询与所有 $L$ 个键进行比对，并将分数作为 $n \times L$ 的表格写入显存。然后 top-$k$ 内核读取该表并保留 $k$ 个条目。在 1M token 的步骤中，单层的表格在每个 GPU 上占用 $32 \mathrm{GiB}$，而选中的索引仅占 $64 \mathrm{MiB}$。由于 top-$k$ 仅保留 $0.2\%$ 的分数，实体化它们浪费了大量显存和内存流量，且该成本随上下文长度增长。

低效点 2：重复选择。在全激活重计算下，选择在每个训练步骤中运行两次（图 2 中的 ②）。在每一层的反向传播之前，前向传播的重运行会再次对所有 $L$ 个键打分并重新选择 top-$k$。这占据了 1M token 步骤中选择总耗时的一半。

图2：在 CP 和全激活重计算下，Megatron-LM 中一个学习型稀疏注意力层的执行过程
图2：在 CP 和全激活重计算下，Megatron-LM 中一个学习型稀疏注意力层的执行过程

核心洞察与设计动机。这两种低效都源于将选择的输出 $\mathcal{T}_t$ 视为常规的连续激活。我们的核心洞察是，选择的输出是一个离散的索引集合。稀疏注意力只读取选中的条目，索引器损失也只读取它们的分数。此外，top-$k$ 不可微，反向传播不需要分数表。这启发了我们的设计：首先，选择不需要实体化分数表，打分内核可以在计算时即时保留或丢弃分数；其次，重计算不需要重复选择，它可以复用前向传播的索引，仅对选中的条目重新打分。

图3：HyScale 中的学习型稀疏注意力层
图3：HyScale 中的学习型稀疏注意力层

三、 方法细节

3.1 架构概览

HyScale 构建在 Megatron-LM 之上，仅改变了学习型稀疏注意力层执行选择的方式，保持训练栈的其他部分不变。图 3 展示了 HyScale 中的一层。遵循上述洞察，两项技术消除了选择的低效：HyTopK 在打分内核内部为每个查询选择 top-$k$，使得前向传播不再将 $n \times L$ 分数表写入显存（①）；选择重放（selection replay）保留前向传播选出的索引，使得重计算仅对这 $k$ 个条目重新打分，而不是所有条目（②）。此外，HyScale 结合了特定的系统优化，将 CP 通信隐藏在计算背后，并提升了内核效率。

3.2 HyTopK 设计

基于采样的 top-$k$ 选择面临的挑战。现有系统将选择作为两个内核运行。基于采样的 top-$k$ 选择将二者合并，通过在键的样本中取第 $k$ 大的分数作为阈值。打分内核在芯片上将每个分数与该阈值比较，仅将通过的键写入小型的候选缓冲区。LiteTopK（引用：Yin 等人，2026）针对推理首创了此设计，它通过附近查询已选择的键来估计阈值。然而，将其应用于训练会面临三大挑战：

  • 挑战 1 (C1)：现有采样在 CP 下失效。LiteTopK 的下界紧凑是因为它利用了附近查询已选择的键。在推理中这是可行的，但在 CP 环境下，各 GPU 同时对其查询打分，使用附近查询的选择意味着必须等待持有它们的 GPU。
  • 挑战 2 (C2)：缓冲区溢出代价高昂。如果阈值过低，通过的键会超过固定大小的缓冲区。如果退回到对所有 $n$ 个查询进行密集的重新打分，即使只有一个查询溢出，也会付出完全的密集计算成本。
  • 挑战 3 (C3)：训练需要确定性的选择。训练系统要求内核具备按位可复现性（bitwise reproducible），以便在出现损失尖峰或硬件故障时重放运行（引用：DeepSeek-AI 等人，2026；Sul 等人，2026）。基于采样的选择使用原子操作将通过的键追加到缓冲区，顺序在每次运行时都会改变。这会影响平局时保留哪些键，以及浮点数累加的舍入结果。对候选者进行排序可以解决此问题，但成本远高于仅对保留的 $k$ 个键进行排序。

HyTopK 的核心洞察。为了弥合采样选择与训练之间的差距，我们提出 HyTopK。其核心洞察是：阈值不需要是一个有保证的下界，只要失败极其罕见且恢复成本低廉即可。基于此，HyTopK 从每个查询自身的上下文中估计阈值（解决 C1），在 GPU 上进行无需主机同步的重试以处理失败查询（解决 C2），并确保选择是确定性的（解决 C3）。

乐观跨步采样（解决 C1）。HyTopK 放弃了严格的下界保证，而是从查询自身的上下文中进行跨步采样。首先，从查询 $t$ 的 $L_t$ 个上下文键中均匀采样 $M = \min(\max(4096, L_t/128), 8192)$ 个键。为了加速，每几个连续的查询共享一个其组合上下文的样本，并在一次矩阵乘法中完成打分。其次，将阈值设定为第 $\lceil 3kM/L_t \rceil$ 好的样本分数，这大约代表整体中第 $3k$ 好的键。凭借覆盖整个上下文的样本和三倍的余量，阈值极少会高于真实的第 $k$ 个分数从而漏掉真正的 top-$k$ 键。

轻量级无主机恢复（解决 C2）。恢复操作在选择之后立即作为 GPU 上的一组短内核链运行。首先，一个检查内核会标记候选数量少于 $k$ 或溢出缓冲区的查询，并在显存中记录失败的查询及其计数。接着，三个恢复内核接管最多 64 个失败的查询，在整个上下文中对每个查询重新打分，精确选择其 top-$k$，并将结果原地写回。这些内核在每次选择后以固定的网格启动，如果计数为零则立即返回。因此，没有失败的选择仅付出几次空启动的代价，而有失败的选择仅为失败的查询买单。在我们的实验中，从未出现过超过 64 个失败查询的情况。

确定性选择（解决 C3）。HyTopK 在其内核中保留了快速的原子写入，并在之后恢复确定性。它仅在影响输出的最小集合内对键进行排序。首先，如果共享第 $k$ 高分数的键多于 $\mathcal{T}_t$ 中剩余的位置，选择内核会对这些平局键的位置运行基数选择（radix select）（引用：Alabi 等人，2012），并保留位置最小的键。其次，使用基数排序（radix sort）（引用：Merrill & Grimshaw，2011）按分数对每个 $\mathcal{T}_t$ 的 $k$ 个条目进行排序，遇到平局则按位置打破。这样，相同输入的两次运行将产生完全一致的 $\mathcal{T}_t$。

3.3 选择重放 (Selection Replay)

重计算中的重复选择。在激活重计算下，现有系统在反向传播前会重新运行选择。它们对所有 $L$ 个键打分并再次选择 top-$k$，尽管层输入没有改变且选择结果相同，这使得选择成本翻倍。

重放的洞察。由于 top-$k$ 不可微，梯度只会到达 $\mathcal{T}_t$ 选中的 $k$ 个条目，而永远不会穿过选择本身。反向传播只需要 $\mathcal{T}_t$ 的值和这 $k$ 个分数。因此，我们可以保留前向传播的索引，在重计算时仅对这些索引指向的条目重新打分，从而将选择重计算的复杂度从每个查询的 $O(L)$ 降至 $O(k)$。

现有策略的挑战。现有的选择性重计算策略（引用：Korthikanti 等人，2023；PyTorch，2026）决定每个操作符是保留输出还是完全重新运行。它们无法表达“只对 $k$ 个键重新打分”的逻辑，因此要么浪费内存（保留所有激活或分数表），要么浪费计算（保留索引但仍需对所有 $L$ 个键打分以获取分数）。

选择重放设计。我们设计了选择重放（selection replay），这是一种决定每个条目（而非操作符）保留什么的重计算策略。它有两种变体：

  • 索引重放（Index replay）：仅保留每个查询的 $k$ 个索引，并仅对这 $k$ 个条目重新打分。
  • 分数重放（Score replay）：同时保留 $k$ 个索引和 $k$ 个分数，完全不重新打分。
图4：表 1 中策略在一个 GLM-5.1 层上的成本 (H20 GPU, 1M tokens)
图4：表 1 中策略在一个 GLM-5.1 层上的成本 (H20 GPU, 1M tokens)

成本比较与实现细节。如图 4 所示，两种重放变体的成本均远低于操作符级别的策略。我们默认使用索引重放，因为其重新打分的成本最多仅占训练步长的 $1.6\%$，且在 GLM-5.1 中比分数重放节省约 5 GiB 的显存。
具体实现步骤如下：前向传播打分并选择，保留 $k$ 个选中的索引；重计算不打所有键的分数，而是加载保留的索引，使用 gather-and-score 内核仅对这 $k$ 个条目打分；反向传播按常规运行。为了节省显存，保留的索引可以在选择后立即通过独立的流异步卸载到固定的主机内存（pinned host memory），并在重计算进入注意力层时异步拷贝回 GPU。

通用性。选择重放适用于任何通过 top-$k$ 选择的学习型稀疏注意力（如 DSA、CSA、QSA），并与现有的模块级重计算策略接口保持兼容。

3.4 整体系统实现

除了上述优化，HyScale 还针对学习型稀疏注意力中另外两个成本高昂的环节进行了系统级优化：同步的 CP 通信和稀疏内核中昂贵的全局内存访问。

3.4.1 针对选择和注意力的解耦 CP 重叠

CP 通信瓶颈。在 CP 下，每个 GPU 需要所有较早 token 的索引器键和 KV 状态。现有系统在选择前同步 all-gather 这两者，并在反向传播时 reduce-scatter 它们的梯度。这些集合通信庞大且同步。
解耦 CP 重叠设计。为全注意力设计的现有 CP 优化（如 Ring Attention 和 DeepSpeed-Ulysses）不适用。我们的洞察是：选择只读取索引器键，而稀疏注意力只读取 KV 状态。因此，HyScale 将两者的 CP 通信解耦，分别传输，并将每次传输与不消耗该数据的计算重叠。前向和重计算中，索引器键的 gather 与查询投影重叠，KV 状态的 gather 与查询投影和选择重叠。反向传播中，KV 梯度的 reduce-scatter 与索引器反向重叠，索引器键梯度的 reduce-scatter 与查询投影反向重叠。

3.4.2 局部性感知的稀疏内核

稀疏内核瓶颈。现有的稀疏内核实现（引用：Li 等人，2025；NVIDIA，2026；Wang 等人，2025）为每个线程块分配单个查询（如图 5a）。这导致严重的 L2 缓存流量：前向内核单独加载每个查询的条目；反向内核不仅单独加载，还使用独立的原子操作累加梯度。
选择局部性的洞察。先前的研究（引用：Yin 等人，2026；Liu 等人，2026a；Meng & Fang，2026）观察到：相邻查询选择的键在很大程度上是相同的。在 GLM-5.1 中，相邻查询平均共享 $71\%$ 的 top-2048 键。

图5：局部性感知的稀疏内核
图5：局部性感知的稀疏内核

内核设计。HyScale 提出了局部性感知的稀疏内核（如图 5b）。每层执行一次轻量级的预处理，在一个批次中构建所有相邻查询对的并集。每个线程块处理两个相邻查询，读取它们选择的并集。共享的键只被加载一次；在反向传播中，内核先将两个查询对共享键的梯度相加，然后再用一次原子操作更新。每个查询会屏蔽掉（mask out）并集中不属于自己的部分，利用原本闲置的 Tensor Core 周期来节省共享键上的 L2 流量，确保输出和梯度精确无误。

3.4.3 实现细节

HyScale 包含约 3.4K 行 Python 和 14.7K 行 GPU 内核代码。此外，它还包含四项端到端技术以消除 Megatron-LM 中的主机同步和 GPU 冗余工作：

  1. CP 布局缓存（CP layout cache）：Megatron-LM 在每层都重建 CP 布局，引发大量主机同步。HyScale 每个微批次（micro-batch）仅构建一次布局供所有层复用。
  2. 序列边界缓存（Sequence boundaries cache）：类似地，HyScale 每个微批次将 RoPE 的序列边界拷贝到主机一次。
  3. 内核融合（Kernel fusion）：融合了索引器的 KL 损失、归一化和 RoPE 等逐元素操作链。
  4. 无排列 KV 收集（Permutation-free KV gather）：Megatron-LM 每次都将收集的 KV 状态重排为 token 顺序。HyScale 则直接调整索引，让选中的索引直接指向收集缓冲区中的原始行，避免了数据重排。

四、 实验环境

  • 硬件配置：实验在一个 H20 集群和一个 H800 集群上运行。每个 H20 节点包含 8 个 NVIDIA H20 GPU (96 GB)，2 个 96 核 AMD EPYC 9K84 CPU，以及 2.2 TiB DRAM。每个 H800 节点包含 8 个 NVIDIA H800 GPU (80 GB)，2 个 48 核 Intel Xeon Platinum 8476C CPU，以及 2 TiB DRAM。节点间通过 8 个 400 Gb/s RoCE NICs (NVIDIA ConnectX-7) 连接。
  • 软件配置：基线系统为运行 Megatron Core 0.19.2 的 Megatron-LM，使用 cuDNN DSA 后端。HyScale 构建在相同代码库上。所有运行默认使用 BF16、zigzag all-gather CP 以及全激活重计算。HyScale 默认开启索引重放并将其卸载到主机。
  • 模型与数据集：从发布的权重训练完整的 GLM-5.1 模型。全局批次大小固定为 4M tokens，每个 GPU 8K tokens。每个样本是 256K、512K 或 1M tokens 的单条序列。后续也评估了 DeepSeek-V4-Flash (CSA) 和 Qwen3.8-Flash-Next (QSA)。

五、 实验结果

1. 端到端训练性能
在 512 个 H20 GPU 上训练全量 GLM-5.1 时，HyScale 的速度最高达到 Megatron-LM 的 $2.07\times$（如图 6a）。这种加速随着上下文长度的增加而扩大（从 256K 的 $1.23\times$ 增长到 1M 的 $2.07\times$）。HyScale 在实现这一加速的同时，保持了几乎相同的显存峰值占用（图 6c）。从图 7 的 1M 步骤耗时分解可以看出，HyTopK 将前向选择时间从 35.9s 削减到 20.8s，而选择重放将重计算选择时间从 37.0s 剧减到 1.0s。更快的注意力层同时也使流水线等待时间减半。

图6：在 512 个 H20 GPU 上从发布权重训练完整的 GLM-5.1
图6：在 512 个 H20 GPU 上从发布权重训练完整的 GLM-5.1
图7：完整 GLM-5.1 在 1M 上下文时一个训练步骤的分解
图7：完整 GLM-5.1 在 1M 上下文时一个训练步骤的分解

2. 上下文并行可扩展性
在强扩展（固定上下文，增加 GPU）测试中，HyScale 的加速比随着 GPU 数量的增加而持续增长。在弱扩展（固定每 GPU 查询数，增加上下文）测试中，HyScale 的层时间增长极其缓慢。如图 8c、8d 所示，在每 GPU 16K 查询的设置下，从 128K 扩展到 8M 上下文，HyScale 的层时间仅增长 $11\times$，而 Megatron-LM 增长了 $190\times$。在 8M 上下文时，HyScale 的速度是基线的 $22.3\times$。这得益于 HyScale 的 CP 布局缓存消除了随 GPU 数量线性增长的主机同步开销。

图8：GLM-5.1 注意力层在 H20 上的 CP 扩展性
图8：GLM-5.1 注意力层在 H20 上的 CP 扩展性

3. 消融实验
在 128 个 GPU 上测试 1M 上下文的单层 GLM-5.1（图 9b）。HyScale 总共节省了 Megatron-LM $66\%$ 的层时间。其中，HyTopK 贡献了 $37\%$ 的节省（因为它消除了前向和重计算中的分数表），选择重放贡献了 $18\%$（跳过了重计算选择），实现优化贡献了剩余的 $11\%$。

图9：H20 上单层 GLM-5.1 的累积消融
图9：H20 上单层 GLM-5.1 的累积消融

4. 架构通用性
在 32 个 H800 GPU 上，针对 1M 上下文测试块级架构。HyScale 使得 DeepSeek-V4-Flash 的 CSA 层加速 $2.01\times$，使得 Qwen3.8-Flash-Next 的 QSA 层加速 $1.57\times$（图 10）。这证明 HyScale 不仅适用于 token 级的 DSA，也能无缝泛化到块级稀疏注意力架构。

图10：32 个 H800 GPU 上单个 CSA 和 QSA 层的累积消融
图10：32 个 H800 GPU 上单个 CSA 和 QSA 层的累积消融

5. HyTopK 效率
图 11a 表明，HyTopK 消除了选择阶段的内存-时间权衡。其峰值显存仅为 1.5–1.9 GiB，且几乎不随上下文增长，运行速度快于占用 12-48 GiB 显存的 Megatron-LM 全分数表方案。图 11b 显示确定性排序仅增加 $0.8\%$ 的时间，即使有 64 个查询失败，无主机恢复也仅使开销增加到 $4.3\%$。图 12 验证了采样效率：HyTopK 的乐观跨步采样在所有缓冲区大小下，失败的查询比例远低于 LiteTopK 和 DoSA（在随机初始化权重下，HyTopK 甚至能做到 0 失败，而基线需要大近两倍的缓冲区）。

图11：HyTopK top-k 选择的端到端效率
图11：HyTopK top-k 选择的端到端效率
图12：采样效率：每种缓冲区大小下失败查询的比例
图12：采样效率：每种缓冲区大小下失败查询的比例

六、 补充细节（相关工作）

学习型稀疏注意力的选择加速。HPC-Ops Top-K（引用：Liu 等人，2026b）、GVR（引用：Cheng 等人，2026）、DeepSelect（引用：Qian 等人，2026）和 Prof-K（引用：Dziarmaga 等人，2026）加速了 top-$k$ 步骤。然而，打分仍然作为一个独立的内核运行，因此分数表仍然被写入显存并读回。为了消除分数表，LiteTopK（引用：Yin 等人，2026）首创了基于采样的选择，将打分与 top-$k$ 选择融合，并利用从附近查询的选择中估计的阈值来过滤分数。DoSA（引用：Meng & Fang，2026）进一步通过低成本地限制每个键的分数来跳过部分打分。这两者都是为推理设计的，在移植到训练时面临重大挑战（见 3.2 节）。本文中的 HyTopK 专为训练设计：它对每个查询自身的上下文进行采样以获得乐观阈值，在没有主机同步的情况下恢复失败的查询，并保持选择的确定性。与 HyScale 正交的工作包括 PIVOT（引用：Liu 等人，2026a）、IndexCache（引用：Bai 等人，2026）、HISA（引用：Xu 等人，2026）和 MISA（引用：Zhou 等人，2026），它们通过在查询或层之间重用选择，或通过对更少的块或头进行打分来减少选择工作。

七、 结论

本文提出了 HyScale，一个用于学习型稀疏注意力的长上下文训练系统。其核心洞察是选择阶段输出的仅是索引，而非连续激活。HyScale 利用这一洞察提出了两项新颖技术：HyTopK 在不写入分数表的情况下选择 top-$k$；选择重放（selection replay）跳过了重计算中重复的选择操作。在 512 个 H20 GPU 上，HyScale 将 GLM-5.1 在 1M token 序列上的训练速度相比 Megatron-LM 提升了 $2.07\times$。此外，它还能泛化到各种学习型稀疏注意力架构中。未来的工作可以进一步探索与其他正交稀疏注意力优化技术的结合。

八、 附录：学习型稀疏注意力的反向传播

本附录以 DSA 和 MLA 为例，使用 2.1 节的符号推导了一个学习型稀疏注意力层的反向传播，并证明选择和稀疏注意力的反向传播仅使用选中的条目。

选择阶段的反向传播。设 $z_{t,h,s} = \mathbf{q}_{t,h}^{I\top} \mathbf{k}_s^I$，公式 (4) 可写为 $I_{t,s} = \sum_h w_{t,h} \mathrm{ReLU}(z_{t,h,s})$。设 $\hat{\mathbf{p}}_t = \mathrm{Softmax}(\mathbf{I}_{t,\mathcal{T}_t})$ 为索引器在选择上的分布。由于 $\mathbf{p}_t$ 被分离（detached）且 $\mathcal{T}_t$ 在分数上是分段常数，公式 (3) 中索引器损失的分数梯度为：

$$\begin{aligned} g_{t,s} = \frac{\partial \mathcal{L}^I}{\partial I_{t,s}} = \begin{cases} \lambda \left( \hat{p}_{t,s} - p_{t,s} \right) & s \in \mathcal{T}_t, \\ 0 & \mathrm{otherwise}. \end{cases} \end{aligned}$$


设 $\delta_{t,h,s} = g_{t,s} w_{t,h} \mathbf{1}[z_{t,h,s} > 0]$，根据链式法则可得：

$$\begin{aligned} \begin{aligned} \frac{\partial \mathcal{L}^I}{\partial w_{t,h}} &= \sum_{s \in \mathcal{T}_t} g_{t,s} \mathrm{ReLU}(z_{t,h,s}), \\ \frac{\partial \mathcal{L}^I}{\partial \mathbf{q}_{t,h}^I} &= \sum_{s \in \mathcal{T}_t} \delta_{t,h,s} \mathbf{k}_s^I, \\ \frac{\partial \mathcal{L}^I}{\partial \mathbf{k}_s^I} &= \sum_{t : s \in \mathcal{T}_t} \sum_{h=1}^{H_I} \delta_{t,h,s} \mathbf{q}_{t,h}^I. \end{aligned} \end{aligned}$$
由于每一个求和都在查询的 $k$ 个选中条目上进行，反向传播仅需要 $s \in \mathcal{T}_t$ 对应的 $z_{t,h,s}$。也就是说，每个查询仅需要 $k H_I$ 个每头分数，而不是前向扫描计算的 $L H_I$ 个。这正是选择重放（见 3.3 节）重新计算的分数。由于索引器的输入被分离，梯度仅到达索引器的投影权重。

稀疏注意力的反向传播。在 MLA 中，查询 $t$ 的头 $h$ 通过 logits $\ell_{t,h,s} = \mathbf{q}_{t,h}^\top \mathbf{c}_s / \sqrt{d}$（对于 $s \in \mathcal{T}_t$）关注选中的潜在变量，其中 $\mathbf{c}_s$ 与其位置部分一起作为键，其潜在部分 $\mathbf{c}_s^V$ 作为值。设权重 $a_{t,h,s} = \mathrm{softmax}_{s \in \mathcal{T}_t}(\ell_{t,h,s})$，输出为 $\mathbf{o}_{t,h} = \sum_{s \in \mathcal{T}_t} a_{t,h,s} \mathbf{c}_s^V$。限制在 $\mathcal{T}_t$ 上的标准注意力反向传播（引用：Dao 等人，2022）给出：

$$\begin{aligned} \begin{aligned} \mathrm{d}\ell_{t,h,s} &= a_{t,h,s} (\mathrm{d}\mathbf{o}_{t,h}^\top \mathbf{c}_s^V - \mathrm{d}\mathbf{o}_{t,h}^\top \mathbf{o}_{t,h}), \\ \mathrm{d}\mathbf{q}_{t,h} &= \frac{1}{\sqrt{d}} \sum_{s \in \mathcal{T}_t} \mathrm{d}\ell_{t,h,s} \mathbf{c}_s, \\ \mathrm{d}\mathbf{c}_s &= \sum_{t : s \in \mathcal{T}_t} \sum_{h=1}^H \Big( \frac{1}{\sqrt{d}} \mathrm{d}\ell_{t,h,s} \mathbf{q}_{t,h} + a_{t,h,s} \mathrm{d}\mathbf{o}_{t,h} \Big), \end{aligned} \end{aligned}$$


其中 $\mathrm{d}\mathbf{c}_s$ 的第一项来自作为键的 $\mathbf{c}_s$，第二项（落在 $\mathbf{c}_s^V$ 上）来自作为值的 $\mathbf{c}_s$。目标 $\mathbf{p}_t$ 是跨头平均的 $a_{t,h,s}$ 且被分离，因此索引器损失不会向稀疏注意力发送梯度。

上下文并行。在 CP 下，$\mathbf{k}_s^I$ 和 $\mathbf{c}_s$ 被 all-gather，因此每个 GPU 从其自身的 $n$ 个查询中计算每个位置 $s$ 的 $\partial \mathcal{L}^I / \partial \mathbf{k}_s^I$ 和 $\mathrm{d}\mathbf{c}_s$。对这些部分和进行 reduce-scatter 操作会将属于每个 GPU 的位置的梯度返回给它们，这正是 2.2 节中提到的反向 reduce-scatter 过程。

方法细节中的参考文献汇总

  • [Alabi 等人, 2012] Alabi, T., Blanchard, J. D., Gordon, B., and Steinbach, R. Fast k-selection algorithms for graphics processing units. ACM Journal of Experimental Algorithmics, 17, 2012. 引用段落：3.2 节（确定性选择），用于说明基数选择算法。
  • [Chen 等人, 2016] Chen, T., Xu, B., Zhang, C., and Guestrin, C. Training deep nets with sublinear memory cost. 引用段落：2.2 节（激活重计算），用于说明激活重计算的背景。
  • [Dao 等人, 2022] Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. FlashAttention: Fast and Memory-Efficient exact attention with IO-Awareness. 引用段落：附录 A（稀疏注意力的反向传播），用于说明标准注意力反向传播。
  • [DeepSeek-AI 等人, 2026] DeepSeek-V4: Towards highly efficient Million-Token context intelligence. 引用段落：3.2 节（训练需要确定性的选择），用于说明训练系统需要按位可复现性。
  • [Korthikanti 等人, 2023] Korthikanti, V. A., Casper, J., Lym, S., McAfee, L., Andersch, M., Shoeybi, M., and Catanzaro, B. Reducing activation recomputation in large transformer models. 引用段落：2.2 节与 3.3 节，用于说明现有的操作符级选择性重计算策略。
  • [Li 等人, 2025] Li, J., Liu, S., and Sun, Y. FlashMLA: Efficient multihead latent attention kernels. 引用段落：3.4.2 节，用于说明现有的稀疏内核实现。
  • [Liu 等人, 2026a] PIVOT: Efficient Query-Group indexing for Token-Level sparse attention. 引用段落：3.4.2 节，用于支撑相邻查询选择键具有高度局部性的观察。
  • [Meng & Fang, 2026] DoSA: Dots sparse attention. 引用段落：3.4.2 节，用于支撑选择局部性的观察。
  • [Merrill & Grimshaw, 2011] High performance and scalable radix sorting: A case study of implementing dynamic parallelism for GPU computing. 引用段落：3.2 节（确定性选择），用于说明基数排序。
  • [NVIDIA, 2026] DeepSeek sparse attention (DSA). NVIDIA cuDNN. 引用段落：3.4.2 节，用于说明现有稀疏内核的单查询分配机制。
  • [PyTorch, 2026] torch.utils.checkpoint. 引用段落：3.3 节，用于说明现有的选择性重计算策略。
  • [Sul 等人, 2026] Mixture-of-Kittens: MoE megakernel for NVL72s. 引用段落：3.2 节（训练需要确定性的选择），用于说明内核确定性的需求。
  • [Wang 等人, 2025] TileLang: A composable tiled programming model for AI systems. 引用段落：3.4.2 节，用于说明现有的稀疏内核分配机制。
  • [Yin 等人, 2026] LiteTopK: Exploiting the curse of dimensionality for a fused Indexer-TopK kernel in Long-Context sparse attention. 引用段落：3.2 节与 3.4.2 节，用于说明 LiteTopK 针对推理首创的采样选择设计，以及选择局部性的观察。