Hierarchical Global Attention (HGA)

# 分层全局注意力：预训练长上下文Transformer的即插即用精确令牌路由
来源：https://arxiv.org/html/2606.30709
## 分层全局注意力：预训练长上下文Transformer的即插即用精确令牌路由

Woernle Frank Vladimir Fedosov Artemiy Grinenko BMW Group项目监督与工业研究方向。主要研究、算法设计与实现。

（2026年6月）

###### 摘要

分层全局注意力（Hierarchical Global Attention, HGA）是预训练长上下文Transformer中密集因果注意力的即插即用替代方案。HGA保留原始检查点参数：预训练的W_Q、W_K、W_V和W_O投影保持不变，不引入任何校准参数，也无需重新训练。

应用于Qwen3-30B-A3B-Instruct-2507-FP8（单块RTX 5090，32GB）时，修补后的模型开箱即可运行于64K令牌上下文，而在此硬件上按令牌级存储K/V是不可能的。

与以往的稀疏注意力方法不同，HGA执行分层两级路由。它首先使用紧凑的RoPE感知摘要检索相关块，然后通过仅路由最相关的组来精炼选择，最后执行精确的令牌级注意力。这种分层检索显著减少了拉取的令牌数量，同时保持对检索到的令牌集的精确注意力，使得基于RAM和NVMe的存储变得可行。

完整的历史令牌K/V位于主机RAM或NVMe存储中，而只有少量路由工作集在注意力期间传输到GPU内存。因此，GPU内存消耗主要取决于模型权重和路由工作集，而不是总上下文长度。

在所有测试的上下文长度（4K - 64K令牌）下，路由注意力与密集注意力之间的差异大约在0.01–0.02 nats范围内，而使用的稀疏度仅为约3%。这些结果表明分层路由引入的近似误差很小，剩余的质量差距很可能由长上下文位置编码主导，而非路由算法本身。

## 1 引言

应用于Qwen3-30B-A3B-Instruct-2507-FP8（单块RTX 5090）时，修补后的模型开箱即可运行于32K令牌上下文，而在此硬件上密集K/V存储根本不可能。在40M SmallLM上，直接复制权重仅导致与密集注意力在8K令牌时的+0.018 nat损失差距，而Triton融合实现对于训练快2.72倍，对于12K令牌的预填充快2.43倍。在Qwen3-30B-A3B-Instruct-2507-FP8上，64K令牌上下文的“大海捞针”检索测试在1.9%稀疏度下实现了100%通过率（3/3深度）。预训练长上下文LLM的实际限制往往不仅是O(n²)的注意力计算，还包括必须留在加速器内存中的密集K/V缓存。这一点在量化大型模型中尤为明显。在Qwen3-30B-A3B-Instruct-2507-FP8部署案例中，FP8模型权重几乎占满了32GB GPU的全部内存预算；在VRAM中存储每个历史令牌键和值使得长上下文几乎无法运行。

HGA正是为此场景设计的系统级补丁。它既不重新训练也不替换检查点的注意力投影。它保留原始的Q/K/V/O投影和注意力归一化，仅更改每个处理块拉取的历史键和值。因此，GPU上保存着预训练模型、紧凑的块摘要表、当前块、少数始终热门的块以及路由工作集。其余的历史令牌K/V位于主机RAM中。

核心观察是，连续的令牌通常不需要完全独立的长距离检索。一个64令牌的块通常讨论一个局部想法、函数、段落或对话轮次，该块中的许多令牌受益于相同的远距离块。因此，HGA首先从当前块路由到之前的块和组。只有在完成这个粗粒度检索步骤后，它才对所选区域的精确令牌键和值执行普通的softmax注意力。简而言之，长距离问题从密集的令牌到令牌的物化转变为块到块的检索，再加上对有限拉取集的精确令牌注意力。

这种设计对预训练模型特别有用。路由摘要由检查点自身的投影键构建；它们是带有RoPE感知相位规则的求和或均值，而不是新学习线性层的输出。不引入任何校准权重。最终注意力输出在当前方法中从不使用合成摘要值：摘要决定拉取什么，而原始令牌K/V计算输出。

实现和基准测试脚本见 https://github.com/vfedosov77/HierarchicalGlobalAttention。我们的目标不是重新设计Transformer注意力，而是提供一个即插即用的分层路由机制，可以直接应用于现有的预训练检查点。HGA并非近似注意力输出本身，而是仅近似检索阶段。一旦选择了相关区域，注意力输出就使用原始令牌级键和值精确计算。这种设计允许现有预训练模型在不修改其学习表示的情况下受益于稀疏长上下文注意力。

