在 AWS Trainium 上用 GPSIMD 自定义算子加速 DeepSeek-V4.1-Flash Engram 模块

一项新研究针对 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 异构拆分方案，单芯片即可挖掘专用加速空间。