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 attentionn_heads ≤ 16per rank — after TP=4 sharding: 24/4 = 6 ✓- Input layout must be
[B, n_heads, S, d_head]
Flux2-klein applicability
| Block type | S | Fits S≤512? | Notes |
|---|---|---|---|
| Single-stream | img_S + txt_S = 256 + 512 = 768 | No | RoPE is applied to the concatenated image+text sequence |
| Double-stream (image) | img_S = 256 | Yes | But double-stream blocks apply RoPE inside FluxAttnProcessor after separate Q/K projections — hooks into NKI require custom processor |
| Double-stream (text) | txt_S = 512 | Yes (boundary) |
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)
| Mode | Steps | Warm avg (s/step) | vs baseline |
|---|---|---|---|
| Eager, baseline | 20 | 0.824 s/step | 1× |
Eager, --fused-qkv |
20 | 14.86 s/step | 18× slower |
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:
- Adds kernel launch overhead (PCIe round-trips for each NKI call)
- Prevents data reuse that XLA would achieve within the fused NEFF
- 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)
| Run | Type | step01 | step02 | step03 | step04 | total |
|---|---|---|---|---|---|---|
| 1 | COLD | 533.449s | 3.868s | 3.868s | 3.868s | 545.053s |
| 2–4 | WARM | 3.868s | 3.869s | 3.869s | 3.869s | 15.475s |
Cold: 533.4s · Warm avg: 3.869 s/step · Throughput: 0.258 steps/s
Compile + --fused-qkv
| Run | Type | step01 | step02 | step03 | step04 | total |
|---|---|---|---|---|---|---|
| 1 | COLD | 651.147s | 19.874s | 3.859s | 3.859s | 678.740s |
| 2–4 | WARM | 3.859s | 3.859s | 3.860s | 3.860s | 15.438s |
Cold: 651.1s · Warm avg: 3.859 s/step · Throughput: 0.259 steps/s
Compile + --flash-attn
| Run | Type | step01 | step02 | step03 | step04 | total |
|---|---|---|---|---|---|---|
| 1 | COLD | 862.344s | 19.601s | 4.159s | 4.159s | 890.263s |
| 2–4 | WARM | 4.159s | 4.159s | 4.159s | 4.159s | 16.636s |
Cold: 862.3s · Warm avg: 4.159 s/step · Throughput: 0.240 steps/s
Compile + --fused-qkv --flash-attn (combined)
| Run | Type | step01 | step02 | step03 | step04 | total |
|---|---|---|---|---|---|---|
| 1 | COLD | 830.249s | 19.558s | 4.149s | 4.149s | 858.105s |
| 2–4 | WARM | 4.149s | 4.149s | 4.149s | 4.149s | 16.597s |
Cold: 830.2s · Warm avg: 4.149 s/step · Throughput: 0.241 steps/s
Summary table
| Mode | Cold (s) | Warm avg/step | Throughput | vs vanilla compile |
|---|---|---|---|---|
| Eager, baseline | 9.3s | 0.835 s/step | 1.198 steps/s | 4.6× faster |
| Compile, vanilla | 533.4s | 3.869 s/step | 0.258 steps/s | 1× (baseline) |
Compile, --fused-qkv |
651.1s | 3.859 s/step | 0.259 steps/s | −0.3% (noise) |
Compile, --flash-attn |
862.3s | 4.159 s/step | 0.240 steps/s | +7.5% slower |
Compile, --fused-qkv --flash-attn |
830.2s | 4.149 s/step | 0.241 steps/s | +7.2% slower |
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
| Kernel | Correct | Practical for eager | Practical for compile |
|---|---|---|---|
| NKI RoPE | — | No (S > 512 for single-stream) | No (same constraint) |
| NKI QKV | Yes | No — breaks XLA fusion (18× slower) | Negligible effect (within noise) |
| NKI Flash Attention | Yes (cosine=0.9999) | TBD | No — 7% slower than vanilla, +62% compile time |
| NKI QKV + Flash Attention | Yes | No | Same as flash-attn alone |