在大模型全面迈向长上下文的今天，Transformer 架构的算力瓶颈愈发凸显。在标准的 Transformer 中，MLP 的计算量随序列长度线性增长，而 Attention 机制的计算量与显存消耗却随着序列长度的增加呈 O(L²) 平方级爆炸。

当上下文从 8K 扩展至 128K 时，计算量足足暴涨了 256 倍。长文档阅读、长对话系统以及复杂的 Agent 场景，皆被死死卡在这个“算力黑洞”里。

为了打破这一魔咒，业界探索出了三条应对路线：一是以 FlashAttention 为代表的“更快地算完”，通过融合与分块把常数项压到极致；二是“算得更省”，通过量化与 KV Cache 压缩降低精度或缓存体积；三是“少算一些”，采用稀疏或线性 Attention 改变复杂度量级。然而，正如工程界的一句名言：“算法设计得再好，底层 Kernel 跑不满硬件，就拿不到任何收益。”

 

基于此，腾讯混元提出的 HiLS-Attention （层次化稀疏注意力）提供了一个极具参考价值的实践样本。围绕其高频使用的滑窗注意力前向算子，北京智源人工智能研究院高性能算子工程师深度复盘了将底层实现分析并重写为 Triton 的完整历程，在 Nvidia H20 芯片上达成性能提升。

一、长上下文时代底层开发的“三座大山”

在真正动手写出一行高性能 Kernel 代码之前，开发者往往要面临工程环境、性能测量与硬件压榨上的三重困境：

1. 环境搭建之难

本次优化的目标，是 HiLS-Attention 架构中极其高频的“滑窗注意力前向”算子。然而，其原生仓库现状令人望而生畏：6个核心 Kernel 全部由 TileLang 手写，代码量高达约 5500 行，且无任何 Triton 实现可供参考。

TileLang 虽然底层控制力强，但安装成本极高。开发者需要从源码编译整个 TVM（高达 641 个编译单元），在这个过程中，不仅遭遇了 Cython 3.3 与 Python Limited API 的严重不兼容，还面临 TVM 子模块 Git clone 仅 46KB/s 的龟速，以及代理节点频繁断连的折磨。相比之下，随 PyTorch 直接内置的 Triton 显得无比轻盈。

2. 测量失真之痛

在开展任何工作之前，都需要明确自己的目标，算子优化也不例外。在优化之前，要明确自己要优化的场景，究竟是kernel时长？还是包含host端和device端的总时长？或者是其他优化目标？

在着手优化前，最大的陷阱往往是基准测试（Benchmark）。仓库自带的 Benchmark 报告显示该算子耗时 370 μs，但当开发者使用 Nsys 探针深入追踪时，却发现纯 Kernel 的执行时间仅为 118 μs，那消失的 252 μs 去哪了？

答案是 Python 包装层的固定开销。每次调用时的 JIT 缓存重新查找、在设备上新建 kv_start 张量、LSE（对数域求和指数）的 log2 到自然对数的逐元素转换，以及推理路径的布局归一化，都在疯狂吞噬时间。在 Decode 等短序列场景，这种包装开销导致 TileLang 的 Kernel 仅占端到端时延的 15%。如果基线是错的，后续所有的优化判断都将建立在错误的沙丘上。

3. 硬件算力之痛

原生的 TileLang 实现，在特定 GPU 与形状下，仅能压榨出硬件约 59% 的峰值算力。大量的时钟周期消耗在了调度等待上，距离逼近硬件极限还有很长的路要走。

 

二、HiLS-Attention 架构揭秘与 Triton 的降维打击

要解决痛点，首先要搞清优化的对象是什么。

1. 算法层：HiLS-Attention 究竟在算什么？

腾讯混元提出的 HiLS-Attention 核心主张是：让“选哪些块（Chunk）”的过程本身可以端到端训练。它的执行分为三步：

• 局部滑窗（Local Branch）：处理邻近 token，同时输出局部的 LSE 作为后续的 logit。

• Chunk 摘要：为了选出最重要的 k 个块，需要计算每个块的 mass。但精确计算 mass 需要完整的 QK 矩阵（这恰恰是我们要避免的）。于是，HiLS 采用了一阶 Taylor 展开，在先验 query μ 处近似计算，得出了代理 key k' 与 bias b'，使得梯度可以顺利回传。

• 选择与组合：通过低秩旁路（如 7B 模型使用 rank 256 的 lmk_q_proj）产生的 Query 进行 Top-k 检索，并将选中的 chunk 分数与局部窗口的 LSE 拼在同一个层次化 Softmax 中竞争。

 

工程细节的暗坑：算法在训练时，每个 chunk 会插入一个特殊的 Landmark token 作为块摘要。这就要求底层算子在实现时，必须在局部分支里精确编写掩码逻辑（mask_lmk）将其屏蔽。

2. 工具层：为什么是 Triton？

Triton 采用块级（block-level）抽象，以 Python 风格写 kernel，编译器负责完成向量化与流水线调度。相比于 PyTorch 框架原生的算子拼接（会产生海量中间张量、访存开销无法避免），Triton 能够通过“融合（Fusion）”把多个算子合成单个 kernel，将中间结果留在寄存器/共享内存中，让显存占用从随序列长度平方增长保持为线性。

