custom-nki-kernels

NKI Kernel Experiments — Flux2-klein-4B on Neuron

Hardware: AWS Trn1.32xlarge (32 NeuronCores), TP=4, bfloat16 Model: black-forest-labs/FLUX.2-klein-4B Shapes: B=1, 512×512 → img_S=256 (2× patchify), txt_S=512, inner_dim H=3072, n_heads=24, head_dim=128

---

1. RoPE kernel (nkilib.core.embeddings.rope)

Kernel constraints

• d_head ∈ {64, 128} (Flux2-klein: 128 ✓)
• S ≤ 512 — applied to sequence-length dimension before attention
• n_heads ≤ 16 per rank — after TP=4 sharding: 24/4 = 6 ✓
• Input layout must be [B, n_heads, S, d_head]

Flux2-klein applicability

Verdict: Not practical. Single-stream blocks (20/25 total) exceed S=512. Double-stream (5/25) would require custom processors. The XLA compiler already fuses RoPE with the surrounding matmuls in the same NEFF — a standalone NKI kernel would break that fusion (see §3).

---

2. Pipeline integration — --fused-qkv flag

Implementation: Flux2AttnProcessorFusedQKV in pipeline.py, activated by --fused-qkv. Replaces 3 separate to_q / to_k / to_v ColwisePar linear calls in double-stream blocks with a single NKI nki_qkv kernel call.

Timing (warm steps, avg of last 3/20 steps)

Output correctness: identical pixel range, mean, and std at every step — the kernel produces correct results.

Root cause: XLA whole-block fusion

In eager (lazy-XLA) mode, the XLA compiler traces the entire transformer block as one HLO program and compiles it into a single NEFF (neff_cache/{hash}.neff). This fuses:

• All Q/K/V projections
• RoPE embeddings
• Flash attention (via custom prim decomposition)
• Output projection + MLP
• Layer norms

Inserting a standalone NKI kernel (@nki.jit) creates opaque tensor boundaries — XLA cannot inline or fuse across NKI kernel calls. The compiler sees:

[XLA subgraph] → NKI qkv kernel → [XLA subgraph]

instead of one monolithic NEFF. This fragmentation:

1. Adds kernel launch overhead (PCIe round-trips for each NKI call)
2. Prevents data reuse that XLA would achieve within the fused NEFF
3. Defeats the cache: the fragmented graphs generate different, smaller NEFFs with no sharing benefit

The 18× slowdown is consistent with this — the baseline fused NEFF is highly optimised; the fragmented version is not.

---

3. Compile mode + fused QKV (--mode compile --fused-qkv) — bug fix note

---

4. Flash attention kernel (flux2_flash_attn)

Script: examples/flux2-klein/nki_flash_attn.py Run: torchrun --nproc_per_node=4 flux2-klein/nki_flash_attn.py

Two-pass online softmax, BLOCK_Q=128, BLOCK_K=128, bidirectional (no causal mask). Uses the older NKI ISA API (sbuf.view / psum.view / hbm.view / nisa.*).

Algorithm

For each head (looped over N=6 sequentially in one kernel instance): For each Q tile (q_idx = 0..5): Pass 1 of online softmax (here collapsed into single-pass via exp-only): For each K tile (ks = 0..5): score_T = k_tile.T @ q_tile (BLOCK_K, BLOCK_Q) via nc_matmul transposed trick probs_T = exp(score_T * scale) out_psum += probs_T.T @ v_tile (BLOCK_Q, D) row_sum += probs_T.T @ ones_v (BLOCK_Q, 1) out = out_psum / row_sum → bf16 → HBM

Note: this is an unnormalized (non-numerically-stable) softmax — no row_max subtraction. Suitable for correctness test; may overflow for long sequences or large activations.

5. Compile mode full comparison

Hardware: trn2.3xlarge, TP=4, bfloat16, 512×512, 4 steps, random weights, 4 runs (1 cold + 3 warm) Date: 2026-03-31 | neuronxcc: 2.0.236418.0a0+9af338ad

All four compile-mode variants measured on the same neuronxcc build for a fair apples-to-apples comparison.

Vanilla compile (no custom kernels)

Cold: 533.4s · Warm avg: 3.869 s/step · Throughput: 0.258 steps/s

Compile + --fused-qkv

Cold: 651.1s · Warm avg: 3.859 s/step · Throughput: 0.259 steps/s

Compile + --flash-attn

Cold: 862.3s · Warm avg: 4.159 s/step · Throughput: 0.240 steps/s

Compile + --fused-qkv --flash-attn (combined)

Cold: 830.2s · Warm avg: 4.149 s/step · Throughput: 0.241 steps/s

Summary table

Interpretation

• Fused-QKV has no measurable effect in compile mode (3.859 vs 3.869 — within run-to-run noise). The Dynamo+NEFF compiler already fuses QKV projections at the HLO level; the explicit NKI kernel neither helps nor hurts, but adds 118s to cold compilation.
• Flash-attn is ~7% slower than vanilla regardless of whether fused-QKV is also enabled. The unnormalized single-pass softmax and sequential head loop are less efficient than the compiler's built-in attention decomposition (two-pass numerically stable, better SPMD utilisation).
• Combining both kernels gives the same result as flash-attn alone (4.149 vs 4.159 — within noise). fused-QKV contributes nothing additional in compile mode.
• Cold compilation time grows with NKI kernel count: vanilla (533s) → fused-qkv (651s, +22%) → combined (830s, +56%) → flash-attn alone (862s, +62%). Each NKI kernel adds a separate KLIR compilation pass inside neuronxcc.

---

8. Conclusions