AWS 自家嘅 Machine Learning Blog 刊出一篇技術文章，講 AI 研究公司 i14 嘅 Laplace 3B 點樣喺 AWS Trainium2 晶片上面做訓練同推理。呢個係廠商自家文章，以下數據全部來自 AWS／i14 嘅聲稱，未見獨立第三方實測。

KV cache 點解會成問題

主流大語言模型（LLM）用 transformer 架構，配 softmax attention。喺推理階段，模型處理愈多 token，key-value（KV）cache 就越大。呢個會拉低並發、推高服務成本，仲會逼到 agent 類型嘅應用要丟棄有用嘅歷史。

i14 講，佢哋設計 Laplace 3B 就係做一個 linear attention 嘅替代方案：recurrent state（遞迴狀態）固定喺每條序列 29.5 MB，同 context 長短無關係。

Linear attention 做咗乜

傳統 softmax attention 喺訓練時計算量隨序列長度呈二次方（O(T²)）增長。Linear attention 改圍住一個每個 token 都更新嘅 recurrent state 來重組計算。Laplace 3B 用 GLA（Gated Linear Attention，一種門控 linear attention）將先前 context 壓縮進一個固定大小嘅矩陣狀態，令訓練 attention 由 O(T²) 變 O(T)，推理狀態維持喺每序列 29.5 MB。

原文提醒：linear attention 係壓縮歷史，唔係保留每個 key 同 value，因此各 workload 都要自評檢索品質。固定狀態對長程 agent、大文件分析同需要持久記憶嘅應用可能有幫助。

點樣改到 Trainium2 上面跑

Laplace 3B 原本用 GPU-optimised 嘅 Triton kernel。AWS Generative AI Innovation Center（GenAIIC）團隊先用標準 PyTorch operators 重寫 GLA kernel，再同 TorchNeuron Native（AWS Trainium 嘅原生 PyTorch backend）比較，最後揀後者，因為效能較好。實現上改咗三個位：

  • GLA attention kernel 用原生 PyTorch，取代 GPU 專屬嘅 fused_chunk_simple_gla Triton kernel，用 16 個 chunk（各 128 tokens）去對齊 NeuronCore 片上 buffer 嘅 128 partitions。
  • 重新設計 Muon optimizer 嘅通訊，兩個 all_to_all collective 將每 rank 嘅通訊量減咗 24 倍，throughput 升咗 1.63 倍。
  • 將 sequential backward pass 編譯，將 GLA recurrence 嘅 16 chunk 步驟靜態展開（static unroll），去貼合編譯器嘅圖大小同片上記憶限制。

數字（全部係廠商聲稱）

喺單一 trn2.48xlarge 實例上：

  • 訓練：約 120,000 tokens 每秒
  • 推理：11,544 輸出 tokens 每秒，跨 1,280 條並發序列

阿調搞兩下

呢篇文章值得睇，因為佢同時示範兩樣嘢：新型 linear attention 架構點樣將長 context 嘅狀態成本由「隨長度無限漲」變成「固定值」，以及 AWS 自家 Trainium2 晶片喺訓練同推理全鏈嘅角色。但要留意：成個數嘅效能數字都出自 AWS 自家文章，TorchNeuron Native 仲係 beta 版，i14 嘅 linear attention 檢索品質有冇獨立驗證過，原文自己都要你用 workload 自評。廠商示範唔等於本站實測，呢點先係關鍵。

阿調｜深度調教