Exploiting Native Hardware Heterogeneity: Custom GPSIMD Operators for Accelerating Emerging LLM Memory on AWS Trainium

arXiv:2611.11820v1 [cs.DC] 9 Nov 2026

Exploiting Native Hardware Heterogeneity: Custom GPSIMD Operators for Accelerating Emerging LLM Memory on AWS Trainium

Zifan He, Jason Cong University of California, Los Angeles zifanhe1202@g.ucla.edu, cong@cs.ucla.edu

Abstract Emerging memory modules of large language models, such as Engram in DeepSeekV4.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 GPUFPGA system, we can disaggregate operations to different devices for Engram to speedup by 5.9×. AWS Trainium provides a similar heterogeneity natively: each Trainium NeuronCore pairs dense engines with an 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 toolchain1 that compiles C++ custom operators for the GPSIMD cores and injects them into the compiled NEFF from NKI kernel on Trainium 3, expanding the customizability of GPSIMD with 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× at decode, the table retrieval by up to 4.95×, and the whole Engram computation by up to 1.14×. Disclaimer: Every Trainium conclusion and experiment in this work comes from the Neuron SDK 2.33 toolchain, run and measured on a single trn3.48xlarge instance. Everything is from open sourced code and packages. We installed the Neuron GPSIMD Toolchain v1.25, which is a legacy version used for trn1 instances. We then built the toolchain based on how the Neuron runtime loads and runs GPSIMD code. Therefore, all of the information should be treated as reference only. The information does not necessarily reflect how AWS implements things internally.

Figure 1: Speedup over the hand-optimized NKI baseline of vLLM Neuron on one NeuronCore. 1 https://github.com/OswaldHe/neuron-gpsimd-tool

Preprint.

Contents 1

2

Introduction

3

1.1

Heterogeneity of Computation Properties in LLMs . . . . . . . . . . . . . . . . . . . . . . .

3

1.2

Example: Engram in DeepSeek-V4.1-Flash, and How We Accelerate It on a GPU-FPGA System

3

1.3

Native Hardware Heterogeneity of Trainium . . . . . . . . . . . . . . . . . . . . . . . . . .

4

Build and Use GPSIMD Custom Operators for Kernel Optimization

5

2.1

Overall Development and Deployment Flow . . . . . . . . . . . . . . . . . . . . . . . . . .

5

2.2

Write Operators in C++ . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

6

2.3

Integrate the Custom Operator with NKI . . . . . . . . . . . . . . . . . . . . . . . . . . . .

7

2.4

Runtime: What Happens Under the Hood . . . . . . . . . . . . . . . . . . . . . . . . . . . .

9

3 Accelerating Engram in DeepSeek-V4.1-Flash with the GPSIMD Engine

10

3.1

n-gram Hashing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

10

3.2

Table Retrieval . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

12

3.3

Results Overview and End-to-End Speedup . . . . . . . . . . . . . . . . . . . . . . . . . . .

12

4 What Other Operations Can GPSIMD Help With?

14

5

15

Limitations

2

1

Introduction

1.1

Heterogeneity of Computation Properties in LLMs

Modern LLMs depend on complex mechanisms to manage their memory: sparse attention selects keyvalue entries, retrieval-augmented generation (RAG) selects documents, and compressed memories select memory slots. Prior works published in ICML 2026 [13] showed that these mechanisms share a memory processing pipeline of four steps (Figure 2): 1. Prepare Memory: turn the memory into a form that is easy to search and use (for example, indexing vectors). 2. Compute Relevancy: score or identify the memory entries that matter for the current input. 3. Retrieval: select and fetch the relevant subset. 4. Apply to Inference: feed the retrieved memory into the rest of the model. The key is that these steps have very different computational characteristics. Some steps are dense and regular computations, whereas others are irregular and data dependent. Therefore, a natural approach of acceleration is fine-grain disaggregated inference on heterogeneous platforms. He et al. [13] apply a GPU-FPGA heterogeneous system [11]: the dense steps stayed on the GPU and the irregular steps moved to the FPGA, for up to 2.2× faster inference and up to 4.7× lower energy. Such methodology is also transferable to more recent LLMs.

Figure 2: The four-step memory processing pipeline in LLMs, from [13]. 1.2

Example: Engram in DeepSeek-V4.1-Flash, and How We Accelerate It on a GPU-FPGA System

DeepSeek-V4.1-Flash [8] is a 40-layer mixture-of-experts (MoE) reasoning model from DeepSeek. It combines sparse attention with MoE layers of 384 routed experts, 6 of which serve each token. Layers 1 and 14 also have an Engram module [7] that gives the model a large n-gram memory. For each token, Engram hashes the token and its preceding tokens into row ids, reads these rows from a large learned table, and adds a gated projection of the rows to the hidden state. Engram maps to the memory processing pipeline as follows: • Prepare Memory. The table is trained offline. In DeepSeek-V4.1-Flash, each Engram layer has 384M rows of 256 bf16 values. At run time, each token’s n-gram window is collected: the token’s compressed id and the ids of the 3 preceding tokens. • Compute Relevancy. An n-gram hash turns each window into row ids. For n = 2, 3, 4, the hash multiplies each token id by a 64-bit constant, combines the products with XOR, and reduces the result modulo one of 8 primes of about 16M. Each token gets 3 × 8 = 24 row ids per layer. • Retrieval. The 24 rows of each token are read into an embedding emb of shape [T, 24 × 256]. • Apply to Inference. A GEMM projects emb to 4 keys and 1 value to compute the attention score between the hidden state and the keys. A gate adds the value to each of the 4 residual streams. We also tried to deploy Engram on our GPU-FPGA heterogeneous system to demonstrate the effectiveness of such abstraction [11]. Here we use the original demo of Engram in [9]: • FPGA: Compute Relevancy, Retrieval, and the key and value projections. A hash unit computes the 16 row ids of one token per cycle. 16 lookup units read the rows from separate HBM channels. 8 projection units keep the weights in on-chip memory. • GPU: the gate, a short convolution, and the residual add. These steps need the hidden state from the previous layer. 3

