基于Triton-TLE优化Tencent HPC-Ops Attention：H100上 Dynamic Decode平均达到 CUDA 的 1.05×、Prefill 达 97.3%

FlagOS智算系统软件栈

32人浏览 · 2026-11-11 18:02:20

FlagOS智算系统软件栈 · 2026-11-11 18:02:20 发布

腾讯混元 HPC-Ops 是覆盖 Attention、Fused MoE 等关键路径的高性能推理算子库。本文基于其 Attention 算子，介绍众智 FlagOS 团队使用 Triton-TLE 在 Hopper 架构上的重构实践：在 H100 上 Dynamic Decode 平均达到官方 CUDA 实现的 1.05×，Prefill 达 97.3%，部分场景加速超过 1.8×。

大模型在线推理请求动态进出、序列不断变长，KV Cache 开销急剧放大，Attention 的访存与调度效率直接决定推理延迟、吞吐与 GPU 利用率。极致性能通常需深入 CUDA 控制 TMA、WGMMA 等底层细节，而 Triton 可编程性高却难以稳定表达。

针对这一矛盾，众智 FlagOS 的 Triton-TLE 提供了分层解法：Triton 表达算法，TLE 映射 Hopper 原生数据流，TLE-Raw 控制关键热路径。基于该方案重构了腾讯混元 HPC-Ops Attention 后，H100 上 Dynamic Decode 平均达到了官方 CUDA 实现的 1.05×，FP8 Prefill 达 0.973×，部分场景加速超过 1.8×。下面文章将详细讲述。

一、Hy3 HPC-Ops Attention

1.Hy3 与 HPC-Ops

真实在线服务中，请求持续进入和退出，序列长度随生成过程不断变化，长上下文显著放大 KV Cache 的容量与访问开销，如何在上述动态负载下同时保障推理延迟、系统吞吐与 GPU 资源利用率，已经成为了大模型推理系统优化的核心问题。

Tencent Hy3 正是这类复杂推理负载的代表。作为腾讯混元面向智能体执行、代码生成和长程推理等场景构建的 Mixture-of-Experts 模型，Hy3 总参数量为 295B、单 Token 激活参数量为 21B，采用 192 个专家和 Top-8 路由，支持最长 256K 上下文与 Multi-Token Prediction（MTP，多 Token 预测），已接入WorkBuddy/CodeBuddy、元宝、ima、Marvis等应用平台，API 也已在腾讯云 TokenHub 上线。

为支撑 Hy3 复杂模型结构与大规模在线服务中的高性能推理，腾讯混元 AI Infra 团队构建高性能推理算子库 HPC-Ops，覆盖 Attention、Fused MoE、Grouped GEMM 以及通信计算融合等关键路径。其核心实现面向 NVIDIA Hopper 架构，使用 CUDA 与 CuTe 显式控制 TMA、WGMMA、Shared Memory、Warpgroup 和设备端任务调度，并已接入 vLLM、SGLang 等主流推理框架

本文聚焦 HPC-Ops Attention 算子，基于 Triton-TLE（Triton Language Extensions，众智 FlagOS 基于 Triton 语言的扩展）重构 Hopper 计算、访存、调度与归约数据流，整体获得了逼近、大部分场景超过官方 CUDA 的性能：

Decode 在动态调度、FP8/BF16 及不同 MTP 场景和 9 类 Workload 上，几何平均达到官方 CUDA 实现的 1.05×，其中，BF16 Dynamic 的 MTP1/2/3 分别达到 1.088×、1.067× 和 1.101×，FP8 Dynamic 的 MTP1/2/4 分别达到 1.075×、1.055× 和 0.964×；
FP8 Prefill 覆盖 NHD/HND、两种量化方式和稀疏/非稀疏执行，整体几何平均达到官方 CUDA 的 0.973×，其中，Mask=False 的完整 KV 路径达到 0.994×，具有 73.3%–74.6% Tile Skip Ratio 的有效稀疏路径达到 0.886×，并相较自身非稀疏执行获得 1.84×–2.96× 的加速。

