北京
三個模型規模，三種讓官方工具鏈卡殼的架構沖突，以及——三個完全不同的“混合頭”定義。這是我們將谷歌Gemma-4家族（2B、4B、12B）搬上AWS Inferentia2加速器時踩過最深的坑。
事情的開頭很直接：Gemma-4不是普通的解碼器模型。它在TPU和XLA編譯路徑上如魚得水，但一組為TPU定制的特性，個體性能拉滿，撞上AWS的推理棧后卻逐個崩潰。逐層嵌入、跨層KV共享、分組查詢注意力搭配極少的KV頭數、滑動窗口與全局注意力交織混合、還有對26.2萬詞表的logit軟截斷——每一條單獨拎出來都是優雅的工程優化，放到一起卻讓面向靜態權重映射的neuronx-distributed完全失去了方向。
官方提供的optimum-neuron加Neuron vLLM后端根本找不到Gemma-4的模型類，它的圖構建器也表達不了“某層直接復用另一層的KV緩存”這種動態依賴。我們第一次啟動時，它確實載入了模型，輸出的文本流暢通順，但全是胡言亂語。從那句亂碼開始，我們進入了長達數周的排雷之旅。
排雷的核心只有一個詞：混合頭。但這個詞在不同大小的模型里完全是三副面孔。在E2B和E4B這兩個變體上，問題出在跨層KV共享。模型中有許多層根本不自己計算鍵和值，它們直接從“宿主層”讀取KV張量。TPU把它看作同一個計算圖的自然視圖，但NxD的ModelBuilder要求每層權重到緩沖區的靜態映射，壓根沒有地方聲明“那一層的K和V就是這一層的K和V”。
我們放棄了和框架較勁，轉而直接對HuggingFace的前向過程做圖追蹤。KV共享變成了被追蹤圖中一個活生生的數據依賴關系。只需要提前枚舉出真正寫入緩存的層（在E2B上只有15層），讓這些層分配KV緩沖區，其余層直接讀取即可。這個改動讓KV緩存大小遠小于“一層一個”的樸素方案，也徹底掃清了E2B和E4B的推理歧途。
12B尺寸上，混合頭的內涵完全不一樣——它指的是同一模型內交錯使用滑動窗口注意力和全局注意力，且兩種注意力類型的KV頭數還不一致。編譯器neuronx-cc在處理這種動態變化時觸及了內部限制。我們通過調整圖分割策略，將不同注意力模式的子圖分別編譯，再在運行時拼接，才繞過了這一障礙。最終，三個模型大小都能在Inferentia2上輸出連貫、正確的結果。
這個過程讓我們看清了一件事：硬件加速的生態差距，往往不體現在算力峰值，而體現在對異構模型結構的容納能力。Gemma-4在TPU上的先天自由，到了其他平臺就變成需要硬開的鎖。而鎖的鑰匙，就藏在那些“共享”“混合”“動態”的細節里。
特別聲明：以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布，本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.