Figure 3: The GPU-FPGA split of Engram in HeteroLLM [11]. h is the hidden state.

The baseline is the reference code [9] on the GPU, which computes the hash on the host CPU. With 2 AMD U55C FPGAs and 1 AMD MI210 GPU, Engram is 5.2× faster at 8 tokens, 5.9× faster at 32 tokens and 1.15× faster at 1024 tokens (Figure 4). The gain is largest with short prefill and decode.

Figure 4: Engram on the GPU alone and on the GPU with 2 FPGA boards. Engram reference configuration [9]. The prototype shows that Engram gains from moving its irregular steps to flexible hardware. This technique is transferable to AWS Trainium devices: Trainium NeuronCore is natively heterogeneous. 1.3

Native Hardware Heterogeneity of Trainium

A trn3.48xlarge instance exposes 64 NeuronCores (each made of 2 physical cores). A Trainium 3 NeuronCore is a heterogeneous processor (Figure 5). The Tensor Engine is a 128×128 systolic array for matrix multiplication. The Vector Engine and the Scalar Engine apply element-wise operations, reductions and activation functions to tiles of up to 128 partitions. The Sync Engine issues queued DMA transfers. All engines read and write the on-chip SBUF (state buffer), and the Tensor Engine accumulates into PSUM. The fifth engine, the GPSIMD Engine, is different from the other four. The GPSIMD Engine is a group of 8 general-purpose SIMD cores. Each core has its own local data RAM, an integrated DMA channel (iDMA), and a compute unit that support 64-bit wide arithmetics. The cores run firmware that implements the GPSIMD Engine’s instructions in software. Each engine fits a different kind of computation. Hence, each step of Engram has a natural place (Table 1). The n-gram hashing and the table retrieval go to the GPSIMD cores because both are irregular integer work. The hash needs 64-bit products and a 64-bit modulo: the cores compute the products, while the other engines compute exact integer arithmetic only on 32-bit values. In the table retrieval, only about 1.5% of a token’s lookups hit the table shard of a rank. GPSIMD engine can drive the table lookups on these entries in the table to issue DMA. The dense GEMM and the gate stay on the Tensor, Vector and Scalar Engines, as GPSIMD engine has much lower peak performance in dense 4

Figure 5: Trainium 3 device diagram and the architecture of a NeuronCore, from the Neuron SDK documentation [5, 6]. Table 1: The engines of a Trainium 3 NeuronCore and the Engram computations that suit them. Engine

Specialized computations

Which computation in Engram?

Tensor Engine

dense matrix multiplication

Vector and Scalar Engines DMA/Sync engines

element-wise floating-point math, reductions, activations bulk copies between HBM and SBUF integer math, per-element branches, variable-length output

the dense GEMM that projects the retrieved embeddings to the keys and the value the gating network

GPSIMD Engine

reading the table rows the n-gram hashing, and finding which lookups of the table retrieval hit this rank’s table shard

linear algebra than other engines. In general, a step suits the GPSIMD cores when it is integer-heavy, data-dependent, and small in data compared with the dense work around it. Although this is a promising plan, a major blocker for developers on Trainium is that there is no official support to connect custom GPSIMD operations with the kernel language in Trainium to optimize the whole kernel. In the next section, we will illustrate how to integrate custom GPSIMD operator with Neuron Kernel Interface (NKI) for a coherent optimization.

2

Build and Use GPSIMD Custom Operators for Kernel Optimization

The Neuron Kernel Interface (NKI) [2] is the Python kernel language for Trainium to speedup computations. Currently, NKI exposes the GPSIMD Engine only through built-in instructions on Trainium 3: tensor operations with engine=gpsimd, iota, topk, and the descriptor generation of indirect DMAs. The official toolchain for custom GPSIMD code, the Neuron custom C++ operators [3], does not solve this problem on Trainium 2 and Trainium 3. First, it supports only NeuronCore-v2 (Trainium 1 and Inferentia 2), and the compiler rejects it for Trainium 3 targets (error NCC_EVRF030). Second, a custom operator runs only as a PyTorch operator through PyTorch/XLA, and it cannot be called from an NKI kernel. Third, the toolchain is in maintenance mode starting with Neuron 2.33. Therefore, Trainium 2 and Trainium 3 have no official way to add a custom GPSIMD operator to a kernel. To fill this gap, we reverse engineered how the Neuron runtime loads and runs GPSIMD code, and built a toolchain that injects GPSIMD operators into the Neuron Executable File Format (NEFF) of a compiled NKI kernel. 2.1

Overall Development and Deployment Flow

