SwiftQK：当归一化层成为张量并行的瓶颈

笔记日期： 2026-08-15 作者： Zhongzhu Zhou 评述论文： SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization 论文作者： Gyudong Kim, Wonjun Han, Young Geun Kim（高丽大学） arXiv： 2608.09160 发表/状态： IEEE Computer Architecture Letters，2026 年 8 月预印本

1. 一篇只有四页正文的短论文，为什么值得写一篇完整笔记

这篇论文很短——含参考文献一共五页——但它是一个在真实 LLM 推理系统中反复出现的典型样例：一个单独看起来微不足道的组件（一个每个元素只有几次浮点运算的归一化层），一旦把模型切分到多张 GPU 上，就可能变成主导延迟的瓶颈。SwiftQK 正是针对这样一个组件：张量并行（Tensor Parallelism, TP）下的 Query-Key 归一化（QK-Norm）。论文在提出任何修复方案之前给出的核心实证发现本身就很惊人：对同时使用 QK-Norm 的 OLMo2 和 OLMo3 而言，TP 从 2 卡扩展到 4 卡时，仅 QK-Norm 本身的开销就从总延迟的约 20% 增长到约 30%，而这部分开销里有 85% 是纯粹的跨 GPU 同步，而非计算本身。这意味着对一个”现代”开源 LLM 来说，增加更多 TP 并行的 GPU 带来的收益正在递减——不是因为注意力或 MLP 块扩展性差，而是因为某一个特定的归一化步骤扩展性很差，并且随着设备数增加而愈发糟糕。

论文提出的修复方案 SwiftQK 在概念上很简单：与其交换完整的 Query/Key 激活向量让每个设备本地计算归一化统计量，不如只交换那个标量统计量本身，并把（此时已经便宜得多的）同步过程与有用的计算在一个持久化的 CUDA kernel 内重叠起来。本篇笔记会花较多篇幅先讲清楚这个问题为什么会存在——张量并行到底切分了什么、RMSNorm 的结构为什么会在你最不希望出现同步点的地方强行插入一个同步点、以及为什么已有的隐藏通信的 kernel 融合技术在这里并不奏效——然后再详细展开 SwiftQK 的三阶段设计、其避免死锁的论证，以及它的数值精度分析。

前置知识：读懂这篇论文之前需要了解的背景

张量并行（TP）切分的是什么。 当一个语言模型层太大——或者延迟要求太苛刻——以至于无法在一张 GPU 上运行时，张量并行会把单层内部的计算切分到多张 GPU 上，这与流水线并行（Pipeline Parallelism, PP）不同——PP 是按层切分模型，不同层放在不同设备上顺序执行。在经典的 Megatron 式 TP 布局中，Attention 或 MLP 子层的第一个线性投影沿输出维度切分（因此每张 GPU 只计算投影激活的一个切片），紧随其后的投影则沿输入维度切分，这样它可以直接消费本地那一份切片，中间不需要任何通信。只有在最后需要把该子层的投影输出跨 GPU 合并时才需要一次通信——通常通过 All-Reduce 实现。这正是 TP 在拥有高速互联（NVLink、NVSwitch）的机器上对延迟敏感的推理场景很有吸引力的原因：一个子层内的两次矩阵乘法之间不需要任何同步点，唯一需要的通信就是每个子层末尾的那一次归约。

RMSNorm，以及它为什么需要整个向量。 均方根归一化（Root Mean Square Normalization）是现代开源 LLM 里几乎普遍使用的归一化函数，它用输入向量 $x \in \mathbb{R}^H$ 的均方根幅值的倒数对其重新缩放：

$$\mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{H}\sum_{j=1}^{H} x_j^2 + \epsilon}} \odot \gamma \tag{1}$$

其中 $H$ 是隐藏维度，$\epsilon$ 是为数值稳定性加入的一个很小的常数（防止 $x$ 接近零时除以零），$\gamma \in \mathbb{R}^H$ 是一个可学习的逐元素缩放参数。式 (1) 中最关键的结构性事实——也是这篇论文之所以存在的全部原因——在于分母：$\sum_{j=1}^{H} x_j^2$ 对 $x$ 的每一个元素求和。如果 $x$ 的 $H$ 个元素被切分到多张 GPU 上——正如 TP 场景下每张 GPU 只持有投影后 Query 或 Key 向量的一个切片——那么任何单张 GPU 都无法仅凭自己那份切片计算出正确的分母。跨 GPU 的某种通信是不可避免的；唯一的问题是需要多少通信量，以及这些通信有多昂贵。

逐层（layerwise）与逐头（headwise）QK-Norm 的区别。 QK-Norm 最初是为了稳定超大模型或超长上下文模型的注意力 logits 而提出的（最早出现在 ViT-22B 中），它在 Scaled Dot-Product Attention 计算之前对投影后的 Query 和 Key 张量应用 RMSNorm。它可以在两种粒度上应用：逐头（headwise）QK-Norm 为每个注意力头独立计算一个归一化因子（因此如果某个头恰好完整地落在一张 GPU 的切片内，式 (1) 中的归约就只需覆盖该头的维度，完全不需要通信），而逐层（layerwise）QK-Norm 则在整个投影后的 Q 或 K 维度上计算单一归一化因子，跨越所有头。论文的调研表明，近期的开源模型系列——OLMo2、OLMo3 等——都收敛到了逐层这一形式。而正是这种版本在 TP 切分下彻底”破功”：因为一个 token 的整个 Q/K 向量被切分到多张 GPU 上，逐层 RMSNorm 唯一的全局分母强行在原本（投影 → 归一化 → 注意力）无需通信的一段计算中间插入了一个同步点。

为什么标准解法（All-Gather）代价高昂，以及为什么简单的 kernel 融合重叠技巧救不了它。 让每张 GPU 都获得计算式 (1) 分母所需信息的默认做法，是在归一化之前对所有 TP rank 的完整 Q/K 激活向量做 All-Gather——也就是在每一对 GPU 之间物理搬运 $O(H)$ 个元素，在每台设备上重建完整向量，然后本地计算平方和。这就是论文自己给出的”改造后的注意力块”示意图（下方图 1(b)）中标注的”Q/K Sync (All-Gather)“步骤，它在 Q/K 投影和注意力计算之间插入了一整个和激活张量同等大小的通信步骤——恰好是标准（不带 QK-Norm 的）TP 中完全没有同步点的位置。已有的用于隐藏 TP 通信的 kernel 融合技术，如 FLUX 和 FlashOverlap，其工作原理是把一次通信与独立的、足够大的一块计算（通常是流水线中的下一个 GEMM）重叠起来，这样等到 GEMM 计算完成时，被传输的数据也差不多到齐了。这种重叠策略对 QK-Norm 来说恰恰失效，因为 RMSNorm 本身的计算极其轻量（每个元素只有几次平方、加法、一次开方、一次除法），相对于All-Gather的通信量而言，本地根本没有足够的独立计算量可以用来隐藏这次通信的开销。论文的表述很精确：仅靠重叠的技术假设”通信可以被隐藏在足够的独立计算背后”——这个假设对 GEMM+All-Reduce 模式成立，但对 RMSNorm+All-Gather 不成立。

2. 在提出修复方案之前，先量化问题有多严重

在给出 SwiftQK 之前，论文花了完整的一节来刻画逐层 QK-Norm 瓶颈到底有多严重，使用了三个几乎共享一切（除了是否、以及如何使用 QK-Norm）的模型家族做受控对比：OLMo（不用 QK-Norm），以及 OLMo2 / OLMo3（都用逐层 QK-Norm）。在 2 卡和 4 卡 A100 上，让基于 ShareGPT 的请求流量跑过 vLLM 的连续批处理引擎，分别在流水线并行和张量并行下运行，论文把 prompt 处理延迟拆成三类：Q/K 同步、Q/K 归一化计算、以及其他一切。

图 1（对应论文 Fig. 1）：带 QK-Norm 的注意力块、其 TP 切分模式，以及由此产生的延迟拆解。

图 1(a)（对应论文 Fig. 1(a)）：Pre-LN Transformer 块，QK-Norm 的两个 RMSNorm 层被直接插入到投影后的 Q 和 K 张量上，紧接在 Q/K/V 线性投影之后、Scaled Dot-Product Attention 之前。

