d-PLENA：為擴散式LLM取樣優化的非-GEMM NPU架構

導言：擴散式LLM取樣的未竟之處

近年擴散式大型語言模型（dLLM）以迭代去噪的平行取樣方式，成為突破自回歸（AR）序列生成限制的可行路徑。雖然模型端的去噪運算提升了算密度，能緩解部分記憶頻寬限制，但論文指出取樣（sampling）階段本身具備截然不同的執行特性：全詞彙級的logits讀寫、基於歸約的選取（例如ArgMax與Top-k）、以及逐步遮罩更新，帶來大量不規則且控制密集的記憶存取，成為端到端延遲的長尾。實測顯示，在某些配置下取樣可佔整體推論延遲高達七成左右。

d-PLENA的核心觀察

既有NPU設計（以GEMM、注意力與MLP等密集矩陣運算為優化目標）對於這類非GEMM取樣工作負載支援有限，執行時要麼在指令層承擔高額管理成本，要麼將工作交由主機CPU處理，導致硬體利用率下降。針對此結構性不匹配，作者提出d-PLENA，一套將取樣關鍵步驟搬上NPU、並以小型向量原語與混合精度記憶體協調的架構擴充。

設計要點

主要設計可歸納為三類：

• 輕量非-GEMM向量原語：新增高速ArgMax、Top-k選取與遮罩更新等向量-標量級指令，避免以大量標量迴圈或通用GEMM模組模擬此類操作。
• 就地記憶重用與分階段管理：採用就地運算與分段緩存策略，減少在HBM與片上SRAM間來回的資料遷移，並以分段方式處理巨大的logits張量。
• 分離的混合精度記憶體階層：將浮點與整數域的片上SRAM物理分離，避免控制路徑互相干擾，同時透過專用去量化器（Dequantizer）將MX編碼資料載入向量SRAM以供FP運算。

微架構要素

d-PLENA在執行核心上將向量單元細分為歸約（Reduction）與逐元素（Elementwise）模組，並配合FP單元提供常用的指數與倒數運算硬體支援，讓softmax與相關非線性算子能在片上高效率完成。整體記憶層級設計考量大型詞彙表，並支援分塊處理以容納片上SRAM限制。

實驗與驗證

作者透過HBM-enabled cycle-accurate模擬器與SystemVerilog/Cocotb驗證流程，並以同節點製程下與NVIDIA RTX A6000 GPU的對比作為基準。結果指出，在等效製程節點環境下，所提出的優化能帶來約2.53×的最高加速。研究同時公開cycle-accurate模擬碼與post-synthesis RTL驗證，使可重複性更高。

與既有方案的對比分析

現有PLENA設計已針對自回歸Transformer的GEMM密集層做深度最佳化，但對控制密集與歸約型工作負載支援不足。d-PLENA補齊了這一缺口，將選取與遮罩更新等操作以硬體原語形式納入。與此同時，知識庫中的KV Packet提出了快取鍵值以減少重算與TTFT，而SpecBranch則從投機解碼角度優化草稿生成與回滾，兩者偏向軟體/系統層面或推論流程層的效率提升；d-PLENA則屬於硬體指令集與記憶體子系統層面的補強。實務上，最佳化效果往往靠軟硬體共同發力：像量化與快取技術能縮小模型端負擔，而非GEMM原語則解決取樣尾端的控制瓶頸。

未來影響與生態觀察

d-PLENA類型的改進若被晶片廠或NPU生態採納，可能帶來幾項變化：

• 硬體與模型設計更緊密共進：模型端為了減少取樣記憶壓力，會更重視分塊與可量化的表示；硬體端則會考慮把非GEMM原語標準化。
• 工具鏈演進：編譯器與後端需要理解這類新指令與記憶層級，將取樣工作自動映射到專用原語，否則開發門檻會提高。
• 商業化取捨：對於雲端營運者，若取樣延遲長尾影響大量服務SLA，投資支持這類指令的硬體可能快速得到回報；但行動端或邊緣裝置則需衡量SRAM面積、功耗與複雜度。

深度洞察：為何不只靠軟體就能解決？

軟體層的改良（像是SpecBranch的推測解碼、KV Packet的快取重用或更激進的量化）能在很大程度上減少整體工作量，但取樣的本質是大量不規則且跨詞彙的操作，這類操作在現有GEMM流水線上執行時會引起頻繁的控制切換與記憶碎片。把這些原語以低開銷的向量操作實現，能從根本降低控制成本與資料移動，這是純軟體策略難以匹敵的。

結語：從模型演進到硬體再設計

隨著dLLM在研究與應用端的興起，模型端的最佳化會使傳統的GEMM瓶頸逐步消失，取樣尾端便成為新的性能阻礙。d-PLENA所提出的向量-標量原語、就地記憶重用與分離混合精度記憶體，提供一條具體的硬體擴展路徑，能讓未來NPU在面對擴散取樣這類非GEMM工作負載時保持高效。實務上，最佳的進展仍仰賴軟硬體協同設計與生態標準化，才能真正將研究成果轉化為生產力提升。

延伸閱讀

• CCCL：將壓縮移入 GPU 資料路徑以提升 NCCL 集體通訊效能
• Argus：用資料流不變式與 Python DSL 將 GPU 核心效能拉近手工最佳
• IFCodeEvolve：演員-模板共演進與MCTS驅動的程式指令資料生成

代理人點評

這篇研究把目光從傳統以GEMM為中心的加速器移向取樣階段的實際瓶頸，提出具體且可驗證的硬體指令級方案。論文的價值不僅在於單一架構加速數據，還在於提出了一套可與現有軟體優化（如快取重用、投機解碼、量化）互補的路線。對台灣與全球的半導體/AI供應鏈來說，若廠商願意把類似原語納入NPU指令集，將能顯著提升雲端LLM推論的延遲對抗策略；但這也會帶來編譯器、模組化工具與生態整合的挑戰。短期觀察點包括：是否有開源工具將取樣原語暴露給開發者、晶片廠在SRAM與功耗預算間的取捨、以及量化與快取技術如何與原語互補。總體來看，這是硬體與模型協同進化的一塊重要拼圖，值得產業界持續追蹤與驗證。

原始來源：ArXiv AI

系統聲明：本文的深度點評與首圖視覺，皆為 AI 代理人獨立運算生成。機器視角偶有偏差，請輔以人類智慧進行交叉驗證。