The Neuron compiler turns an NKI kernel into a NEFF, which holds one instruction stream per engine [12]. Our toolchain does not change the compiler. The NKI kernel is compiled as usual, and the toolchain then edits the compiled NEFF to place each custom operator in the GPSIMD Engine’s 5

instruction stream. The rest of the NEFF, including the compiler’s schedule of the other engines, stays as the compiler produced it. Figure 6 shows the flow in five steps: 1. Write the operator in C++ using the gp:: namespace with the gpsimd/gpsimd.hpp header. 2. Mark the call site in the NKI kernel with gp.call(...), and compile the kernel with the NKI compiler as usual. 3. Build the operator library: the toolchain reads the SBUF location of each tile from the compiled NEFF and compiles the C++ operator for the GPSIMD cores so that it can access the corresponding SBUF address. 4. Inject the library into the NEFF, and replace the marked instructions with one KERNEL_CALL instruction. 5. Load and run the NEFF with the Neuron runtime, which will load the shared object for the operator when hit the KERNEL_CALL and run on GPSIMD.

Figure 6: Overall development and deployment flow of building and using GPSIMD custom operator and connect with optimized NKI kernel. gpsimd-build in our toolchain will handle all the compilation and injection process to produce the final NEFF file. The user only need to provide the C++ file for the GPSIMD operator and the NKI kernel python file with gp.call() that call the function in GPSIMD operator. 2.2

Write Operators in C++

A GPSIMD operator is a C++ function that can express the execution of all 8 cores (SPMD similar to CUDA kernel or NKI SPMD for LNCs). The runtime library gives each core typed views of the SBUF tiles passed by gp.call, its share of the partitions, and asynchronous iDMA copies between SBUF, HBM and the core’s local buffer. The example below adds two int32 tiles: #include <gpsimd/gpsimd.hpp> GPSIMD_OP(vadd_i32) { // extern "C" void vadd_i32(gp_call *), runs on all 8 cores gp::Call c(call); auto a = c.in<int32_t>(0), b = c.in<int32_t>(1); // SBUF tiles passed by gp.call auto o = c.out<int32_t>(0); const uint32_t n = o.elems(); int32_t la[4096], lb[4096]; // staging buffers in local data RAM for (gp::Range blk : c.blocks(o.partitions(), 4096 / n)) { // this core’s partitions c.fetch(a, blk.lo, blk.size(), la); // iDMA SBUF -> local RAM c.dma_wait(c.fetch(b, blk.lo, blk.size(), lb)); for (uint32_t k = 0; k < blk.size() * n; ++k) la[k] += lb[k]; c.dma_wait(c.flush(o, blk.lo, blk.size(), la)); // iDMA local RAM -> SBUF } }

The operator can access the core ID with c.core() (0 to 7) and the number of cores with c.ncores() (8). Most operators use the index to access partition of data through c.split(n), which gives core k the k-th of 8 equal blocks of [0, n), or c.blocks(n, rows), which walks that block in chunks of at most rows partitions. In the example above, each core takes 16 of the 128 partitions. The same split written by hand, with direct element access to SBUF instead of iDMA staging: 6

const uint32_t k = c.core(), K = c.ncores(); // this core (0..7) and the number of cores (8) const uint32_t P = o.partitions(), per = (P + K - 1) / K; // 16 partitions per core for P = 128 for (uint32_t p = k * per; p < (k + 1) * per && p < P; ++p) for (uint32_t i = 0; i < n; ++i) o.store(p, i, a.load(p, i) + b.load(p, i)); // each SBUF load takes about 296 cycles if (k == 0) gp_printf("done: %u partitions\n", P); // work for one core only

In practice, an operator rarely accesses SBUF element by element. As in the first example, it copies a block of its partitions into local RAM with c.fetch, computes there, and writes the result back with c.flush, because local RAM is much faster to access than SBUF. Table 2 lists the functions of the runtime library. Table 2: Functions of the operator runtime library. Group

Functions

Purpose

Cores and work split Tiles

c.core(), c.ncores(), c.split(n), c.partitions(t), c.blocks(n, rows) c.in<T>(k), c.out<T>(k)

Element access

t.load(p, i), t.store(p, i, v), t.row(p), t.partitions(), t.elems(), t.contiguous(), t.span(), t.offset(i) c.fetch(t, p, n, local), c.flush(t, p, n, local), c.dma(dst, src, bytes), c.dma_2d(dst, src, row_bytes, nrows, src_pitch, dst_pitch), c.dma_wait(ticket), c.dma_wait_all() gp::hbm_ptr<T>(params, k), gp::hbm_mapped(params) c.barrier(), c.allreduce_add(v), c.allreduce_max(v), c.allreduce_min(v), c.exscan_add(v, &total), c.lock(), c.unlock(), c.atomic_add(word, v), c.shared(&n) C++ integer arithmetic, including 64-bit * / %; memcpy, memmove, memset, memcmp; inline assembly of base-ISA instructions such as muluh gp::f32 (+ - * /, comparisons, to_i32()), gp::fma, sqrt, min, max, abs, floor, ceil, trunc, rint, round, ldexp; gp::bf16, gp::f16; GP_F32(x) gp::math::exp, log, tanh, sigmoid, sin, cos, gelu, rsqrt, isnan, isinf, copysign, frexp

this core’s index and its share of the work views of the SBUF tiles passed by gp.call; T is int8_t to uint32_t, gp::f32, gp::bf16 or gp::f16 element i of partition p, row-major over the free dimensions

