Kernels:

---

aws-neuron

/

vjepa2-neuron-kernels

Like 2

Follow

AWS Inferentia and Trainium 196

Reworks the fused MLP kernel and the residual + LayerNorm kernel. K2, the weight folds and the
layer are unchanged apart from one line in layers.py that packs the folded fc1 bias.

Results (trn2.3xlarge, LNC=2, all outputs torch.equal to v1.1)

Full encoder (torch.compile(model.encoder, backend="neuron"), 24 layers, one 16-frame clip):

Standalone at the encoder shape (2 interleaved rounds each):

K1/K3 do no matmuls, so the profiler reports MFU 0 for them. Achieved HBM bandwidth
(bytes moved / device time) is their roofline metric.

Absolute times are on torch-neuronx 2.14.3 + nki 0.7.0b1. The same v1.1 code measures 26.63 ms/clip
there, against 23.94 in the README, so compare the relative change.

What changed in K4 (found from the device profile)

1. x before weights. v1.1 issued b1 -> all of W1 -> all of W2 -> x, so the first matmul waited
   for ~16 MB of weights it didn't need yet (first matmul at ~57 us). Order is now
   b1 -> x(rb0) -> W1 -> x(rb1) -> W2.
2. b1 in one DMA. v1.1 loaded fc1 bias as 32 DMAs of [128, 1] fp32 (4 B per partition), which
   trickled in until ~69 us. Each gelu waited on its bias tile, and through its PSUM bank the fc1
   matmuls behind it. b1 is now host-packed to [128, H/128] (pack_b1) and lands in one 16 KB DMA on
   the Scalar engine's HW-DGE queue, so gelu's wait no longer counts W1/W2 traffic.
3. Row block 0's x transposed on the Tensor engine, which is idle in the prologue.
4. Rewrite: 3D SBUF tiles instead of tile lists, nl.tile_size.pmax / gemm_moving_fmax
   instead of literals, and exact-multiple shapes asserted instead of min() remainder handling.
   b2 stays a per-(row block, output chunk) partition-broadcast load; hoisting it into one load was
   neutral standalone but cost ~0.7 ms/clip in the encoder.

What changed in K1/K3

1. dge_mode.none on every DMA. In v1.1 each of the ~40 DMAs had its descriptors generated by
   GpSimd (software DGE) at ~0.75 us apiece, one after another; the first add waited ~8 us for them.
   With dge_mode.none the runtime generates descriptors before execution. Hardware DGE on the Sync
   queue (dge_mode.hwdge, engine=sync) was as fast standalone but 0.1 ms/clip slower in the
   encoder, so it is not used.
2. All token tiles resident in SBUF, every load issued before any compute, so a tile's loads
   never wait behind the previous tile's stores.

The compute instructions and their order are unchanged, so results are bitwise identical.

Compatibility

• mlp still accepts b1 as [H] (packed on the fly) as well as the pre-packed [128, H/128];
  pack_b1 is exported. Op names unchanged.
• New K4 requirement: K % 512 == 0, H % 128 == 0, R % 1024 == 0. The encoder shape (R=2048, K=1024,
  H=4096) satisfies these.
• K1/K3 SBUF use is ~116 KB/partition at H=1024 with 2 branches (all tiles resident).
• Toolchain: all numbers above are torch-neuronx 2.14.3 + nki 0.7.0b1. Standalone K4 and K1/K3
  were also verified on the latest public NKI release, nki 0.6.0 + neuronx-cc 2.27.5334: K4 248.6 /
  247.7 us vs 291.6 / 294.0 us for v1.1, K1 42.6 us, K3 48.1 us, all bitwise identical. The full
  encoder could not be run on nki 0.6.0 with torch-neuronx 2.14.3: v1.1 itself fails there
  (CPU fallback failed for 'model_default'). nki 0.6.0 rejects list comprehensions inside kernels,
  so the kernels use explicit loops.