不可不知小技巧｜H800 上基于 Triton-TLE 优化 FlashKDA，长文本推理性能提升超 40%

编者按线性注意力具有更低的长序列计算复杂度但理论优势并不会自动转化为 GPU 上的实际性能。即使执行同一套 KDA 算法线程是否长时间等待、状态保存在什么位置以及数据需要搬运多少次都会直接影响最终延迟。本文通过对比 FlashKDA 与 Triton-TLE展示如何让并行准备阶段减少同步空转并让顺序递推阶段避免主状态的反复搬运。在本文汇总的 12 组 H800 配置中这些优化带来了 1.380× 的几何平均加速在两个重点长文本配置上的加速约为 1.4×。在大模型推理场景里长文本越长标准全注意力需要处理的数据就越多计算量和显存占用也随之增长。Kimi Linear 采用 KDAKimi Delta Attention与全注意力层 3:1 的混合结构在模型效果与推理效率之间取得平衡。KDA 不必保存全部历史信息而是将其压缩到一份固定大小的状态中并通过更细粒度的门控决定哪些信息保留、哪些信息遗忘。它的状态更新可以拆成“按通道衰减”和“一次小规模校正”从而把多个 token 合并成一个 chunk 处理为 GPU 并行计算创造条件。FlashKDA 将这套计算拆成两个阶段K1 并行准备各个 chunk 所需的数据K2 再按时间顺序更新状态。我们用 Triton-TLE 沿用这一框架但分别优化了两条路径K1 通过减少线程等待让 GPU 计算单元更持续地工作K2 则把反复使用的主状态留在寄存器中减少它在不同存储层级之间的搬运。在本文汇总的 12 组 H800 配置中TLE 相对 FlashKDA 取得 1.380× 的几何平均加速在两个重点长文本配置上的加速约为 1.4×。01FlashKDA 是什么Kimi Delta Attention 可以看成一条带门控的递归状态链。每个 token 先用 key 从旧状态中读出预测再用 value 与预测之间的误差更新状态最后由 query 读取输出。它绕开了标准 attention 的完整T × T矩阵代价是引入了串行依赖后一个 token 的状态必须等前一个 token 计算完成。FlashKDA 的做法是固定BT16把序列切成num_chunks ceil(seq_len / 16)个 chunk。chunk 内的 16 个 token 被整体改写为向量扫描和小矩阵运算不再逐 token 推进完整状态不同 chunk 的局部算子可以并行构造只有携带状态的 chunk 间递推必须保持时间顺序。忽略部分缩放和尾块处理后完整结构可以写成BT 16num_chunks ceil_div(seq_len, BT)# K1所有 (sequence, head, chunk) 相互独立可以同时执行。parallel_for (sequence, head, chunk):token_range chunk * BT : (chunk 1) * BTq_c, k_c, g_c, beta_c slice_inputs(token_range) # [16, D]# 以下 scan 和 16×16 矩阵运算在 chunk 内并行处理 16 个 token。g_prefix prefix_scan(safe_gate(g_c))Aqk[chunk] build_causal_qk_matrix(q_c, k_c, g_prefix)Akk_inv[chunk] invert_strict_lower_k_matrix(k_c, g_prefix, beta_c)w[chunk], qg[chunk], kg[chunk], decay[chunk] build_state_terms(q_c, k_c, g_prefix, beta_c)# K2不同 (sequence, head) 状态链可以并行同一条链内的 chunk 必须串行。parallel_for (sequence, head):state initial_state_or_zero()for chunk in range(num_chunks): # sequentialtoken_range chunk * BT : (chunk 1) * BTv_c, beta_c slice_value_and_beta(token_range)v_new Akk_inv[chunk] (sigmoid(beta_c) * v_c- w[chunk] state)output[token_range] qg[chunk] state Aqk[chunk] v_newstate state * decay[chunk] transpose(kg[chunk]) v_new因此K1 的启动网格覆盖所有(序列, head, chunk)用充足的 chunk 数量提供并行度K2 的启动网格只覆盖(序列, head)每个 CTA 持有一条状态链并依次消费 K1 生成的工作区。这里的分块并未消除状态依赖而是把依赖从逐 token 降低为每 16 个 token 一次并把依赖之外的计算提前并行完成。代码用H表示 Q/K head 数、HV表示 value head 数。这里讨论两种后端共同支持的HV H路径所以下文统一写作H。Chunk KDA 两阶段结构02原始设计FlashKDA 的两阶段拆分与优化FlashKDA 首先解决了两类工作不适合共享同一资源配置的问题。早期原型把 chunk 内计算和状态递推融合在一个内核中前者的高并行度被后者的串行依赖限制导致大量 streaming multiprocessor SM缺少可调度的工作。拆成 K1/K2 后两个阶段可以分别选择启动网格、线程数和数据驻留策略根据 FlashKDA v1 深入解析https://github.com/MoonshotAI/FlashKDA/blob/d2ff19a/docs/20260420-flashkda-v1-deep-dive.md这次拆分带来了至少 15% 的端到端收益。两个阶段都以 BT16 为基本计算粒度。这个尺寸既控制了安全门控变换safe gate在 BF16 下的数值范围也使 chunk 内的L成为16 × 16严格下三角矩阵。由于L^160单位下三角矩阵(IL)的逆可通过有限 Neumann 展开计算同时16 × 16的矩阵尺寸也与 Tensor Core 计算块自然对齐。K1以 chunk 为单位并行准备K1 面向大量相互独立的 chunk优化目标是降低单个 Cooperative Thread ArrayCTA 的资源占用让更多 CTA 并发执行。每个 CTA 负责一个(序列, head, chunk)用 Tensor Memory AcceleratorTMA 一次性将 q、k、g、beta 和门控偏置装入 shared memory再依次完成归一化、门控前缀和、衰减变换、小矩阵乘和 Neumann 展开。为贴近源码下面沿用k_decayed/q_decayed/k_restored/g_total/INV/Mqk这些工作区名称其中INV和Mqk分别对应前文简式中的Akk_inv和Aqk其余张量共同编码 K2 读取旧状态和更新状态所需的变换。省略具体布局和尾块处理后执行流程如下// 一个 8-warp CTA 处理一个 chunk。flashkda_k1(chunk):TMA.load(q, k, g, beta, dt_bias - shared_memory)cta_barrier()parallel_all_warps: normalize_qk_in_place()cta_barrier()parallel_threads: g_prefix, g_total safe_gate_and_prefix_sum(g, dt_bias)cta_barrier()qk_gprefix_regs load_from(smem.q, smem.k, smem.g_prefix)cta_barrier()// 两组数据生命周期不重叠复用同一片 shared memory。write_qk_variants_to(reused_smem, qk_gprefix_regs, g_total)cta_barrier()warp_0: L mma(k_decayed, k_inv)warp_1: Mqk mma(q_decayed, k_inv)cta_barrier()apply_causal_masks_and_beta(L, Mqk, beta)cta_barrier()warp_0: INV neumann_inverse(L)cta_barrier()TMA.store(k_decayed, q_decayed, k_restored,g_total, INV, Mqk - workspace)这里的 union 是资源控制的关键q/k/g 完成衰减变换后不再使用原来的存储空间随即交给k_decayed/q_decayed/k_inv/L/INV/Mqk等中间量仅这一处生命周期复用就节省约 14 KB shared memory。K1 固定使用包含 8 个 warp 的 CTA配合__launch_bounds__(256, 8)寄存器用量被压到每线程 32 个。使用 FP16 执行 Neumann 展开求逆、以 2 为底的ex2.approx和基于tanh.approx的 sigmoid则进一步减少了 chunk 内小型计算的开销。K2用流水线顺序推进状态K2 的约束完全不同每个 CTA 对应一个(序列, head)必须沿 chunk 顺序推进同一份状态。FlashKDA 将 CTA 划分为 1 个加载 warp、4 个 MMA warp 和 1 个写回 warp并设置 3 级输入流水线和 2 级输出流水线。加载 warp 用 TMA 将 K1 工作区和当前 value 搬入 shared memoryMMA warp 等待当前输入槽位就绪完成递推后提交输出槽位写回 warp 再用 TMA 将结果写回全局内存。省略布局和边界处理后三个角色的数据流如下// 一个 6-warp CTA 处理一条状态链。shared BF16 state[D, D]input_pipe pipeline(stages3)output_pipe pipeline(stages2)state load_initial_state_or_zero().to(BF16)load_warp:for chunk in chunks:slot input_pipe.acquire(chunk)TMA.load(v, beta, k_decayed, q_decayed, k_restored,g_total, INV, Mqk - slot)input_pipe.commit(chunk)mma_warps:for chunk in chunks:tile input_pipe.wait(chunk)out_slot output_pipe.acquire(chunk)prediction k_decayed stateU INV (sigmoid(beta) * (v - prediction))output q_decayed state Mqk Ustate BF16(state * g_total transpose(k_restored) U)store(out_slot, output)output_pipe.commit(chunk)input_pipe.release(chunk)store_warp:for chunk in chunks:out_slot output_pipe.wait(chunk)TMA.store(out_slot - global_output)output_pipe.release(chunk)完整状态以 BF16 保存在 shared memory 中因此 4 个 MMA warp 每处理一个 chunk 都从中读取旧状态并将更新后的状态写回原处如果接口要求 FP32 初始或最终状态则只在循环前后进行格式转换。中间量 U 则用 MOVM_T 在寄存器内转换 MMA 分片布局避免仅为布局转换而往返 shared memory。输入加载、状态递推和输出写回由两条流水线并行推进但状态本身仍严格沿 chunk 顺序更新。综上FlashKDA 形成两种差异化资源配置策略K1 依靠较小的单 CTA 资源占用扩大并发K2 则以共享内存中的状态控制递推阶段的资源占用。这套设计兼顾了占用率与不同配置下的扩展性。Triton-TLE 沿用两阶段结构但针对 H800 的目标负载在关键路径上采用了不同的资源取舍。03瓶颈剖析从 FlashKDA 到 Triton-TLE 的设计动机在 H800 的目标负载上FlashKDA 的两个阶段呈现出不同的资源特征。以定长 N1, H96, T1024 为例K1 已有接近满载的占用率却没有转化为同等水平的指令发射K2 的启动网格小于 SM 数量继续压低单 CTA 的资源占用也无法增加可并行的状态链。Triton-TLE 因而没有为两个阶段套用同一种优化策略。这两个阶段的优化思路不一样K1 依靠大量独立 chunk 提供并行度需要让已有 warp 更持续地发射有效工作Triton-TLE 用显式共享内存张量和异步加载组织跨阶段的数据复用再配合张量级矩阵运算与资源自动调优选择更紧凑的 CTA。K2 无法并行展开相邻 chunk需要把更多片上资源集中到单条状态链上这里输入加载和输出写回仍可与 MMA 重叠因此 Triton-TLE 用 TensorDescriptor、tle.gpu.copy、流水线和 warp 特化warp specialization把外围搬运组织在寄存器递推周围。需要强调的是这不是一次把整个 K2 迁移到 WGMMA 的优化。两种实现都使用 TMA在 Hopper 上Triton-TLE 只有kg^T v_new这一步状态更新使用 WGMMA。设计核心始终围绕 K1 的有效发射率和 K2 的状态驻留方式。04为什么选择 Triton-TLE将数据流与执行角色显式编码Triton 本身已经能够用tl.dot表达张量计算用triton.autotune搜索内核配置也能通过 TensorDescriptor 使用 TMA。Triton-TLE 的价值并不是提供 Triton 无法访问的硬件指令而是在这些能力之上把共享内存对象、异步搬运、缓冲槽位、执行角色及其资源分配组织成可以组合的一等编程抽象。因此这里所说的“相对于 Triton 的独有能力”指的是 Triton-TLE 补充的程序结构而不是对 TMA 或 WGMMA 等硬件能力的独占。对于 KDA 这样同时包含高并行准备阶段和串行状态链的算子这些抽象使数据放在哪里、由谁搬运、何时可被下一角色消费都能直接体现在程序结构中。本文实现用到的主要能力如下这些能力在两个阶段承担的职责并不相同。K1 借助显式共享内存张量和异步加载组织数据复用再与 Triton 的tl.dot和自动调优配合选择更适合 chunk 内计算的 CTA 配置。K2 使用tle.pipe与tle.gpu.warp_specializepipe 把w/v/qg/kg/Aqk/Akk/gk等多组输入作为一个有明确生命周期的数据槽位传递warp 专门化则把外围搬运与顺序递推分给不同角色。FP32 主状态跨 chunk 驻留在寄存器中本身是一项数据驻留决策并非单独的 Triton-TLE原语。Triton-TLE 的作用是让整个递推循环留在 MMA 消费者中同时用结构化的数据流把加载和写回围绕它组织起来。下面分别说明这些抽象如何映射到 K1 和 K2。05Triton-TLE 实现把资源用在关键路径上K1从高占用率转向高有效发射率K1 的任务是把每个 chunk 的原始q/k/g/beta转换为一组 K2 可以直接使用的小矩阵和向量。不同 chunk 之间没有状态依赖所有(序列, head, chunk)独立并行。K1chunk 内算子构造简化后的算法如下for each chunk in parallel:q, k l2_normalize(q), l2_normalize(k)g_prefix prefix_scan(safe_gate(g))Aqk causal_lower((q * exp(g_prefix)) (k * exp(-g_prefix))^T)L strict_lower((k * exp(g_prefix)) (k * exp(-g_prefix))^T)Akk_inv neumann_inverse(L * sigmoid(beta))emit w, qg, kg, g_last, Aqk, Akk_inv其中Aqk描述 chunk 内 value 对输出的直接贡献Akk_inv用来修正 valuew/qg/kg/g_last则把旧状态的读取、更新和衰减整理成 K2 需要的形式。K1 完成后K2 不再处理门控扫描或三角求逆只需执行形状规则的矩阵乘和状态递推。FlashKDA K1 的 8 个 warp 并非在每个阶段都有等量工作。归一化和衰减阶段可以利用较多线程L/Mqk的 MMA 和 Neumann 逆矩阵计算却只激活少数 warp阶段之间的全 CTA 屏障又让其余常驻 warp 等待。结果是占用率很高但每个周期的有效指令发射仍然有限。K1 不需要 K2 那样的生产者—消费者流水线。Triton-TLE 在这里的直接作用是显式分配原始 BF16q[16,128]、原始 BF16k[16,128]和 FP32 门控前缀和三个共享内存张量并用异步加载填充它们供后续阶段重复读取。归一化系数保存在寄存器中Aqk/Akk和逆矩阵计算则用张量级tl.dot表达。省略边界处理和具体布局后实现思路可以写成下面的伪代码它只描述数据驻留和计算组织不对应完整的 API 签名triton.autotune(configsk1_resource_configs, ...)def k1(...):q_buf tle.gpu.alloc([BT, K], dtypebf16, scopetle.gpu.smem)k_buf tle.gpu.alloc([BT, K], dtypebf16, scopetle.gpu.smem)gc_buf tle.gpu.alloc([BT, K], dtypefp32, scopetle.gpu.smem)q_ptr tle.gpu.local_ptr(q_buf, tile_indices)k_ptr tle.gpu.local_ptr(k_buf, tile_indices)gc_ptr tle.gpu.local_ptr(gc_buf, tile_indices)q tle.load(q_block, is_asyncTrue)k tle.load(k_block, is_asyncTrue)g tle.load(g_block, is_asyncTrue)tl.store(q_ptr, q)tl.store(k_ptr, k)tl.store(gc_ptr, prefix_scan(safe_gate(g)))q_rstd, k_rstd l2_rstd(q), l2_rstd(k)Aqk tl.dot(q_with_decay, transposed_k_with_decay)Akk tl.dot(k_with_decay, transposed_k_with_decay)Akk_inv neumann_inverse_with_tl_dot(Akk)# 再次读取共享内存中的 q/k/g构造 K2 所需的工作区。emit_workspace(q_ptr, k_ptr, gc_ptr, q_rstd, k_rstd)store_outputs(Aqk, Akk_inv)这里tle.gpu.alloc/local_ptr和异步tle.load是 Triton-TLE 提供的数据驻留与搬运原语tl.dot和自动调优来自 Triton。二者结合后K1 不必沿用 FlashKDA 固定的 8-warp 阶段划分而可以在包含 2/4/8 个 warp 的 CTA 配置间选择这里讨论的负载最终采用包含 4 个 warp 的 CTA。这种选择并没有减少指令数代表配置中Triton-TLE K1 执行38.51M条指令高于 FlashKDA 的29.99M。它的收益来自更少的全 CTA 同步和角色空转使issue active从39.45%提升到63.69%barrier stall/issue从20.21降到0.89。因此包含 4 个 warp 的 CTA 虽然使用更多寄存器、理论占用率也更低Tensor Core 工作却更连续。Aqk/Akk采用[batch, head, chunk, 16, 16]布局使 K1 写出的每个16 × 16矩阵与 K2 读入的对应矩阵都在内存中形成连续的数据块。K2让 FP32 状态留在寄存器里K2 中每个(序列, value head)对应一条状态链。输入加载和输出写回可以通过流水线重叠但 chunkc1必须等 chunkc产生新状态——循环内部的递推路径才是真正要缩短的目标。K2state recurrence 与流水线Triton-TLE K2 保留了生产者—消费者流水线角色配置为 4 个加载 warp、4 个 MMA warp 和 1 个写回 warp输入和输出各使用四级流水线。最关键的变化是完整的 FP32 主状态在整个 chunk 循环中始终留在 MMA 消费者的寄存器上下文里。这里需要区分逻辑角色与实际启动配置。加载/MMA/写回的逻辑划分是4/4/1但硬件以分区粒度分配 warp最终 CTA 包含 384 个线程即 12 个物理 warp。代码中外层num_warps4配置默认的加载分区warp_specialize的[4, 1]分别指定 MMA 和写回工作分区。实现层面Triton-TLE 不只是换一种方式调用tl.dot。TensorDescriptor 和tle.gpu.copy负责搬运数据块两个容量为 4 的tle.pipe管理输入与输出tle.gpu.warp_specialize再把加载、MMA 和写回工作分给不同 warp。省略参数和边界处理后这套生产者—消费者模型可以概括为下面的角色级伪代码load_pipe tle.pipe(..., capacity4)store_pipe tle.pipe(..., capacity4)def load_producer():for chunk in chunks:slot load_pipe.acquire(chunk)tle.gpu.copy(input_descriptors, slot)load_pipe.commit(chunk)def mma_consumer():state load_initial_state().to(tl.float32)for chunk in chunks:tile load_pipe.wait(chunk)output, state recurrence(tile, state)out_slot store_pipe.acquire(chunk)store(out_slot, output)store_pipe.commit(chunk)load_pipe.release(chunk)def store_consumer():for chunk in chunks:output store_pipe.wait(chunk)tle.gpu.copy(output, output_descriptor)store_pipe.release(chunk)tle.gpu.warp_specialize(load(load_producer, 4),mma(mma_consumer, 4),store(store_consumer, 1),)加载生产者最多提前准备 4 个 chunk写回消费者也能独立写回已完成的输出。它们可以和 MMA 发生重叠但不会改变状态递推的依赖关系只有持有 FP32 主状态的 MMA 消费者必须严格沿 chunk 顺序推进。Triton-TLE 的作用是把可并行的搬运和写回组织在这条串行路径周围。在这个执行框架内MMA 消费者的核心递推如下# b_h is the FP32 master state and lives across the whole loop.b_h load_initial_state().to(tl.float32)for chunk in chunks:w, qg, kg, Akk_inv, Aqk, g_last, v_beta load_pipe.wait(chunk)# Stage a BF16 operand in shared memory for HMMA.b_h_bf b_h.to(tl.bfloat16)# Correct the value with the old state.kh tl.dot(w, b_h_bf).to(tl.float32)v_new tl.dot(Akk_inv, (v_beta - kh).to(tl.bfloat16)).to(tl.float32)# Read the old state and add the chunk-local output.output scale * tl.dot(qg, b_h_bf)output tl.dot(Aqk, v_new.to(tl.bfloat16))# Keep the updated master state in FP32 registers.b_h b_h * exp2(g_last)[:, None]b_h tl.dot(tl.trans(kg), v_new.to(tl.bfloat16)).to(tl.float32)这份 FP32 主状态始终跨 chunk 驻留在 MMA 消费者的寄存器里。为了供 HMMA 读取每个 chunk 会生成一份128 × 128的 BF16 操作数暂存在共享内存中参与w h和qg hkg^T v_new则由 WGMMA 以 BF16 输入、FP32 累加的方式直接更新寄存器中的主状态。Triton-TLE 主要避免了将更新后的主状态反复写回共享内存再重新读出。FlashKDA 与 Triton-TLE 的数据驻留差异相应地K2 每线程使用 168 个寄存器每个 SM 最多只能驻留一个 CTA。不过在 H800 上定长N1, H64/96的主要负载中这并不会增加调度波次H800 有 132 个 SM64 或 96 个 K2 CTA 都能在一个调度波次内启动。此时更多寄存器是避免主状态反复往返共享内存的代价不能脱离上下文直接判定为性能损失。K2 的实际启动网格为N × H变长或批处理场景还要考虑序列数N。06性能结果与适用边界性能对比主要在 132 个 SM 的 NVIDIA H800 上跑输入 BF16KV128、BT16。端到端延迟由 CUDA 事件测量NCU 重放耗时replay duration只用于区分 K1/K2 的阶段贡献。完整测试可见文末源码仓库。在汇总的 12 个 H800 配置中Triton-TLE 均快于 FlashKDA几何平均加速为1.380×范围为1.272×–1.439×。当T8192时重点关注的H64/96负载分别达到1.419×和1.401×H800 上 Triton-TLE 相对 FlashKDA 的加速H96T1024的阶段数据进一步说明两个内核的收益来源并不相同在这个配置上K1 自身加速1.264×减少22.240 usK2 自身加速1.474×减少47.424 us。按两段 NCU 重放耗时的减少量计算K1 和 K2 分别贡献约32%和68%。这一比例仅用于说明 K1 和 K2 各自对阶段耗时下降的贡献不应视为对端到端加速的精确拆分。H20 的 78 个 SM 则揭示了这套资源选择的边界。对于 Triton-TLE 生产版本后端在定长N1、逐步增加 head 数的测试中H64时仍有1.136×加速到H96每个 SM 只能驻留一个 CTA 的 Triton-TLE K2 至少需要第二个调度波次加速随之降到1.035×。在汇总的 12 个配置中Triton-TLE 有 11 个领先几何平均加速为1.050×。H20 上的加速与调度边界07结语回顾一下整条优化路径FlashKDA 已经完成了最关键的一步用 BT16 在数值稳定性和计算效率之间找到了平衡点并把 token 并行的准备阶段和 head 并行的递推阶段拆开了。Triton-TLE 沿 着这套结构继续优化K1 用显式 shared memory 张量、异步加载和张量级计算撑起更紧凑的 CTA让已有的 warp 能持续发有效指令K2 则把 FP32 主状态跨 chunk 留在寄存器里用生产者-消费者流水线把加载和写回安排在串行递推的周围。这次优化最值得记住的经验是寄存器、占用率、屏障等待这些指标都不能脱离具体负载去解读。在 H800 的重点场景里拿到了约 1.4× 的端到端加速H96,T1024 的阶段数据也说明 K1 和 K2 都有贡献——K1 靠减少同步空转改善了 chunk 内计算的有效发射K2 靠寄存器驻留主状态减少了每轮递推的指令和 shared memory 往返生产者-消费者模型负责把外围搬运重叠起来。对定长 N1, H64/96K2 的资源选择没有增加调度波次但在只有 78 个 SM 的 H20 上同一选择从 H96 开始就会引入额外波次。性能判断终究要回到目标配置的关键路径结合资源配置、调度波次和消融结果来解释性能剖析数据。本次优化源码地址https://github.com/flagos-ai/FlagGems-vllm/blob/main/src/flaggems_vllm/ops/FLA/chunk_kda.pyFlashKDA v1 深度解析https://github.com/MoonshotAI/FlashKDA/blob/d2ff19a/docs/20260420-flashkda-v1-deep-dive.md