2. Prefill 与 Decode：同一个 Attention，两类性能问题

Prefill 阶段一次处理 Prompt 中的大量 Query Token，矩阵的 M 维较大，QK 和 PV 具有较高的计算密度，核心是计算吞吐与流水线效率问题；Decode 每一步只处理一个或少量 Query Token，却需要读取此前积累的完整 KV Cache，所以，KV 访存、Kernel 固定开销和序列不均衡的问题变得更加突出。

Prefill：面向高计算密度的 Block-Sparse Attention

HPC-Ops Prefill Attention 利用 Prefill 阶段较大的 Query 规模，提高 QK 与 PV 的计算密度，从而充分发挥 Hopper Tensor Core 的吞吐能力。不同任务规模下，Kernel 对并行度与流水效率的侧重有所不同：当 Query Tile 数量较少时，需要保留更多独立 CTA 以充分利用 GPU；当 Query Tile 数量充足时，则可以通过 Persistent CTA 和 Warp Specialization，将数据搬运与 Tensor Core 计算更充分地重叠。为此，HPC-Ops 在 BF16 路径中根据 Query Tile 数量选择 Multi-stage 或 Persistent Warp-specialized Kernel；对于需要额外处理 Paged KV、量化 Scale 和 Block Mask 的 FP8 路径，则直接采用 Warp-specialized 结构。其核心执行过程可以概括为以下三个阶段：

Tile Scheduling：按 Batch、Query Head 和 Query Token 划分计算任务，并根据有效 Query Tile 数量选择 Multi-stage 或 Warp-specialized 执行路径。
Pipelined Attention Compute：使用 TMA 将 Q/K/V 异步搬入多级 Shared Memory Buffer，通过 Barrier 协调加载与计算 Warpgroup，流水执行 WGMMA QK、Causal Mask、FP32 Online Softmax 和 WGMMA PV。
Epilogue：遍历有效 KV Tile 后完成最终 Softmax 归一化与精度转换，并通过 Shared Memory 和 TMA Store 写回 Attention 结果。

对于 Paged KV Cache 和 Block-sparse 场景，计算阶段还会根据 Block Table 定位物理 K/V Page，并依据 Block Mask 仅处理有效 KV Tile，从源头减少无效的数据搬运与矩阵计算。

Decode：面向混合长度请求的动态 Split-K Attention

HPC-Ops Decode Attention 的核心目标，是解决 Continuous Batching 下不同请求 KV 长度差异造成的 GPU 负载不均。固定 Split 数会让长请求成为拖尾任务；固定 Chunk 大小并按照最长请求构造规则 Grid，又会为短请求启动大量没有有效工作的 CTA。为此，HPC-Ops Dynamic Decode Attention 将 Decode 拆分为 Assign、Split-K Compute 和 Combine 三个阶段：

Assign：将各 KV 序列切成 Tile，并按 CTA 工作量生成负载均衡的 Task Map。
Split-K Compute：Persistent CTA 按 Task Map 并行处理 KV Chunk，利用 TMA + WGMMA 完成 QK、Softmax 和 PV，并输出局部结果与 LSE。
Combine：根据各 Chunk 的 LSE，对局部输出重新归一化加权合并，得到与完整 Attention 数学等价的结果。

除上述 Dynamic 路径外，官方实现还提供 Static Split-K，即按照固定切分数量拆分 KV 序列，不运行动态 Assign，启动链路更短，但无法针对 Batch 内实时变化的长度分布重新平衡任务。HPC-Ops Decode 官方 Test 覆盖参数如下：

二、Triton-TLE Prefill

在覆盖 Q/KV 长度、NHD/HND 布局、两种量化方式以及稠密/稀疏 Mask 的 24 个测试配置上，Triton+TLE FP8 Prefill 相对 HPC-Ops CUDA 的总体几何平均为 0.973×，性能覆盖 0.803×–1.816×。其中，无稀疏 Mask 的 Dense 路径已整体逼近 CUDA，几何平均达到 0.994×；quant_type=0 在全部布局与 Mask 配置上的几何平均达到 1.024×，为CUDA的 0.945×–1.816×。短序列场景中，NHD、quant_type=0 的 Dense 路径达到 1.816×，HND、quant_type=1 的短序列 Mask 路径达到 1.510×。