三、实战路径：从对齐、测准到压榨极限

将 HiLS 的滑窗注意力前向算子用 Triton 重写，不仅是一场性能的较量，更是一场对语义等价性与硬件特性的深度探索。这场实战经历了三个关键阶段：

阶段一：测准基线，跨越语义差异的鸿沟

首先要确保“测得公平”。我们使用 CUDA Graph 回放彻底消除 Host 端的发射开销，统一双方的 block_M=128 配置，并覆盖 11 组不同形状进行对比。跨越两种语言的语义鸿沟是最大的耗时点，必须确保严格对齐：

• LSE 对数域：TileLang 内部使用 log2，包装层需乘 ln2；而 Triton 直接输出自然对数。

• 输出布局：必须处理从 (B, Lq, Hq, Dv) 到 (B, Hq, Lq, Dv) 的转置对齐。

• 掩码等价性：expand_to_chunk 的边界向下对齐逻辑，以及 masklmk 的组合判定（(j + 1) mod C != 0），必须在两个版本中做到数学上的一致。

阶段二：初步战果，Triton 的全线碾压

在纯 Kernel 的较量下，Triton 实现了性能的提升：

• 训练与预填充：稳定领先原版 1.34x ~ 1.83x（长序列下稳定领先 30% ~ 60%）。

• 解码（Decode）场景：经过后续的 Split-KV 优化，加速比飙升至 6.48x。

• 硬件利用率：从初期的 79%，经过两轮优化后达到 84%（预填充 4K/8K 口径），而 TileLang 仅为 59%。

阶段三：打破直觉，直击 Tensor Core 发射瓶颈

Triton 已经领先，但极限在哪？NCU 工具显示计算受限达 73.6%。

最初怀疑是复杂的稀疏掩码拖慢了速度。然而，真实的消融实验推翻了直觉：去掉整个掩码运算后，性能反而变慢了 3.9%；而去掉第二个矩阵乘（PV）后，性能直接快了 37.3%。真正的瓶颈是 Tensor Core 的发射效率，两个 MMA（矩阵乘加）占据了约 45% 的时间。 这说明掩码运算完全被延迟掩盖，是非 MMA 逻辑中极其廉价的部分。

 

行之有效的调优策略：

• 分块调参（BN=32/stages=3）：将共享内存占用从 80KB 降至 48KB，使每个 SM 可并行 block 数从 2 翻倍到 4，完美掩盖了更多 MMA 的延迟。

• 两段式 Split-KV（专治 Decode）：针对 Decode 阶段 Lq=1 导致 Grid 仅有 8 个 Block（78个SM中70个空闲）的问题，将 KV 切分给多个 Block，写出部分 online-softmax 状态后再合并。

• 启发式谓词：Split-KV 是否开启，绝不能按序列长度 Lq 简单判定，必须基于 Grid 饱和度（如 cdiv(lq,64) * b * hq ≤ 32）。否则批量预填充场景会遭遇严重的性能回退。

• 分支补全：针对 D>128 与 fp32 分支重新调参，例如在 fp32 下使用更窄的 BN。

四、写在最后

面对大模型序列规模持续扩张带来的底层算力挑战，算法层面的创新若无法转化为 Kernel 级的极致性能，便难以真正落地。本次针对 HiLS-Attention 滑窗算子的 Triton 重写，在 NVIDIA H20（SM90架构）、bf16 精度、基于 CUDA Graph 计时的测试环境下，在11 组形状上全面超越 TileLang，最高实现 6.48x 加速，并将硬件利用率从 59% 提升至 84%。

这一工程实践的核心启示在于：精准的测量基准（剔除包装层假开销）是优化的前提，而硬核的消融数据打破了“掩码最耗时”的主观直觉，将真正的瓶颈（Tensor Core 发射）暴露无遗。唯有坚持数据驱动、尊重硬件规律，才能真正将优秀的算法设计转化为逼近极限的澎湃算力。

 

关于众智FlagOS社区

为解决不同 AI 芯片大规模落地应用，北京智源研究院联合众多科研机构、芯片企业、系统厂商、算法和软件相关单位等国内外机构共同发起并创立了众智 FlagOS 社区，目前已经有 100多家成员单位。

FlagOS 是一款专为异构 AI 芯片打造的开源、跨芯系统软件栈，支持 AI 模型一次开发即可无缝移植至各类硬件平台，大幅降低迁移与适配成本。它包括大型算子库、跨芯AI 编译器、并行训推框架、跨芯通信库等核心开源项目，致力于构建「模型-系统-芯片」三层贯通的开放技术生态，通过"一次开发跨芯迁移"释放硬件计算潜力，打破不同芯片软件栈之间生态隔离。

官网：https://flagos.io

GitHub 项目地址：https://github.com/flagos-ai

GitCode 项目地址：https://gitcode.com/flagos-ai

SkillHub: https://skillhub.flagos.io