前面幾天的 SGLang RVV 後端走手寫路線：自己決定 operator boundary，用 C++ / RVV intrinsic 寫 kernel，再透過 torch.ops.sgl_kernel.* 接回 Python。今天開始看讓 torch.compile 從一般 PyTorch 程式擷取計算圖，再由 TorchInductor 產生 CPU kernel。
今天我主要參考 ASPLOS 2024 的 PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation。今天的文章可以配合論文一起看，然後論文描述 PyTorch 2 初期的設計與 2023 年 nightly build 實驗，後來版本的細部 heuristic 和 API 可能已調整，這裡把它當成理解 compiler architecture 的主要來源。
linear + relu 開始linear + relu 開始import torch
def f(x, w, b):
return torch.nn.functional.linear(x, w, b).relu()
opt_f = torch.compile(f, backend="inductor")
x = torch.randn(64, 2048)
w = torch.randn(3072, 2048)
b = torch.randn(3072)
y = opt_f(x, w, b)
torch.compile 傳回一個包裝後的 callable，第一次執行 opt_f(x, w, b) 時，compiler stack 才根據當下的 Python 路徑與 tensor metadata 開始 graph capture 和 JIT compilation。
在 CPU 情境下，這次呼叫可以分成三段
Capture
Python bytecode -> FX graph + guards
Compile
AOTAutograd -> decompositions -> Inductor IR -> scheduling/codegen
Run
C++ compiler -> .so -> Python wrapper 載入並呼叫 kernel
後續呼叫會先檢查 guards。條件符合時可以重用已編譯的 artifact；條件不符時，Dynamo 會嘗試使用其他 cache entry，或重新 capture 與 compile。
PyTorch eager mode 採用 define-by-run，Python 程式執行到一個 tensor operation，dispatcher 就送出對應的 operator。這個模式保留了 Python 的 control flow、class、list、dictionary、logging 與第三方套件，也能直接用 print 或 pdb 除錯。
compiler 只看到單一 operator 時，很難做跨 operator 的最佳化。例如 linear 後面緊接 relu，eager mode 可能產生中間 tensor，再啟動下一個 kernel。看到整段 graph 的 compiler 才有機會融合 epilogue、減少中間 buffer，或重排 loop 與 memory access。
早期 graph capture 方法常要求使用者改寫程式，或只記錄某次 execution path。論文用 torch.jit.trace、torch.jit.script、Lazy Tensor 與 torch.fx.symbolic_trace 說明這個困難：PyTorch 現實中的模型已經大量使用 Python 語意，graph capture 必須處理這些程式，不能只限定使用者寫成純 graph language。
TorchDynamo 的做法是在 CPython 執行 bytecode 之前分析它，擷取可編譯的 PyTorch operation sequence，然後把 compiled callable 寫回新的 Python bytecode。程式可以同時包含 compiled graph 和原本的 Python execution。
上半圖的 graph capture 與 lowering 是共用路徑。到 backend 才會根據 device、operator、shape、dtype 與 layout 建立可用實作。CPU 路徑主要產生 C++ / OpenMP。GPU 的 pointwise 與 reduction 常由 Triton codegen 處理，GEMM 類 operator 還可以保留 ATen external kernel，並在啟用 max-autotune 後與 Triton 或其他可用 template 比較，選擇最快的有效的實作，之後 wrapper 呼叫選中的 kernel 或 external library。
TorchDynamo 是 Python-level JIT compiler，它透過 PEP 523 提供的 CPython frame evaluation hook，在 function frame 執行前介入。
Dynamo 的 symbolic evaluator 會逐條讀取 Python bytecode，同時維護
遇到 tensor operation 時，Dynamo 在 FX graph 加入 node，遇到 Python constant、list、dictionary、module attribute 或 user-defined function 時，它會用對應的 symbolic representation 追蹤。能夠安全 inline 的 function 會被展開，依賴 tensor shape 或 Python type 的 branch 可能被特化成當前路徑。
完成分析後，Dynamo 會產生新的 Python bytecode，新 bytecode 會呼叫 compiled FX graph，恢復 local / stack state，執行必要的 side effects，最後回到原本 Python control flow。
JIT compiler 會根據擷取當下看到的資訊做 specialization
x.dtype == torch.float32
x.device == cpu
x.size() == (64, 2048)
x.stride() == (2048, 1)
module.training == False
Guards 將這些假設轉成 runtime checks。每個 compiled cache entry 都會搭配 guard function；下次呼叫時，guard 通過才能使用對應的轉換後 bytecode 與 compiled artifact。
Guards 不只來自 TorchDynamo。AOTAutograd 與 TorchInductor 也可以為自己的 specialization 加入條件。所以重新編譯的原因可能來自 shape、stride、dtype、Python value、module state，或 backend 做的 codegen 選擇。
Dynamo 遇到無法安全處理的 operation 時，會在當下位置產生 graph break。常見來源包含第三方 Python library、tensor 轉成 Python value，以及依賴 tensor data 的 control flow。
發生 graph break 時，Dynamo 會先編譯已經擷取的 partial FX graph，將無法擷取的 bytecode 交回 CPython，然後透過 continuation function 從中斷點後面繼續分析。論文用 resume_at_X 表示這類 continuation，它會接收跨過 graph break 仍然活著的 variables，恢復 stack / exception state，再跳回原本程式位置。
compiled graph 0
-> CPython 執行不支援的運算
-> resume_at_X(...live variables...)
-> compiled graph 1
這個機制讓 partial graph 仍然可用，Graph break 太多時，kernel fusion 的範圍會變小，Python 與 dispatcher overhead 也會回來，因此它同時是功能性邊界與效能診斷訊號。
TorchDynamo 擷取的 tensor operation 會用 FX GraphModule 表示，對前面的例子，graph 大致可能包含
placeholder x
placeholder w
placeholder b
call_function aten.linear
call_function aten.relu
output
FX Graph 是 Dynamo 與 backend compiler 之間的一個介面，還沒有完成硬體特定scheduling。後續 pass 會繼續做 decomposition、functionalization、shape propagation、lowering 與 fusion
backend= 參數也表示 Dynamo 和 Inductor 是可分開的組件，Dynamo 負責 capture 和重寫 Python，預設 backend 是 Inductor，使用者也可以提供其他 backend 來接收 FX graph。
TorchDynamo 先擷取 model forward，訓練時還需要 backward graph，AOTAutograd 會用 fake tensor inputs 執行 eager autograd engine，記錄 joint forward / backward graph，再用 min-cut algorithm 切成獨立的 forward 和 backward graph。這個切分還會考慮 activation 要保留還是在 backward 重新計算，藉此調整記憶體使用。
這裡的 fake tensor 只保留 shape、dtype、device、stride 等 metadata，不需要配置實際tensor data。compiler 可以用這些 metadata 做 shape propagation 和 operator analysis。依賴實際資料的 operation 無法只靠 fake tensor 決定結果，可能使 Dynamo 切開 graph。
AOTAutograd 還會做兩類對 backend 很重要的轉換
這個系列接下來主要看 inference，但 AOTAutograd 仍然會參與 decomposition 和 functionalization。
TorchInductor 是 torch.compile 的預設 backend。論文強調它的幾個設計選擇，用與 PyTorch eager 接近的 tensor / storage abstraction 處理 stride、view、alias 和 mutation，compiler 主體用 Python 實作，TorchInductor 也是很多後端，根據後端產生程式碼，像是 GPU Triton 和 CPU C++ / OpenMP。
Inductor 不需要為每個 PyTorch operator 寫一份專用 codegen，一部分 operator 會先通過 decomposition 展開，再將簡化後的 FX node lowering 到 Inductor IR。如果遇到沒有 lowering 的 operator，Inductor 可以產生 fallback kernel node，呼叫原本 PyTorch implementation。
論文將 Inductor IR 稱為 define-by-run loop-level IR。Loop body 由可執行的 Python function 定義，並透過一組小型 primitive operations 表示
ops.load
ops.store
ops.reduction
ops.index_expr
ops.indirect_indexing
ops.masked
Tensor shape 和 stride 可以用 SymPy symbol 表示，loop body 只需描述某一個 index 上應該執行的運算。同一份 IR 可以切換 analysis handler 來收集 memory access，也可以切換 codegen handler 產生 Triton 或 C++。
Scheduler 會把 IR buffer 轉成 scheduler nodes，建立 memory read / write dependency，然後決定
Fusion 先要通過 dependency 與 indexing 檢查。寫入與讀取順序不相容時，把兩個 loop 合在一起可能改變結果。通過合法性檢查後，Inductor 再根據 fusion category、預估可節省的 memory traffic，以及 nodes 在原 graph 的距離排序候選項目。
這段解釋了 compiler 為什麼需要 graph。將 pointwise、reduction 與 scatter 合併成較少 kernels，可以讓中間值留在 register 或 cache 附近，減少寫回主記憶體後再讀取的次數。
CPU backend 會產生 C++ / OpenMP，分成 vectorized 與 non-vectorized 兩類路徑。
Vectorized 路徑會做 tiling，並把多數 operation 對應到 PyTorch at::vec::Vectorized abstraction。這個 abstraction 能在不同 CPU SIMD ISA 上提供共用介面。Non-vectorized 路徑則產生比較標準的 C++ / STL code。兩條路徑都可以用 #pragma omp for 平行化 loop，reduction 則根據 loop 是否平行來選擇 OpenMP reduction 或 C++ accumulator。
一個大幅簡化後的 pointwise kernel 可能長這樣：
extern "C" void kernel(const float* x, float* y, int64_t n) {
#pragma omp parallel for
for (int64_t i = 0; i < n; ++i) {
y[i] = x[i] > 0.0f ? x[i] : 0.0f;
}
}
實際 generated code 還要處理 tensor pointer、size、stride、dtype、temporary buffer、parallel scheduling 與 external kernel call。Wrapper codegen 會負責 tensor size calculation、memory allocation / deallocation，並呼叫 C++、Triton 或 external source 的 kernels。所以一個 compiled graph 可以同時包含 generated kernel 和 library call。
CPU C++ source 會交給 C++ compiler，再編譯成 Linux shared object（.so）。GCC 或 Clang 的選擇受 PyTorch / Inductor 設定、環境變數與目標平台影響，常見來源包含
CXX / CC
torch._inductor.config.cpp.cxx
torch._inductor.config.cpp.march
g++ 或 clang++
對 x86 或 Arm CPU，target flags 會決定 compiler 可以使用的 SIMD ISA。RISC-V 情境下，rv64gc 不包含 V extension，開啟 RVV 需要對應的 -march 與 CPU capability。Generated C++ 使用 RVV intrinsic 時，還要確認 target compiler 版本能正確編譯。
我目前的 Banana Pi 實驗將 Inductor JIT compiler 固定為 GCC/G++ 15。這是本實驗環境的工具鏈選擇，主要用來避開早期 compiler 的 RVV intrinsic 相容性問題，不代表所有 RISC-V CPU 必須使用同一個 compiler 版本。
LLM 推論會讓 batch size、sequence length 與 KV cache 狀態持續變化。如果 compiler 只支援 static shape，每種 shape combination 都可能需要新 kernel，編譯成本與 cache artifacts 會迅速增加。
論文的 dynamic shape 設計使用 symbolic integers 表示 size 和 stride。每個 symbolic size 同時保留第一次 capture 看到的 concrete hint。遇到 shape-dependent branch 時，Dynamo 用 hint 選擇當前路徑，再產生 guard 保護這個 specialization，中間 tensor 的 shape 則透過 meta function 傳播，不需要執行真實運算。
例如
def f(x, y):
z = torch.cat([x, y])
if z.size(0) > 2:
return z * 2
return z + 2
compiler 會根據當次 input 選擇其中一條 branch，並把 z.size(0) 改寫成與 input shapes 有關的 symbolic expression。當後續 input 無法滿足 guard，就需要使用其他 cache entry 或重新編譯。
這是 serving 環境要檢查 recompilation 的原因。dynamic=True 不等於一份 kernel 自動涵蓋所有控制路徑；symbolic shape 可以擴大 artifact 可重用的範圍，仍然受 guards、data-dependent control flow 與 backend codegen capability 限制。而且 dynamic shape 推理本身也有成本，精確 heuristic 會隨 PyTorch 版本調整。
torch.compile 可以減少 Python 與 dispatcher overhead，但論文的消融實驗指向更大的來源，Inductor 的 inlining 與 fusion 將 pointwise、reduction 與 scatter 組合成更少 kernels，藉此減少 memory traffic。論文關閉 fusion 與 inlining 後，整體結果會落到 eager 以下。所以只看 Python 呼叫次數，無法解釋 compiler 的全部收益。
論文實驗在 NVIDIA A100 上 180 多個模型的實驗，有 2.27× inference 與 1.41× training geometric mean 加速
CPU 的第一次執行可能包含
Dynamo bytecode analysis
-> guards / graph breaks
-> AOTAutograd + decompositions
-> Inductor lowering + scheduling
-> C++ source generation
-> GCC / Clang compile + link
-> .so load
-> first kernel execution
後續執行在 guards 通過時才能重用 compiled artifact。因此 benchmark 至少要分開
LLM serving 還要記錄重新編譯次數、cache 命中率、每種 batch / sequence shape 的 artifact 數量，以及 weight prepack 所增加的時間與記憶體。
讀完流程後，可以先看 graph break、guards 與 recompilation
TORCH_LOGS="graph_breaks" python example.py
TORCH_LOGS="guards" python example.py
TORCH_LOGS="recompiles" python example.py
graph_breaks 幫助找到 Python 和 compiled graph 的分界；guards 列出 artifact 的使用條件，recompiles 會顯示哪個 guard failure 觸發重新編譯。這三組訊息能把第一次很慢或shape 變化後又 compile轉成可查驗的原因。
torch.compile 和 SGLang RVV 的關係前面幾篇的 SGLang RVV backend 是手寫 operator 路線
SGLang Python
-> torch.ops.sgl_kernel.*
-> 手寫 C++ / RVV intrinsic kernel
PyTorch Inductor RVV 是 compiler 路線
SGLang / PyTorch model code
-> torch.compile
-> FX / ATen graph
-> TorchInductor CPU backend
-> generated C++
-> RVV micro-kernel
手寫路線要自己維護 operator boundary、tensor layout、custom op schema、dispatch 與 kernel。
Compiler 路線則要處理 graph capture、guards、fallback、ISA selection、template routing、weight prepack 與 artifact cache。
對目前的 PyTorch Inductor RVV 實驗，需要追蹤的位置
Inductor CPU target selection
-> RVV capability / compiler flags
-> C++ template routing
-> BF16 GEMV / GEMM micro-kernel
-> generated C++ artifact
-> SGLang workload
明天會進入 PyTorch Inductor RVV，具體看 VecRVV 如何被選到、RVV BF16 micro-GEMM 如何產生，以及 blocked weight layout 為什麼會影響 SGLang 的 prefill 與 decode。
iThome鐵人賽
看影片追技術