Triton 如何用 FpSan 比较两个内核的浮点语义是否一致-育师

Triton 如何用 FpSan 比较两个内核的浮点语义是否一致

【免费下载链接】tritonDevelopment repository for the Triton language and compiler项目地址: https://gitcode.com/GitHub_Trending/tri/triton

优化一个 Triton 内核之后（优化版对照参考版、融合版对照未融合的组合、或两个应该数学等价的调度变体），需要确认改写后的内核仍然实现着与原版相同的符号计算。Triton 仓库中的 FpSan 就是为这个任务设计的编译器插桩模式：它把选定的浮点 Triton IR 操作改写为基于整数位模式的确定性 "payload algebra"，其目标不是近似 IEEE 浮点算术，而是保留选定的代数结构——在消毒语义下符号等价的两个内核会继续保持结果一致，而错误改写、错误操作数、错误数据流或缺失同步则倾向于扰动结果。这篇文章给出在 FpSan 模式下完成一次"两个内核浮点语义一致性"比较的操作路径。

适用前提（均来自仓库文档）：

FpSan 是编译器特性，interpreter 模式下不生效；
目前支持所有 NVIDIA 硬件，以及 AMD 的gfx942、gfx950、gfx1250。

先理解 FpSan 输出代表什么

FpSan 比较的是 payload，不是 IEEE 数值。对每个浮点宽度w，FpSan 定义浮点位模式与w位整数 payload 之间的一一映射，算术在模2^w意义下回绕。概念上：

embed(x)把浮点位模式映射为整数 payload；
unembed(u)把整数 payload 映射回浮点位模式；
被消毒的浮点操作实现为unembed(F(embed(...)))。

嵌入方式刻意让几个关键常数保持稳定：embed(+0.0) = 0、embed(+1.0) = 1、embed(-1.0) = all-ones。这正是x + 0 = x、x * 1 = x这类恒等式在 FpSan 下自然成立的原因。

FpSan 精确保留的性质包括：

add、subtract、一元取负、multiply、FMA 及 dot 类累加的环恒等式（如x - x = 0、-(-x) = x、加法与乘法的结合/交换律、乘法对加法的分配律）；
exp/exp2的选定指数恒等式（如exp2(x + y) = exp2(x) * exp2(y)、exp2(0) = 1）；
sin/cos的角度恒等式（和差公式、cos(x)^2 + sin(x)^2 = 1）；
经过 cast、load、store、copy 时的 payload 相等性；
对尚无更丰富代数模型的一元函数，提供确定性的 op-distinguishing tag。

因此文档给出的核心判读是：如果两个内核在保留性质下应当是同一个符号计算，它们就应当产生相同的 payloads（文档注明这一点基于超越数论中一个被普遍相信的猜想 Schanuel's conjecture）。同时文档明确：FpSan 结果只能与其他 FpSan 结果比较，不能与普通浮点输出比较。

启用 FpSan

在你要插桩的那次编译或运行之前开启 FpSan。

从 Python 侧：

import triton triton.knobs.compilation.instrumentation_mode = "fpsan" # compile and run kernels here triton.knobs.compilation.instrumentation_mode = ""

从 shell 侧：

TRITON_INSTRUMENTATION_MODE=fpsan python your_script.py

TRITON_INSTRUMENTATION_MODE对应 knobs.py 中的instrumentation_mode（默认空字符串，即关闭），your_script.py换成你要检查的脚本。

可选参数：如果两个内核中包含浮点缩窄 cast，可以设置TRITON_FPSAN_HOMOMORPHIC_CASTS=1，使缩窄 cast 改用简单的 payload 截断；文档说明该选项 "preserves addition across independently downcast partial reductions"，即跨独立降精度局部归约的加法仍保持一致。只在比较对象涉及窄类型局部累加、且需要加法在 cast 两侧可比时才需要它。

比较两个内核的操作路径

文档指出 FpSan 最有效的用法是：在同一个 FpSan 模式下比较两个内核，或同一个内核的两个版本。文档列出的典型用途包括：

优化内核对照简单参考内核；
融合内核对照未融合的组合；
两个应该数学等价的调度变体；
检查 accumulator 选择、predication 或 TMEM pipeline 是否保持了预期的 payload flow。

操作步骤：

准备两个待比较的内核（例如优化版与参考版），保证两者消费同一份输入、产出相同形状的输出。
在 fpsan 模式下（上面两种方式任选其一，两次运行必须处于同一 FpSan 模式），分别运行两个内核，各自得到一份输出缓冲。
对两份输出做逐元素位级精确比较。FpSan 模式下输出是 unembed 回的 payload，等价内核应当 payload 完全相同。

仓库里有一份可直接照读的可执行参照：test_fpsan.py。其中test_binops_payload_semantics展示了完整模式——把 int32 存储缓冲用triton.TensorWrapper(out, dtype=torch.float32)包装后交给内核，设置fresh_knobs.compilation.instrumentation_mode = "fpsan"运行，最后用下面的断言把输出与期望的 payload 逐位比较：

def _assert_payload_equal(actual, expected) -> None: np.testing.assert_array_equal(_as_payload_np_unsigned(actual), _as_payload_np_unsigned(expected))

同一文件中的test_expect_zero_upper_triangle_exp则演示了"普通模式跑一遍 + fpsan 模式跑一遍"的对照写法，并断言两种模式的输出不相等（assert not torch.equal(regular_out, fpsan_out)）——这正是"FpSan 输出不可与普通浮点输出直接比较"这条规则的具体体现。

结果怎么判断

两份输出 payload 逐位相同：两个内核对 FpSan 保留的性质而言是同一个符号计算（受前述 Schanuel's conjecture 前提约束）。
存在不同：文档指出 wrong rewrites、wrong operands、wrong dataflow 或 missing synchronization 都会扰动结果，所以差异说明两个实现在 FpSan 语义下不等价。文档给出的排查方向是检查 accumulator 选择、predication、TMEM pipeline 是否仍保持预期的 payload flow。
无论结果如何，只拿 FpSan 结果对 FpSan 结果下结论；把 FpSan 输出与普通模式输出、或与常规数值容差比较，都不可作为一致性依据。

哪些检查不该交给 FpSan

FpSan is not an IEEE simulator。文档明确列出不依赖它的情形：

真实浮点排序、舍入、NaN 传播、无穷、次正规数或异常；
log、sqrt、erf、floor、ceil、rsqrt等 tagged unary op 的真实超越语义；
期望的浮点位模式（即在内核中 bitcast 浮点与整数的场景）。

tagged unary op（tl.log/tl.log2/tl.sqrt/tl.rsqrt/tl.erf/tl.floor/tl.ceil及 precise square root variants）的改写是可逆 payload 标签：同一 op 下相同输入给出相同 payload、不同一元 op 有不同 tag，但刻意不保留sqrt(x)^2 = x、log(x*y) = log(x) + log(y)这类真实数学恒等式。

文档对适用边界的总结：当你想检查 IEEE 边界情况、真实超越精度、NaN/inf 行为、低精度硬件格式的解码语义时，FpSan 是 poor fit，应改用常规数值测试。此外，对 Gluon MMA 与 tensor memory 的支持保留的是 payload 数据流，"do not make FpSan a substitute for race or synchronization checking"——数据竞争检测不在它的职责内。