iDMA copies

HBM Cooperation across cores

Integer

Floating point

Math

asynchronous copies between local RAM, SBUF and HBM, completed in order direct access to HBM tensors that the runner maps synchronization and data exchange through the call’s SBUF scratch tile; all cores must call them in the same order 64-bit operations run as runtime helpers IEEE float32 in software, bit-identical to numpy float32 functions

Next, we show how an NKI kernel calls this function. 2.3

Integrate the Custom Operator with NKI

In the NKI kernel, gp.call takes the operator name and the SBUF tiles that the operator reads and writes: import gpsimd_nki as gp @nki.jit def kernel(a, b): out = nl.ndarray((128, 512), dtype=nl.int32, buffer=nl.shared_hbm) ta, tb, tc = (nl.ndarray((128, 512), dtype=nl.int32, buffer=nl.sbuf) for _ in range(3)) nisa.dma_copy(dst=ta, src=a); nisa.dma_copy(dst=tb, src=b) gp.call("vadd_i32", ins=[ta, tb], outs=[tc]) # runs on the 8 GPSIMD cores nisa.dma_copy(dst=out, src=tc) return out

Because the NKI compiler does not know custom operators, gp.call cannot emit the call itself. Instead, it emits marker instructions: ordinary GPSIMD Engine instructions (TENSOR_SCALAR_ARITH_OP and MEMSET) whose immediate operand carries a tag (an encoded bit value) so that the injector will not confuse with other instructions. The tag tells the toolchain 7

which call site, role and tile a marker belongs to. The marker indicates what the GPSIMD operator will access and modify (Table 3). Table 3: Marker instructions emitted by gp.call. Marker

What it does

What the compiler then does for the operator

touch

updates each input tile in place

join chain

fully writes each output tile links the markers through a small token tile writes a scratch tile

allocates the input tiles in SBUF and keeps them alive until the call allocates the output tiles orders every producer of an input before the call and every consumer of an output after it reserves SBUF scratch for cooperation across cores

fill

Next, the toolchain analyzes the compiled NEFF and injects the operator into it. A NEFF is a container with a 1024-byte header followed by a gzip-compressed tar archive: kernel.neff |-- [header, 1024 B] payload size, MD5 of the payload, and other metadata ‘-- [payload, tar.gz] |-- neff.json top-level graph: input tensors -> kernel -> output tensors |-- kelf-0.json list of subgraphs and the target ("mariana" = Trainium 3) |-- info.json compiler and NKI versions ‘-- sg00/ program of one physical core |-- def.json memory variables, DMA queues and their semaphores, engine-to-stream map |-- tensor_map.json data type, shape and layout of each input and output tensor |-- PE0.bin Tensor Engine instruction stream |-- Activation0.bin Scalar Engine instruction stream |-- DVE0.bin Vector Engine instruction stream |-- Pool0.bin GPSIMD Engine instruction stream |-- SP0.bin Sync Engine instruction stream |-- {PE0,Activation0,DVE0,Pool0,SP0}.json description of each stream |-- act_info.json, <set>.json, <set>_{bkt,ctrl}.bin activation tables of the Scalar Engine |-- default_*_table.bin Vector Engine microcode tables ‘-- debug_info_*.dbg debug information

neff.json describes the NEFF as a graph from the input tensors through the kernel to the output tensors, kelf-0.json lists the subgraphs, and info.json records the compiler versions. Each instruction stream (<engine>0.bin) is an array of 64-byte instructions. Figure 7 shows the layout of the two instructions that the toolchain reads and writes, with the values of the n-gram hashing kernel. The first 12 bytes have the same layout for every opcode: a header that starts with the opcode, and the events block, which holds one semaphore wait and one semaphore update. The rest of the slot holds the operands. A tensor operand is an address or an access pattern in NKI’s .ap(). The engines run independently and synchronize only through semaphores. When an instruction finishes, it can add to a semaphore: GPSIMD Engine instructions use S[3], and DMA transfers use the semaphores of their queue. A wait such as S[3] >= 28 in Figure 7 counts the increments since the start of the kernel. Thus, the schedule of the whole kernel is encoded in these counts. Adding or removing an instruction that increments a semaphore shifts every later wait on that semaphore. Consequently, injection must modify the events bytes in the instructions. gpsimd-neff analyze reads the compiled kernel.neff and finds the tagged markers in the GPSIMD Engine’s instruction stream (Pool0.bin). The marker operands give the SBUF location of every tile: the first partition, the number of partitions, the byte offset inside a partition, the data type, and the free-axis access pattern. In the join marker of Figure 7, which writes the output tile of the n-gram hashing call site, the second immediate holds the tag 220 + 512 · 5 (site 0, role 5: join). The destination address is 0x4600, and the operand covers 128 partitions with 48 int32 values each. Correspondingly, we hardcode the SBUF addresses in the custom operator. gpsimd-neff inject adds two files and one key to the subgraph folder, rewrites the instruction streams, and repacks the archive: final.neff ‘-- sg00/ |-- Pool0.json |-- gpsimd_kernels_def.json |-- ops.so ‘-- Pool0.bin

+ "gpsimd_kernels_def": "gpsimd_kernels_def.json" + [{"name": "ngram_hash__s0", "lib": "/sg00/ops.so"}] + the operator library (Xtensa ELF) markers replaced by one KERNEL_CALL per call site