图 1(b)（对应论文 Fig. 1(b)）：在 TP 下，投影后的 Q 和 K 张量被切分到多张 GPU 上（图中双卡场景每张 GPU 分别持有 $Q_1/K_1/V_1$ 或 $Q_2/K_2/V_2$）。蓝色箭头标出的是必须在归一化后的 $Q$/$K$ 送入注意力计算之前完成的 All-Gather Q/K 同步步骤；黄色菱形标出的是 TP 本身在每个子层末尾就已经需要的、与 QK-Norm 无关的另一次 Layer Sync（All-Reduce）。

图 1(c)（对应论文 Fig. 1(c)）：OLMo（无 QK-Norm）、OLMo2、OLMo3（均为逐层 QK-Norm）在 PP*（一个 PP 基线）与 TP 下、2 卡与 4 卡下的归一化总 prompt 处理延迟。延迟在每个模型/GPU 数量组合内部相对 PP* 值归一化；TP 的堆叠柱状图把总延迟拆成 Q/K Sync（蓝色）、Q/K Norm Computation（棕黄色）与 Others（灰色）。

图 1(c) 背后的数字是论文最核心的动机性结果。对于没有 QK-Norm 的 OLMo，从 PP 切换到 TP（4 卡）能把延迟降低 36.2%，这正是我们期望 TP 通过避免流水线气泡带来的”正常”收益。但对于 OLMo2 和 OLMo3，同样是 4 卡下 PP→TP 的切换，延迟却只分别降低了 7.5% 和 12.0%——TP 通常带来的优势有很大一部分被 QK-Norm 的开销吃掉了。深挖为什么会这样：在 PP 下，QK-Norm 的同步加计算开销几乎可以忽略不计，OLMo2 和 OLMo3 分别平均只占总执行时间的 1.4% 和 1.3%（因为 PP 根本不切分隐藏维度，所以不需要做任何跨 GPU 归约）。而在 TP 下，这一开销膨胀到 2 卡下的 20.0%/19.0%，并在 4 卡下进一步增长到 30.1%/29.7%。这种”随 GPU 数增长而增长”的模式正是问题所在的证据：当你把隐藏维度切分得越细、分布到越多 GPU 上时，每张 GPU本地需要做的那部分 RMSNorm 计算量在缩小（每卡工作量更少），但同步成本——等待每个对端的偏和变得可见——并不会随之缩小，反而在（此时已经更小的）总延迟中占据了更大的比例。论文精确地定位了驱动因素：Q/K 同步在 2 卡下占 QK-Norm 开销的 73.6%/72.8%，在 4 卡下上升到 85.6%/85.5%。换句话说，RMSNorm 在做的计算本身几乎无关紧要；开销几乎全部来自搬运数据（或者，正如我们接下来会看到的，仅仅是等待几个标量变得可见）在 GPU 之间的成本。

3. SwiftQK 的设计：先缩减通信量，再隐藏剩余的延迟

SwiftQK 的策略有两个相互独立的组成部分，值得明确指出两者都是必需的——论文的消融式基线对比（第 4 节讨论）清楚表明，单独使用其中任何一个都无法获得完整收益：

1. 缩减必须跨 GPU 传输的数据量。 与其 All-Gather 整个 $O(H)$ 元素的 Q/K 激活向量让每张 GPU 都能本地重新计算 $\sum_j x_j^2$，不如让每张 GPU 只计算自己本地切片的平方和，然后只与对端交换这一个标量。这把一个 $O(H)$ 的通信模式变成了 $O(1)$ 的模式（每张 GPU 一个浮点数，而不是每张 GPU 一整个向量）。
2. 隐藏剩余部分的延迟。 即使把负载缩减成了一个标量，GPU 仍然要等待每个对端的标量到达，才能计算出全局的平方和——这是一个同步延迟问题，而不是容量问题，它不会仅仅因为负载变小了就消失。SwiftQK 把这段残余延迟隐藏在不依赖于归约结果的其他有用计算背后。

算法 1：融合的多 GPU 持久化 RMSNorm kernel

论文把这两个想法都打包进了一个单一的融合 CUDA kernel 里，从不启动独立的通信和归一化 kernel，逐 token 执行三个阶段：

算法 1 —— 融合的多 GPU 持久化 RMS-Norm kernel（复现自论文，附带展开说明）

输入：本地输入切片 X_local，本地权重 W_local，
      总隐藏维度 H_total，GPU 数量 N，
      当前 rank r，IPC 缓冲区 B_IPC
输出：归一化后的本地切片 Y_local

对每个分配到的 token t（持久化的 block 级步进循环）：

    # ---- 阶段 A：本地平方和 ----
    S_local <- BlockReduceSum( X_local[t]^2 )      # 只对本地切片做归约
    如果 thread_id == 0:
        B_IPC[r].sumsq[t] <- S_local                # 把本地标量发布到 IPC 缓冲区
    __syncthreads()                                  # 屏障：阶段 A 结束

    # ---- 阶段 B：通信与计算重叠 ----
    如果 warp_id == 0:
        # --- 通信路径（仅 Warp 0）---
        向对端 GPU RemoteWrite(flag)                 # 通知"我的标量已就绪"
        SpinWait(local_flag)                          # 等待所有对端发出信号
        S_global <- sum_{k=0}^{N-1} B_IPC[k].sumsq[t]  # 点对点标量归约
        RMS_inv  <- 1 / sqrt(S_global / H_total + eps) # 全局 RMS 因子
    否则:
        # --- 计算路径（其余 warp，与上面并发执行）---
        对每个分配到的隐藏维度 j:
            X_local[t][j] <- X_local[t][j] * W_local[j]  # 权重乘法（与 S_global 无关）
    __syncthreads()                                  # 屏障：阶段 B 结束

    # ---- 阶段 C：最终归一化 ----
    对每个分配到的隐藏维度 j:
        Y_local[t][j] <- X_local[t][j] * RMS_inv     # 应用全局 RMS 缩放

逐步说明每个阶段为什么存在：

阶段 A —— 把 $O(H)$ 的交换变成 $O(1)$ 的交换。 关键的观察是，式 (1) 的分母只需要标量 $\sum_j x_j^2$，并不需要原始向量 $x$ 本身。每张 GPU 只对它本地持有的元素用标准的块级并行归约（BlockReduceSum）计算这个和，然后把这一个浮点数写入一个进程间通信（IPC）缓冲区，对端 GPU 可以通过 NVLink 直接读取这个缓冲区，无需经过 host 端或某个集合通信库的 All-Gather 原语。这是设计中缩减数据量的那一半：与其为每一对 GPU 移动 $H/N$ 个浮点数（All-Gather 下的切片大小），SwiftQK 为每一对 GPU 只移动恰好 1 个浮点数。

阶段 B —— 用独立的工作隐藏等待。 即使只需要交换一个标量，也总会有某张 GPU”最后一个到”，而其他每张 GPU 都必须等待这个标量变得可见后才能算出 $S_\text{global}$。这是一个延迟问题——一个不可避免的同步屏障——无论怎么缩减负载量都无法消除它。SwiftQK 的应对方式是让 GPU 在等待的同时做一些有用的事：它把每个 CUDA block 的 warp 拆成两个互不相交的角色。Warp 0 完全专注于通信路径——向对端发信号、自旋等待对端的信号、执行那一次微小的标量点对点归约——而 block 中的其余每个 warp 则对本地切片执行 RMSNorm 的权重乘法（式 (1) 中的 $x_j \times \gamma_j$），这一步在数学上与 $S_\text{global}$ 无关，因此可以在全局和尚未算出之前安全地提前进行。这是一个真正巧妙的技巧：它之所以可能，是因为 RMSNorm 恰好可以分解为一个需要全局归约的归一化尺度项，和一个不需要全局归约的可学习权重乘法项，二者可以以任意顺序甚至并行计算。

阶段 C —— 应用最终结果。 一旦阶段 B 的两条路径都完成（由 __syncthreads() 屏障强制保证——block 中的所有 warp 都必须到达这一点，意味着归约和权重乘法都已完成），最后一个阶段把已经完成权重缩放的本地切片乘以刚刚算出的全局 RMS 因子，并直接把结果写回 HBM。值得注意的是，只有最终归一化后的切片会被写出——论文明确指出这是除了 IPC 标量交换之外唯一的内存层级交互，因此这个 kernel 完全避免了物化任何大的中间张量（比如 All-Gather 版本 QK-Norm 需要的那个被聚集起来的完整 Q/K 向量）。