对于具有实际稀疏性的场景，Block Mask 可跳过 73.3%–74.6% 的 KV Tile，且由于避免了无效 KV Tile 对应的数据搬运和矩阵计算，Triton+TLE 的延迟相较对应 Dense 路径获得 1.84×–2.96× 的加速，总体相对 CUDA 的几何平均为 0.886×，覆盖 0.803×–0.984×；其中 quant_type=0 达到 0.964×，覆盖 0.945×–0.984×，已较为接近 CUDA，而 quant_type=1 的几何平均为 0.814×，覆盖 0.803×–0.823×。

就优化思路方面，Triton+TLE 的 Prefill 实现围绕 Hopper 的计算与访存流水展开：一方面将 Block Mask 压缩为连续的 Active-Tile List，使无效 KV Tile 在进入 TMA 和 WGMMA 前即被过滤；另一方面通过 Warp Specialization 分离数据搬运与 Tensor Core 计算，并以 Page-aware TMA Pipeline 重叠页表解析、Paged KV 加载、QK、Online Softmax 和 PV；同时针对高层编译器难以稳定生成的稀疏压缩与 V 转置路径，进一步使用 TLE-Raw 构造可控的 Hopper 原生指令序列。

核心优化包括：

1.TLE-Raw：稀疏 Active-Tile Compaction 与原生 V 转置

TLE 能够表达 TMA、WGMMA、Shared Memory 和异步 Barrier，但对于 Warp Ballot/Popcount 驱动的动态任务压缩，以及具有特定字节排列和 Swizzle 的 FP8 V 转置，高层布局系统难以稳定生成与 CUDA 基线一致的指令序列。为此，@dialect(name="cuda") 将两个底层 CUDA Device 微内核接入 TLE，并使用 tle_raw.call_smem 直接传递 Shared-Memory MemDesc，接管编译器难以高效表达的局部关键路径。

TLE-Raw路径1——稀疏任务压缩：

Producer Warp 使用 Warp Ballot 汇总32个 Tile 的 Mask Predicate，通过 Popcount 计算有效 Tile数，再根据每个线程在有效 Predicate 中的前缀计数结果，计算其写入位置，将有效 Tile Index 连续写入 Shared Memory，生成稳定的 Active-KV-Tile List。Q TMA 在压缩开始前发起，Mask 扫描能够与 Q 搬运并行。后续 Producer 和两个 Consumer 共享同一份列表，只有有效 Tile 才会进入 Page Table 解析、TMA 加载和 WGMMA 计算。否则，高层实现通常只能让 Producer 和 Consumer 分别遍历完整 Mask，并在循环内部通过 Predicate 判断 Tile 是否有效，虽然能够保证功能正确，但会重复执行 Mask Load、索引计算和条件分支，稀疏 Tile 仍然占据循环迭代；更重要的是，编译器难以将这一过程稳定收敛为 Warp-level Stream Compaction，稀疏率无法充分转化为调度开销的下降。

TLE-Raw 路径2——PV WGMMA 的 $V^T$ 构造：

Prefill 的概率矩 P 在 Online Softmax 后已经位于寄存器中，因此 $PV$ 采用P为 Register Operand、 $V^T$ 为 Shared-Memory Operand 的 FP8 RS-WGMMA。原始 V Tile 的 Token-major Layout 无法直接满足 RS-WGMMA 对转置方向、FP8 字节顺序与 Shared-Memory Swizzle 的要求，因此需要在 QK 与 PV 之间完成一次 V Layout Transformation。