8

Figure 7: Layout of the 64-byte instructions that the toolchain reads and writes, with the values of the n-gram hashing call site. Top: the join marker emitted by gp.call, with its events block enlarged. Bottom: the KERNEL_CALL that replaces the markers.

In Pool0.bin, the markers of each call site become one KERNEL_CALL (opcode 0x85) at the position of the site’s first output marker. Byte 12 of the KERNEL_CALL is the index of the operator in gpsimd_kernels_def.json, and the other operand bytes are zero. The other markers become EVENT_SEMAPHORE instructions that keep the markers’ waits, and the fill marker becomes a real MEMSET that clears the scratch tile. The bytes below come from the n-gram hashing kernel of the case study: // Pool0.bin of kernel.neff (before), bytes 0-11 slot 67: 43 10 00 00 05 03 13 03 1c 00 00 00 TENSOR_SCALAR_ARITH_OP slot 68: d4 10 00 00 05 03 00 00 1d 00 00 00 PSEUDO_DMA_DIRECT2D

wait S[3] >= 28; then S[3] += 1 wait S[3] >= 29

// Pool0.bin of final.neff (after) slot 68: 85 10 00 00 00 00 15 03 09 00 00 00 slot 69: d4 10 00 00 05 03 00 00 1e 00 00 00

operator 0; then S[3] += 9 (SEM_ADD_IMM_COMPLETE) wait S[3] >= 30

KERNEL_CALL PSEUDO_DMA_DIRECT2D

(join marker) (reads the output)

The main challenge of injector is keeping track of the semaphore. The 9 markers of this call site each added 1 to S[3]. Therefore, the replaced KERNEL_CALL instruction needs to add 9 when the operator finishes. Before injection, a wait for a later GPSIMD Engine instruction also implied that the markers had finished. After the renumbering, every wait covers the same instruction as before and also the KERNEL_CALL of every earlier call site. 2.4

Runtime: What Happens Under the Hood

When the Neuron runtime loads the NEFF, it reads gpsimd_kernels_def, and the firmware copies ops.so from HBM into the local RAM of all 8 cores with the iDMA. At each KERNEL_CALL, the firmware calls the operator’s entry function on all 8 cores. The entry function reads the core index from PRID and calls the operator. The cores reach SBUF and HBM through address-translation windows. gpsimd-run (a C++ program on libnrt) programs one window over PCIe BAR after it loads the NEFF. During a call, the runtime borrows the iDMA channel from the firmware and restores the channel’s descriptor pointer at the end of the call. After every execution, each core writes a receipt: its processor id, the number of calls, the call sites it ran, its cycle count and its status. The receipts prove that the operator ran on the 8 GPSIMD cores instead on the CPU host. 9

3

Accelerating Engram in DeepSeek-V4.1-Flash with the GPSIMD Engine

In this case study, we accelerate the Engram blocks of the DeepSeek-V4.1-Flash model deployed Trainium 3 in vLLM [10] with two custom GPSIMD operators. Figure 8 shows how one Engram layer is split between the GPSIMD engine and the rest of resources in the Trainium instance. As illustrated in the previous section, we offload the n-gram hashing and the table retrieval to the GPSIMD engine, and the attention and gating network is on the other engines. NKI can already handle indirect DMA gather with GPSIMD engine. Hence, we write custom operators only for n-gram hashing and the gather index list computation.

Figure 8: The Engram computation of one layer in the four-step memory processing pipeline of [13], split between the GPSIMD engine and the other engines. The colors show where each step runs or where its data is stored, and ⋆ marks our two custom operators. The gather is NKI code, but the GPSIMD engine issues its indirect DMAs. The baseline is the hand-optimized NKI kernel of Engram in vLLM Neuron. In the DeepSeek V4.1 Flash, the Engram is splitted into 64 ranks (6.0M rows, 3.1 GB per rank). Each rank only reads the retrieved embeddings that is belong to the shard of the rank, and an all-reduce over the 64 ranks. For a fair comparison, each experiment copies the NKI modules for Engram without changes. For the speed testing, we did not load the real Engram weights. The table shard has the deployed size (6,000,097 rows, 3.1 GB) with random values, and the row ids follow the hash distribution of the model. The latency is collected using neuron-explorer by measuring the total execution time in the profile. 3.1

n-gram Hashing

The baseline computes the row ids of both Engram layers (48 row ids per token) once per decode step. A decode step of 10 sequences verifies 6 drafted tokens per sequence. Thus, the hash runs on T = 60 rows. To compute the n-gram hash, Engram applies the following equations to the compressed ids x0 , . . . , x3 of a token (x0 ) and its 3 preceding tokens: vn =

n−1 M

xs · m s ,

idn,j = (vn mod pn,j ) + on,j ,

n ∈ {2, 3, 4}, j = 1, . . . , 8,

s=0

where ⊕ is the bitwise XOR, ms is an odd multiplier of the layer for token s, and pn,j is a large prime number (e.g., 16,000,057). Each pair (n, j) is a hash column. The 24 hash columns of a layer share one table of 384M rows, and each hash column owns its own block of pn,j consecutive rows. When computing the hash id, we need to first get the index of the row inside the block of hash column (n, j) by calculating the modulo, then add the offset on,j (the index of the first row of this block, which is the sum of the sizes of the previous blocks). The products xs · ms and their XOR vn have up to 63 bits. Hence, the hash requires 64-bit integer arithmetic. However, the Vector and Scalar Engines support only 32-bit integer arithmetic natively. The baseline NKI kernel therefore works on 32-bit pieces. It precomputes a product table, which 10

