jburtoft/qwen35-deltanet-neuron-kernels

Qwen3.5 Gated Delta Net NKI Kernel for AWS Neuron

Notice (2026-07-16): This repo accelerates the prefill (CTE / context encoding) path only. Its NeuronGatedDeltaNet module ignores cache_params, so calling model.generate() with it as-is recomputes the full recurrent state on every decoded token — correct but quadratic.

For accelerating both prefill AND decode — including proper cache_params (conv state + recurrent state) handling for model.generate() — use jburtoft/qwen35-deltanet-tkg-full instead. That repo is a strict superset: it includes this repo's sequential CTE kernel plus a batched TKG single-token kernel for decode, and it delivers 1.63x end-to-end model.generate() speedup (verified on Qwen3.5-0.8B, trn2.3xlarge, PyTorch Native Beta 3).

The performance numbers in this README (5-9x on isolated single-layer prefill) are still accurate for the specific single-layer-prefill benchmark harness they were measured with, but they should not be interpreted as end-to-end model.generate() speedup.

---

NeuronCore-optimized NKI kernel that replaces the Qwen3_5GatedDeltaNet module in Qwen3.5 models with a hardware-accelerated implementation that keeps the recurrent state matrix (128×128 float32) entirely in on-chip SBUF memory — eliminating HBM round-trips per token.

Supported Models

All Qwen3.5 models with Qwen3_5GatedDeltaNet linear attention layers. head_dim must be exactly 128 (SBUF partition width) and head_k_dim == head_v_dim (both checked at forward time).

Model	Key Heads	Value Heads	GQA Ratio	Tested
Qwen3.5-0.8B-Base	16	16	1:1	✅ 9.05x speedup at S=16, 5.4x at S=51 (isolated layer, Beta 3)
Qwen3.5-2B-Base	16	16	1:1	-
Qwen3.5-4B-Base	16	32	1:2	✅ 5.37x speedup at S=15 (isolated layer, Beta 3)
Qwen3.5-9B-Base	16	32+	1:2+	-
Qwen3.5-27B	16	48	1:3	-
Qwen3.5-35B-A3B (MoE)	16	32	1:2	-

Requirements

• AWS Neuron SDK 2.29+ (NKI 0.3.0+)
• PyTorch Native (torch-neuronx) with NKI support -- tested on Beta 3 (torch-neuronx 2.11.3, NKI 0.4.0b4, neuronx-cc 2.25)
• transformers >= 5.10.0.dev0 with KernelConfig support (PR #46339, merged 2026-06-09)
• kernels library (HuggingFace kernel loading)
• head_dim == 128 and head_k_dim == head_v_dim (asserted at forward time)

Usage

from transformers import AutoModelForCausalLM, KernelConfig

model_id = "Qwen/Qwen3.5-0.8B-Base"

kernel_config = KernelConfig({
    "Qwen3_5GatedDeltaNet": "jburtoft/qwen35-deltanet-neuron-kernels:NeuronGatedDeltaNet",
})

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    dtype="bfloat16",
    kernel_config=kernel_config,
    device_map="neuron",
    trust_remote_code=True,
)

Performance

Benchmarked on trn2.3xlarge (LNC=2, 4 logical NeuronCores) with PyTorch Native Beta 3 on 2026-07-15. Isolated single-layer forward pass with matched random weights vs the HF Qwen3_5GatedDeltaNet chunk fallback (torch_chunk_gated_delta_rule) reference.

Qwen3.5-0.8B-Base (GQA 1:1)

Seq Length	Baseline (chunk fallback)	NKI Kernel	Speedup
2 tokens	92.1 ms	10.9 ms	8.41x
10 tokens	94.7 ms	10.6 ms	8.97x
15 tokens	94.6 ms	10.5 ms	9.05x
16 tokens	94.9 ms	10.4 ms	9.16x
51 tokens	99.3 ms	18.2 ms	5.45x

Qwen3.5-4B-Base (GQA 1:2)

