快速後端

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。

前四個 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/參考容差是 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 的對比。

倉庫檢查

所有命令都使用 /home/ckl/projects/S/laya/.venv/bin/python；ruff 和 zensical 來自那個 virtualenv 的 bin 目錄。

新測試套件在打包測試既有的豁免項裡登記為需要可選的 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 裡。

復現的失敗發生在打包階段，匯出已成功：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 圖、分詞、匯出或編譯。

這裡的 AMP 指 fp32 參數配 bf16 autocast。顯式 bf16 指權重和殘差流都是 bf16、不帶 autocast；要跟包做最接近的執行對比就用那幾列。顯式 bf16 eager/compiled 和 AOTI 在這次修復之前都不可用。

在兩個 AOTI 對比裡，決策和動作的 argmax 都在 7/7 行上一致。機率是原始 softmax 輸出，未經校準。動作分佈在這些輸入上已經飽和，所以它的機率增量為零並不意味著底層 logits 與 AMP 相同。顯式 bf16 compiled 前向也與顯式 bf16 eager 不同（最大決策 logit 增量 0.5，機率增量 0.00659859）。降精度編譯不是逐位精確的。

GPU 與桌面應用和其它 Python 作業共用；CPU 編譯也是共用的，且沒有鎖定頻率。這些是觀測到的耗時，不是隔離環境下的加速宣告。包住這些執行的 nvidia-smi 快照：

仍存在的限制

• 該產物特定於這個 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 核心；不涉及任何託管服務。