让模型写出一段能编译的 Triton 或 CUDA kernel 并不算困难。真正困难的是回答后续问题：它实现的是不是同一个算子？在 ragged sequence、非对齐 shape、FP8 和随机采样下还正确吗？快的是哪个 GPU、哪组输入？当线上请求没有匹配的实现时，系统会怎样退回已知正确路径？

FlashInfer-Bench 试图把这些问题串成一条闭环：

生产调用
  └─> 记录 Definition 与真实 Workload
        └─> 人或 AI 生成 Solution
              └─> 隔离编译与正确性检查
                    └─> 可复现性能评测
                          └─> 形成 Trace
                                └─> apply() 按输入选择实现
                                      └─> 无匹配时 fallback

因此它既不是一个单纯的 kernel 小测验，也不是“AI 生成代码后自动替换生产”的无条件开关。它的核心价值是用统一契约连接生成、评测和部署；而生产可信度仍取决于契约是否完备、workload 是否代表真实流量、验证是否覆盖风险，以及上线策略是否保留控制面。

一个 kernel 为什么不能只用函数名描述

名字同为 attention 的两个调用，可能在这些方面完全不同：

  • prefill 还是 decode；
  • paged KV 还是 ragged KV；
  • MHA、GQA、MLA 或 DSA；
  • head dimension、query/KV head 数、page size；
  • causal、window、logits soft cap；
  • BF16、FP16、FP8 或量化 KV；
  • sequence length 是否动态；
  • plan() 阶段准备了哪些 metadata。

最优实现也随 shape 改变。隐藏维度为 128 与 4096 的 GEMM 不应强行使用同一个 tile；短 decode 和长 prefill 对 occupancy、并行切分和 HBM 流量的要求也不同。

如果只告诉生成模型“优化这个 attention”，它可能得到一个在单一样例上很快、却悄悄忽略 page table、mask 或输出 LSE 的实现。FlashInfer Trace 首先要做的，就是把“同一个算子”定义成可核验的合同。

FlashInfer Trace 的四个组成部分

一个 Trace 是某个 Solution 在某个 Definition 和具体 Workload 上的不可变评测记录。顶层关系为：

Definition ─┐
Solution   ─┼─> Trace ─> Evaluation + environment
Workload   ─┘

Definition：算子合同

Definition 描述：

  • 输入、输出的名称、shape、dtype 和 layout；
  • axes 以及每个 axis 是常量还是变量；
  • axes 与张量 shape 的约束；
  • 运算语义和正确但不必很快的 reference；
  • op type、标签和说明。

系统用一条严格规则判断两个调用能否属于同一 Definition：axes 集合相同、每个 axis 的 const/var 角色相同、所有 const axis 的值相同。

例如 head_dim=128 若被视为常量，就已经成为 kernel 身份的一部分；batch_size 若为变量，则由每个 workload 实例给出。

FlashInfer-Bench 倾向于更具体的 Definition，并避免 optional input 或一个 flag 切换多种行为。原因是明确性：当 causal 与 non-causal 的语义不同，建立两个 definition 比把分支藏在运行时参数中更容易生成、验证和无歧义 dispatch。

一个概念化定义可以写成：

name: fused_add_rmsnorm_h4096
op_type: rmsnorm
axes:
  batch: {type: var}
  hidden: {type: const, value: 4096}
inputs:
  hidden_states: {shape: [batch, hidden], dtype: bfloat16}
  residual:      {shape: [batch, hidden], dtype: bfloat16}
  weight:        {shape: [hidden], dtype: bfloat16}
outputs:
  output:         {shape: [batch, hidden], dtype: bfloat16}
  residual_out:   {shape: [batch, hidden], dtype: bfloat16}
reference: reference.py

这只是说明字段关系，具体 schema 应以当前官方文档为准。

Workload：合同的一次真实实例

Workload 给 variable axes 填入具体值，并提供真实或可复现生成的输入。相同 definition 可以有很多 workload：batch 为 1、17、256，sequence length 为 128、8192，paged KV 的 indptr 也可呈现完全不同的 ragged 分布。

输入值是否要保存，取决于它会不会影响正确性和性能：

  • 普通 dense GEMM 常可按固定 seed 随机重建浮点输入；
  • attention 的 page index、sequence length 和 indptr 必须保留真实结构；
  • top-p sampling 的概率分布会影响分支和循环次数，真实 tensor 很有价值；
  • 极端数值、重复概率、空序列等边界不能只靠均匀随机覆盖。

