← 研踪 PaperTrailarXiv:2611.11820 · v1
Exploiting Native Hardware Heterogeneity: Custom GPSIMD Operators for Accelerating Emerging LLM Memory on AWS Trainium
首次提交 2026-11-09 11:28
版本更新 2026-11-09 11:28
雷达观测 2026-11-10 13:18
摘要
Emerging memory modules of large language models, such as Engram in DeepSeek-V4.1-Flash, mix irregular integer work (a 64-bit n-gram hash and a sparse table lookup) with dense computation (a GEMM and a gate). When deploying on a GPU-FPGA system, we can disaggregate operations to different devices for Engram to speed up by 5.9$\times$. AWS Trainium provides a similar heterogeneity natively: each Trainium NeuronCore pairs dense engines with a customizable GPSIMD Engine, and all engines share the on-chip SBUF. However, the Neuron Kernel Interface (NKI) exposes the GPSIMD cores only through fixed built-in instructions, and the official GPSIMD custom-operator toolchain does not support Trainium 2 or 3 and cannot be used with NKI. In this work, we build a toolchain that compiles C++ custom operators for the GPSIMD cores and injects them into the compiled NEFF from the NKI kernel on Trainium 3, expanding the customizability of GPSIMD with the current software stack. Over a hand-optimized NKI baseline of DeepSeek-V4.1-Flash, two GPSIMD operators speed up the n-gram hashing by 1.79$\times$ at decode, the table retrieval by up to 4.95$\times$, and the whole Engram computation by up to 1.14$\times$.
作者备注：https://github.com/OswaldHe/neuron-gpsimd-tool
观测到的版本
相邻记录可能跳过中间版本；首次发现 v2 及以上也会标为版本更新。
v1 · 2026-11-09 11:28 · 最新观测版本
Exploiting Native Hardware Heterogeneity: Custom GPSIMD Operators for Accelerating Emerging LLM Memory on AWS Trainium
Emerging memory modules of large language models, such as Engram in DeepSeek-V4.1-Flash, mix irregular integer work (a 64-bit n-gram hash and a sparse table lookup) with dense computation (a GEMM and a gate). When deploying on a GPU-FPGA system, we can disaggregate operations to different devices for Engram to speed up by 5.9$\times$. AWS Trainium provides a similar heterogeneity natively: each Trainium NeuronCore pairs dense engines with a customizable GPSIMD Engine, and all engines share the on-chip SBUF. However, the Neuron Kernel Interface (NKI) exposes the GPSIMD cores only through fixed built-in instructions, and the official GPSIMD custom-operator toolchain does not support Trainium 2 or 3 and cannot be used with NKI. In this work, we build a toolchain that compiles C++ custom operators for the GPSIMD cores and injects them into the compiled NEFF from the NKI kernel on Trainium 3, expanding the customizability of GPSIMD with the current software stack. Over a hand-optimized NKI baseline of DeepSeek-V4.1-Flash, two GPSIMD operators speed up the n-gram hashing by 1.79$\times$ at decode, the table retrieval by up to 4.95$\times$, and the whole Engram computation by up to 1.14$\times$.
查看此版本解读 →查看此版本原文 ↗