如果使用通用 BlockedLayout 描述这一转换，会由于编译器无法从中完整推导 WGMMA Fragment 所需的线程与寄存器分布，导致生成了大量细粒度 Shared Load、Register Permutation 和 Shared Store；如果使用 WGMMA-aware 静态 Layout ，所列指令总数由 232 条减少至 160 条，但仍只能引导、不能控制底层代码生成：写回侧依然退化为标量 STS.U8，无法稳定形成 Hopper 原生的 LDSM/STSM 转置通路:

v_copy_layout = gl.BlockedLayout(...)
              ↓
v_ldsm_stsm_layout = _wgmma_layout(...)
vt_smem.store(v_stage.load(v_ldsm_stsm_layout))
# 人为指定 WGMMA/MMA 的线程和寄存器分布
128 × LDS.U8            32 × LDS.U16
96 × PRMT        →      64 × PRMT
8 × STS.128             64 × STS.U8
# 触发向量化访存升级， 指令数减少

因此，在 FP8 字节排列、Shared Swizzle 和 RS-WGMMA Operand Layout 的组合约束下，仅依赖通用 load → transpose → store 难以保证与优化 CUDA 微内核一致的代码生成，从而需要通过TLE-Raw显式执行 ${LDSM}^T\rightarrow \mathrm{PRMT} \rightarrow {STSM}^N\rightarrow {Shared }V^T$

${LDSM}^T$ ：面向 Tensor Core 的转置式 Shared-to-Register 读取；

$PRMT$ ：重排相邻 FP8 字节；

${STSM}^N$ ：将结果写入 WGMMA 可直接消费的Shared Layout；

LDSM(next) → PRMT/STSM(current) 软件流水：下一子 Tile 的读取提前覆盖当前子 Tile 的字节重排与写回；

tle.gpu.wgmma ：负责发射 FP8 RS-WGMMA。

TLE-Raw 作为接管高层布局系统难以稳定优化的 V 转置微内核，其核心作用是将不可控的编译器 Layout Conversion，替换为指令序列、字节排列和内存可见性均可验证的 Hopper 原生微内核。否则，依赖通用 load → transpose → store Layout Conversion，容易生成大量细粒度 Shared Load、Permutation 和 Store，且无法稳定触发原生 LDSM/STSM 指令；两个 Consumer 还可能重复执行相同的 V 转置，进一步放大指令量、Bank Conflict 和 Barrier 等待。

2.Warp-specialized Page-aware TMA Pipeline

Paged KV Cache 的数据分散在多个物理 Page 中。每处理一个 KV Tile，Kernel 都需要先读取 Block Table，将逻辑位置转换为物理 Page ID，再搬运对应的 K、V 和量化 Scale。为了避免这些不规则访存阻塞 Tensor Core，tle.gpu.warp_specialize 将一个 CTA 划分为一个 Producer Warpgroup 和两个 Consumer Warpgroup：

Producer 负责 Active-Tile List 消费、Page Table 解析以及 Q/K/V/Scale 的 TMA 搬运；

两个 Consumer 分别处理 128-row Query Tile 的两个 64-row 分块，执行 FP8 QK、Online Softmax 和 FP8 PV。

这一角色划分将数据准备与 Tensor Core 计算分配给不同的 Warpgroup。Producer 可以在 Consumer 计算当前 Tile 时准备下一 Tile，两个Consumer 则共享 Active-Tile List、Page Metadata 和 K/V 数据，减少重复的索引解析与访存操作。

此外，流水线进一步按照数据的实际使用顺序来组织。Producer 在等待 Shared Memory Buffer 可写之前预取物理 Page ID，使 Page Table 访问与上一Tile 的计算重叠；获得 Buffer 后，依次发射 K、K Scale 和 V 的 TMA 请求。K 是 QK 阶段最先使用的数据，因此全部 K Page 会优先加载。Consumer 在 K 就绪后即可启动 SS-WGMMA，而 V 的搬运可以继续与 QK 和 Online Softmax 重叠，直到 PV 阶段才需要完成。