官方工具支持在运行服务时 tracing，并按 shape 或自定义 key 去重。大张量可只存必要的整数 metadata，tensor payload 则可以 safetensors blob 保存。其目标不是无差别录下所有流量，而是得到体积可控又保留性能多样性的 workload 集。

Solution：某个合同的具体实现

Solution 可以来自 CUDA、Triton、CUTLASS、PyTorch 或其他支持的 backend。除源代码外，还要记录：

  • 目标 GPU architecture；
  • 软件和库版本约束；
  • 作者是人、模型还是 agent pipeline；
  • 编译配置和依赖；
  • 对应 Definition；
  • 必要的构建入口。

Solution 的函数签名必须与 Definition 对齐，输出也必须是合同声明的 tuple。调用一个已有高性能库并不自动构成违规；是否允许、链接成本如何计算、是否真正生成了自定义 kernel，应由具体比赛或评测规则明确。

Evaluation：这次运行发生了什么

Evaluation 记录某个三元组的：

  • 编译、shape/dtype、数值或 runtime 状态；
  • 正确性统计；
  • warm-up 与测量后的 latency；
  • reference/baseline 结果；
  • GPU、driver、CUDA、框架和库版本；
  • 日志与错误信息。

常见状态会区分 PASSED、INCORRECT_SHAPE、INCORRECT_DTYPE、INCORRECT_NUMERICAL、RUNTIME_ERROR 和 COMPILE_ERROR。这样“0 分”不再混成一个结果：agent 可以知道是代码没有编译，还是运行了但语义错误。

为什么数据要来自服务 trace

合成 benchmark 常把 shape 取成整齐的 2 的幂，输入分布也近似均匀。线上 LLM 服务却充满结构：

  • continuous batching 让每个请求的 context length 不同；
  • paged KV 的 block table 与物理页分布影响访存；
  • MoE 每个 expert 接收的 token 数高度不均；
  • top-p/top-k sampling 的有效候选数随 logits 改变；
  • tensor parallel 会把 head、hidden 和 expert 维度切成特定常量；
  • 模型层之间可能复用同一 op type，却有不同 axis 配置。

FlashInfer-Bench 从运行模型与真实 prompt 轨迹中提取主干算子，覆盖 attention、GEMM、MoE、normalization 和 sampling 等类别。然后沿性能敏感的 axes 与 tensor statistics 去重，保留具有代表性的子集。

这里需要理解一个边界：trace dataset 代表的是采集时选定的模型、引擎配置和流量，不是宇宙中所有输入。若线上从 TP=8 改为 TP=4、启用另一种量化或请求长度分布发生变化，就要重新采集或补充 workload，而不能把旧榜单当成永久最优表。

正确性不是一个统一的 allclose

不同类型 kernel 的“正确”含义不同。FlashInfer-Bench 将验证器按语义区分。

确定性算子：逐元素误差界

GEMM、normalization 与多数 attention 输出可与 reference 逐元素比较。对每个元素要求：

\[|y_{sol}-y_{ref}| \le \epsilon_{abs} + \epsilon_{rel}|y_{ref}|\]

此外任何 NaN 或 Inf 都直接拒绝。atol/rtol 必须按 dtype、累积顺序和算子条件数设置；为让某个候选通过而事后放宽全局阈值，会把正确性门槛变成可操纵参数。

低精度算子：matched ratio

FP8 等低精度 kernel 相对高精度 reference 可能系统性出现少量较大误差。仅把全局 tolerance 放得很宽，会容忍所有元素都不够准确。

FlashInfer-Bench 可采用 matched-ratio 规则：仍使用较紧的标准误差界，但要求至少比例 $\rho$ 的输出元素满足它。这样允许少量 outlier，同时保留对主体数值的约束。

Matched ratio 也不是质量证明。若错误集中在少量但语义关键的位置，服务层仍可能退化，因此生产验证还应比较 logits、模型输出或端到端任务。

随机算子：比较分布而不是样本

Sampling kernel 每次输出可以不同，拿一次 candidate sample 与一次 reference sample 做相等比较没有意义。

验证过程先根据输入概率和 top-k/top-p mask 得到理论分布 $q$，多次运行候选得到经验分布 $\hat f$，再计算 total variation distance：

\[TVD(q,\hat f) = \frac{1}{2}\sum_i |q_i-\hat f_i|\]

只有 TVD 低于阈值才通过；同时，任何被 mask 排除的 token 一旦被采样，都立即判错。这同时检查了统计分布与硬约束。