持久化 kernel 循环。 “对每个分配到的 token t（持久化的 block 级步进循环）“这一表述很关键：SwiftQK 没有为每个 token 启动一个 CUDA block（那样会为每个 token 引入一次 kernel 启动与销毁的开销，也会让阶段 B 中的对端同步逻辑变得脆弱，因为来自不同启动批次的 block 无法保证同时驻留），而是一次性启动一个有界、固定数量的 block，每个驻留的 block 在一个循环中处理一串 token，对每个 token 重复阶段 A–C。这正是 GPU 编程意义上”持久化 kernel”的含义，也是下一小节死锁安全性论证背后的机制。

为什么死锁安全性需要单独论证

阶段 B 中的 SpinWait 是一种忙等待：一个 warp 原地自旋，反复检查一个标志位，直到对端的写入变得可见。只要需要写入该标志位的对端 block 确实在并发运行，这就没有问题——但 GPU 会以硬件调度器选择的任意顺序把 block 调度到流式多处理器（SM）上，如果你启动的 block 数超过了可同时驻留的数量，一些 block 就会停留在启动队列里，根本还没开始执行。如果已经在运行的 Block A 正在自旋等待还在排队、尚未被调度的 Block B 发出的信号，而 Block B 反过来又需要某种只有 Block A 完成后才会释放的资源——这就是经典的循环等待——kernel 就会死锁：每个 SM 都被一个卡在自旋状态的 block 占据，而没有任何一个 block 能完成从而为排队中的 block 腾出空间。

为什么这里是真实的风险，而不是假设性的。 SwiftQK 的跨 GPU 对端同步本质上要求处理同一 token 的对应 block 在每一张 GPU 上都同时驻留并参与点对点握手——如果 GPU 2 上的”对端” block 还没被调度到某个 SM 上，GPU 1 的 Warp-0 在等待 GPU 2 标志位时的自旋就会永远自旋下去（或者直到硬件超时），因为根本不会有任何东西来设置那个标志位。这正是朴素的”每 token 一个 block”kernel 设计在规模化时容易陷入的失败模式，也正是促使论文选择持久化 kernel 设计的约束条件——而不仅仅是一个性能优化。

修复方法：只启动能保证同时驻留的数量的 block。 令 $B_\text{res}$ 表示给定 kernel 每 block 资源占用（寄存器、共享内存、线程数）情况下，一个 SM 上能同时驻留的最大 block 数，令 $N_\text{SM}$ 表示设备上的 SM 数量。SwiftQK 最多启动 $B_\text{res} \times N_\text{SM}$ 个 block——这个数字可以由 CUDA occupancy calculator 根据编译后 kernel 的资源使用情况提前确定——这保证了每一个被启动的 block 从 kernel 开始运行的那一刻起就是驻留的，因此阶段 B 中的自旋等待永远是在等待一个已经在运行的 block，而不会是一个还在排队中的 block。持久化的逐 block token 循环（每个 block 依次处理多个 token，而不是为每个 token 启动一个新 block）正是让这个固定、有界的网格仍然能够处理任意数量 token 的机制。

设计选择讨论——为什么不干脆使用一个自带死锁保护的库原语？ 一个显而易见的替代方案是依赖 NCCL 或类似的集合通信库，它们已经内置了成熟的集合操作死锁规避逻辑。论文没有走这条路，是因为集合通信库通常围绕 kernel 启动边界的语义设计——一次集合调用通常本身就是一次完整的 kernel 启动（或一系列启动），有自己的调度和同步开销，这恰好重新引入了 SwiftQK 试图通过把通信和计算融合进一个 kernel 来消除的那种启动开销。SwiftQK 接受的权衡是：它必须自己重新实现死锁安全性（通过上述基于 occupancy 的有界网格论证），以换取避免 NCCL 的启动开销，并对单次 kernel 调用内哪些 warp 做哪些工作拥有细粒度控制。值得指出的边界条件是：这个基于 occupancy 的安全性论证依赖于在编译/启动时准确知道 kernel 的实际资源使用情况；如果寄存器或共享内存使用发生变化（例如由于编译器版本变化、不同的数据类型，或者一个 SM 资源预算不同的 GPU 架构），安全的网格大小就必须重新计算，否则死锁自由的保证就不再成立。这是一个基于库的集合通信不会遇到的真实可移植性约束。

4. 评测：设计真的兑现了收益吗？相对于什么基线？

论文用四个对比点评测 SwiftQK，值得拆解一下，因为这正是用来分辨哪一部分设计（缩减数据量 vs. 隐藏延迟）在起作用的关键：

• All-Gather-based QK-Norm —— 未修改的基线：整向量交换，不做任何重叠尝试。
• Comm-Overlap —— 把 All-Gather 通信与 RMSNorm 计算重叠，即在不缩减通信负载的前提下应用通用的 FLUX/FlashOverlap 式重叠思路。
• MiniMax(eager) —— 标量统计量聚合（与 SwiftQK 阶段 A 相同的数据量缩减），但不融合进单一 kernel、也不与独立计算重叠——通信和归一化顺序执行。
• MiniMax(fusion) —— 把标量聚合进一步与 RMSNorm 计算融合进一个优化过的 kernel（论文引用的一个已有 vLLM 实现），但论文自己的结果表明，这个基线仍未达到 SwiftQK 阶段 B 那种 warp 级别的通信-计算重叠。
• SwiftQK —— 数据量缩减和通信-计算重叠都具备，融合进一个持久化 kernel。

这是一个构造得很好的消融阶梯：Comm-Overlap 单独隔离出”仅靠重叠、不缩减数据量是否有帮助”，MiniMax(eager) 单独隔离出”仅靠缩减数据量、不做重叠是否有帮助”，MiniMax(fusion) 在缩减数据量之上加上了融合，而 SwiftQK 又在此基础上加上了 warp 级别的重叠。测试模型包括 OLMoE（7B，混合专家）、OLMo 2（13B，稠密）和 OLMo 3（32B，稠密）——之所以选择它们，是因为它们在模型规模和架构上各不相同，却都使用逐层 QK-Norm，从而隔离出论文关心的那个变量。微架构层面的性能剖析运行在两张通过 NVLink 互联的 RTX 3090 上；端到端服务评测运行在 4 卡和 8 卡通过 NVLink 互联的 A100 服务器上，所有 kernel 都集成进了 vLLM。

图 2（对应论文 Fig. 2）：微架构性能剖析、数值精度对比，以及端到端服务性能。

图 2(a)（对应论文 Fig. 2(a)）：相对延迟、相对 NVLink TX 吞吐量、相对 SM issue rate，均相对 All-Gather 基线归一化，在 64 和 4096 输入 token 数下、跨 OLMoE、OLMo2、OLMo3 测得。SwiftQK 在每一种配置下都展现出最低的相对延迟，同时展现出最高的 SM issue rate——这是 Warp 0 的通信路径与其他 warp 的计算路径确实在并发运行、而不是互相阻塞的证据。

图 2(b)（对应论文 Fig. 2(b)）：All-Gather-based、MiniMax(fusion) 和 SwiftQK 相对于金标准全精度（FP64）RMSNorm 参考值的平均绝对误差与 RMSE，分别针对 BF16 和 FP8（e4m3）激活。SwiftQK 的误差与其他方法相当，而不是更大——也就是说，标量偏和聚合带来的累加顺序变化，并没有在低精度格式本身已带来的误差之上引入额外的数值误差。

图 2(c)（对应论文 Fig. 2(c)）：对比五种方法在 OLMoE/OLMo2/OLMo3 上、4 卡和 8 卡下的每输出 token 时间（TPOT，越低越好）和饱和请求吞吐量（越高越好），在饱和点以下扫描请求速率（RPS）。SwiftQK（⑤，深蓝色）在每一个 TPOT 面板中都稳定地位于其他四条曲线下方，在每一个吞吐量柱状图中都位于其他四者上方。

解读这些数字。 在微架构性能剖析中（图 2(a)），SwiftQK 相对 All-Gather-based QK-Norm 把 QK-Norm 延迟降低了 81.4–93.9%，相对 MiniMax(fusion) 又进一步降低了 29.4–77.0%——后一个对比很重要，因为它隔离出了 SwiftQK 的 warp 级重叠在已经具备数据量缩减和融合的 MiniMax(fusion) 之上额外带来的收益。在 4096 token 下，SwiftQK 的 SM issue rate 比 MiniMax(fusion) 高 2.8–4.6 倍，直接证实 SwiftQK 的 Warp-0/其他 warp 拆分设计确实在通信等待期间让计算单元更忙碌，而不是报告的延迟收益来自某种无关的 kernel 效率差异。

