在 AWS Trainium 上用 GPSIMD 自定义算子加速 DeepSeek-V4.1-Flash Engram 模块
Exploiting Native Hardware Heterogeneity: Custom GPSIMD Operators for Accelerating Emerging LLM Memory on AWS Trainium
把 DeepSeek-V4.1-Flash 的 Engram 记忆模块搬到 Trainium 3 上跑，作者自己写 C++ 算子塞进 GPSIMD 核心，n-gram 哈希快了 1.79 倍，表查找快近 5 倍。在 AWS 上部署记忆增强 LLM 的工程团队可以看看这套绕过官方工具链限制的做法。
一项新研究针对 DeepSeek-V4.1-Flash 中的 Engram 记忆模块，在 AWS Trainium 3 上构建了一套工具链，将 C++ 自定义算子编译到 NeuronCore 的 GPSIMD Engine 并注入 NKI 内核编译生成的 NEFF。此前官方 GPSIMD 自定义算子工具链不支持 Trainium 2 或 3，也无法与 NKI 配合使用，该工作补上了这一空缺。
Engram 这类新兴记忆模块混合了不规则整数运算（64 位 n-gram 哈希和稀疏表查找）与密集计算（GEMM 和门控）。在 GPU-FPGA 系统上通过将操作拆分到不同设备，Engram 可获得 5.9 倍加速；Trainium 的 NeuronCore 原生具备类似的异构结构，密集引擎与可定制 GPSIMD Engine 共享片上 SBUF。基于手写优化的 NKI 基线，两个 GPSIMD 算子使 decode 阶段 n-gram 哈希加速 1.79 倍，表检索最高加速 4.95 倍，整个 Engram 计算最高加速 1.14 倍。
对在 AWS 上部署带有记忆模块的新一代 LLM 的团队来说，这意味着不必依赖 GPU-FPGA 异构拆分方案，单芯片即可挖掘专用加速空间。
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$.
本条情报归属 AITOP「论文」核心栏目，重点聚焦北美前沿实验室与硅谷科技巨头。在跨区域产业观察坐标中，该动态直接关联「AI 开发者工具链、Coding Agent 与工程化实践」研究主线，反映出当前技术演进与商业落地的最新风向。涉及核心实体与技术议题包括：DeepSeek-V4.1-Flash、AWS Trainium、NKI、推理优化，读者可通过下方标签进入对应专题页查阅历史脉络与关联对比。
信源核验说明：本文原始信息发布自一手权威信源「arXiv: DeepSeek」（Tier-1 一手发布），具备高确定性参考价值。
