compile=True 與 TileLang 快路徑：工程筆記

這些筆記講 compile=True 和 fast=True 在 README 之外的行為。它們來自在 #472、#576 和 #718 上工作時的實測，環境是 RTX 4070 Ti SUPER、torch 2.11 和 tilelang 0.1.14。放在這裡，是為了 下一個人不用再測一遍。

後端選擇

Agent(..., backend="auto") 和 laya.load(..., backend="auto") 選擇啟用後端類層。預設仍是 eager。顯式的 backend= 優先於 compile 和 fast；省略它就保留這兩個標誌現有的行為。

• eager：原裝的 PyTorch 前向，在任何受支援的裝置上。
• compile：僅 CUDA 的 torch.compile，帶動態形狀和 reduce-overhead 模式、桶填充、持久化 inductor 快取，並在安裝時預熱。它複用與 compile=True 相同的獨立維度作用域，後者保留其現有的預設模式和 CPU 支援。設 LAYA_COMPILE_WARMUP=0 推遲後端預熱，設 LAYA_INDUCTOR_CACHE_DIR 選擇它的快取目錄（預設 ~/.cache/laya/inductor）。
• tilelang：當前快路徑之上的一個介面卡，使用 agent 的 bf16 或 fp16 dtype。
• auto：當安裝了 TileLang、且有受支援的 ModernBERT 編碼器和 dtype 時，在 CUDA 上用 TileLang；否則在 CUDA 上編譯；其它裝置上用 eager。
• onnx：laya.load(..., backend="onnx", onnx_path="model.onnx") 返回現有的 ONNXAgent。沒有 onnx_path 時，它用 laya.onnx。

一個不可用的後端會發出一個 RuntimeWarning，說出解析到的後端，並回退到 eager。要強制要求某個後端，就用 agent.set_backend("tilelang", strict=True)。切換會等待活躍推理；agent.backend 報告活躍的名稱，agent.backend_object 暴露已安裝的物件。agent.set_backend("compile", warmup=False) 把編譯推遲到推理時，所以編譯錯誤那時會在請求上浮現。agent.warmup() 仍然可用。agent.deaccelerate() 移除通過類層安裝的後端。

Router 通過 Router(agent_kwargs={"backend": "auto"}) 轉發顯式的選擇。它們預設不傳後端參數，保持與現有類 Agent 建構函式的相容性。作用域受限的 CPU OOM 重試會分離後端，並在模型回到其原始裝置時恢復它。

compile=True 實體化注意力掩碼

Eager SDPA 把 ModernBERT 的 (rows, 1, L, L) 注意力掩碼當作一個廣播檢視。在 compile=True 使用的動態形狀下，inductor 無法證明最後一維是對齊的。它把掩碼擴充套件到每個頭，並填充進一個真實的 rows x heads x L x L 緩衝區。在 bf16、12 個頭的配置下，32 行 x 1024 token 時就是 0.8 GB。

• GPU 有餘量。 這個緩衝區消耗頻寬，每個長批次幾十毫秒。
• GPU 接近佔滿。 快取分配器會顛簸，同一次呼叫可能要幾十秒。

如果在繁忙的 GPU 上用長批次編譯，就限制批大小（predict_batch(..., batch_size=)) or use fast=True。TileLang 注意力讀取打包的 QKV 緩衝區，並按序列長度做掩碼， 所以它沒有這樣的緩衝區。

冷啟動

• 首次編譯。 每個圖要幾十秒。compile=True 需要兩個圖：一個用於批，一個用於單行，後者由 torch 特化。compile=True 現在在載入期間呼叫 agent.warmup()。compile_warmup=False 恢復惰性 編譯，而 agent.warmup(shapes=...) 仍可手動使用。Eager 和 TileLang 載入不會自動預熱。這些形狀 覆蓋常見的請求，而不是每一種可能的形狀守衛。
• 預熱失敗。 自動預熱是盡力而為：失敗會發出一個 RuntimeWarning，說出錯誤（包括底層編譯器 錯誤），載入帶著完好的 torch.compile 包裝器和編譯設定返回。例如，沒有 MSVC 的 Windows 可以 在預熱失敗的情況下仍以 compile=True 載入。後續請求仍使用編譯後的模型並浮現編譯失敗；Laya 不會把它們切到 eager 執行。顯式的 agent.warmup() 呼叫也會傳播失敗，包括在一次失敗的自動預熱 之後。所以一次成功的載入並不保證編譯後的推理已經就緒。
• Laya 快取的選擇啟用。 laya.load(..., compile=True, compile_cache=True) 只在程序級 TORCHINDUCTOR_CACHE_DIR 缺失時才設定它，設為 $XDG_CACHE_HOME/laya/torchinductor，或在 XDG 未設定或不是絕對路徑時設為 ~/.cache/laya/torchinductor。已有的設定（包括更早的 PyTorch 編譯 所設的那個）優先。該目錄在載入時建立；檔案系統錯誤會傳播。compile_cache=False（預設）、 eager 和 TileLang 載入不動環境。這不會移動或刪除舊快取。容器仍然需要一個持久的 home/卷。快取 相容性與失效由 PyTorch 管理；一次 GPU、torch、編譯器、模型或輸入守衛的變化可能要求重新編譯。
• 跨重啟。 Inductor 的 FX-graph 快取把編譯好的圖放在 TORCHINDUCTOR_CACHE_DIR 下。預設 在 /tmp 下，重啟機器或重啟容器就沒了。把它設成一個持久目錄，或者容器裡的一個卷，第二個程序 就會載入這些圖，而不是重新編譯。在 #472 的實測裡，這把預熱從約 120 秒降到約 50 秒。