在数值精度方面（图 2(b)），论文的担忧是合理且值得认真对待的：重新排列一个求和的累加方式（每张 GPU 切片的本地和通过点对点标量归约合并，相对于对重建后的完整向量做一次大归约）在原则上可能改变浮点数的舍入行为，因为浮点数加法不满足结合律。SwiftQK 的应对方法是，即使底层激活以 BF16 或 FP8-E4M3 存储，也用 FP32 来累加本地平方和以及跨 GPU 的标量归约。相对 FP64 金标准参考值测得的最大绝对误差、平均绝对误差和 RMSE（BF16 下为 5.0e-2、2.1e-3、3.4e-3；FP8 下为 5.0e-1、2.5e-2、3.9e-2）与其他方法的结果非常接近——也就是说，SwiftQK 的设计并没有在 BF16/FP8 格式本身已经带来的误差之上引入有意义的额外数值漂移。对于一篇修改了一个对数值敏感的归一化操作的累加模式的系统论文来说，这是一项真正重要的检查，论文能够纳入它而不是只报告速度数字，值得肯定。

在端到端服务性能方面（图 2(c)），相对 All-Gather 基线，SwiftQK 把平均 TPOT 降低了 29.5%，把饱和吞吐量提高了 25.4%；相对两个中间基线，它把 TPOT 分别降低了 17.9%（对比 Comm-Overlap）和 28.5%（对比 MiniMax(eager)），吞吐量增益分别为 14.6% 和 20.1%。对论文核心主张最有说服力的一点是，即使相对已经做了标量聚合和融合的 MiniMax(fusion)，SwiftQK 仍然把 TPOT 降低了 14.3%，把吞吐量提高了 8.8%。论文对这最后一个对比的表述是恰当的：它证明了 SwiftQK 的端到端收益”不仅仅来自标量统计量聚合，也来自将缩减的通信量与融合的持久化执行以及 kernel 内的通信-计算重叠相结合”——也就是说，设计的两个部分（第 3 节的两条编号策略）都在独立地做出贡献，而不是其中一个做了全部工作、另一个只是搭便车。

4b. 对算法 1 的数值化走读

只看符号不够直观，我们不妨在一个 2 卡 TP 的玩具场景里，逐 token 走一遍算法 1 的具体数字：总隐藏维度 $H_\text{total} = 8$，均匀切分成两份，每张 GPU 持有该 token 投影后 Query 向量的 $H_\text{local} = 4$ 个元素，取 $\epsilon = 10^{-6}$、$\gamma_j = 1$（为了让算术保持简洁）。

设该 token 未切分前的真实投影 Q 向量为

$$x = [\,1.0,\ 2.0,\ -1.0,\ 0.5,\ 3.0,\ -2.0,\ 1.5,\ 0.5\,] \tag{2}$$

GPU 0 持有前 4 个元素 $x^{(0)} = [1.0, 2.0, -1.0, 0.5]$，GPU 1 持有后 4 个元素 $x^{(1)} = [3.0, -2.0, 1.5, 0.5]$。

每张 GPU 上的阶段 A。 各自计算本地平方和：

$$S_\text{local}^{(0)} = 1.0^2 + 2.0^2 + (-1.0)^2 + 0.5^2 = 1.0 + 4.0 + 1.0 + 0.25 = 6.25 \tag{3}$$ $$S_\text{local}^{(1)} = 3.0^2 + (-2.0)^2 + 1.5^2 + 0.5^2 = 9.0 + 4.0 + 2.25 + 0.25 = 15.5 \tag{4}$$

每张 GPU 只把一个标量（$6.25$ 或 $15.5$）写入该 token 对应的 IPC 缓冲区槽位。作为对比，All-Gather-based QK-Norm 需要把每张 GPU整个4 元素切片搬到对端——每个方向各 4 个浮点数，而不是 1 个。

阶段 B：并发进行的点对点归约。 每张 GPU 的 Warp 0 在两个 IPC 标量都可见后读取它们：

$$S_\text{global} = S_\text{local}^{(0)} + S_\text{local}^{(1)} = 6.25 + 15.5 = 21.75 \tag{5}$$ $$\mathrm{RMS}^{-1} = \frac{1}{\sqrt{S_\text{global}/H_\text{total} + \epsilon}} = \frac{1}{\sqrt{21.75/8 + 10^{-6}}} = \frac{1}{\sqrt{2.71875 + 10^{-6}}} \approx \frac{1}{1.6489} \approx 0.6065 \tag{6}$$

在 Warp 0 计算式 (5)-(6)（这需要等待两个标量都可见，是整个 kernel 里唯一真正的同步延迟）的同时，每张 GPU 的其他 warp 同时在为本地切片计算 $x_j^{(r)} \times \gamma_j$。由于这个玩具例子里 $\gamma_j = 1$，这一步在数值上是空操作，但在一般情形下（学到的 $\gamma \ne 1$），这是与式 (5)-(6) 的等待并行发生的真实、独立的工作。

阶段 C：应用全局因子。 每张 GPU 用共享的 $\mathrm{RMS}^{-1} \approx 0.6065$ 乘以自己（已经完成权重缩放的）本地切片：

$$y^{(0)} = 0.6065 \times [1.0, 2.0, -1.0, 0.5] \approx [0.607, 1.213, -0.607, 0.303] \tag{7}$$ $$y^{(1)} = 0.6065 \times [3.0, -2.0, 1.5, 0.5] \approx [1.820, -1.213, 0.910, 0.303] \tag{8}$$

与未切分计算的一致性核验。 直接对完整的 8 元素向量 $x$（式 (2)）计算 RMSNorm，得到 $\sum x_j^2 = 6.25 + 15.5 = 21.75$——与式 (5) 中的 $S_\text{global}$ 完全一致，证实了这种切分后的两阶段计算与未切分的计算完全数学等价（除浮点舍入顺序外），而不是一种近似。这正是论文所说 SwiftQK 保持了“同样的全局归一化语义”的意思：式 (3)-(8) 把一个 $O(H)=8$ 元素的跨 GPU 交换缩减到了 2 个标量的交换（每个方向每对 GPU 1 个），同时计算出与参考 All-Gather 路径完全相同（除 FP32 累加顺序外）的输出。

通信节省从哪里来。 在这个玩具例子中，节省只是从 4 个浮点数减少到 1 个（4 倍缩减），因为 $H_\text{local}=4$ 很小。而在真实部署中，一个 13B 参数模型在 4 卡 TP 下每张 GPU 的 Q/K 切片可能持有大约 $H_\text{local} \approx 1024$-$2048$ 个元素（取决于头数和 TP 度数）——因此同样的阶段 A 技巧把每对 GPU 超过一千个浮点数的交换缩减成一个标量，是超过 1000 倍的原始负载缩减，这与论文报告的 81-94% 延迟降低主要由数据量而非重叠技巧驱动是一致的（重叠技巧贡献的是相对 MiniMax(fusion)（已经具备数据量缩减）的额外 29-77% 改进）。

4c. SwiftQK 与其他 TP 通信缩减策略的对比

把 SwiftQK 的方法放到一个更宽的谱系上，与更广泛的 TP/服务文献中探索过的其他通信缩减技术对比，有助于看清它组合想法的新颖之处到底在哪里。

技术	缩减通信量？	重叠剩余延迟？	融合进一个 kernel？	适用于 QK-Norm 的全局归约模式？
标准 All-Gather TP	否（完整向量）	否	否（通信与归一化分开）	是（基线）
FLUX / FlashOverlap（GEMM+通信重叠）	否	是，但仅当计算/通信比很大时	部分（通信/GEMM 融合）	否——RMSNorm 计算太轻，无法隐藏 All-Gather
MiniMax(eager) 标量聚合	是（$O(H)\to O(1)$）	否（通信后归一化顺序执行）	否	是
MiniMax(fusion)	是	无 warp 级重叠	是（通信+归一化融合）	是
SwiftQK	是	是（warp 级，kernel 内）	是（单一持久化 kernel）	是
逐头 QK-Norm（架构性回避）	不适用——通过将每个头的归一化保持在单张 GPU 切片内来完全避免归约	不适用	不适用	绕开问题而非解决问题