TLE 则负责表达整套 Hopper 异步执行协议：Tensor Descriptor 描述不同物理 Page 上的 TMA 访问，Shared-Memory MemDesc 表达两级 Ring Buffer，tle.gpu.copy 发起异步 TMA 搬运，Barrier 的 Empty/Full Phase 则管理每个 Buffer 的生产、消费和复用。K、K Scale 与 V 使用独立的逐页完成 Barrier，因此 Consumer 可以在 QK 结束后立即释放 K Buffer，而 V Buffer 保持有效至转置与 PV 完成。通过 Page Table 提前解析、K 优先发射、K/V 生命周期解耦和两级缓冲复用，Pipeline 将页表访问与大部分 KV 搬运隐藏在 WGMMA 计算窗口中，使 Paged KV Cache 的不规则访存尽量不进入 Tensor Core 的关键执行路径。

其余优化如下表所示：

三、Triton-TLE Dynamic Decode

官方 Benchmark 覆盖 9 类 Workload ：

在上述 9 Workload 中，Triton+TLE Dynamic Decode 实现在不同数据精度、KV 布局与 MTP 配置下，形成了接近甚至超过 CUDA 较稳定的性能覆盖。以 MTP=1 为例，BF16 Dynamic 在9类 Workload、NHD/HND 两种布局共18个测试上的几何平均加速比为 1.088×，性能范围为 CUDA 的 1.031×–1.225×，均超过 CUDA；FP8 Dynamic 在9类 Workload、两种布局及两种量化方式共36个测试上的几何平均加速比为 1.075×，性能范围为 CUDA 的 0.919×–1.362×，其中 Q per-token per-head、K/V per-tensor 量化路径在18个测试点上的几何平均加速比达到 1.120×，性能范围为 1.057×–1.362×。

MTP>1 场景同样保持出良好的扩展性，BF16 Dynamic MTP2/MTP3 的几何平均分别达到 1.067×/1.101×，在 9个Workload上达到 CUDA 的1.002×–1.153× 和 1.017×–1.230×；FP8 Dynamic MTP2 的几何平均达到 1.055×，覆盖CUDA实现的 0.902×–1.503×；MTP4 在更高计算密度和资源压力下仍达到 CUDA 的 0.964×，并在部分负载上达到 1.589×。

具体在 9 Workload 上的测试结果，以 BF16 NHD 与 FP8 NHD+Q Per-token Per-head，K/V Per-tensor 为例，分别如下图所示：

在优化思路上，Triton+TLE 实现围绕 Hopper 的 Tensor Core、TMA 和 Cluster DSM 重新组织 Decode 数据流。计算侧针对 Small-M 负载重构 Attention 计算方向，以 Typed RS-WGMMA、Narrow16、Padded32 等专用几何减少 Query Padding；访存侧通过 Projected TMA 统一 NHD/HND，并利用零拷贝 V 转置视图衔接 $V^TP^T$ 的 RS-WGMMA 数据通路，在避免 V 预转置和额外 Workspace的同时构建多级异步流水。调度与归约侧，Dynamic 路径通过 GPU-resident Task Map 按有效 KV 工作量进行紧凑任务调度，Static 路径则采用解析式Split-K 映射；Cluster 内部利用 DSM 优先合并 Online-Softmax Raw State，仅将无法片上闭合的结果提交给跨 Cluster Finalizer。进一步实现通过 Workload-aware Policy、Deferred Normalization、场景特化快速路径和原生 MTP 的联合计算，在负载均衡、片上通信、归约开销与 Tensor Core 利用率之间取得针对不同场景的平衡。

具体而言，核心优化实现包括：

1.Small-M WGMMA 计算布局优化

Decode 阶段的 Query 规模很小，直接套用面向大矩阵设计的 WGMMA Shape 会产生大量无效计算。以 GQA8 MTP1 为例，每个 KV Head 对应 8个 Query Head，常规QKT需要将 8 个有效 Query Row 补齐至 WGMMA 的 (M=64)，其余 56 行不仅产生无效 FLOP，还会扩大Accumulator、寄存器和 Shared Memory 占用。

因此，Triton-TLE 实现将 Attention 计算方向重构为：

$S^T=KQ^T， P^T=\operatorname{softmax}_{KV}(S^T)， O^T=V^TP^T$