Seq Length	Baseline (chunk fallback)	NKI Kernel	Speedup
2 tokens	93.5 ms	18.4 ms	5.09x
10 tokens	94.0 ms	17.7 ms	5.30x
15 tokens	94.4 ms	17.6 ms	5.37x
16 tokens	93.8 ms	18.6 ms	5.05x

The speedup comes from:

1. SBUF state residency: The 128×128 state matrix stays in on-chip SBUF throughout the recurrence (no HBM reads/writes per token)
2. Hardware-native operations: Uses NeuronCore's nc_matmul, nc_transpose, and vector engine for all operations
3. Efficient decay: Leverages hardware exp activation and element-wise multiply
4. PSUM read fusion and in-place state ops (added 2026-07-15): saves 3 tensor_copy operations per token by feeding PSUM directly into tensor_tensor/tensor_scalar and writing state in place. Net -9.5% wall-clock at S=51 vs the pre-optimization version.

Numerical Parity

Verified against fp32 CPU reference (Qwen3_5GatedDeltaNet chunk fallback):

• cosine similarity ≥ 0.9999 across all tested (model, seq_len) combinations
• max abs diff ≤ 4.7e-3 (bfloat16 accumulation noise on signals of norm ~10-17)
• norm ratio 0.9987 - 1.0001 (unbiased -- kernel does not systematically over- or under-estimate)

Tested seq_lens: 2, 10, 15, 16, 51, 100, 127, 129, 200. Both Qwen3.5-0.8B and Qwen3.5-4B.

How It Works

The kernel implements the gated delta rule recurrence:

For each token t:
  1. state *= exp(g_t)              — decay state
  2. kv_mem = state^T @ k_t         — read from memory
  3. delta = (v_t - kv_mem) * β_t   — compute update
  4. state += outer(k_t, delta)     — write to memory
  5. o_t = state^T @ q_t            — produce output

The surrounding PyTorch operations (projections, causal conv1d, gating, RMS norm) run via standard Neuron compilation.

Correctness Contract

• head_v_dim == head_k_dim == 128 (hard requirement, asserted at forward time)
• State grows or shrinks per token via state *= exp(g_t). HF's Qwen3.5 uses g = -A_log.exp() * softplus(a + dt_bias) which is strictly non-positive, so exp(g_t) ∈ (0, 1] and state cannot overflow. If a downstream model changes the gating to positive-drift, the kernel does not clamp.
• No shared B/C across heads. Each (batch, head) invokes the kernel independently.

See the module docstring in build/torch-neuron/__init__.py for the full contract.

Architecture

NeuronGatedDeltaNetLayout  — Weight structure (defines conversion_mapping)
NeuronGatedDeltaNet        — Forward pass (calls NKI kernel per head, includes correctness assertions)
deltanet_recurrent_fwd     — @nki.jit kernel (token-by-token recurrence with in-place state ops)

Development

Kernel optimization work for this repo is tracked in the deltanet-kernel project. Contributions welcome via PR.

Changelog

2026-07-15 (Task 007 -- PSUM fusion + in-place state)

• F1: In-place tensor_scalar(dst=state, data=state, ...) for state decay (saves 1 SBUF alloc + 1 tensor_copy per token)
• F2: Fuse nc_matmul PSUM read directly into the subsequent tensor_tensor(v_t - kv_mem) (saves 1 SBUF alloc + 1 tensor_copy per token)
• F3: In-place tensor_tensor(dst=state, data1=state, data2=outer_prod, add) for state update (saves 1 SBUF alloc + 1 tensor_copy per token)
• Correctness (C1): added runtime assertions for head_v_dim == 128 and head_k_dim == head_v_dim
• Numerical (C2): documented state-decay dynamic-range constraint in the module docstring
• Net wall-clock effect: -9.5% at S=51, negligible at S=16/S=2. Parity: bit-identical (18/18 test cases pass with cosine sim 0.99995-0.99997)
• Also attempted (reverted): bf16 operand matmul (Task 013). Saved 58% of PE instruction-time but the cast overhead (VectorE +25%) canceled the wall-clock win because PE is not on the critical path in this token-serialized kernel. Reverted; findings documented in deltanet-kernel project Task 013.

License

Apache 2.0