这张表明确化了第 4 节消融阶梯已经实证的结论：SwiftQK 并不是单独引入了一个全新的想法——标量统计量聚合已经存在（MiniMax-eager），通信/计算重叠也已经存在（FLUX/FlashOverlap）——它的贡献在于认识到 RMSNorm 的特定结构（可加性分解的全局统计量，加上一个独立的权重乘法项）让这两个已有想法能够以一种任何单一想法都无法单独实现的方式，组合进一个单一、避免死锁、持久化的 kernel 中，而这种特定组合之所以必要，是因为 RMSNorm 对于通用 GEMM 式重叠来说太轻、而对于不带重叠的标量聚合来说又太对延迟敏感。

5. 设计选择讨论：为什么这样选、替代方案是什么、可能在哪里失效

为什么用标量聚合而不是更聪明的部分 All-Gather（比如分层或树形归约）？ 树形或分层的 All-Gather 原则上可以在不缩减原始数据量的前提下，通过更高效地利用 NVLink 拓扑来降低交换完整向量的有效通信成本。SwiftQK 的方法在数据量这个维度上严格更优——无论拓扑如何，它都把需要移动的数据从 $O(H)$ 个元素缩减到 $O(1)$——但代价是这一切之所以可行，仅仅是因为 RMSNorm 的全局统计量可以归约成单一的平方和；如果换成一种全局统计量需要不止一个标量的归一化方案（比如同时需要均值和方差，或者需要跨 token 而非仅仅逐 token 的统计量），这个技巧就不会那么干净地推广过去。这里的边界条件是架构性的：SwiftQK 的技巧特定于那些能够在各切片之间可加性分解的、基于归约的归一化统计量，对于统计量不以这种方式分解的归一化变体，需要重新思考。

为什么按 warp 划分工作，而不是让同一批 warp 在通信和计算之间时分复用？ 把 Warp 0 专门分配给通信路径、其余 warp 分配给计算，是一个在 kernel 设计时就静态决定的划分，而不是运行时决定的。一个显而易见的替代方案——让所有 warp 都在检查通信标志位和做逐元素工作之间交替——可以避免”浪费”一个 warp 的计算能力在主要是自旋等待上，但会让正确性论证复杂得多（你需要仔细推理每个 warp 究竟何时检查标志位、何时做计算，而不是一个干净的两路划分），而且这会带来一种持久化、对死锁敏感的 kernel 最难以承受的那种微妙调度 bug 风险。论文的静态划分是更保守、更容易验证的选择，代价是 Warp 0 本身损失了一些 SM issue rate 的潜力（不过论文报告的相对 MiniMax(fusion) 2.8–4.6 倍的 SM issue rate 优势表明，这个代价在实践中并不大）。边界条件是：当 $H$（一张 GPU 切片下的头维度 × 头数）相对于一个 block 内 warp 数量变得很大时，放弃一整个 warp 的计算能力去做通信，相对代价会变得更小；而对于非常小的隐藏维度切片（比如小模型上激进的 TP 度数），把整整一个 warp 专门用于本质上只是几个标量交换的工作，可能会成为相对更大的固定成本，不过论文没有报告出现这种代价占主导的场景。

为什么专门用 FP32 累加，而不是全程停留在原生（BF16/FP8）精度？ 直接用 BF16 或 FP8 累加平方和会特别严重地放大舍入误差，因为平方运算本身已经压缩了被求和数值的有效动态范围，而低精度格式的尾数位很少，难以吸收累积的舍入误差。FP32 累加是一种标准的缓解手段（在混合精度训练的累加器等场景中也很常见），在这个场景下几乎不增加额外成本，因为归约本身作用于每张 GPU 一个标量，而不是一整个张量，所以把这一个值提升到 FP32 来做累加几乎是免费的。用原生低精度累加的替代方案会严格更差，论文正确地避免了它；这一部分的贡献在数值层面并不新颖，但它是一个必要且被正确应用的工程细节，一个不够仔细的实现很容易漏掉它。

为什么论文的基线对比集只局限于 Megatron 式的层内 TP，没有讨论序列并行或专家并行的交互？ 这与其说是对某个具体选择的批评，不如说是一个范围性的观察，但值得作为设计边界指出来：SwiftQK 完全是在标准张量并行注意力/MLP 切分的语境下评测的。许多生产级推理服务栈会把 TP 与序列并行（对前向传播的部分计算额外切分序列维度）结合使用，或者对于像 OLMoE 这样的 MoE 模型，与专家并行（沿专家维度而非隐藏维度切分）结合使用。论文确实测试了 OLMoE，但仅在与稠密模型相同的 TP 切分假设下测试；它没有讨论 SwiftQK 这种逐 GPU 切片的归约模式，在 Q/K 隐藏维度的切分方式与专家路由维度的切分方式不同的场景下会如何交互——而这正是大规模 MoE 模型实际部署中的一种真实配置。

5b. 将死锁自由论证形式化

第 3 节关于死锁安全性的讨论值得完全严谨化，因为它是整个设计中最关键的正确性保障（一个偶尔会卡住的快 kernel，比一个永远不会卡住的慢 kernel 更糟糕）。

风险本身：作为一个等待图问题来建模。 把每个 CUDA block 的执行建模为等待图中的一个节点：如果 block $u$ 正在自旋等待只有 $v$ 能产生的信号，则从 $u$ 到 $v$ 有一条边。当且仅当这个等待图中存在一个环，且该环上的 block 从未被同时调度时，kernel 才会死锁。具体到 SwiftQK 的阶段 B，处理同一 token 索引 $t$ 的所有 GPU 上的对端 block 在等待图中形成一个强连通分量——每张 GPU 的 Warp 0 都在等待其他所有 GPU 对该 token 的 IPC 写入，然后才能计算 $S_\text{global}$。这是算法本身不可避免的要求（全局和确实需要每个切片的贡献）——因此预防死锁的唯一手段就是保证这个强连通分量的所有成员都同时驻留并执行，永远不要排队。

为什么限制网格大小是充分的（而不仅仅是必要的）。 GPU 硬件调度器保证，一旦一个 block 被分发到某个 SM 上，它会运行到完成而不被其他 block 抢占（忽略罕见的上下文切换场景），并且只要不被外部事件阻塞就会取得进展。如果 SwiftQK 恰好启动 $B_\text{res} \times N_\text{SM}$ 个 block——这是 occupancy calculator 根据编译后 kernel 的寄存器/共享内存足迹确定的、能同时驻留的最大数量——那么构造上所有被启动的 block 都会在 kernel 启动时即被分发到某个 SM（没有任何 block 留在启动队列里等待空位）。由于每张对端 GPU 上对应的 block 都以相同的上限启动，且所有 GPU 同时执行相同的启动（这隐含在论文将单一逻辑 kernel 调用跨所有 TP rank 的结构中），因此一个 token 的跨 GPU 强连通分量的每一个成员都保证从时间零开始就驻留——封死了通向死锁的唯一途径。这就是为什么上限必须是最多 $B_\text{res} \times N_\text{SM}$（少一些也安全，只是并行度低一些；多一些就会重新引入启动队列和风险）。

论文表述中一个被忽略的细节：持久化步进循环的块内顺序。 因为每个驻留 block 会依次处理多个 token（持久化步进循环），不同 block 处理完分配到的 token 的速度可能不一致（由于内存访问延迟等因素），值得确认一下：一个处理 token $t+1$ 的 block 不会需要等待一个仍卡在处理 token $t$ 的阶段 B 的对端 block。由于每个 token 的阶段 B 同步只涉及分配给同一 token 索引的对端 block（而非所有驻留 block 之间的全局屏障），一个提前完成 token $t$ 的 block 可以直接进入 token $t+1$ 的阶段 A，无需等待处理其他 token 的无关 block。这在算法结构中是隐含的，但值得明确指出，因为一个不小心引入了跨 token 全局屏障的持久化 kernel 实现，即使没有技术上死锁，也会重新引入不必要的串行化。

5c. 超出论文评测范围的实际部署考量

一个评估 SwiftQK 生产部署可行性的系统工程师合理会提出的几个运维问题，超出了第 4 节基准测试套件的覆盖范围：

与动态批处理和可变序列长度的交互。 vLLM 式连续批处理意味着每一轮调度步骤中活跃 token 集合都会变化，且批次中不同请求可能处于 prefill 或 decode 的不同阶段。SwiftQK 的持久化 kernel 需要正确处理每次调度步骤 token 数量的变化，因为 $B_\text{res} \times N_\text{SM}$ 个固定 block 组成的网格必须既能优雅处理大量 token 的 prefill 批次（每个 block 在其步进循环中处理多个），也能处理只有少量 token 的纯 decode 批次（此时部分 block 在某次迭代中可能根本没有工作可做）。论文将其集成进 vLLM 暗示了这一点已经处理得当，但四页短文格式没有空间详细描述这个调度集成逻辑。