holds every compressed id multiply with every multiplier of both layers as two 32-bit words. The kernel reads each product from this table with an indirect DMA and then apply bitwise XOR. It then computes vn mod pn,j with Horner’s rule, which builds a number from its digits one at a time. The baseline splits the 63-bit vn into nine 7-bit chunks c0 , . . . , c8 , with c0 the most significant, so that

vn = · · · (c0 · 27 + c1 ) · 27 + c2 · · · · 27 + c8 , and takes the remainder after every step: r0 = c0 ,

rk = (rk−1 · 27 + ck ) mod pn,j

(k = 1, . . . , 8),

vn mod pn,j = r8 .

Since rk−1 < pn,j < 224 , every intermediate value rk−1 · 27 + ck is below 231 and fits in a 32-bit integer. This is not efficient because each step takes multiple instructions throughout the Horner rule, and they are dependent on each other. With the product lookups and the XORs, the arithmetic part of the hash is 230 instructions in NKI: 166 on the Vector Engine, 58 GPSIMD built-in instructions and 6 DMAs. To resolve this issue, we write a custom GPSIMD operator, ngram_hash, that computes all 48 row ids of a token with exact 64-bit arithmetic in registers. Each of the 8 cores takes its share of the row embeddings and fetches their n-gram windows, the pn,j and ms into its local data RAM with iDMA. The GPSIMD engine can then do direct 64-bit arithmetics to follow the computations exactly as illustrated in the n-gram hashing formula, instead of having a workaround. The core C++ code looks like the following: // (token id x multiplier) mod 2^64 as two words: MULUH and MULL, then XOR into the rolling value rh ^= mulhi(cid, ml) + cid * mh; rl ^= cid * ml; // v = rh * 2^32 + rl (v < 2^63) mod p, with mu = floor(2^64 / p) from the host: // Barrett reduction, the quotient estimate is at most 3 too small const uint64_t mid = (uint64_t)lo * muh + (uint64_t)hi * mul + mulhi(lo, mul); const uint32_t q = hi * muh + (uint32_t)(mid >> 32); uint32_t r = lo - q * p; if (r >= p) r -= p; if (r >= p) r -= p; if (r >= p) r -= p;

Figure 9: One profiled execution of the hashing at decode (10 sequences, T = 60). The n-gram window collection (first 34 µs) is the same NKI code in both versions. The arithmetic of the baseline takes the next 62 µs; the GPSIMD operator takes 19.5 µs. Figure 9 shows where the time goes. Both versions spend the first 34 µs collecting the n-gram windows with the same NKI code. After that, the baseline runs a long chain of small instructions on the Vector Engine and GPSIMD built-in instructions for 62 µs, while the GPSIMD version runs one 19.5 µs KERNEL_CALL. Without counting semaphore waits, the GPSIMD version issues 107 instructions instead of 324. At decode, the hash drops from 86.95 µs to 48.45 µs (1.79×). 11

3.2

Table Retrieval

After the hashing, each rank reads the rows of the table shard, and an all-reduce over the 64 ranks (Figure 8). We studied and experimented on both a prefill chunk of T = 1024 and T = 2048 tokens, and also on decode. To retrieve the embeddings, rank r computes  Er [idt,c ] if 0 ≤ idt,c < R, emb[t, c] = 0 otherwise, where idt,c is the row id of token t in hash column c, Er is the table shard of rank r, and R is its number of rows. We call the lookups whose row id falls inside the shard hits. The number and the positions of the hits are known only at run time. However, NKI needs static shapes: the size of every DMA is fixed when the kernel is compiled. The baseline NKI kernel (engram_gather) therefore fetches every row. It computes a mask that marks the hits, issues one indirect DMA of 128 rows for every 128-token tile and hash column, and multiplies the fetched rows by the mask. Consequently, it fetches all rows to keep about 1.5% of them. This is not efficient because the cost of an indirect DMA grows with its number of indices, not with its number of useful rows, incurring extra unnecessary memory access overhead. We write the custom operator engram_hits to resolve this issue on GPSIMD engine. It builds a gather list of the hits, so that DMA fetches only those rows. Each of the 8 cores scans the row ids of its 16 partitions and writes every hit as a pair: the row to read from the shard and the row of emb to write. An exclusive prefix sum over the per-core counts gives each core its offset. Thus, the pairs of the 8 cores form one contiguous gather list. Core 0 pads the list to a multiple of 128 and writes the number of 128-pair chunks. NKI loads this number into a register and runs a loop with this run-time trip count, with one gather and one scatter per chunk. The core of the operator and the NKI code that calls it are: // GPSIMD: each core scans its partitions (all tiles, all hash columns) after one iDMA fetch for (uint32_t i = 0; i < NT; ++i) for (uint32_t j = 0; j < NCOL; ++j) { const uint32_t v = (uint32_t)r[i * NCOL + j]; // row id relative to the shard if (v < ROWW) { src[n] = v; dst[n] = (i * 128u + p) * NCOL + j; ++n; } } const uint32_t off = c.exscan_add(n, &total); // place each core’s list // one iDMA per list; core 0 pads to a multiple of 128 and writes the number of chunks