随机验证还必须固定或记录 RNG 语义、样本数和置信区间。样本太少会把统计噪声误判为实现错误，也可能让有偏 kernel 偶然过关。

边界输入比大批随机数更重要

Reference 比较只有在 workload 覆盖错误触发条件时才有效。一个用于生产的验证集还应主动构造：

  • 0、1、warp 边界及非对齐维度；
  • 最大支持 shape 与接近显存上限的输入；
  • 空 batch、空页和最后一页未填满；
  • 极大/极小值、全零、重复值、NaN/Inf 输入策略；
  • 高度倾斜的 MoE routing；
  • top-p 恰好跨候选边界、概率和不精确等情况；
  • 不同 stride、非 contiguous tensor；
  • 多 stream 并发及反复执行。

还需要检查越界写。输出数值看似正确，不代表 kernel 没有踩到相邻 buffer；race condition 也可能只在多次运行或特定 SM 调度下出现。计算 sanitizer、guard region、重复执行与独立进程都能补足单次数值比较看不到的风险。

性能测量怎样避免自欺

正确性通过后，才有资格谈速度。FlashInfer-Bench 为每张 GPU 设置跨进程 device lock，避免两个 benchmark 同时占用设备。候选先执行若干次 untimed warm-up，再用 CUDA event 记录多轮 device-side 时间。

仍需明确几个选择：

测的是 warm 还是 cold

JIT 编译应从单次 kernel latency 中剔除，但首次生产请求是否会承担编译成本，是部署问题。L2 cache 是否 flush、数据能否复用、CUDA Graph 是否启用，也会显著影响结果。报告必须说明所测场景。

用均值还是分位数

官方核心 benchmark 可用多轮均值形成稳定排序；生产压测还要看 P50/P95/P99，因为频率变化、资源竞争和动态分支可能制造长尾。

按 workload 平均是否合理

如果线上 90% 是小 batch、10% 是大 batch，给所有 workload 相同权重会选出错误的全局赢家。更合理的调度表按具体 feature key 选实现，或用生产频率加权评估整体收益。

Kernel 加速不等于服务加速

若某 kernel 只占请求时延的 5%，即使快 2 倍，Amdahl 定律给出的端到端上限也很小：

\[S_{total} = \frac{1}{(1-f)+f/S_{kernel}}\]

其中 $f$ 是原始时间占比。替换还会引入 dispatch、layout conversion、workspace 和同步；最终必须在完整模型服务中测 TTFT、TPOT、吞吐与显存。

隔离执行是在防什么

AI 生成代码可以编译失败、非法访问、死循环、耗尽显存，也可能通过读取残留显存或缓存 reference output 来“刷”正确性分数。FlashInfer-Bench 提供两类运行方式：

  • 完全隔离模式：每个 solution 在独立 subprocess/CUDA context 中运行，完成或超时后销毁；
  • persistent worker：每 GPU 保留长生命周期进程与预热备用 worker，提高大规模 sweep 吞吐，失败时切换/恢复。

隔离模式减少跨 solution 状态残留，persistent 模式减少 context 初始化与重复编译。可疑或反复失败的候选应升级到更严格隔离，而不是为了跑得快继续污染共享 worker。

在生产供应链中还应再加一层：生成与评测环境不应持有生产凭据；编译容器限制网络与文件访问；artifact 经过源码审计、依赖锁定、签名和 provenance 记录。Benchmark runtime isolation 不能代替整个软件供应链安全。

排名应该回答两个问题

一个 agent 的结果至少有两个维度：

  1. Resolved/Correctness：多少定义与 workload 能产生通过验证的实现；
  2. Performance：通过的实现相对 baseline 快多少、覆盖多少 workload。

只统计最快成功样例，会奖励“多数失败、偶尔撞中一个 shape”的策略；只统计通过率，又会让所有 solution 都调用保守 reference。性能曲线、speedup threshold 下的覆盖率、各 workload latency 和失败类别组合起来，才能描述能力。

而榜单是随 dataset、硬件、validator 与参赛 solution 更新的动态状态。文章不应把某一天的模型名次写成长期事实；可复现结论应指向冻结 snapshot、版本与评测环境。

apply() 怎样把 Trace 变成运行时选择

FlashInfer-Bench 的部署连接点是 apply()。启动前，它可以从本地 Trace 数据库构建索引：

  1. 过滤没有通过正确性或误差阈值的 evaluation；
  2. 从 workload 提取 shape 等 dispatch feature；
  3. 为每个 key 选择最快 solution；
  4. 对高频 solution 提前编译；
  5. 在线调用时用少量索引查找完成选择。