故障和掉队处理。 论文中描述的 SpinWait 原语没有明确的超时或故障检测机制。在一个生产环境的多 GPU 服务集群中，一个对端 GPU 可能变慢（热节流、无关进程暂时占用 SM、ECC 引发的计算重试），或在最糟糕的情况下完全失效（驱动崩溃、Xid 错误）。无超时的自旋等待在故障情况下会使整个流水线阶段无限期挂起，而不是优雅降级——这是任何建立在这个原语之上的系统都会面对的真实运维风险，而这篇短文的范围并未触及，需要一个未被描述的监控层（例如能够检测并重启卡住 kernel 的看门狗机制）来处理。

跨 GPU 架构的可移植性。 报告的结果特定于通过 NVLink 互联的 RTX 3090（消费级，Ampere）和 A100（数据中心级，同样 Ampere）GPU。更新的架构（Hopper、Blackwell）有不同的 SM 资源预算、不同代的 NVLink、在某些情况下还有专门针对这种小消息对端同步的硬件级原语（比如 NVLink Switch 多播），这可能使 得 SwiftQK 自建的 IPC 缓冲区方法变得没那么必要，或者可以被利用来让它变得更快——论文没有讨论 $B_\text{res} \times N_\text{SM}$ 的 occupancy 计算和持久化 kernel 结构在更新的硬件代中是否需要重新调优（或甚至是否仍然最优）。

5d. 为什么这个问题值得一篇独立的短文？对“小开销”系统优化的一般性反思

SwiftQK 值得从一个更广的角度看待：它是一个关于“小开销变大开销”现象的典型案例。在单 GPU 训练/推理中，RMSNorm 的计算量小到可以忽略不计。但一旦引入分布式执行、尤其是引入需要跨设备归约的操作，原本微不足道的计算就可能因为同步开销而变得不成比例地昂贵。这种现象并不限于 QK-Norm：任何需要在分布式执行中计算一个全局统计量的操作——比如混合专家模型中的负载均衡统计、长序列推理中的全局 KV 缓存大小统计、或者分布式训练中的梯度范数裁剪——都可能面临同样的模式：计算微不足道，同步开销却随并行度增长而增长。SwiftQK 的两步方法论（先将充分统计量缩减到最小，再将剩余延迟与独立计算重叠）作为一种设计模式，对这一整类问题都有参考价值，而不仅仅局限于 RMSNorm 本身。这也是这篇仅四页正文的 IEEE CAL 短文值得词头词尾细读的原因：它解决的虽然是一个看似小众的问题，但解决思路具有较强的可迁移性。

5e. 与其他归一化方案的广义对比：LayerNorm 与 GroupNorm 会遇到同样的问题吗？

值得单列一个问题：SwiftQK 的技巧能否推广到其他常见的归一化方案上？答案取决于该归一化方案的全局统计量是否仍然可加性分解。

LayerNorm（层归一化）： LayerNorm 需要同时计算均值 $\mu = \frac{1}{H}\sum_j x_j$ 和方差 $\sigma^2 = \frac{1}{H}\sum_j (x_j - \mu)^2$，两者都需要对整个隐藏维度求和。但关键在于，方差可以重写为 $\sigma^2 = \frac{1}{H}\sum_j x_j^2 - \mu^2$，也就是说只需要两个可加性分解的量：$\sum_j x_j$ 和 $\sum_j x_j^2$。这意味着 SwiftQK 的阶段 A 思路可以直接推广到 LayerNorm：每张 GPU 本地计算两个标量而非一个（本地和与本地平方和），写入 IPC 缓冲区，对端聚合后得到全局均值和全局方差，同样可以在阶段 B 中与权重乘法重叠。通信量从 $O(1)$ 变成 $O(2)$，但仍然远远小于 $O(H)$。

GroupNorm / 分组归一化： 如果分组方式恰好与 TP 的切分轴对齐（每个分组完整地落在一张 GPU 上），则完全不需要跨 GPU 通信，这正是论文对逐头 QK-Norm 的概念性描述（每个头自成一组）。但如果分组轴与切分轴不对齐（例如一个分组跨越了两张 GPU 的切分边界），则仍然需要跨 GPU 归约，SwiftQK 的技巧同样适用。

RMSNorm 以外需要非可加性统计量的方案： 相比之下，若一种归一化方案需要的全局统计量本身不可加性分解（比如需要对每个元素取绝对值后取中位数而非均值的“平伸归一化”，或者需要跨 token 而非仅逐 token 的统计量，如 BatchNorm），那么 SwiftQK 的整套思路就无法直接搬过去——中位数无法从各分片的局部中位数重建，必须重新交换更多信息（甚至可能需要部分或完整的排序信息）。这是本文第 5 节已经指出的边界条件的一个具体例子：SwiftQK 的适用范围精确地匹配于“全局统计量可加性分解”这一数学性质，而不是任何归一化操作的通用解。

6. 局限性：论文自己承认的部分（以及一些没有明说的部分）

论文对自身范围的局限性相当坦诚。它明确指出 Transformer 块内 QK-Norm 的放置方式和定义本身仍是一个正在演进的设计空间（引用了对替代归一化策略的持续探索），并明确指出其贡献只能推广到”跨越 TP 切分的激活并需要跨 GPU 同步”的归一化方案——也就是说，这个技巧不会自动迁移到未来某个模型可能采用的每一种归一化变体上。论文的评测也局限于特定的硬件基础（通过 NVLink 互联的 RTX 3090 和 A100），以及三个特定的 OLMo 系列模型，没有测试使用 QK-Norm 的其他模型家族（比如已经采用类似归一化的 Gemma 或 Qwen 变体），这就留下了一个悬而未决的问题：观察到的开销百分比和加速比是否能推广到具有不同头维度/总隐藏维度比例的架构上。

6b. 一个容易被忽略的细节：为什么这篇短文选择发在 IEEE CAL 而非会议长文

SwiftQK 发表在 IEEE Computer Architecture Letters（CAL）而非 MLSys、OSDI 或 ASPLOS 这类会议长文，这个选择本身就值得片刻注意。CAL 作为一个专门发表短、快的体系结构类创新的期刊，适合容纳这种“单一、集中、技术深度高但评价范围窄”的工作：一个具体的单 kernel 优化，专注于一个特定的性能问题，而非一个完整的系统设计。相应地，这篇短文也就自然地不会包含会议长文常见的内容——例如对相关工作的全面文献综述、与多种归一化方案的系统性对比，或者跨多种集群拓扑结构的可扩展性实验。对读者而言，这意味着本篇笔记中指出的一些“缺口”（比如缺乏跨互联类型的评测、缺乏误差范围统计），与其说是论文的疏忽，不如说更多是格式约束下的取舍——四页正文的篇幅注定了论文必须在“讲清楚一个想法”和“覆盖所有边缘情况”之间做出选择，而作者选择了前者。这也是为什么本篇笔记第 5b、5c、5e 节要花额外篇幅去补充论文本身没有空间展开的形式化论证、部署考量和横向对比：论文的短小精悍是它的优点，但也意味着读者需要自己补上这些工程实践中真正会遇到的问题。

这篇论文特有的弱点和缺陷。 首先，评测完全局限于 NVLink 级别的互联，点对点带宽和延迟都很理想；论文没有给出阶段 B 中持久化自旋等待模式在纯 PCIe 多 GPU 配置（在更廉价或更旧的服务部署中很常见）下，或者在跨节点配置（GPU 间通信必须经过网络结构而非 NVLink）下会表现如何——在 NVLink 上很廉价的自旋等待策略（微秒级往返）在更高延迟的链路上可能变成真正浪费性的忙等待，而论文对这种情况没有提供任何回退方案或讨论。第二，死锁安全性论证虽然逻辑上成立，但完全依赖于 occupancy 计算（$B_\text{res} \times N_\text{SM}$）被正确计算且在运行时保持有效；论文没有讨论如果 kernel 与同一 GPU 上的其他并发 kernel 一起启动会发生什么（在真实的多租户服务系统中这是常见场景，某个模型的 QK-Norm 计算可能与来自另一个请求或模型的无关工作共享 SM），这可能会改变实际同时驻留的 block 数量与静态 occupancy 计算所预测的不同，从而悄悄重新引入论文声称已经消除的死锁风险。第三，消融对比集（All-Gather、Comm-Overlap、MiniMax(eager)、MiniMax(fusion)、SwiftQK）是一个合理的阶梯，但论文没有报告任何延迟或吞吐量数字在多次重复运行之间的方差或置信区间——对于一篇核心主张建立在共享硬件上测得的 10-30% 区间百分比差异之上的论文而言（GPU 存在因热节流、驱动层调度抖动等因素造成的逐次运行方差），缺乏误差范围或重复试验统计是一个真实的严谨性缺口。