经过转置后，64-token KV Tile 被映射到 WGMMA 的 (M=64)，8 个 Query Head 则映射到原生 (N=8)，从计算几何上消除了 Query 维的Padding。PV 阶段进一步使用 TLE Typed RS-WGMMA，其中逻辑转置后的 $V^T$ 作为 Register Operand， $P^T$ 作为 Shared-Memory Operand，通过 tle.gpu.wgmma 和静态 Operand Layout 构造完整的数据通路。进一步地，针对不同精度和 MTP 规模，Triton-TLE 实现进一步根据有效 Query Row 数选择 N=8、N=16 或 N=32 的合法 WGMMA 几何；当有效行数无法直接匹配硬件支持的矩阵形状时，仅补齐至相邻的最小合法尺寸，使计算规模尽可能贴近实际 Query 规模。该设计减少了无效计算和寄存器占用，并改善了 CTA Occupancy，使 Small-M Decode 也能高效利用 Hopper Tensor Core。

2.零拷贝 Paged-KV 后端

Paged KV Cache 的物理地址不连续，NHD 与 HND 又具有不同的存储顺序； $O^T=V^TP^T$ 计算还要求同一份 V Tile 以 [dim,token] 参与 $V^TP^T$ 。如果在 Global Memory 中额外维护 $V^T$ ，或者在 Shared Memory 中执行物理转置，不仅需要额外 Workspace，还会增加 HBM 往返、同步和流水停顿。

基于此，Triton-TLE 实现首先将 K/V Descriptor 的逻辑坐标重排为 [block,head,token,dim]，再带维度投影的 TMA（Projected TMA）发起 [1,1,64,128] 的异步搬运。写入 Shared Memory 时，前两个大小为 1 的维度被投影消除，直接形成了 [64,128] 的二维 K/V Tile；而NHD 与 HND 的差异则由 DescriptorStride 吸收，因此两种物理布局可以复用统一的计算主循环。

对于 $V^TP^T$ 计算，FP8 实现先为 TMA 写入 shared memory 的 [64,128] V tile 构造零拷贝转置 memdesc；随后从该视图读取 V，并在寄存器中组织为逻辑 [128,64] 的 FP8 操作数，供 RS-WGMMA 与 shared memory 中的PT 相乘，避免了额外的 shared-memory 转置缓冲区；BF16 进一步利用 TLE 编译器的布局传播和冗余搬运消除能力，能够将 local_load → transpose → local_alloc → WGMMA 数据流收敛为WGMMA的零拷贝转置视图，同一份 [64,128] V Tile 无需物理搬移，即可按逻辑 [128,64] 组织为 RS-WGMMA 的 Typed Register Operand。

@builtin
def memdesc_transpose_2d(value, _semantic=None):
  order = [1, 0]
  handle = _semantic.builder.create_memdesc_trans(
      value.handle, order
  )
  layout = value.type.layout.make_permute(order)
  # 返回共享同一底层存储的转置 MemDesc
  return buffered_tensor(handle, layout, ...)

v_page = v_smem.slot(buf)                 # [token, dim]
v_page_t = memdesc_transpose_2d(v_page)   # logical [dim, token]
v_reg_t = tl.load(tle.gpu.local_ptr(v_page_t))

pv = tle.gpu.wgmma(
  v_reg_t, p_smem,
  trans_b=True,
  out_dtype=tl.float32,
)

在此基础上，TLE MemDesc、异步 TMA Copy 和 Barrier 共同构成多级流水，并将 V 的 Barrier Wait 延迟至 PV 真正发射之前，结合 Shared-Memory 缓冲区复用与物理页号预取，使下一 Tile 的 K/V 搬运与当前 Tile 的 QK、Online Softmax 和 PV 计算尽可能重叠。该路径既避免了 V 预转置和额外 Workspace，也降低了页表解析、地址计算与数据准备的串行依赖。

3.Task-aware 调度与 Cluster DSM 分层归约

