---
license: mit
---


# 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

| 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:
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)

| 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 |