作者可能低估或省略的局限性。 论文把逐头 QK-Norm 基本描述为无需通信（“如果某个头恰好完整地落在一张 GPU 的切片内就不需要任何通信”是本文作者从论文自身表述中的推断，而非论文的明确断言，但论文确实用逐头形式与逐层形式作对比，来说明后者是更难处理的情形）——但论文从未真正在一个使用逐头 QK-Norm 的模型上做基准测试来实证验证这一基线假设；第 2 节/图 1 中整个动机性对比都是围绕清一色使用逐层 QK-Norm 的模型构建的，因此读者不得不凭信任接受”逐头基本免费”这一表述，而无法看到它在实际使用逐头 QK-Norm 变体的模型（比如 ViT-22B 或任何采用逐头变体的 LLM）上得到验证。另外，SwiftQK 的持久化 kernel 设计必然要求在每张 GPU 上、在每一次使用 QK-Norm 的层的整个前向传播期间，都保持一组固定、有界的 CUDA block 常驻——这是一项真实的占用税（GPU 在持久化 kernel 驻留期间无法把这些 SM 槽位用于其他任何事情），论文没有对此进行量化；对于一个拥有数十个 QK-Norm 层的模型而言，一个必须保持驻留以服务 RMSNorm 调用的持久化 kernel，可能会实质性地限制周围那些 GEMM 密集型注意力和 MLP 计算可用的 SM 容量，而这项成本在评测中的任何地方都没有被报告。

具体的、可操作的改进建议。 (1) 报告尾延迟百分位数（p99 TPOT），而不仅仅是平均值——一个带有忙等待成分的、受通信约束的操作正是那种在竞争压力下（比如某个对端 GPU 因热节流或被无关 kernel 占用 SM 而暂时变慢）尾部行为可能与均值大幅偏离的代码路径，而论文当前的 TPOT 数字对此毫无可见性。(2) 为非 NVLink 互联（PCIe、跨节点 RDMA）添加明确的回退模式，或者至少给出讨论，因为自旋等待设计的成本收益权衡高度依赖于点对点延迟，而许多真实部署（尤其是成本敏感的推理服务集群）并非每一对 GPU 都有 NVLink 可用。(3) 直接量化持久化 kernel 的占用税——针对一个有代表性的模型，报告整个前向传播过程中，有多大比例的 SM-秒被 SwiftQK 的常驻 block 消耗、又有多少可供周围的 GEMM 计算使用，这样评估这一权衡的系统工程师就能定量而非定性地进行推理。(4) 把评测扩展到至少一个使用逐头 QK-Norm 的模型上，以实证验证（而非假设）逐头变体在 TP 下确实无需通信，这将强化论文把逐层 QK-Norm 定性为唯一有问题情形的论述。

6c. 将论文数字与现有开源实现对照阅读的建议

根据论文引用的 rms_norm_tp.py（提交哈希 99a8561）是 MiniMax(fusion) 基线的具体实现，对于想要深入理解的读者而言，把这个已有实现与论文描述的算法 1 并排阅读，是一个非常有效的理解方式。具体而言，可以重点关注三个地方：第一，已有实现的标量聚合部分是否也使用了 IPC 缓冲区还是依赖于更重的集合通信库进行标量交换（前者开销更小但需要自己实现对端刷新逻辑，后者更简单但开销更大）；第二，它是否真的实现了阶段 A 和阶段 C 的融合，还是仍然分成两个独立的 kernel 启动（后者就会引入额外的 kernel 启动开销，这正是 SwiftQK 想要消除的目标之一）；第三，它是否具备 SwiftQK 阶段 B 那种 warp 级别的通信-计算分工，还是把所有 warp 都用于顺序执行同步和归一化。图 2(a) 中 SwiftQK 相对 MiniMax(fusion) 额外 29.4-77.0% 的延迟降低，实际上就是这个第三点差异的直接体现。

7b. 从部署成本角度重新看待 29.5% 的 TPOT 降低

数字本身很容易变得抽象，不妨把 29.5% 的平均 TPOT 降低放到一个具体的运维情境中。假设一个推理服务集群在 4 卡 A100 上以 TP 方式服务 OLMo2-13B，在未采用 SwiftQK 之前，饱和请求吞吐量为某个基准值 $T_0$。根据图 2(c)，SwiftQK 将平均 TPOT 降低 29.5%、将饱和吞吐量提高 25.4%。对于一个每天 24 小时持续服务、以固定 GPU 集群规模运行的推理集群而言，25.4% 的吞吐量提升大致等价于在不新增任何 GPU 的情况下，把可服务的有效请求容量提升近四分之一——换句话说，对于一个需要 4 个 4 卡节点才能满足高峰请求量的部署，同样的负载在采用 SwiftQK 后可能只需要约 3.2 个节点就能达到。对于一个长期运行、A100 节点按小时计费的云上推理服务而言，这相当于在不降低 SLO 的前提下直接节省了一个接近 20% 的 GPU 小时开销，而这一切仅仅来自对一个单一归一化 kernel 的重写，无需修改模型权重、训练流程或服务调度逻辑。这个对比也解释了为什么这类看似微小的 kernel 级优化在大规模部署中往往比看起来更“宏观”的系统重设计更具性价比——它不需要重新验证模型效果，不需要重新训练，也不需要开发新的调度策略，只需要替换一个 kernel 实现。

7c. 一个完整的回顾：从一个开源模型的视角重读图 1(c)

回到图 1(c) 的数据，把它从一个开源模型开发者的角度重新读一遍也很有启发性。OLMo 系列是一个很好的对照组，因为同一家开发团队在短时间内迭代了三代模型，其中两代引入了 QK-Norm。这意味着开发者在选择采用 QK-Norm 时，很可能并没有充分评估它在 TP 服务下的代价——QK-Norm 的采用决策很可能是基于训练稳定性的考量作出的（这正是 QK-Norm 最初被引入的动机，见前述的 ViT-22B），而非推理服务成本。这个观察本身就指向了一个更深层的系统设计教训：训练时做出的架构决定（为了提升训练稳定性或模型质量）很少会在设计阶段就同时考虑到推理时的分布式部署成本，而训练和推理在并行化方式上的不对称性（训练常用数据并行+流水线并行，推理常用张量并行以降低延迟）意味着训练阶段看不到、也不会监测到的开销，可能在推理阶段突然变得非常突出。这也是为什么这类跨训练-推理阶段的性能特征化工作（就像论文第 2 节做的那样）具有独立于具体优化方案的价值——它们能帮助模型设计者提前意识到自己的架构选择将在下游部署中带来什么样的隐藏成本。

8. 可复现性说明

论文对其基线的引用具体到足以在精神上独立复现：All-Gather-based QK-Norm 被描述为标准 TP 实现；MiniMax(fusion) 被明确引用为一个已有的、公开可用的 vLLM kernel（rms_norm_tp.py，在 vLLM GitHub 仓库中以提交哈希 99a8561 引用），这为任何想复现对比基线的人提供了一个具体、可核查的工件。评测所用的模型（OLMoE-7B、OLMo2-13B、OLMo3-32B）都是艾伦人工智能研究所 OLMo 系列公开发布的模型检查点，请求负载（ShareGPT）是一个标准的公开可用数据集。这篇短文（可以理解，受限于 IEEE CAL 四页正文的格式约束）没有完全说明的是：报告数字所用的精确 CUDA/驱动/vLLM 版本锁定，每种 GPU 架构所用的具体 occupancy calculator 输出（$B_\text{res}$ 值），或者用于产出图 2(a) token 数扫描的具体微基准测试框架——想要精确复现论文报告的百分比改进的读者，需要要么联系作者获取这些细节，要么根据论文引用的 vLLM 集成自行重建。

8b. 术语快查表