HPC-Ops 官方 Dynamic 的均分调度策略能够避免按请求分配 CTA 时长序列占用时间过长的问题，使不同 CTA 承担相近的 KV 计算量，对 Ragged Batch 具有较好的负载均衡效果。但这种调度不能同时保证属于同一 Attention 输出的 Split 具有归约局部性：一个(batch,q_token,q_head) 的 KV 区间可能被拆分并分配给多个相互独立的 CTA，每个 CTA 都需要将局部 Output 和 LSE 状态写入 Global Workspace，再由 Finalizer 回读并完成全局 Softmax 合并。KV 切分越细，虽然 CTA 负载越均衡，产生的 Partial 数量也越多，由此增加：

Global Workspace 写入与回读；

Completion Counter 和同步竞争；

Finalizer 的输入规模与归约延迟；

短任务进入完整 Split-K 流程所需的固定开销。

因此，优化的关键不只是“将 Tile 分得更平均”，而是同时解决计算负载均衡与 Partial 归约局部性。Triton-TLE 引入面向任务特征与 Cluster 拓扑的动态调度：Dynamic 路径先根据负载选择 KV chunk 的目标长度，再按各请求的实际 KV 长度生成任务；对于同一逻辑输出，连续的 chunk 按 group = chunk_id // cluster_size 分组，组内的 rank = chunk_id % cluster_size 决定执行该 chunk 的 CTA；这样，同组 CTA 位于同一个硬件 Cluster，可通过 DSM 先合并 Partial；只有任务跨越多个 Cluster 时，才需要进一步合并各 Cluster 的结果。调度器则按任务规模选择归约范围：单 CTA 任务直接写回；较小的 Split-K 任务可在 4-CTA Cluster 内由两个 CTA 组成子组完成合并；更大的任务使用完整的 2、4 或 8-CTA Cluster。特别是，Dynamic Task Map 按实际任务构建紧凑 Grid，仅保留 Cluster 对齐所需的少量占位 CTA，减少短请求带来的空任务；Static 路径则由 Program ID 解析规则 Split-K 映射，省去 Task Map 的构建和读取。由此，Cluster 既提供 KV 并行计算能力，也为同一输出的 Partial 提供片上归约范围。

具体而言，Cluster归约中每个计算 CTA 输出尚未完成全局归一化的局部 Online Softmax Raw状态，并通过 tle.device_mesh建立硬件 Cluster，利用 tle.remote构造对其他 CTA Shared Memory 的远程访问试图，并通过tle.distributed_barrier 保证 Partial 的发布与读取之间的内存可见性和执行顺序：

# Example: a hardware Cluster composed of two CTAs
mesh = tle.device_mesh({
  "block_cluster": [("cluster_x", 2)]
})
cluster_rank = tle.shard_id(mesh, "cluster_x")

# Rank 0 owns the DSM reduction buffers.
owner_o = tle.remote(o_smem, 0, scope=mesh)
owner_m = tle.remote(m_smem, 0, scope=mesh)
owner_l = tle.remote(l_smem, 0, scope=mesh)

# Rank 1 publishes its raw state to Rank 0.
if cluster_rank == 1:
  tl.store(tle.gpu.local_ptr(owner_o), local_o)
  tl.store(tle.gpu.local_ptr(owner_m), local_m)
  tl.store(tle.gpu.local_ptr(owner_l), local_l)

tle.distributed_barrier(mesh)

# Rank 0 performs the Softmax-aware merge.
if cluster_rank == 0:
  peer_o = tl.load(tle.gpu.local_ptr(o_smem))
  peer_m = tl.load(tle.gpu.local_ptr(m_smem))
  peer_l = tl.load(tle.gpu.local_ptr(l_smem))

  global_m = tl.maximum(local_m, peer_m)
  w0 = tl.exp2(local_m - global_m)
  w1 = tl.exp2(peer_m - global_m)

  global_l = local_l * w0 + peer_l * w1
  global_o = local_o * w0 + peer_o * w1

