TileLang Mamba2 Chunk Scan 前向算子：自动调优实现、H800 实测吞吐与 Triton/Helion 基线对比
项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang 本文以仓库中 benchmark/mamba2/README.md 记录的 Mamba2 chunk_scan 前向算子基准测试为主体，完整继承其评测参数、复现命令与实测吞吐数据，并结合 benchmark/mamba2/benchmark_mamba_chunk_scan.py 的源码深入讲解该 TileLang kernel 的分块策略、指数衰减实现与 @autotune 搜索空间设计，帮助读者掌握在 Mamba2 状态空间模型关键算子上编写、调优并跨编译器评测 TileLang kernel 的完整流程。
一、评测目标与固定参数配置
该基准聚焦 Mamba2（状态空间模型 SSM 的第二代）推理/训练管线中的核心前向算子 chunk scan：给定按 chunk 组织的输入 cb、投影序列 x、时间衰减系数 dt、累加衰减 dA_cumsum、输出投影 C、上一 chunk 的隐状态 prev_states 与跳跃连接系数 D，计算输出序列。
文档固定了如下评测形状（这也是 README 中所有吞吐数据的适用前提）：
| 参数 | 取值 | 含义 |
|---|---|---|
batch | 8 | 批大小 |
heads | 80 | 注意力头数（nheads） |
groups | 1 | 头分组数（ngroups，GQA 式共享） |
chunk_size | 256 | 每个 chunk 的时间步数 |
dim | 64 | 每头维度（headdim） |
dstate | 128 | 状态空间维度 |
seq_len | 1024 ~ 32768 | 扫描的序列长度（nchunks = ceil(seq_len / chunk_size)） |
计算精度为 T.float16 输入、T.float32 累加（见 kernel 源码中 dtype = T.float16、accum_dtype = T.float32），并开启 TL_ENABLE_FAST_MATH 快速数学 pass。
二、如何复现：完整命令
README 给出的官方复现方式如下（在仓库根目录执行，进入 benchmark/mamba2 后通过 stdin 调用 chunk_scan_fwd）：
cd benchmark/mamba2
python - <<'PY'
from benchmark_mamba_chunk_scan import chunk_scan_fwd
batch = 8
heads = 80
groups = 1
chunk_size = 256
dim = 64
dstate = 128
for seq_len in [1024, 2048, 4096, 8192, 16384, 32768]:
res = chunk_scan_fwd(
batch,
seq_len,
chunk_size,
groups,
heads,
dim,
dstate)
tflops = (2 * batch * seq_len * chunk_size * heads * dim * 0.5 + 2 * batch * seq_len * heads * dim * dstate) / res.latency * 1e-9
print(f"seq_len={seq_len:5d} latency={res.latency:.6f}ms TFlops={tflops:.3f}")
PY
几点说明（均可在源码中印证）：
chunk_scan_fwd被@autotune(configs=get_configs(), warmup=10, rep=10)装饰（见 benchmark_mamba_chunk_scan.py），调用后返回的是自动调优结果对象，其.latency属性即最优配置的测得延迟（ms），这正是上式中res.latency的来源；脚本__main__分支还会进一步打印best_config与ref_latency。- FLOPs 公式为
2*batch*seq_len*chunk_size*heads*dim*0.5 + 2*batch*seq_len*heads*dim*dstate，与 benchmark_mamba_chunk_scan.py 中total_flops的定义一致：前一项是 chunk 内 CB 与 x 的因果 matmul（0.5因子对应下三角因果掩码只有一半有效），后一项是C与prev_states的跨 chunk 状态 matmul。 - 若希望与 Triton、Helion 基线同场对比，可直接运行脚本本体，它提供命令行参数
--batch、--heads、--groups、--seq_len、--chunk_size、--dim、--dstate、--tune（见 benchmark_mamba_chunk_scan.py），默认值与 README 的固定参数完全一致。注意该脚本依赖mamba_ssm（Triton 版_chunk_scan_fwd）与helion两个第三方包，缺失时会抛出带安装提示的ImportError。
三、实测结果：H800 上随序列长度扩展的吞吐
README 记录的测试环境为：仓库 commit 8a5eb569704bfea64478c29adcfe3a09e3c2b12c，GPU 为 NVIDIA H800 SXM（驱动 560.35.05），使用默认自动调优搜索空间。实测数据如下表：
| Seq_len | Latency (ms) | Throughput (TFLOPs) |
|---|---|---|
| 1024 | 0.169 | 126.477 |
| 2048 | 0.329 | 130.195 |
| 4096 | 0.645 | 133.054 |
| 8192 | 1.278 | 134.362 |
| 16384 | 2.531 | 135.711 |
| 32768 | 5.076 | 135.379 |
从数据形态看：延迟随 seq_len 近似线性增长（1024→32768 时延迟从 0.169 ms 增至 5.076 ms，约 30 倍，与 32 倍长度基本吻合）；吞吐随序列变长从 126.5 TFLOPs 爬升到约 135 TFLOPs 后趋于平稳——短序列下 kernel 启动与访存占比更高，长序列下计算密度摊薄了固定开销。复现结果会因驱动、CUDA 版本与调优随机采样而异，上表应视为该特定环境的参考基线。
四、与 Triton、Helion 基线的对比
README 声明的基线版本为：Triton v3.5.0 + mamba-ssm v2.2.6.post3，以及 Helion v0.2.1。三者在同一组输入上的吞吐对比见下图（图注为 NVIDIA H100）：
- Triton 基线直接调用
mamba_ssm.ops.triton.ssd_chunk_scan._chunk_scan_fwd，计时用do_bench(..., _n_warmup=10, _n_repeat=10)； - Helion 基线是在脚本内用
@helion.kernel()现写的等价 kernel（按[nheads, chunk_size, headdim, batch, nchunks]分块，先做C @ prev_states.T再循环累加因果衰减后的cb @ x，见 benchmark_mamba_chunk_scan.py），并通过run_example与 PyTorch 参考实现ref_program对拍。
五、源码深读：chunk_scan_fwd TileLang Kernel
chunk_scan_fwd 的完整实现位于 benchmark_mamba_chunk_scan.py，核心结构如下。
5.1 张量布局与 Kernel 网格
八个张量签名为：
cb: (batch, nchunks, ngroups, chunk_size, chunk_size)—— 每个 chunk 内 CB 预计算结果；x: (batch, seqlen, nheads, headdim)、Output: (batch, seqlen, nheads, headdim)；dt,dA_cumsum: (batch, nheads, nchunks, chunk_size)—— 逐时间步衰减系数及其累计；C: (batch, seqlen, ngroups, dstate)；prev_states: (batch, nchunks, nheads, headdim, dstate)；D: (nheads)。
网格划分为 T.Kernel(nheads, ceildiv(chunk_size, block_M) * ceildiv(headdim, block_N), batch * nchunks)，即：bz 遍历头（头到组的映射用 bz // (nheads // ngroups) 实现 GQA 共享，groups=1 时退化为每头独立），by 遍历 (batch, chunk)，bx 把单个 chunk 的 [chunk_size, headdim] 输出块按 m_idx（时间维）/n_idx（维度维）二维切分（见 benchmark_mamba_chunk_scan.py）。
5.2 两阶段累加：状态贡献 + chunk 内扫描
- 跨 chunk 状态项：将
C块与prev_states块载入 shared（C_shared: (block_M, block_Dstate)、prev_state_shared: (block_N, block_Dstate)），执行T.gemm(C_shared, prev_state_shared, acc_o, transpose_B=True)，再按行乘上scale_m_local[i] = exp2(dA_cumsum_m * p)（p = 1.44269504是 log2(e) 常数，把自然指数转换为exp2以利用硬件快速指令）。 - chunk 内因果扫描：
for k in T.Pipelined(loop_range, num_stages=num_stages)流水化 K 维循环。每一拍内：- 载入
cb块并逐元素乘exp2(dA_m * p - dA_k * p)得到跨时间步的相对衰减； - 再乘
dt列向量； - 用
T.if_then_else(m 起始行 >= k 起始列)施加下三角因果掩码； - 载入
x块并执行T.gemm(cb_local, x_shared, acc_o)累加。 这一整段与 PyTorch 参考实现ref_program中scores_decay = cb * decay; masked_fill(~tril)再 einsum 的逻辑一一对应（见 benchmark_mamba_chunk_scan.py 与 benchmark_mamba_chunk_scan.py）。
- 载入
- 跳跃连接与写回：累加
x_residual * D[bz]，经acc_o_shared写回Output对应切片。cb_shared与x_residual_shared通过T.annotate_layout绑定tilelang.layout.make_swizzled_layout以减少 shared memory bank conflict。
5.3 自动调优搜索空间
get_configs() 定义的默认搜索空间为（见 benchmark_mamba_chunk_scan.py）：
block_M ∈ {64, 128, 256}（chunk 时间维切块）block_N ∈ {32, 64}（headdim 维切块）block_K ∈ {64, 128, 256}（K 维流水步长）block_Dstate ∈ {128}（固定为 dstate）num_stages ∈ {1, 2, 3, 4, 5}（T.Pipelined流水级数）
笛卡尔积共 3×2×3×1×5 = 90 个候选配置。@autotune 装饰器（实现在 tilelang/autotuner/tuner.py）会把每个配置映射为 chunk_scan_fwd 的同名形参，多线程并行编译后逐个做正确性校验与计时（warmup=10, rep=10），计时底层走 tilelang/profiler/bench.py 的 do_bench（基于 CUDA event，含缓存刷新与多种回读模式）。调优结果带缓存：以函数源码、形状参数、配置集合与编译/剖析参数哈希为 key 落盘，同形状重复调用可直接命中缓存。若不想走调优流程，也可按 autotune 文档提示直接在调用时显式传入 block_M/block_N/... 等可整参数跳过搜索。
六、适用前提与限制
- 本 kernel 为 fp16 输入 / fp32 累加 的 CUDA 实现，评测数据仅对 H800 SXM（图注 H100）环境有效，其他 GPU 需重新实测；
- 复现依赖
mamba_ssm与helion（对比基线时）及einops（参考实现中rearrange/repeat），安装方式见脚本头部的 ImportError 提示； seq_len需能被chunk_size整除（参考实现中有assert seqlen == nchunks * chunk_size）；- 文中所有吞吐数字仅反映该 commit 与该环境下的测量，不应外推为跨平台的性能承诺。
小结
这份基准展示了 TileLang 在 Mamba2 这类"chunk 内矩阵乘 + 状态递推"混合算子上的标准工作流：用 T.prim_func 显式描述 shared/fragment 数据流与 T.Pipelined 流水，用 @autotune 对分块尺寸与流水级数做穷举搜索，再与 Triton、Helion 生态的同类 kernel 在同一 do_bench 计时口径下对拍。完整的可运行入口是 benchmark/mamba2/benchmark_mamba_chunk_scan.py，评测记录与数据表见 benchmark/mamba2/README.md，调优器机制可继续深入 tilelang/autotuner/tuner.py。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
