所有標籤Google 的 Srinath Mandalapu 用兩講、兩百多頁,把一個 attention 從 Python 一路追到 TPU 的 VLIW 指令。L12 講 JAX 生態、TPU Ironwood 的記憶體與運算單元,以及 XLA 怎麼把 attention 編成三個融合 kernel。L13 講 XLA 做不到的部分:用 Pallas 自己控制 HBM 與 VMEM 之間的搬運,寫出 FlashAttention,再加上區塊稀疏變成 Splash Attention。思路跟 GPU 相同,差別在於:TPU 上排程主要交給編譯器,要介入就得用 Pallas 把迴圈與區塊大小攬回自己手上。
L13 問的是:寫快程式的專家太少,怎麼辦?投影片給三個答案。一是提高抽象層級:Halide 把「算什麼」(演算法)和「怎麼算」(排程)拆成兩種語言,同一段模糊濾波可以只改一行排程就換成分塊、向量化、多核版本。二是智慧搜尋:因為排程空間定義清楚,可以用搜尋加學出來的成本模型自動產生排程。三是新興的 LLM agent:讓模型寫 CUDA、執行、看 profiler、反思、再改,並用範例資料庫或 prompt 最佳化讓它自我改進。投影片最後把問題留給你:真正的價值在 DSL 設計,還是 LLM agent?
L11 問的是:硬體為 AI 專用化之後,程式設計師要付出什麼?投影片以 H100 為例:要吃滿 Tensor Core,就得用 16×16 tile、讓 TMA 非同步搬資料、讓 producer 與 consumer warp 管線化,寫起來很複雜,所以有了 ThunderKittens 這類 DSL。另一條路是資料流架構(SambaNova SN40L):用 map/reduce/zip 等平行模式描述計算,編譯器做分塊、metapipelining 與佈局,投影片說它能把 Llama 3.1 8B 的整個 decoder 融成一個 kernel。
CS107 第 14 講用 sum_array 的 10 條 x86-64 指令拆開反組譯輸出:左側是位址與機器碼,右側是 AT&T assembly;讀者的任務是由 opcode、operand、register 與控制流程還原 C,而不是手寫組合語言。
TypeScript 7 把 compiler 與 language service 原生重寫為 Go,官方與大型使用者回報約 10 倍級的改善;收益來自共享記憶體平行化,代價是舊 API 與 compiler option 必須清理。