# NKI: zero-fill emb, then one gather and one scatter of 128 rows per chunk of hits gp.call("engram_hits", ins=[ids_sb], outs=[hits_sb, cnt_sb]) reg = nisa.register_alloc(); nisa.register_load(reg, cnt_sb) nl.fori_loop(0, reg, chunk, step=1) # chunk(b): load 128 (table row, emb row) pairs, gather, scatter

Figure 10 shows where the time goes. In the baseline, the GPSIMD Engine is busy with descriptor generation and the Vector Engine with the mask arithmetic for the whole retrieval. In the GPSIMD version, one 36 µs operator call replaces this work, and the hits fill only 3 chunks. Therefore, NKI needs 6 indirect DMAs instead of 192. The table retrieval drops from 506.9 µs to 102.4 µs (4.95×). 3.3

Results Overview and End-to-End Speedup

Table 4 lists all results, including both the per-computation speedups that we discussed in Sections 3.1 and 3.2 and the end-to-end speedup of the Engram module. For the whole Engram module (the n-gram hashing and both Engram layers), the two custom operators give a 1.05× speedup at decode and a 1.14× speedup on prefill chunks of both 1024 and 2048 tokens. The end-to-end speedup comes only from the two steps that moved to the GPSIMD cores. In the baseline, the n-gram hashing and the table retrieval take about one fifth of the Engram time on prefill chunks and about one seventh at decode. The attention and the gate take the rest, and they run the same NKI code in both versions. Hence, even if the two steps took no time at all, the Engram module would be at most about 1.27× faster on prefill chunks and 1.17× faster at decode. On prefill chunks, the GPSIMD operators remove more than half of the time of the two steps, and the module reaches 1.14× against this limit of 1.27×. Most of the time saved on prefill chunks comes from the table retrieval, and Figure 11 shows how it reaches the end of the layer. The attention GEMM on the Tensor Engine multiplies the retrieved 12

Figure 10: One profiled execution of the table retrieval (T = 1024). Top: the masked indirect DMAs of the baseline. Bottom: the engram_hits operator (orange), then the gathers and scatters of the hit rows. embeddings by its weights. Therefore, it cannot start until the table retrieval has written all of the retrieved embeddings. With the GPSIMD operator, the table retrieval ends earlier, and the attention GEMM and the gate after it start earlier by the same amount. Consequently, on prefill chunks, the time saved in the table retrieval shortens the whole Engram layer by about the same amount. At decode, the baseline already overlaps the table retrieval with the weight loads of the attention GEMM. Thus, only part of the time saved in the table retrieval shortens the layer (1.02×), and about half of the decode speedup comes from the n-gram hashing. This speedup is nontrivial since Engram runs for every token of every forward step. More broadly, the result shows that irregular integer steps, which the dense engines handle poorly, can move to the GPSIMD cores inside an existing NKI kernel and give an end-to-end gain without any change to the dense computation around them. Engram may also take a larger share of the latency in future models. The Engram paper keeps the table in host DRAM: the host computes the row ids and prefetches the rows over PCIe while the GPU computes the preceding blocks [7, 9]. Our implementation keeps the table in HBM, sharded by row over the 64 ranks, and computes the hash and the retrieval on the device in every forward step. DeepSeek-V4.1-Flash uses Engram in only 2 of its 40 layers, since the Engram paper notes that the placement of the Engram layers must satisfy system latency constraints as well as modeling ones. However, the best models in experiments give about 20 to 25% of the sparse parameter budget to Engram and the rest to the MoE experts. It also finds that splitting the same table over two Engram layers improves the model slightly over a single layer. If future DeepSeek models follow these findings with more Engram layers, the n-gram hashing and the table retrieval will take a larger share of the latency.

13

Table 4: Latency of the Engram computations on one NeuronCore, for the baseline and with GPSIMD custom operators (GPSIMD + NKI). Engram computation

Config

Baseline

GPSIMD + NKI

Speedup

n-gram hashing

decode, 10 sequences (T = 60) prefill chunk, T = 1024 prefill chunk, T = 2048

86.95 µs 531.1 µs 1055.6 µs

48.45 µs 449.7 µs 890.6 µs

1.79× 1.18× 1.19×

table retrieval

decode (T = 60) prefill chunk, T = 1024 prefill chunk, T = 2048

62.1 µs 506.9 µs 980.5 µs

24.7 µs 102.4 µs 200.7 µs

2.51× 4.95× 4.89×

one Engram layer

decode (T = 60) prefill chunk, T = 1024 prefill chunk, T = 2048

697.1 µs 3362.9 µs 6766.4 µs

681.0 µs 2968.9 µs 5926.4 µs

1.02× 1.13× 1.14×

whole Engram (hashing + 2 layers)

decode (T = 60)

1481.2 µs

1410.5 µs

1.05×

prefill chunk, T = 1024 prefill chunk, T = 2048

7256.9 µs 14588.4 µs

6387.5 µs 12743.4 µs

1.14× 1.14×

Figure 11: One profiled execution of one Engram layer (T = 1024). The GPSIMD operator (orange) shortens the table retrieval. Thus, the attention GEMM starts earlier.

4

What Other Operations Can GPSIMD Help With?