我们的贡献是：

- • 为预训练GQA/MoE Transformer提供即插即用的精确令牌注意力替代方案，重用原始Q/K/V/O权重和归一化，无需任何校准参数；
- • 一个Qwen3-30B-A3B-Instruct-2507-FP8集成，其中FP8模型在RTX 5090上，无需微调，开箱即可运行于32K令牌上下文的RAM支持路由K/V；
- • 一种块共享路由方案：邻近查询重用一个小型拉取工作集，例如在当前Qwen3配置中，对于64令牌块约1K路由中间令牌，加上固定的汇点/局部/当前块；
- • 一种两级块-组路由策略，结合确定性注意力汇点和近期局部块与上下文中间部分的内容路由检索；
- • 仓库基准测试显示，在40M SmallLM上，8K令牌时复制权重的损失差距仅为+0.01828 nats，12K令牌时训练步加速2.72倍；
- • 一种用于RAM支持推理的分层K/V存储抽象，其中完整的历史令牌K/V位于主机RAM中，而只有有限的工作集占用VRAM，将GPU内存消耗与上下文长度解耦；
- • 在Qwen3-30B-A3B-Instruct-2507-FP8上的“大海捞针”评估显示，无需任何微调，64K令牌上下文的检索准确率达100%。

## 2 背景

#### 因果注意力。

对于输入 X ∈ ℝ^{n×d}，一个标准注意力头计算

Attn(Q,K,V)_i = ∑_{j≤i} softmax_j ( q_i k_j^⊤ / √(d_h) ) v_j。 (1)

精确的令牌级感受野随 i 增长，因此查询-键比较的总数为 O(n²)，且 K/V 缓存与所有之前令牌成比例。

#### 分组查询注意力。

分组查询注意力（GQA）在查询头组之间共享键/值头[1]。HGA保留这种布局。源检查点的 W_Q、W_K、W_V 和 W_O 张量保持其形状和名称。在 Qwen3 封装器中，原始投影和 Q/K RMSNorm 模块通过引用保留。

#### RoPE。

旋转位置编码（RoPE）通过位置相关角度旋转查询和键维度[12]。只有当摘要键的相位与查询的 RoPE 几何兼容时才有用。因此 HGA 在构建块和组摘要时使用第 3 节所述的混合规则。

## 3 方法

### 3.1 分区

序列被划分为固定长度 C 的块；当前实现使用 C=64。每个块再细分为大小为 g_s 的组。40M 基准使用 g_s=16，路由块数 K_c=20，路由组数 K_g=32。Qwen3-30B 精确令牌封装器直接打开 topk_chunks=16 个选定块，并保持固定的第一个和最近窗口。

封闭块维护三种数据：

- • 用于第一阶段路由的块键摘要；
- • 用于第二阶段路由（当启用组路由时）的组键摘要；
- • 用于路由后精确注意力的原始令牌级键和值。

只有最后一项在最终注意力输出中使用。

### 3.2 混合RoPE路由摘要

块摘要必须保持与 RoPE 旋转后的查询的可比性。对已经旋转的键进行平均对于快速变化的高频对效果良好，但可能过度强调令牌级相位噪声。在原始键空间平均然后在一个代表性位置旋转结果，对于缓慢变化的低频对效果更好。

因此 HGA 使用一种混合策略。对于每个键维度对：

- • 高频 RoPE 对先按令牌旋转然后平均；
- • 低频 RoPE 对在原始键空间平均，然后在块或组的中间位置旋转。

相同的原则用于块摘要和组摘要。这些摘要是已投影键的和或均值；没有单独的摘要线性层。它们仅作为路由键使用，在当前输出 softmax 中从不作为值向量使用。

### 3.3 路由策略

对于活动块中的查询，HGA 从确定性和路由部分构建候选上下文。

#### 确定性可见性。

当前块始终可见，并带有因果掩码。可配置数量的第一个块始终可见，作为注意力汇点。可配置数量的最近封闭块始终可见，作为局部上下文。如果紧邻的前一个块不在该最近窗口中，则强制包含。

#### 基于内容的中间路由。

上下文的中间部分（排除始终可见的第一个和最后一个窗口）竞争路由预算。查询根据块键摘要进行评分。在块-组变体中，对选定块内的组摘要进行评分，打开前 K_g 个组以获取精确令牌级 K/V。在 Qwen3 精确令牌变体中，选定块直接按 KV 头作为令牌 K/V 拉取。