Cluster 中不负责最终归约的 CTA 将局部状态发布到指定归约 Owner CTA 的 Distributed Shared Memory（DSM），Owner CTA 在 Cluster Barrier 后完成 Softmax-aware Merge，并结合 Head Sharding 将不同 Head 的归约工作分散到多个 CTA。对于包含 S 个 Split 的逻辑输出，传统路径最多需要向 Global Workspace 写入 S 份 Partial；采用大小为 C 的 Cluster 后，Cluster 内部先完成局部合并，跨 Cluster Finalizer 面对的输入数量可降至约ceil（S/C）；特别的，如果全部 Split 能够在单个 Cluster 内闭合，则可直接完成 Deferred Normalization 并写回最终结果，无需进入跨 Cluster Finalizer，只有不能片上闭合的 Cluster-level State 才写入 Global Workspace。

由此，Task-aware 调度保留了官方 Tile 级平均分配在 Ragged Batch 上的负载均衡优势，同时根据任务规模避免过度 Split；Cluster DSM则进一步将原本缺乏局部性的 CTA Partial 组织为片上归约组。二者共同缓解了“切分越细、计算越均衡，但全局归约越昂贵”的矛盾，减少Global Workspace 流量、Finalizer Fan-in 和归约关键路径。

Triton-TLE Dynamic Decode 实现的其余优化如下表所示：

朴素 Triton 为什么不够？

Triton 擅长表达张量计算、GPU 调度逻辑和编译期 Kernel 特化，但 Decode 优化还需要精确控制 Hopper 的 TMA 异步搬运、WGMMA Operand Layout、Shared Memory 生命周期以及跨 CTA 的 DSM 通信。这些能力难以仅通过常规的 load/store、dot 和 Program ID 抽象高效表达。

例如：

FP8 PV 所需的VT 必须满足特定的字节排列与 Shared-Memory Swizzle，需要稳定生成LDSMT → PRMT → STSMN等 Hopper 原生指令序列；

Small-M WGMMA计算需要重构计算几何并构造 Typed RS-WGMMA，并将 V Tile 转换为零拷贝逻辑转置视图；

Paged KV Cache 需要通过 Tensor Descriptor 和 Barrier 构建多级 TMA Pipeline，使页表解析、K/V 搬运与 QK、Softmax、PV 计算相互重叠；

Cluster 分层归约还需要访问其他 CTA 的Shared Memory，并严格保证 Partial State 的可见性和同步顺序。

若仅使用朴素 Triton，这些路径通常需要额外的数据重排、Global Workspace或 Kernel 边界，从而增加访存、同步和归约开销。

TLE 为这些优化提供了分层表达能力：通过 TMA、Typed WGMMA、Shared-Memory MemDesc、Warp Specialization 和 Barrier 构建 CTA 内的异步计算流水；通过 Cluster Mesh、Remote MemDesc 和 Distributed Barrier 实现跨 CTA 的 DSM 协作。对于高层布局系统难以稳定表达的局部热路径，Prefill 进一步使用 TLE-Raw 接入 Warp-level Stream Compaction 以及原生 V 转置微内核，从而直接控制指令序列、FP8 字节排列和Shared-Memory Layout。

在此基础上，Prefill 通过 Active-Tile Compaction 将稀疏性转化为实际的访存与计算缩减，并以 Page-aware Warp-specialized Pipeline 提高 TMA/WGMMA 重叠度；Decode 则通过 Small-M WGMMA、原生 MTP、Task-aware 调度和 Cluster DSM 分层归约，分别降低无效计算、重复 KV 访问、SM 空转以及 Global Workspace 流量。最终，Workload-aware Policy 再根据序列形态、布局、精度、量化方式和 MTP 规模选择相应的 Kernel 变体。

Triton 负责 Attention 算法、任务调度和编译期特化，TLE 负责将这些设计映射为 Hopper 原生的数据流、执行布局与同步协议，TLE-Raw 则补足少数必须进行指令级控制的关键路径。三者协同，使 Prefill 和 Decode 在保持高层可编程性的同时，获得逼近乃至超过优化 CUDA 实现的性能。