可選的 CUDA 圖

agent = laya.load("convaiinnovations/laya", compile=True,
                  compile_cache=True, compile_mode="reduce-overhead")

compile_mode 預設是 "default"；在活躍的編譯路徑上只接受 "default" 和 "reduce-overhead"。Eager 和 TileLang 載入忽略編譯選項。CPU 編譯仍然可用，但 CUDA 圖記錄只適用於 CUDA。CUDA 模式需要 PyTorch 的 torch.compiler.cudagraph_mark_step_begin API；沒有它的更舊構建會丟擲一個明確的錯誤。

動態的 Dynamo 圖並不意味著形狀無關的 CUDA 圖：新的具體形狀可能需要再次預熱和記錄，而不需要新的 Dynamo 圖。兩個預設的合成預熱形狀不會預記錄每一種請求形狀。重複的形狀可以受益，但變化的形狀可能要付出額外的延遲並保留圖池。PyTorch 可能對不受支援的操作或配置跳過 CUDA 圖；設定這個模式並不保證捕獲。

Laya 把每次編譯後的 CUDA 前向標記為一個新的步驟，在它的各個 agent 之間把這些前向序列化，並在釋放鎖之前把兩個輸出張量克隆到編譯圖之外。這使被保留的輸出在重放之間保持有效，代價是兩次複製和序列化的前向執行。這個鎖不協調無關的、應用自己擁有的編譯模型；共享 CUDA 圖迭代或使用自定義流的呼叫方必須管理它們自己的協調。磁碟快取跨程序複用編譯後的程式碼，而不是即時的 CUDA 圖記錄或它們的裝置記憶體。

用 benchmarks/bench_compile_defaults.py 復現冷/重啟計時、記憶體和快取計數；見記錄的測量。

AOTInductor：還不行

為每個 checkpoint 和 GPU 架構預先編譯一個產物（torch._inductor.aoti_compile_and_package）， 就能完全省掉編譯。在 torch 2.11 上，它卡在打包這一步：

• 匯出沒問題。 對 DecisionModel 做 torch.export 能成功，約 5 秒，帶動態的 rows、 markers 和 tokens。tokens 必須宣告成 16 的倍數（16 * Dim(...)）；普通的 range 過不了 匯出器自己的 L % 8 對齊檢查。這跟上面是同一個掩碼對齊問題。
• 打包失敗。 怎麼失敗，取決於程式是怎麼匯出的：
  • 在 autocast 下，程式帶著 dtype 斷言，AOTI 在 autocast 之外會踩中它： Tensor dtype mismatch! Expected: torch.bfloat16, Got: torch.float32。
  • 從不帶 autocast 的 bf16 副本出發，trace 會在前向裡失敗： mat1 and mat2 must have the same dtype。DecisionModel.forward 在 action head 之前 把池化後的 state 和置信度特徵升到 fp32，而 autocast 通常會把這個協調好。

所以產物這條路需要一個 dtype 明確的 action head：要麼把它的輸入轉成 head 的 dtype，要麼讓 head 跑 fp32。

TileLang 的可移植性：核心只支援 CUDA

tilelang 為 CUDA、HIP、Metal、WebGPU 和一個 C 後端註冊了 target。沒有 AMD 或 Apple 硬體時， 還能回答的問題是：laya/tl_kernels.py 到底能不能為 CPU 降低。在 Linux x86-64 上用 tilelang.compile(kernel.prim_func, target=...) 探測：

每個 Laya 核心在 "c" 上都失敗，原因屬於這三類之一：

C 後端確實接受：

• fp32 逐元素迴圈（T.Parallel）；
• T.Pipelined，它會降低成一個普通迴圈；
• 帶 T.alloc_local 累加器的 T.gemm，它會降低成一個標量三重迴圈。

所以一個 CPU 版本會是第二套核心，而不是一個 target 開關。它的 GEMM 會是不分塊的標量迴圈， 也就沒法跟原裝前向在 CPU 上已經在用的 MKL/oneDNN 路徑競爭。這三樣構造也是在 HIP 和 Metal 上 首先要檢查的：fragment、帶 GemmWarpPolicy 的 T.gemm，以及 bf16/fp16 支援。

要復現第二張表的第一行：tilelang.compile(K.gemm_kernel(768, 768).prim_func, target="c")。