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.

File	What it is
`kernels/deltanet_full_batched_v2_35b.py`	DeltaNet decode kernel. The full fused DeltaNet block (causal conv + SiLU, L2-norm + q-scale, gates, recurrent state update, RMSNormGated) wrapped in an outer `affine_range(B)` so a single custom call processes all B sequences (DMA-coalesced "batch-as-heads"). Self-contained.
`kernels/deltanet_full_batched_35b_ops.py`	`torch.ops.*` registration + fake/meta impl so the decode kernel drops into a `torch.compile` graph without a graph break.
`kernels/deltanet_chunked_prefill_35b.py`	DeltaNet prefill kernel. Processes the whole sequence in chunks of `C` using the closed-form chunked gated-delta-rule. The Woodbury correction `A = (I − A_str)⁻¹` for the strictly-lower nilpotent `A_str` is built with a doubling product `(I+T)(I+T²)…` — `log₂(C)` static matmuls, no data-dependent slicing. Requires `nkilib`.
`kernels/deltanet_chunked_prefill_35b_ops.py`	`torch.ops.*` registration for the prefill kernel.
`kernels/moe_fused_w8_35b.py`	FP8 MoE decode kernel (see below). Fuses the local routed experts (gate/up GEMM → SwiGLU → down GEMM → affinity scale → expert-sum) with official E4M3FN weights decoded to BF16 tile-by-tile. Self-contained (no `nkilib`).
`kernels/moe_fused_w8_35b_ops.py`	`torch.ops.moe_w8.*` registration for the MoE kernel variants (the `_ob` / "block_ob_coalesced" variant is the fastest).
`moe_w8.py`	CPU-side, offline weight prep for the MoE kernel: decode official E4M3FN + block scales, requantize to per-output-block legacy-E4M3, and pack the coalesced operands + `build_local_affinities`. Pure torch.
`st_reader.py`	Dependency-free safetensors reader used by `moe_w8.py` to read the official FP8 checkpoint.
`examples/prepare_fp8_moe_weights.py`	Runnable example: official FP8 checkpoint → the six MoE-kernel operands (one layer, one TP rank) + routing `affinities`; `--run-kernel` calls the op end-to-end on Trn2.

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:

Path	Best batch that loads	Throughput
FP8 `block_ob_coalesced` (this kernel)	BS=128	681.3 tok/s
BF16 full-graph	BS=32 (BS≥64 OOMs on load)	320.6 tok/s

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:

bash

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.

dstair/gated-deltanet-nki-trainium

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

Contents

Design notes

FP8 MoE decode kernel (moe_fused_w8_35b)

Usage

License

dstair/
gated-deltanet-nki-trainium

FP8 MoE decode kernel (`moe_fused_w8_35b`)