概念上类似：

from flashinfer_bench import apply

result = apply(
    def_name_or_resolver=resolve_definition,
    runtime_kwargs=kernel_inputs,
    fallback=known_good_implementation,
)

当没有匹配 definition/workload、solution 不兼容或全局功能关闭时，fallback 仍是行为基准。

对简单 stateless 函数，resolver 可以直接查看输入 shape。FlashInfer 的 attention wrapper 常具有 plan()/run() 两阶段状态，adapter 需要在 plan 时保存 indptr、page size 等上下文，在 run 时取回完整参数才能正确匹配。只拦截最终 tensor 调用，可能已经丢失 definition 所需的调度信息。

当前官方文档列出的 built-in adapter 覆盖是有限集合，并对 page size、causal 等条件有明确约束。启用 FIB_ENABLE_APPLY=1 不代表所有 FlashInfer 调用都会被替换；不满足 predicate 的路径应原样执行。

“可以动态替换”与“应该直接上线”之间

apply() 提供了技术路径，生产发布还应设置 promotion pipeline：

generated
  -> offline correctness
  -> hidden/adversarial workloads
  -> sanitizer + isolation
  -> reproducible benchmark
  -> source review + signed artifact
  -> staging end-to-end
  -> shadow compare
  -> small canary
  -> gradual rollout
  -> default eligible

每个 artifact 至少绑定：

  • source 与 build recipe hash；
  • GPU compute capability；
  • driver/CUDA/compiler/FlashInfer 版本；
  • Definition、Trace dataset snapshot 与 tolerance；
  • 通过的 workload feature 范围；
  • benchmark 和 sanitizer 报告；
  • fallback 与禁用开关。

线上监控则应覆盖 illegal access、NaN/Inf、输出抽样比对、P99 latency、OOM、fallback rate 和 solution 命中分布。某个 driver 或库升级后，应使旧 artifact 重新进入验证状态，而不是只按 kernel 名继续复用。

一个可执行的上线策略

先选择低风险算子

从语义简单、调用边界清楚、可逐元素验证的 fused normalization 或特定 GEMM 开始。Paged attention、MoE routing 和 stochastic sampling 的状态空间更大，适合在验证体系成熟后接入。

将 Definition 当成 API 版本

任何 layout、mask、scale 或 side effect 的变化都应产生新 definition/version。不要悄悄修改 reference 后沿用已有 benchmark，这会让旧 solution 看似兼容。

将 workload 分成公开、隐藏与线上回放

公开集合帮助 agent 调试，隐藏集合防止 hard-code，脱敏线上回放检查真实分布。三者用途不同，不应互相替代。

同时保留两个回退层级

  • dispatch 无匹配时立即回到 FlashInfer 原实现；
  • 已匹配 solution 线上异常时，熔断该 artifact 并在进程/集群级禁用。

GPU illegal memory access 可能破坏当前 CUDA context，仅在同一调用里捕获异常并重试不一定安全；服务进程还需要健康检查和重启策略。

用端到端收益决定是否推广

Kernel 微基准胜出只是 promotion 的必要条件。若 layout adapter、dispatch 或调度交互抵消了收益，就没有上线价值。最终决策应基于真实模型、真实请求分布和业务 SLO。

FlashInfer-Bench 的真正意义

AI 生成 kernel 的最大障碍不是缺少更多代码候选，而是缺少共同语言与可追溯闭环。FlashInfer Trace 把任务、输入、实现和结果锁在一起，使下面的关系可以被复查：

为什么生成这个实现？   -> Definition
它针对什么场景？       -> Workload
代码和目标环境是什么？ -> Solution
它如何证明正确且更快？ -> Evaluation
线上为什么选中它？     -> apply index + compatibility predicate
出问题如何退出？       -> fallback + rollout control

这套结构同样适用于人工 kernel、编译器搜索和已有库实现，并不依赖代码一定由 AI 生成。

小结

FlashInfer-Bench 把 GPU kernel 优化从“贴一段代码和一个加速数字”推进为一组可复现对象：Definition 定义语义，Workload 表示真实调用，Solution 保存实现与环境，Evaluation 记录正确性和性能，apply() 再按输入选择已经验证的候选。

它确实提供了动态替换能力，但安全上线仍需隐藏边界测试、隔离执行、artifact provenance、端到端回放、灰度、监控和 fallback。只有这些环节都存在，“AI 写出的更快 kernel”才可能变成“生产系统可以控制风险地使用的更快 kernel”。

参考资料