因此该方法不是纯粹固定的稀疏模式。它故意将固定汇点/局部窗口与中间上下文的内容路由检索结合起来。这与 MInference 使用的 A 形思想有关，其中初始汇点令牌和局部令牌始终可见，但 HGA 在剩余块上使用预训练键空间添加了一个 top-k 召回路径[5]。

### 3.4 精确令牌输出注意力

路由后，注意力模块仅拼接来自以下部分的真实令牌级键和值：

1. 1.当前块的因果可见部分；
2. 2.始终可见的第一个和最近块；
3. 3.根据实现变体，从路由中间上下文打开的组或选定块。

softmax 在此选定令牌集上精确计算。没有摘要输出令牌、摘要门控或每头标量校准参数。

### 3.5 检查点兼容性

该模块与密集注意力参数兼容。在 40M 模型中，密集权重直接复制到路由模型中。在 Qwen3 FP8 实现中，路由注意力封装器通过引用保留原始的量化投影和归一化。唯一改变的操作是每个查询允许关注的键和值集。

这是与需要从头训练或专门稀疏检查点的方法的主要区别。HGA 可作为预训练模型的系统级补丁使用。微调可以改善模型对新路由模式的使用，但对于 Qwen3-30B RAM 缓存演示或 40M 仅复制损失结果而言并非必需。

### 3.6 成本与内存

对于固定的块大小、固定汇点/局部窗口和固定路由预算，每个处理块关注并传输有界数量的真实令牌。因此，昂贵的令牌级注意力和令牌 K/V 传输与处理的令牌数呈线性关系：

T_{token-attn} = O( n ⋅ ( C + (F + L)C + B_{route} ) ) = O(n)， (2)

其中 F 和 L 分别是第一个和最近块的数量，B_{route} 是固定的路由令牌预算，例如打开组的 K_g g_s 或打开块的 K_c C。在 Qwen3 精确令牌配置中，K_c=16，C=64，因此每个块的路由中间预算为 1024 令牌，尚未加上确定性汇点、局部和当前块。

紧凑的路由表远小于令牌 K/V 缓存：每个块一个块键，每个块几个组键。参考实现可以直接扫描此表；对于非常长的上下文，此表可以限制、缓存或索引。重要的系统特性是完整的令牌 K/V 不需要驻留在 VRAM 中。VRAM 主要由模型权重、热的第一/最近窗口、块摘要和当前路由工作集决定。

## 4 实现

### 4.1 Qwen3-30B-A3B FP8 封装器

大型模型路径将 Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 的每个注意力模块替换为精确模式下的 QwenRoutedAttention。官方模型卡将此模型描述为 305 亿参数的 MoE，其中 33 亿参数被激活，48 层，GQA 包含 32 个查询头和 4 个 KV 头，FP8 量化，原生支持 262K 上下文[10]。HGA 封装器不修改模型权重。

当前的 RAM 缓存配置使用：

完整的历史令牌 K/V 记录存储在主机 RAM 中。每一步，路由器选择一个小型块集，存储仅将这些块拉到计算设备。连续的令牌常常请求重叠的块，因此一个有界 LRU 缓存保持频繁重用的块驻留。这与密集缓存不同，密集缓存中所有之前的令牌 K/V 必须始终物化在加速器上。

### 4.2 40M SmallLM 实现

SmallLM 实验使用 8 个解码器层，隐藏层大小 384，6 个查询头，2 个键/值头，FFN 大小 2048，RoPE 和 GQA。路由注意力使用 C=64，g_s=16，K_c=20，K_g=32。融合 CUDA 路径使用 Triton 和 torch.compile；密集基线是禁用路由的相同模块，回退到对 RoPE 查询和键的因果缩放点积注意力。

40M 测试路径有意不使用任何校准参数。零样本质量结果将相同的密集权重复制到密集和路由模型中，因此测量的差距隔离了稀疏路由而非训练差异。

### 4.3 分层 K/V 存储

KvRouter 实现将路由与存储分离。一个存储暴露四个操作：追加新封闭的块、读取块摘要用于路由、获取第二路由级别的组摘要、以及收集选定块或组的精确令牌 K/V。

基于 RAM 的实现（RamKVCacheStore）按温度分区数据：

- • 热点：块摘要以及始终可见的第一个和最近块始终驻留在计算设备上。在训练期间，这些张量保留梯度。