快速后端
TileLang 可移植性
laya/tl_kernels.py 里的五个内核都能通过 compile_cpu 在 TileLang 的 CPU（c）目标上执行。这是一种显式的标量 fp32 特化，不是 Agent.accelerate 的 CPU 加速后端。它需要 TileLang 和一个本地 C++ 编译器。它不引入任何运行时服务依赖。
import torch
from laya import tl_kernels as K
A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)compile_cpu(factory, *args, **kwargs) 接受与五个 GPU 工厂相同的形状与操作选项。它选择 cpu=True、dtype="float32"、目标 c，并关闭向量化。所有浮点输入和输出都必须是连续的 CPU float32 张量；注意力长度仍为 int32。调用它之前，用 .cpu().float().contiguous() 转换 16 位激活。LayerNorm 仍然原地更新它的残差张量，RoPE 仍然原地更新 Q/K 列。保留编译好的内核，以复用它的动态维度。
CPU 特化把 fragment 和 shared 分配替换为局部缓冲区，使用串行的 T.grid 循环，并省略注意力的 GPU swizzle 标注。GEMM 使用 TileLang 的标量 CPU 实现，归约使用局部缓冲区。原本的分块算法、填充谓词、注意力掩码和在线 softmax 仍与 GPU 实现共用。CPU 中间结果保持 fp32，包括注意力概率；GPU 中间结果保留其原始 dtype。这里不宣称任何 CPU 加速或全模型 CPU 后端。
在 TileLang 0.1.14 上重新探测
观测环境：Linux、Python 3.12.13、torch 2.11.0+cu130、TileLang 0.1.14 和一块 RTX 4070 Ti SUPER。早先的可移植性顾虑针对的是编译未经改动的 GPU 特化，而不是 CPU 降低的可能性。在基线提交 fa9a2a7 上，checkout 里没有这个文档页。
测试用 target="c" 和 GPU 的 FAST pass 配置编译 factory.get_tir(...)。以下是确切的诊断文本（省略了源码位置和堆栈踪迹）。每个内核都探测了 bf16 和 fp16。
| 内核与探测维度 | 首个 bf16 失败 | 首个 fp16 失败 |
|---|---|---|
gemm_kernel(128, 64) |
Check failed: layout_map.count(buffer) != 0 (0 vs. 0) : The layout for fragment C_l can not be inferred correctly. |
同上 |
gemm_geglu_kernel(64, 64) |
CPU fill only supports local and global buffers, but got dst scope `local.fragment`. |
同上 |
add_ln_kernel(128) |
CPU reduce only supports local src and local/local.var dst buffers, got src scope `local.fragment` and dst scope `local.fragment`. |
同上 |
rope_kernel(2, 64) |
Cannot convert type bfloat16 to C type |
C++ 编译器：error: no matching function for call to ‘vec_type<float, 4>::vec_type(half4&)’ |
attn_kernel(1, 64, 2, 64) |
Check failed: layout_map.count(buffer) != 0 (0 vs. 0) : The layout for fragment s_c can not be inferred correctly. |
同上 |
前四个 fragment/归约失败是 tvm.error.InternalError。bf16 代码生成失败也是 InternalError。FP16 RoPE 抛出 RuntimeError: Compilation Failed!，后面跟着编译器调用和源码；上面的诊断文本输出到 stderr。它生成的向量转换在把 float 向量转回 half 向量时也失败了。
为了把 dtype 支持与 fragment 支持分开，测试用 cpu=True（局部缓冲区和串行循环）再次编译每一个内核，保留 bf16 并关闭向量化。此时五个内核都以完全相同的方式失败：
Cannot convert type bfloat16 to C type因此，fragment 是 GEMM、GEGLU 和注意力的第一道阻碍，fragment 归约是 LayerNorm 的第一道阻碍，而 bf16 独立地成为全部五个内核的阻碍。RoPE 没有 fragment 分配或归约；GEMM 和 GEGLU 没有显式的 T.reduce_* 操作。注意力的归约起初被它的布局失败掩盖了。仅针对 fp32 的单独 fragment 探测，在没有任何 GEMM 或 bf16 的情况下，为 reduce_sum 和 reduce_max 都复现了填充错误和归约错误。只替换 scope 也产生了 GEMM 缓冲区的这个语义错误（其它受影响的缓冲区是 Ci、x 和 s）：
[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.串行的 fp32 特化移除了这些阻碍。诊断断言固定在 0.1.14 版本，在其它版本上跳过，那些版本上应当重新探测这些错误。数值 CPU 测试在其它版本上继续运行。
数值校验
运行 python -m pytest tests/test_fast_cpu.py -q -s。在上述环境上：58 通过。仅 CPU 的检查不需要 CUDA；只有 GPU 对比在没有 CUDA 时跳过。覆盖范围包括不均匀的 GEMM M/N/K 分块、所有 GEMM 尾声（epilogue）、GEGLU、所有残差/偏置组合、大的残差值、RoPE 位置回绕与未被触碰的 V 列、静态/动态注意力形状、滑动窗口、部分分块、不相等长度、空序列，以及有限的填充输出。
随机种子是 1234。GPU 对比使用相同的输入值，先舍入到 bf16 或 fp16，再升到 fp32 供 CPU 执行。这些是有界夹具的绝对容差，不构成对任意量级或模型深度的保证。注意力对比使用有效的 query 行，与 tests/test_fast.py 中一样。
| 内核 | CPU 与 fp32 参考的最大差 | CPU 与 GPU 的最大差（两种 dtype） | CPU/GPU 容差 |
|---|---|---|---|
| GEMM | 2.38419e-7 | 0.00770831 | 0.05 |
| GEGLU | 2.98023e-8 | 0.000208303 | 0.05 |
| LayerNorm | 1.07288e-6 | 0.0156183 | 0.05 |
| RoPE | 0 | 0.0130053 | 0.05 |
| Attention | 5.96046e-7 | 0.00377572 | 0.02 |
CPU/参考容差是 2e-5（RoPE 为 2e-6）。残差流更新完全相等，包括无残差的情形。
GPU 保持的证据
cpu 关键字默认为 false。既有的 bf16/fp16 默认值、GPU 分配 scope、并行循环、swizzle 和 FAST 选项均未改变。默认的 fp32 GPU 调用仍然抛出 ValueError。
前后对比运行分别使用了通过 git show fa9a2a7:laya/tl_kernels.py 取出的原始模块和改动后的模块。基线运行时，原始模块通过 importlib 以 laya.tl_kernels 载入；其它 Laya 代码和 Python 环境保持不变。
python -m pytest tests/test_fast.py -q，其中全前向测试被配置为加载下面标识的缓存英文 checkpoint：改动前 13 通过；改动后 13 通过。最初没有 checkpoint 时是 11 通过 / 2 跳过。两次完整运行都输出了该 checkpoint 既有的温度钳制警告（choice:11+=0.10058280825614929 -> 0.5）。- 对既有内核测试做带种子的捕获：24 个输出张量逐位相同，包括残差更新；前后最大差值为 0。
- 上述五个探测形状、两种 dtype 生成的 CUDA 源码：10/10 逐字节相同。这也覆盖了 RoPE，它不在原来的 fast 测试套件里。
python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ...以及等价的 fp16 命令：每个 dtype 288 个问题、60 个 state。所有前后 JSON 记录（fp32、原装和 fast 概率）比较完全相等；前后最大概率差值为 0。
MODEL 是缓存的 convaiinnovations/laya 英文快照 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851；运行时使用 HF_HUB_OFFLINE=1。这里不宣称做过多语言或类型化决策 checkpoint 的对比。
| dtype | type | n | fast 与原装的最大差，前 = 后 | fast/原装的 argmax 一致率，前 = 后 |
|---|---|---|---|---|
| bf16 | choice | 48 | 0.0310 | 47/48 |
| bf16 | noul | 180 | 0.0756 | 180/180 |
| bf16 | score | 60 | 0.0152 | 60/60 |
| fp16 | choice | 48 | 0.0069 | 48/48 |
| fp16 | noul | 180 | 0.0092 | 180/180 |
| fp16 | score | 60 | 0.0040 | 60/60 |
仓库检查
所有命令都使用 /home/ckl/projects/S/laya/.venv/bin/python；ruff 和 zensical 来自那个 virtualenv 的 bin 目录。
| 命令 | 结果 |
|---|---|
ruff check laya/ --select=E9,F63,F7,F82,F401,F811 --line-length=120 |
All checks passed! |
python -m compileall -q laya/ tests/ |
退出码 0，无输出 |
python tests/test_router.py |
703 通过，0 失败 |
python tests/test_criteria.py |
198 通过，0 失败 |
python tests/test_hooks.py |
240 通过，0 失败 |
python tests/test_hooks_api.py |
415 通过，0 失败 |
python tests/test_packaging.py |
131 通过，0 失败 |
uv pip install --python /home/ckl/projects/S/laya/.venv/bin/python -r requirements-docs.txt |
检查了 3 个包（已安装）；该 virtualenv 没有 pip |
zensical build --strict --clean |
No issues found；没有 griffe: 行 |
新测试套件在打包测试既有的豁免项里登记为需要可选的 TileLang extra 和一个 C++ 编译器。API 契约测试固定了这个新增的仅关键字 CPU 选择器、未改变的 GPU dtype 默认值和 compile_cpu 签名，且不把 TileLang 引入基础 CI 环境。
AOTInductor
DecisionModel 可以导出、编译进一个 .pt2 包，并用 torch._inductor.aoti_load_package 加载。这个包同时返回决策 logits 和动作 logits。分词、填充、温度校准和答案格式化仍然由调用方负责；这不会新增 Agent 后端，也不改变它默认的执行路径。
已关闭的 PR #472 记录了早先的 dtype 阻碍。动作头的池化 state 和置信度特征以 fp32 计算，即使模型的权重显式是 bf16。没有 autocast 时，第一个动作线性层因而收到 fp32 输入和 bf16 权重。现在前向把拼接后的输入转成 head 的权重 dtype，只在 autocast 之外。Softmax、熵和返回的决策 logits 保留它们的 fp32 计算。既有的 fp32 参数 eager 推理，包括 fp16/bf16 AMP，保持其数值；混合的权重/autocast dtype 也保留 autocast 原本的转换。
把一份单独的评估副本转成 bf16 导出，在 autocast 之外。把 token 轴声明为 16 * Dim("tokens16", ...)，以满足注意力对齐防护。行数和标记数可以各自独立地动态。不要假设在 autocast 下捕获的导出能在该上下文之外打包：这个配方要使用显式的权重 dtype。
离线复现
基于断言的检查默认使用一个随机初始化的微型 ModernBERT，不访问网络。要测量真实模型就传入一个本地 checkpoint 目录。缺少 CUDA、AOTInductor API 或 C++ 编译器时会产生显式的 SKIP；在受支持的安装上，编译和一致性失败会让检查失败。
python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
python scripts/check_aoti.py --model /path/to/local/multilingual \
--output-dir /tmp/laya-aoti输出目录里有 decision.pt2、results.json、inputs.pt 和 eager.pt。JSON 包含完整的 nvidia-smi 快照、导出/打包/加载耗时、包字节数、延迟，以及两个头的最大绝对 logit/概率增量。二进制产物和编译器缓存被有意排除在仓库之外。
在 RTX 4070 Ti SUPER 上实测
2026-10-03，Linux x86-64，Python 3.12.13，torch 2.11.0+cu130，transformers 5.17.0，NVIDIA 驱动 615.71.09，16,376 MiB 显存。Checkpoint：convaiinnovations/laya，multilingual 子目录，revision 为 1c5edc17a7acd8701df6fc341c0d179f1c62c982。基线是 main fa9a2a7。完整的机器快照和未舍入的测量在 aoti_multilingual_rtx4070.json 里。
| 测量 | 前 | 后 |
|---|---|---|
| 导出 | 5.00 s | 3.89 s |
| 打包 | 16.59 s 后失败 | 60.02 s |
| 包大小 | 没有产物 | 645,729,621 字节（615.82 MiB） |
| 在已初始化的进程中加载 | 不可用 | 0.447 s |
| 在空缓存的崭新进程中加载 | 不可用 | 4.867 s |
复现的失败发生在打包阶段，导出已成功：mat1 and mat2 must have the same dtype, but got Float and BFloat16。每次运行都使用一份单独且初始为空的 Inductor 缓存；AMP 编译在打包之前运行，所以打包耗时不是一个全新解释器冷启动的测量。新进程检查只加载包和已保存的输入（没有 checkpoint），并屏蔽了 torch.compile 和打包入口点。它复现了两个头的答案。PyTorch 仍然把一个小的 CPU AVX 能力探测编译进了最初为空的缓存；没有编译任何模型图或 CUDA 内核。因此，在这套运行时上，「加载时没有任何编译器活动」这种笼统说法是不准确的。
同一批已保存的输入包含跨三个批次的七行决策，带有真实分词后的账单/退款文本、choice/score/noul 问题、填充和不等的有效标记数。每个延迟都是十次预热后五组、每组 30 次前向的中位数，使用同步的墙钟计时。torch.compile(dynamic=True) 使用默认模式。这些延迟数字不包含显式的 CUDA 图、分词、导出或编译。
| 行 × token × 标记槽 | AMP eager 前 → 后（ms） | AMP compiled 前 → 后（ms） | 显式 bf16 eager（ms） | 显式 bf16 compiled（ms） | AOTI bf16（ms） |
|---|---|---|---|---|---|
| 2 × 128 × 3 | 15.1173 → 14.1200 | 6.7117 → 6.9076 | 13.9635 | 4.8835 | 2.2674 |
| 1 × 256 × 3 | 15.4687 → 14.4156 | 6.7274 → 5.4772 | 14.7994 | 4.7250 | 2.4154 |
| 4 × 160 × 5 | 14.8452 → 14.5384 | 7.3872 → 5.5934 | 14.7115 | 5.1393 | 3.6018 |
这里的 AMP 指 fp32 参数配 bf16 autocast。显式 bf16 指权重和残差流都是 bf16、不带 autocast；要跟包做最接近的执行对比就用那几列。显式 bf16 eager/compiled 和 AOTI 在这次修复之前都不可用。
| 对比，取全部七行的最大值 | 决策 logits | 决策概率 | 动作 logits | 动作概率 |
|---|---|---|---|---|
| 既有 eager 前后对比，fp32 与 fp16/bf16 AMP | 0 | 0 | 0 | 0 |
| AOTI 与显式 bf16 eager | 0.5625 | 0.00659859 | 0 | 0 |
| AOTI 与既有 bf16 AMP eager | 0.4375 | 0.00769910 | 8.0 | 0 |
在两个 AOTI 对比里，决策和动作的 argmax 都在 7/7 行上一致。概率是原始 softmax 输出，未经校准。动作分布在这些输入上已经饱和，所以它的概率增量为零并不意味着底层 logits 与 AMP 相同。显式 bf16 compiled 前向也与显式 bf16 eager 不同（最大决策 logit 增量 0.5，概率增量 0.00659859）。降精度编译不是逐位精确的。
GPU 与桌面应用和其它 Python 作业共用；CPU 编译也是共用的，且没有锁定频率。这些是观测到的耗时，不是隔离环境下的加速声明。包住这些运行的 nvidia-smi 快照：
| 运行 / 快照 | GPU 利用率 | 已用显存 | 功耗 | 温度 / 状态 |
|---|---|---|---|---|
| 前 / 开始 | 28% | 2,817 MiB | 12 W | 33°C / P8 |
| 前 / 结束 | 10% | 8,560 MiB | 23 W | 36°C / P3 |
| 后 / 开始 | 0% | 2,634 MiB | 12 W | 33°C / P8 |
| 后 / 结束 | 73% | 6,476 MiB | 160 W | 42°C / P2 |
仍存在的限制
- 该产物特定于这个 checkpoint、精度、PyTorch/运行时栈和 GPU 目标；没有测试到其它硬件或 PyTorch 版本的可移植性。
- 这项检查导出 rows 1–8、tokens 32–512（16 的倍数）、marker slots 2–8。它覆盖三种形状（包括与导出示例不同的形状），而不是这些范围内的每一个点。一个有效选项可能使用带填充的 marker slots；真正的一槽张量走单独的单选项分支，需要单独导出。任意 token 长度和生产环境的分桶/派发层不在这次改动范围内。
- 七行验证的是打包这一回归，而不是广泛的 checkpoint 准确率或校准后的置信度一致性。把导出副本转换成 bf16，与在 autocast 下保留 fp32 权重是不同的；既有 eager 路径本身保持不变。
- 该脚本需要一套支持 CUDA 的 PyTorch 构建和本地编译器工具链才能创建产物。
.pt2内嵌模型和 CUDA 内核；不涉及任何托管服务。