从编译器能力到公开算子
这个项目围绕一个具体目标展开：让 TIRx 在 Hopper SM90a 上表达并执行 Gated Delta Net prefill。工作包含 TVM 编译器扩展、packed 变长 GDN 实现，以及与固定 CuTeDSL 和 FLA 版本的对比。
这里最有代表性的工程问题，是如何把 WGMMA、TMA、共享内存布局、寄存器累加器和同步协议连成一条真正可调用的路径。只有指令能够生成，还不足以证明算子能够正确工作。
SM90a 支持具体补了什么
实现以一个完整的 128-thread warpgroup 为 WGMMA 执行范围，接受 FP16/BF16 输入与 FP32 寄存器累加。共享/共享（SS）和寄存器/共享（RS）两条操作数路径都需要布局与生成指令保持一致。
| 开发部分 | 解决的问题 |
|---|---|
| WGMMA 分派与 lowering | 限定 sm_90a、M=64、N∈{16,32,64,128}，K≥16 且为 16 的倍数；将合法算子映射到原生指令。 |
| 寄存器 fragment 布局 | 表达 A 操作数、FP32 累加器，以及较大累加器中对齐的输出列切片。 |
| 共享内存 descriptor | 让 MN-major / K-major 与 swizzle 布局按同一含义被 WGMMA 消费。 |
| TMA 与 Host ABI | 处理显式 global stride、TensorMap 的存储对齐和 dtype，使 host/device 两端传递一致。 |
| 尾块 copy | 对逻辑有效前缀加掩码，在需要时重新计算 swizzle 后的共享内存偏移。 |
commit/wait 协议保持对调用方可见，方便将多次异步操作组织进流水线。一个关键取舍是把合法条件写进分派检查：通用 sm_90 或不满足 warpgroup 条件的请求不能误入 WGMMA 路径。
用状态递推解释 GDN
GDN 的输出依赖不断更新的状态。以公开的 [V,K] 状态矩阵 H 为例，每个 token 先衰减状态，再计算预测误差，将更新写回，最后得到输出。状态方向、门控和舍入顺序都属于实现需要保留的语义。
H_decay = alpha × H_previous
prediction = H_decay @ k
update = beta × (v − prediction)
H = H_decay + update ⊗ k
output = scale × (H @ q)GPU 实现将递推组织成 64-token chunk，通过矩阵运算生成 QK、KK 和 WY 数据，再执行状态扫描与输出。逆矩阵过程包含精度可见的 8 → 16 → 32 → 64 阶段；数学上的等价变形不一定保持这些数值行为。
packed 输入让多个独立序列共享连续 token 存储。每个序列必须拥有自己的状态与尾块：无效 Q/K/V 置零、alpha 取一、beta 取零，重放不能越过 cu_seqlens 边界。公开接口支持 MHA、GQA、GVA 三类头关系；本页性能矩阵覆盖的是 MHA 和 GVA。
为什么保留三种调度
同一个算子的短 packed 序列与长序列，未必适合相同的状态驻留和重算方式。公开 wrapper 只使用 host 可见的形状与状态选项选择路径，不读取设备上的 cu_seqlens 数值来决定路由。
通用 pipeline
256-thread prepare 生成 chunk 数据，128-thread scan/output 完成递推。未命中特化条件的有效输入走这条 TIRx 路径。
寄存器 replay
单个 512-thread 融合 kernel，针对 4096 tokens、10 条序列、8 heads 且带初始与最终状态的精确组合。
前驱 chunk replay
256-thread prepare 与 256-thread replay/output，两个 value warpgroups 共驻；只重放一个前驱 chunk，配合双缓冲与重算。
特化路由使用精确元数据键，并要求显式 alpha、beta。near miss 返回通用 TIRx 实现，因此测试需要覆盖公开分派和相邻输入。内部 kernel 的快，不足以代表用户调用路径的快。
公开版本复测：六个场景
采用仓库中 2026 年 7 月 29 日公开 tag 的独立复测：一套 H20-3e 环境、sm_90a、BF16 Q/K/V、head dimension=128、scale=0.73。下表是公开 callable 的 CUDA-event 时间，每项取独立进程平均延迟的中位数。
| 场景 | TIRx ms | CuTeDSL ms | FLA ms | TIRx / CuTeDSL |
|---|---|---|---|---|
| single-t512-h8-mha-zero | 0.102550 | 0.121454 | 0.373282 | 0.844350 |
| single-t1024-h8-mha-state | 0.157430 | 0.174915 | 0.389517 | 0.900037 |
| single-t1024-h8-hv16-gva-state | 0.143251 | 0.177971 | 0.394927 | 0.804915 |
| single-t4096-h16-mha-zero | 0.263054 | 0.451661 | 0.442887 | 0.582414 |
| packed-n10-t4096-h8-mha-state | 0.173384 | 0.170712 | 0.392756 | 1.015654 |
| packed-n20-t8192-h8-hv16-gva-state | 0.520107 | 0.555996 | 0.629735 | 0.935451 |
六行的 TIRx/CuTeDSL 几何平均比值为 0.834812，TIRx/FLA 为 0.453700。五行快于 CuTeDSL；packed-10 的比值为 1.015654，约慢 1.57%，在预注册 ±2% 噪声带内，不能算作领先。
计时包含 20 次预热和每份记录 100 次迭代，通常每实现、每行运行 3 个独立进程；packed-10 按预注册规则扩大到 7 个，总计 66 份计时记录。比较使用固定 CuTeDSL GDN 快照与官方 FLA 实现，不混用 GDN2 comparator。
验证如何支撑结论
公开 tag 复测确认了源码与运行身份、逐记录正确性及六行公开调用计时，结论等级为 CHARACTERIZATION。它为公开版本的可测量行为提供了证据。
更早的历史封存另有同步、内存安全与代码生成检查：9 次公开调用的 host-visible 同步审计未发现违规，packed redzone 检查通过，Compute Sanitizer 报告零 race hazards 与零 sync errors，六个 case、11 个阶段的代码生成记录中 spill 与 stack 为零。这些检查属于历史版本，不能自动转移到复测版本。