• TP（张量并行）： 在单层内部将计算切分到多张 GPU，与 PP 相对。
• PP（流水线并行）： 按层将模型切分到多张 GPU，不同层在不同设备上依次执行。
• QK-Norm： 在注意力计算之前对投影后的 Query 和/或 Key 张量应用 RMSNorm，以稳定注意力 logits。
• 逐层 QK-Norm： 在整个投影后 Q/K 维度上（所有头一起）计算单一归一化因子。
• 逐头 QK-Norm： 为每个注意力头独立计算一个归一化因子。
• All-Gather： 一种集合通信原语，通过交换每个设备的本地切片，在每个参与设备上重建完整张量。
• IPC 缓冲区： 对端 GPU 可以直接访问的一块内存区域（进程间通信，通过 NVLink），无需经过 host 中转的数据传输。
• TPOT（每输出 token 时间）： 解码过程中相邻生成 token 之间的平均延迟，一个标准的 LLM 服务延迟指标。
• RPS（每秒请求数）： 用于在饱和点以下扫描服务负载的投放请求速率。
• $B_\text{res}$： 具有 SwiftQK 特定寄存器/共享内存足迹的 CUDA block 能在一个 SM 上同时驻留的最大数量。
• $N_\text{SM}$： 目标 GPU 上的流式多处理器（SM）数量。

8c. 一个方法论上的总结：什么时候值得为一个小组件写一个专用 kernel

回顾整篇论文，它实际上给出了一个可以提炼成通用气含训则的诊断流程：首先，确认该组件的全局统计量是否真的需要跨分布式分片的完整数据（而非仅仅需要一个可加性分解的标量摘要）；其次，测量该组件的本地计算量相对其需要的跨设备通信量是否真的很小（如果本地计算量已经足够大，通用的 GEMM+通信重叠技术就已经足够，根本不需要专门写一个融合 kernel）；最后，确认该组件在目标并行度下的开销占比确实随并行度增长而增长（否则专门优化它的性价比有限）。QK-Norm 同时满足这三个条件：全局统计量可加性分解为单一平方和，本地计算相对 All-Gather 通信量微不足道，开销随 TP 度数增长而增长——这三个条件的同时成立，正是这类专用 kernel 优化能够带来巨大收益、且值得专门投入工程资源去实现的信号。

8d. 与预先训练好的模型集成 SwiftQK 时需要注意什么

对于已经在训练时使用逐层 QK-Norm 的现有开源检查点（如 OLMo2、OLMo3），部署时接入 SwiftQK 原则上不需要任何重新训练或微调——因为它计算的是与原始 All-Gather 路径数学上等价的同一个归一化函数，只是换了一种计算它的方式。这是这类系统级 kernel 优化的一个重要优势：它们可以作为推理引擎层面的替换部件接入，而不需要触及模型训练流水线。但需要注意的是，图 2(b) 的数值精度对比只针对 BF16 和 FP8-E4M3 两种激活格式做了验证；如果一个推理服务栈使用其他低精度格式（例如 INT8 量化激活），需要重新验证 FP32 累加路径是否仍然能把误差控制在可接受范围内，而不能直接假设论文报告的误差数字会推广到未测试过的数据类型上。

8e. 小结：一张图看清 SwiftQK 的设计逻辑链条

在进入最终的实践意义讨论之前，把前面所有讨论串起来回顾一遍会很有帮助。这篇论文的整个逻辑链条可以概括为五步：第一步，观察到现代 LLM 普遍采用逐层 QK-Norm，而它的全局归一化因子需要整个 Q/K 向量（式 1）。第二步，指出 TP 下这个向量被切分到多张 GPU 上，因此必须引入某种跨 GPU 同步，而标准做法（All-Gather）需要搬运整个向量，且现有的通信-计算重叠技术无法隐藏这个开销（因为 RMSNorm 自身的计算量太小）。第三步，利用 RMSNorm 分母只需要一个可加性分解的标量（平方和）这一数学性质，把通信量从 $O(H)$ 压缩到 $O(1)$。第四步，利用 RMSNorm 可以分解为归一化尺度项（需要全局归约）和权重乘法项（不需要）这一事实，将剩余的同步延迟与独立计算在 warp 级别重叠。第五步，为了让这一切在实际硬件上安全可靠地运行，引入基于 occupancy 计算的有界持久化网格，确保每个需要相互等待的 block 都从一开始就同时驻留，从而彻底排除死锁风险。每一步都直接回应了前一步暴露出的具体问题，这种环环相扣的设计逻辑链条正是一篇优秀系统论文应有的样子，也是本篇笔记花了大篇幅逐阶段展开的原因。

8f. 对 vLLM 类推理引擎开发者的可操作清单

如果你是一个推理引擎（vLLM、SGLang 或类似系统）的开发者，想把这篇论文的思路落地，一个具体的自核清单可以是：（1）在你自己的 TP 服务栈中插入细粒度计时器，分别测量 QK-Norm 相关的 All-Gather 通信时间与它的归一化计算时间，确认是否存在图 1(c) 那种随 TP 度数增长而增长的开销占比；（2）如果确认开销显著，优先考虑集成论文引用的已有 vLLM 实现（rms_norm_tp.py，提交哈希 99a8561）作为中间步骤，它已经能提供标量聚合带来的大部分收益；（3）如果需要进一步扳耐剩余的 warp 级重叠收益，再考虑自行实现类似 SwiftQK 的持久化融合 kernel，但必须同时实现本笔记第 5b 节形式化过的基于 occupancy 的有界网格论证，以及第 5c 节讨论的失败/超时处理机制，不能直接搬运论文代码而不做任何适配。

8g. 对硬件供应商和互联标准制定者的一点延伸思考

虽然 SwiftQK 本身是一个纯软件（CUDA kernel）层面的解决方案，但它所解决的问题的根源——小数据、高频率的跨 GPU 标量同步——也值得硬件设计者思考。目前的 NVLink 互联主要为大批量数据传输（如模型权重、激活张量）优化，而 SwiftQK 需要的是频繁、小批量的标量交换及其对应的带外信号机制。若未来的 GPU 互联标准（或 NVLink 的后继代）能为这类极小负载、延迟敏感的对端同步提供硬件级原语支持（例如一个专用的、不占用常规数据通道带宽的信号总线），那么这类 kernel 级软件变通就可能不再需要，或者可以在那个硬件原语之上获得进一步的性能提升。这也是为什么本篇笔记在第 5c 节特意指出硬件可移植性是一个未充分讨论的边界条件——这不仅仅是论文自身的局限，也是整个领域值得长期跟进的方向。

9. 如果你在构建或运维一个 TP 服务栈，这意味着什么

如果你今天正在为使用 QK-Norm 的模型运维一个张量并行的服务系统，论文第 2 节的问题刻画本身就是可以立即付诸行动的，无论你是否采用 SwiftQK 这个具体方案：值得去剖析一下，在你自己的 TP 度数下，自己系统里有多少 TP 延迟可以归因于 QK-Norm 同步而非其他一切，尤其是如果你已经观察到增加更多 TP 并行 GPU 带来的收益在递减。如果这部分开销确实显著，SwiftQK 的核心思路——把跨 GPU 交换缩减到最小的充分统计量，并在一个融合的 kernel 内把剩余的同步延迟与独立计算重叠——是一个值得应用的模式，即使超出这个具体的 RMSNorm 场景也是如此：任何在 TP 下逐 token、逐切片的基于归约的归一化或统计量计算，只要其全局统计量像平方和那样能够在各切片之间可加性分解，都是同样的”标量聚合加重叠”处理方式的候选对象。

更广义地看，这篇短文也是一个提醒：在评估一个分布式系统的性能时，不能只看单个组件在单卡上的计算复杂度，而必须同时考虑它在分布式切分下需要什么样的跨设备同步。一个单卡上微不足道的操作，在切分后可能因为引入了必要的同步而变成系统的真正瓶颈——SwiftQK 正是这个普遍现象的一个清晰、可量化的注脚，也是本篇笔记想要传达的最核心的一点启发。

对于想要开始自己实验的读者，一个合理的上手顺序可以是：先在本笔记第 4b 节那个玩具例子的基础上，写一个单 GPU 参考实现验证式 (3)-(8) 的数值一致性；再在两卡环境下用标准 NCCL All-Gather 实现一个最基础的逐层 QK-Norm TP 基线，测量它在你自己硬件上的开销占比；最后再实现 MiniMax(eager) 风格的标量聚合版本作为中间步骤。这样一步步地复现论文的消融阶梯，比直接跳到实现完整的 SwiftQK 持久化融合 kernel 更容易定位问题，也更易于在本地硬件上验证第 5b 节形式化过的死锁安全性论证是否真的成立。