2026-09-13
2026-09-13
AttnFuse 将 RoPE 作为一等注意力编译操作
AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels
AttnFuse 解决了灵活注意力中的一个特定编译器接口限制。PyTorch 的 flex_attention 允许研究人员在 Python 中描述自定义注意力模式并将其编译为融合内核,但其设计是在中心矩阵乘法之后应用修改,从而将旋转位置嵌入(RoPE)排除在外。AttnFuse 则改变了注意力语言:RoPE 和其他在乘法之前的变换成为显式且可组合的操作。
该 DSL 有意保持狭窄。研究人员将十个高级构建块组合在一起,包括 RoPE、因果性和 softmax,编译器将每个格式良好的程序转换为一个融合的 GPU 内核。其四阶段流程——fusing、tiling、lowering 和 code generation——因此针对受限的注意力中间表示而不是任意 GPU 程序。实际的研究问题是,这种限制是否足以揭示有用的变体,而不要求每位研究者都实现内核。
报告的性能是有条件的但幅度显著。在 RTX 3090 上,论文报告对 RoPE+causal 模式相比 flex_attention 有 2.10× 的加速;而在 H100 上,论文报告一个完整的 Llama-3-8B 训练步在性能上与 PyTorch 的手工调优后端相差不超过 5%。作者在系统方面的最强教训是,这些结果并不意味着应始终选择融合的 RoPE。他们的“Rotation Calculus”将预旋转视为内存带宽成本,而将内核内旋转视为计算与划块成本。推导出的 Hopper 交叉点约为 N*=5,500,与在序列长度 4K 与 8K 之间的测量一致。
边界条件会改变解释。当前 Hopper 优化覆盖因果 MHA/GQA、head dimension 64 或 128、fp16/bf16、序列长度至少为 2,048,以及仅前向执行。在 H100 上当 N≥8K 时,论文报告预旋转胜出,这使得当前的融合 RoPE 路径在计算丰富的硬件上对极长上下文不那么有利。此外,论文报道的 HuggingFace 端到端图在注册的注意力函数之前应用 RoPE,因此该实验追踪的是普通因果注意力而非融合的 RoPE。
一个可迁移的操作是将每个提议的注意力变换视为一次编译器选择实验:实现融合路径和预变换路径两者,然后在 GPU 架构、序列长度、head dimension 和精度维度上进行对比测量。要带入新项目的问题不是简单地“该操作能否被融合?”,而是“在何种硬件与形状条件下,融合能击败额外的内存流量?”。
阅读此文以研究一个具体的编译器决策:何时应将诸如 RoPE 的注意力变换融合进内核,何时应将其保留为独立操作?
Abstract; S1.SS0.SSS0.Px3.p2; S3.SS1.p3; S5.SS3; S9.SS1.p1-p2; S6.SS5.SSS0.Px1.p1
mKernel 将块级融合内核扩展到 NVLink 域和 RDMA 节点之间
mKernel: Fast Multi-GPU, Multi-Node Fused Kernels
mKernel 针对分布式 GPU 执行中的一个具体边界:论文指出现有的融合内核大多局限于单个 NVLink 域,而在独立流上重叠通信与计算仅减少了部分通信成本。其方法通过在块(tile)粒度上重叠计算、节点内 NVLink 通信与节点间 RDMA,从而将融合扩展到多节点。由此该改动具有架构性:跨节点的数据移动成为持久内核块调度的一部分,而不是独立的集合通信阶段。
每个持久内核将流式多处理器划分为计算角色和通信角色。一个在 GPU 上的控制器在运行时调整该划分,因为有用的分割随内核和输入形状而变化。数据移动是分层的:mKernel 使用 NVSwitch 进行本地规约与广播,然后在节点间网络上交换每节点的聚合结果。对于 RDMA 提交,GPU 端命令队列会通知一个直接基于 libibverbs 实现的主机代理;论文对 InfiniBand 和 AWS EFA 报告了相同的方法。该库实现了五种模式:AllGather+GEMM、GEMM+ReduceScatter、GEMM+AllReduce、Ring Attention 和 MoE Dispatch+GEMM。
论文报告在两个 16-GPU H200 集群上,对于以未融合 cuBLAS 或 FlashAttention 后接 NCCL 集合通信作为基线的情形,GEMM+AllReduce 的最高加速为 1.72×,Ring Attention 的最高加速为 1.88×。这些为作者在声明测试平台上的最大值,并不证明每个内核或拓扑都能同等受益。在作者的 ConnectX-7 实验中,GPUDirect Async 据报道相较于主机辅助的 GPU 发起通信收益很小。自适应控制器在 21 个配置上的几何平均相比静态最佳固定分区获得了 1.18× 的结果,但该自适应实验因测试平台限制仅限于单节点。
对于构建分布式内核的研究者来说,可迁移的做法是从拓扑向内设计通信调度:首先识别何者可以在 NVLink 域内规约或广播,然后仅跨节点发送必要的聚合结果,同时为通信保留显式的 SM 容量。一个有价值的后续问题是,当节点间 RDMA 争用、节点数和拓扑发生变化时,同一控制器是否仍然受益。该问题应与融合收益分开测试,在相同形状和基线下比较静态与自适应的 SM 划分;论文并未为该自适应组件提供多节点证据。
阅读此文以学习一种具体的、面向拓扑的配方:如何将计算与多节点通信融合,同时明确报告的加速值及自适应 SM 划分仅限单节点的限制。
abstract1.1; S3.SS1.SSS0.Px1.p1; S4.p1; S7.p1; S5.SS5.SSS0.Px2.p1; S5.SS5.p1
SAS 使用语言建模损失训练稀疏注意力上下文排序
SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking
后训练注意力稀疏化要求一个选择器为每个查询仅保留一小组历史 token 或块。实用约束是固定的注意力预算:每个被保留的上下文单元都必须证明其存在的价值。论文识别了可训练选择器中的不匹配:它们为上下文单元打分、应用硬 Top-K 选择,从而阻断了语言建模损失直接更新选择器。相反,许多方法对层级稠密注意力分布进行蒸馏。这可以模仿原始模型的注意力位置,但并不直接优化在稀疏预算下哪些单元对预测最有用。
SAS 改变了训练接口而不是替换 Top-K 推理。其选择器为历史块输出分数,将它们转换为归一化的 softmax 门,并将每个门的对数注入注意力 logits 中的 attention softmax 之前。语言建模损失因此可以通过普通反向传播直接更新选择器。前向传递保留连续的门值,而推理仍使用硬 Top-K 选择。一个 Triton 内核将每块门注入整合进 FlashAttention 风格的计算,以实现内存高效的长上下文训练。
论文的消融将四个实现选择置于核心位置:将门放在注意力内作为对数空间偏置;用 softmax 对门归一化;保留连续分数而非将其降为二值掩码;以及主要在被选定的稀疏范围上训练。这些选择在受控消融中进行评估,使用 AttnGate 选择器,块大小 block size 为 64、Top-K 为 32,基线为在 OpenR1-Math-220k 上训练的 Qwen3-4B。
报告的评估涵盖推理、长上下文理解和 agentic 任务。摘要称 SAS 在各种注意力预算下稳定地优于可训练稀疏注意力基线,且在严格预算下收益特别显著。由于摘要中未提供数值效应量或上下文长度拆分,该证据支持在所报告设置下关于排序方向性的结论,而非方法在每个长上下文情形下均胜出的主张。
一个有用的复现实验是:在保持骨干模型、选择器、块大小和 Top-K 预算不变的情况下,比较对密集注意力进行蒸馏的办法与 SAS 的梯度路径;然后逐一移除一个设计选择——内层与外层门、归一化与非归一化门、连续值与二值、以及密集训练与被选范围训练——并测量任务质量与选择器代价。这可区分收益来自目标对齐、门的参数化,还是训练范围。
如果您正在实现后训练的块稀疏注意力,并需要一种在固定 Top-K 预算下用语言建模目标训练选择器的方法,阅读此文。
Abstract; S1.p4; S4.SS1.SSS0.Px2
Herald 从七特征激活轨迹对有害提示进行分类
早期工作使用单一的均值差拒绝方向,并通过在层与 token 位置上消融残差流激活来测试该方向。对本简报相关的限制是,单一方向的干预本身并不记录安全信号在深度上的出现时间。Herald 针对此可观测性缺口:为每层从带标签的 prefill 最后 token 状态学习一个监督的 LDA 伤害方向,并形成跨层的投影序列。
它将该序列压缩为七个几何特征——包括斜率、曲率、单调性、起始(onset)及相关统计量——并用一个 288 参数的 MLP 进行分类。该方法在方向学习时不需要梯度,报告的方向存储为 262 KB,对于一个 32-layer、d=4096 的模型,prefill FLOPs 约为 2.6×10−6。
在 WildGuardMix 上,Herald 在约 1,000 个每类样本时接近其性能平台点,并在所有测试的骨干模型上于 100 个样本时比潜在基线高出超过 2 F1。在对 Llama-3.1-8B-Instruct 与 OLMo2-7B-Instruct 进行的 200 个提示越狱干预中,将峰值单调性层的伤害方向作零投影将拒绝率从约 94% 降至约 60%,拒绝由 Llama-Guard-3 作为仲裁判定。这为一种低成本的检测与干预配方提供了有用证据,但并不证明单层普遍解释拒绝行为。
报告的优势有界。Herald 在 ToxicChat 上落后于最佳 guard 4.3 F1,在 OpenAI Moderation 上落后 8.4 F1。WildGuardMix 的训练数据与 WildJailbreak 的测试数据共享 WildChat 作为来源语料,因此这些分布并非完全独立。由于 HPD 在 prefill 时测量,该方法可能会错过仅在延伸推理后才可解析的有害意图。
一个具体的后续是:将整个源语料留出,比较七特征轨迹与单层分数,并为生成时添加后-prefill 有害性的探测器。该实验将检验轨迹是否为可迁移信号,而不仅仅是训练分布的紧凑编码。
阅读此文以学习如何将每层隐藏态探针转为低成本的审查器,并测试其表观越狱优势是否能在源语料层面的留出与仅测 prefill 的盲点下存活。
S1.p6 / S4.SS3.SSS0.Px1 / A1.p1; S6.T2 / S6.SS2.p1; Appendix N (A14.SS0.SSS0.Px1 / A14.T17); S6.p1 / S8.SS0.SSS0.Px2.p1.1; Appendix S (A19.SS0.SSS0.Px1.p1.1); Appendix U (A21.SS0.SSS0.Px3.p1.1) · §3.1 (S3.SS1.p1–p2) and Figure 1 (S1.F1)
Internal-DW 在不改变预测展开的情况下重新加权长时滞反向路径
在长时滞自回归预测中,BPTT 将每个未来损失通过重复的自回归雅可比乘积传递。这里检验的具体限制是:被传输的梯度幅度并不等同于可靠的学习信号——在受控系统中,远端梯度可能增大而其信噪比下降,因为可预测信号与不可预测变动被一并传输。
该方法 Internal Dual-Wiener routing(Internal-DW)仅改变反向路由。它保留完整的前向展开和所有视界损失,同时对内部梯度路由按可靠性加权。在每个残差块处,有界的 Wiener 增益对恒等通路与非线性通路进行加权;这些增益由路由级梯度统计量与显式噪声模型估计。该干预因此是仅在反向时进行,而非改变前向展开。
决定性证据分两阶段给出。在已知 SNR 的对照中,Internal-DW 在恢复可预测梯度信号的保持集上降低了泛化误差并改善了预测。在四个以历史主导、弱驱动的测试平台上,相对于完整 BPTT,它将预测误差降低了 5.16%–13.76%;在所有四个上均战胜了梯度截断和雅可比正则化,并在三者上优于按验证选择的 TBPTT。报告的比较使用了共享的模型、训练视界、优化器和其他设置,采用了三条匹配的随机种子;TBPTT 的设置与 DW 采样器通过验证过程选出。拟合的训练视界分析还发现,Internal-DW 的最优视界在每条匹配种子及种子平均拟合中都不比完整 BPTT 的最短,且在所有四个数据集上具有更低的拟合最小误差。
这些结果是条件性的而非保证。当可用历史受限或所选噪声采样器未能表示主导的驱动相关变动时,当前估计器的收益会减弱或反转。供给的示例包括在 NARMA 上约 7.87% 的预测劣化,以及在一次受驱动的 Mackey–Glass 干预后 7.71% 的退化。此外,Internal-DW 最小化的是局部二次的路由风险而非观测到的展开损失,因此它并不保证会产生更好的训练轨迹。
一个可迁移的研究操作是:在应用选择性反向控制前先测试可靠性——使用保持集探针评估可用历史是否占主导,以及候选采样器是否能表示驱动相关变动,然后在匹配视界下比较完整 BPTT、Internal-DW、截断与 TBPTT。记录按视界的梯度幅度、估计的 SNR 和预测损失。关键问题是:估计器预测的路由可靠性是否与实际展开目标的改善相对应。
阅读此文以学习如何检验长时滞梯度幅度是否反映可预测信号,并识别在何种历史与采样器条件下按可靠性加权的路由停止带来好处。
Abstract; S5.SS2.SSS0.Px1.p1–p2; S5.SS2.SSS0.Px2.p1; S6.p1
本期按 2026-09-13 的候选论文事后编制,核验日期为 2026-09-16。