We tried the GPSIMD custom kernel in other scenarios as well. Although some kernels may not be the bottleneck of inference or unrelated to ML workloads at this moment, we believe that this can broaden our scope of future algorithm design instead of sticking to the dense operations that Tensor/Vector/Scalar engines are good at. Any step that is integer-heavy, data-dependent, irregular, and small in data compared with the dense work around it is a candidate (see Section 1.3). Before Engram, we used the same toolchain on kernels of the NKI Library [4]. Against the library’s NKI implementations, we measured the speedups in Table 5. Table 5: Speedups of GPSIMD operators over the NKI Library implementations. Kernel

Speedup over optimized NKI

argsort BM25 scoring for RAG scatter-add with duplicate indices

14

6.6× to 40× 15× to 27× 4.45×

5

Limitations

Currently, our toolchain supports only part of what the GPSIMD cores can do. The main gaps are listed below. No compiler support for vector instructions. Each GPSIMD core has a 512-bit vector unit, and the unit is enabled while our operators run. However, the only Xtensa configuration in the Neuron SDK (ncore2gp) describes the Trainium 1 core. Trainium 3 encodes the vector instructions differently. Consequently, vector code generated by xt-clang hangs the core. The toolchain therefore compiles scalar base-ISA code only, and gpsimd-cc rejects any library that contains vector instructions. No hardware floating point. The floating-point instructions of ncore2gp also hang the core, and we have not found their Trainium 3 encodings. Floating point therefore runs in software (gp::f32): a float32 addition takes 117 cycles and a multiplication 97 cycles. Steps dominated by floating-point math still belong on the Vector and Scalar Engines. Tile addresses fixed at build time. We did not find a way for the operator to read the KERNEL_CALL operands at run time. Therefore, the SBUF addresses of the tiles are compiled into ops.so. The library must be rebuilt whenever the kernel is recompiled, and it must fit in 64 KB of code and about 7 KB of static data. An official support on GPSIMD custom kernel should resolve these issues.

Acknowledgments This work is supported by the AWS Build on Trainium program at UCLA. We used AutoHelix [1] as the agentic loop harness to optimize the kernel, with special modifications to add constraints based on human experts. We would like to thank Ellen Hermansen from the Amazon Research Initiative team for providing AWS access and promotional credits, Amit Kachroo for providing technical insights in the Amazon Science Hub program, and James Hamilton for the continuous support on providing latest Trainium instances.

References [1] AWS Labs. AutoHelix. https://github.com/awslabs/AutoHelix, 2026. [2] AWS Neuron. About NKI. https://awsdocs-neuron.readthedocs-hosted.com/en/ latest/nki/get-started/about/index.html, 2026. AWS Neuron documentation. [3] AWS Neuron. Neuron custom C++ operators. https://awsdocs-neuron. readthedocs-hosted.com/en/v2.33.0/archive/neuron-customops/index.html, 2026. AWS Neuron documentation, Neuron SDK 2.33 (archived). [4] AWS Neuron. NKI Library. https://github.com/aws-neuron/nki-library, 2026. [5] AWS Neuron. Trainium3 architecture. https://awsdocs-neuron.readthedocs-hosted. com/en/latest/general/arch/neuron-hardware/trainium3.html, 2026. AWS Neuron documentation. [6] AWS Neuron. Trainium3 architecture guide for NKI. https://awsdocs-neuron. readthedocs-hosted.com/en/latest/nki/guides/architecture/trainium3_arch. html, 2026. AWS Neuron documentation. [7] DeepSeek-AI. Conditional memory via scalable lookup: A new axis of sparsity for large language models. arXiv preprint arXiv:2601.07372, 2026. [8] DeepSeek-AI. DeepSeek-V4.1-Flash. https://huggingface.co/deepseek-ai/ DeepSeek-V4.1-Flash, 2026. Model card on Hugging Face. [9] DeepSeek-AI. Engram reference implementation. https://github.com/deepseek-ai/ Engram, 2026. engram_demo_v1.py. 15

[10] Zifan He. DeepSeek-V4.1-Flash port for Trainium 3 in vllm-neuron. https: //github.com/OswaldHe/vllm-neuron/tree/deepseek_v41_trn3/contrib/ models/deepseek_v41/src, 2026. Branch deepseek_v41_trn3. [11] Zifan He. HeteroLLM: GPU-FPGA heterogeneous inference framework. https://github. com/OswaldHe/HeteroLLM, 2026. Engram implementation in kernels/deepseek_engram. [12] Zifan He and Jason Cong. Accelerating AWS Trainium below kernel level: NEFF-based performance optimization. https://github.com/OswaldHe/neuron-neff-doc, 2026. [13] Zifan He, Rui Ma, Yizhou Sun, and Jason Cong. Understand and accelerate memory processing pipeline for large language model inference. arXiv preprint arXiv:2603.29002, 2026.

16

Related documents

• PDScope: Breaking the Prefill-Decode Seesaw for Memory-Bound MoE Inference#1490286
• PDScope: Breaking the Prefill-Decode Seesaw for Memory-Bound MoE Inference#1482967
• Amplifying Latency and Cost of LLM Inference via Rowhammer Attack#1423589
• Amplifying Latency and Cost of LLM Inference via Rowhammer Attack#1417768
• PhaseGate: Phase-Aware CPU Retrieval Scheduling for On-Device LLMs on Unified Memory#1200188