gated-deltanet-nki-trainium

Gated DeltaNet (+ FP8 MoE) — AWS Neuron (NKI) kernels for Trainium2

Custom NKI kernel implementations for the Qwen3.5-family hybrid MoE architecture (model_type: qwen3_5_moe), targeting AWS Trainium2 via PyTorch Native (torch.compile(backend="neuron")):

• Gated DeltaNet linear-attention kernels (decode + chunked prefill), and
• a fused block-scaled FP8 MoE decode kernel for the 256-expert / top-8 MoE.

These kernels wrap heavy operators in opaque NKI custom ops to make compilation easier and reduce the number of eager operations that would otherwise run.

⚠️ Hardware / software requirements. These are research / bring-up kernels for AWS Trainium2, built against a pre-GA build of PyTorch Native.

Contents

Design notes

• One static forward → one NEFF. The kernels are shaped for a fixed-shape compiled forward.
• Numerical stability. The chunked prefill recurrence is stable in fp32 only because q, k are L2-normalized internally (the kernel does this) — the doubling product is otherwise ill-conditioned.
• Validated on a CPU oracle first. The chunked math was validated against an HF-exact CPU reference (out/state diff ~1e-7) before any device compile.
• Batch packing. For even batch sizes with C=16, DN_PAIRED_BATCH=1 packs adjacent independent C16 token blocks into block-diagonal C32 solves while keeping private recurrent states.

FP8 MoE decode kernel (moe_fused_w8_35b)

A high-batch decode kernel for the 256-expert / top-8 MoE. FP8 W8A16 on the experts is a capacity → high-batch-throughput lever, not a per-token-latency win: it halves the resident expert weights (~16 → 8 GB/core), which is what lets large batch sizes fit on one device. Measured on trn2.3xlarge, TP=8/LNC=1, 40 layers, seq=256:

BF16 is faster per token at low batch; FP8 wins aggregate throughput by unlocking a batch bf16 cannot reach. Do not expect FP8 to speed up BS=1 decode.

Interface — torch.ops.moe_w8.fused_fp8_block_coalesced_ob(hidden, gate_up, down, gate_up_scales, down_scales, affinities) → [B, H] f32 (local-expert routed sum, pre all-reduce). At TP=8 (H=2048, I=512, 256 experts → 32 local/rank), B ∈ {32, 64, 128}, gate_up=[E,H,2,I] int8 (E4M3 bytes), down=[E,I,H] int8, scales [E,128,·] bf16, affinities=[B,E] f32. I must be 512.

Serving contract (this is a serving-path kernel, not a drop-in layer): the router (softmax → top-8 → norm) and shared expert stay in the surrounding torch.compile(fullgraph=True) graph; the kernel computes only this rank's experts (expert-parallel); an all-reduce sums the partials; the shared expert is added after. Build affinities with moe_w8.build_local_affinities(...). Weight prep is offline/CPU via moe_w8 (needs the official FP8 checkpoint + the BF16 base for router/shared/attention/DeltaNet/embeddings/LM-head). Requires DECODE_FULLGRAPH, decode-only (no prefill path), --optlevel 1, one static compile per batch shape, TP=8/LNC=1.

Validation (trn2.3xlarge, 2026-08-13): kernel vs its quantized CPU reference — cosine 1.0000000, NRMSE 0.00218% at BS=32 with 32 real experts (BS=32/64/128 synthetic all pass, cosine ≈ 1.0). block_ob_coalesced's coarse per-output-block scale requantizes the official FP8 experts at weight cosine 0.99965 / NRMSE 2.67% and single-expert output cosine 0.99907 / NRMSE 4.36%; at the full-model level it generates token-identical output to the more-exact block_pow2_coalesced path (routing + top-8 sum + shared expert + depth wash out the per-expert error). If tighter per-expert numerics are needed, block_pow2_coalesced (also in moe_fused_w8_35b.py) maps each 128×128 block exactly at lower throughput.

Usage

Import the _ops module to register the custom ops, then call them from your compiled forward. See the tensor layouts documented at the top of each kernel file (state / mixed_qkv / conv_state / gates → new_state / new_conv_state / output; and for MoE, the 6-tensor signature above).

For the FP8 MoE kernel, examples/prepare_fp8_moe_weights.py is a runnable end-to-end example — official FP8 checkpoint → kernel operands + affinities, with an optional on-device --run-kernel:

python3 examples/prepare_fp8_moe_weights.py \
    --checkpoint /path/to/Qwen3.5-35B-A3B-FP8 \
    --layer 0 --world-size 8 --rank 0 --run-kernel --batch 128

The kernels were developed as part of a full Qwen3.5-35B-A3B reference implementation on PyTorch Native; that harness (weight sharding, TP, the surrounding model code) lives in the source repository:

• Source: https://github.com/dstair/neuron-pytorch-native (contrib/qwen3.6-35b-a3b)

License

Apache-2.0. See LICENSE and NOTICE. This repository contains original kernel code only and does not redistribute any model weights.