A model-first walkthrough: understand what Llama-3.2-1B inference IS, then how this codebase runs it on AMD NPU2 hardware.
+ + + + + + +Llama-3.2-1B is a 1.24-billion-parameter decoder-only transformer language model from Meta, released in 2024. Given a sequence of input tokens, it produces a probability distribution over the vocabulary for the next token. Repeated autoregressively, this generates text.
+LlamaConfig at llama32_1b_weights.py:36)| Parameter | Value | What it means |
|---|---|---|
n_layers | 16 | Number of stacked transformer blocks |
emb_dim (d_model) | 2048 | Hidden dimension everything flows through |
n_heads | 32 | Number of Q heads in attention |
n_kv_heads | 8 | Number of K/V heads (GQA: 4 Q heads share each KV head) |
head_dim | 64 | Per-head dimension. Note: 32 × 64 = 2048 = emb_dim |
hidden_dim | 8192 | FFN intermediate width (gate/up/down projections expand to this) |
vocab_size | 128256 | Tokenizer vocabulary size; LM Head outputs this many logits |
seq_len | 2048 | Fixed prefill length in this implementation (not a model property) |
| weight dtype | bfloat16 | 16-bit brain-float for all weights and activations |
| RoPE base | 500000 | Rotary Position Embedding base frequency |
| Component | Per layer | × 16 layers | Per-tensor shape |
|---|---|---|---|
| Attention norm weight | 2,048 | 32,768 | (2048,) |
| Q projection | 4.19 M | 67.1 M | (2048, 2048) |
| K projection | 1.05 M | 16.8 M | (2048, 512) |
| V projection | 1.05 M | 16.8 M | (2048, 512) |
| O projection | 4.19 M | 67.1 M | (2048, 2048) |
| FFN norm weight | 2,048 | 32,768 | (2048,) |
| Gate projection | 16.8 M | 268 M | (2048, 8192) |
| Up projection | 16.8 M | 268 M | (2048, 8192) |
| Down projection | 16.8 M | 268 M | (8192, 2048) |
| Per-layer subtotal | 61.0 M | 976 M | ~ 122 MB bf16 |
| Embedding table | — | 263 M | (128256, 2048) |
| Final norm | — | 2,048 | (2048,) |
| LM Head (vocab projection) | — | 263 M | (128256, 2048) |
| Grand total | ≈ 1.50 B | ~ 3.0 GB bf16 |
Note: Llama-3.2-1B uses untied embeddings (LM Head is a separate parameter from the embedding table). That's why total is ~1.50 B not ~1.24 B if you sum just the published parameter count. The embedding table is loaded but the embedding lookup is a host-side numpy index, not an NPU kernel.
+ + +Llama-3.2-1B is just 16 of these blocks stacked, sandwiched between a token embedding lookup at the start and a final RMSNorm + LM Head at the end. (See A3 for the full top-level pipeline.)
+ +One transformer block is a function block(x) → output where both x and output have the same shape [B, S, H]. The block has two sub-blocks (attention and FFN), each with a residual connection. We diagram them separately to keep each readable.
| Symbol | Meaning | Llama-3.2-1B value |
|---|---|---|
B | batch size | 1 (this implementation is single-stream) |
S | sequence length | 2048 (prefill) or 1 (decode) |
H | hidden dim (d_model) | 2048 |
L | number of decoder layers | 16 |
N_h | query head count | 32 |
N_kv | KV head count (GQA) | 8 |
G | GQA group size = N_h / N_kv | 4 |
d_h | per-head dim = H / N_h | 64 |
D_ff | FFN intermediate dim | 8192 |
V | vocab size | 128256 |
Note: H = N_h · d_h = 32 · 64 = 2048, and the K/V projection output is N_kv · d_h = 8 · 64 = 512 (smaller than H because of GQA).
+ +From the block's input x, the attention sub-block produces an updated hidden state with cross-position information mixed in (causally — only earlier positions affect later ones). Three weighted projections (Q, K, V) plus RoPE, attention compute, and an output projection. The output is added to a saved copy of x (residual).
[B, S, H] → [B, S, H], weight γ: [H]y = x · rsqrt(mean(x², dim=-1) + ε) · γ(b, s) position is normalized independently along the hidden dim. No mean subtraction (unlike LayerNorm), no bias. The mean is over 2048 elements per row.[B, S, H] → [B, S, N_h·d_h] (= [B, S, H] since H = N_h · d_h), weight W_q: [H, N_h·d_h]Y = X @ W_q (no bias)(b, s) row maps independently; B · S can be flattened into the M dim for batching. In our impl: prefill is a GEMM at M=2048; decode is a GEMV at M=1.[B, S, H] → [B, S, N_kv·d_h], weight W_k: [H, N_kv·d_h]Y = X @ W_k (no bias)[B, S, H] → [B, S, N_kv·d_h], weight W_v: [H, N_kv·d_h]Y = X @ W_v[B, S, N_h, d_h] → [B, S, N_h, d_h] (unchanged), reads cos/sin LUT of shape [S, d_h](b, s, h) head's d_h-vector by the angle determined by position s. Q_roped[b,s,h,i] = Q[b,s,h,i]·cos[s,i] − Q[b,s,h,i+d_h/2]·sin[s,i] (half-split convention)generate_rope_lut). Pure data movement + multiplies; no reductions.[B, S, N_kv, d_h] → [B, S, N_kv, d_h] (unchanged)[B, S, N_kv, d_h] unchangedq_roped: [B, S, N_h, d_h], k_roped: [B, S, N_kv, d_h], v: [B, S, N_kv, d_h] → out: [B, S, N_h, d_h]K^T swaps the seq and d_h dims.scores[b,h,s,t] = Q[b,s,h,:] · K[b,t,h//G,:] / √d_h — note the GQA index h//G shares one KV head across G query heads.scores[b,h,s,t] = −∞ for t > s so query position s only attends to positions 0..s.P[b,h,s,t] = softmax(scores[b,h,s,:]) — normalized over the LAST dim (key positions). Row-wise per query.out[b,s,h,:] = Σ_t P[b,h,s,t] · V[b,t,h//G,:]S × S). FlashAttention fuses all 5 sub-steps into a tiled kernel that never materializes the full S × S matrix in memory. No learnable weights. Memory-bound for large S, compute-bound for small S.[B, S, H] → [B, S, H], weight W_o: [H, H]proj = attn_out @ W_o (no bias)x: [B, S, H] + proj: [B, S, H] → [B, S, H]res1 = x + projx to the projection output. No reduction, no broadcast (both inputs same shape). Output is the input to the FFN sub-block.Takes the attention sub-block's output (call it res1) and applies a 3-projection feed-forward network with SwiGLU activation. Like the attention sub-block, the result is added to a saved copy of the input.
[B, S, H] → [B, S, H], weight γ: [H]γ (called ffn_norm).[B, S, H] → [B, S, D_ff], weight W_gate: [H, D_ff]gate = X @ W_gate[B, S, H] → [B, S, D_ff], weight W_up: [H, D_ff]up = X @ W_up[H, 2·D_ff] GEMM (some implementations do this); ours keeps them separate.[B, S, D_ff] → [B, S, D_ff] (unchanged)SiLU(x) = x · σ(x) = x / (1 + e^{−x})[B, S, D_ff] × [B, S, D_ff] → [B, S, D_ff]swiglu[i] = SiLU(gate[i]) · up[i] — Hadamard product.silu_and_mul.cc), saving one full pass over the 8192-wide tensor.[B, S, D_ff] → [B, S, H], weight W_down: [D_ff, H]down = swiglu @ W_downres1: [B, S, H] + down: [B, S, H] → [B, S, H]out = res1 + downsilu_and_mul.cc). Saved per-pass over the 8192-wide tensor.rope_halfsplit.cc) uses the half-split convention: (d[i], d[i + d_h/2]) are paired for rotation. llama.cpp and the original RoPE paper use interleaved (d[2i], d[2i+1]). The two produce DIFFERENT outputs for the same input — they are not interchangeable. Our LUT layout is [cos_0..cos_{d_h/2-1}, sin_0..sin_{d_h/2-1}] (concatenated, not interleaved), matching the half-split rotation.causal=True and never materializes a mask matrix; it just skips attending to t > s in the inner loop.x · rsqrt(mean(x²) + ε) · γ. ε is a small constant (1e-5 typically) for numerical stability.W_lm is a separate parameter from the embedding table W_emb. (Some smaller models tie them to save parameters.) Both are [V, H]; together they account for ~526 M of the model's 1.5 B parameters.
+The 14 ops above map to the production NPU kernels as follows:
+ +| Sub-block | Model ops | NPU realization |
|---|---|---|
| Attention | +RMSNorm + Q proj + K proj + V proj + RoPE Q + RoPE K | +rms_gemms_rope.elf — 6 sub-launches stitched into one ELF |
| Scaled dot-product attention | +flash_attn.elf — 1 launch (separate ELF; un-mergeable) | |
| (boundary) | +O proj + Residual #1 | +First 2 sub-launches of o_ffn.elf |
| FFN | +RMSNorm + Gate proj + Up proj + SiLU·mul + Down proj + Residual #2 | +Remaining 6 sub-launches of o_ffn.elf |
So one transformer block = 3 NPU calls (rms_gemms_rope + flash_attn + o_ffn) wrapping a total of 15 sub-launches (6 + 1 + 8). The grouping is not the natural "attention sub-block / FFN sub-block" boundary — instead, the cut is "before FlashAttention" vs "after FlashAttention", because FA must be its own ELF (compile-time scaling issue documented in docs/explain.md). Why this exact grouping is best — and why all 15 sub-launches don't go into one ELF — is the topic of Part B.
Below is one Llama-3.2-1B layer written as plain NumPy — useful as a reference for the math, independent of NPU plumbing. (The actual production NPU pipeline is described in Part B; numerical correctness is gated by make verify against HF transformers bf16 — see VERIFICATION.html.)
def transformer_block(x, lw, rope_lut, config):
+ # Attention sub-block
+ normed = rms_norm(x, lw.attn_norm)
+ q = normed @ lw.wq
+ k = normed @ lw.wk
+ v = normed @ lw.wv
+ q_roped = apply_rope(q, rope_lut)
+ k_roped = apply_rope(k, rope_lut)
+ attn_out = attention(q_roped, k_roped, v, config) # GQA, causal mask
+ res1 = x + attn_out @ lw.wo
+
+ # FFN sub-block
+ normed2 = rms_norm(res1, lw.ffn_norm)
+ gate = normed2 @ lw.w_gate
+ up = normed2 @ lw.w_up
+ swiglu_out = silu(gate) * up
+ output = res1 + swiglu_out @ lw.w_down
+ return outputThe diagram below shows the whole inference call as 6 stages. The decoder block is collapsed (×L) — its internals are diagrammed in A2.
+ + + +[B, S] integer indices → [B, S, H] bf16, weight W_emb: [V, H]x[b, s, :] = W_emb[token_ids[b, s], :] (table lookup)[B, S, H] → [B, S, H] per block, repeated L timesx ← block_i(x, layer_weights[i], rope_lut) for i in 0..L-1[B, S, H] → [B, S, H], weight γ_final: [H]γ (called final_norm).[B, S, H] → [B, S, V], weight W_lm: [V, H] (untied — separate from W_emb)logits = X @ W_lm.T (no bias)pred_pos), as a 1×V GEMV partitioned 8 ways. See A7 for why this is sufficient.[B, S, V] → [B, S] integer indicesnext_token = argmax(logits, dim=-1)pred_pos to get the next token. CPU operation in our impl (cheap — V=128256 single argmax).The forward pass above works for ANY input length. But there are two common usage patterns:
+ +| Mode | Input | What we do | Output | Cost |
|---|---|---|---|---|
| Prefill | +The full prompt: token_ids of length S = prompt_len |
+ One forward pass with seq=S. Save K, V at every layer for every position into a "KV cache" — we'll need them for decode. Argmax at position S-1 gives the first generated token. |
+ 1 token + populated KV cache | +~1.27 s for S=2048 | +
| Decode | +One token at a time: x of shape (1, 2048) — embedding of the previous output token |
+ One forward pass with seq=1. Use the KV cache in attention — the new K, V for this position get appended. Argmax gives the next token. | +1 new token + KV cache extended by 1 position | +~92 ms per token | +
To generate N tokens of text from a prompt: 1 prefill call + N decode calls. The KV cache is built once during prefill and grows by one row per decode step.
+ + +For a sequence of length T, attention computes:
+ +Q = X @ Wq # shape (T, n_heads, head_dim)
+K = X @ Wk # shape (T, n_kv_heads, head_dim)
+V = X @ Wv # shape (T, n_kv_heads, head_dim)
+attn = softmax(Q @ K.T / √d) @ V # causal maskedDuring decode, position T+1 only adds one new query Q[T+1]. But that query needs to attend to all previous K[0..T] and V[0..T]. If we threw those away after the prefill and recomputed them, we'd redo O(T) work per decode step.
+ +Once K[i] and V[i] are computed for any position i, they never change again (they only depend on x[i] and weights, not on later tokens). So we store them in a per-layer cache and append a new entry per decode step.
+ +Allocated in llama32_1b_inference.py:369:
+ +k_cache = np.zeros(
+ (config.n_layers, n_kv_heads, max_seq, head_dim),
+ dtype=bfloat16,
+)
+v_cache = np.zeros((config.n_layers, n_kv_heads, max_seq, head_dim), dtype=bfloat16)| Dimension | Size | Why |
|---|---|---|
n_layers | 16 | Each layer has its own K, V (different transformations of x) |
n_kv_heads | 8 | GQA — only 8 distinct heads (vs 32 Q heads) |
max_seq | prompt_len + n_tokens | Enough room for the prompt + every generated token |
head_dim | 64 | Per-head dimension |
Total memory: 16 × 8 × max_seq × 64 × 2 bytes = 16,384 × max_seq bytes ≈ 32 MB at max_seq=2048. Tiny compared to the 3 GB of weights — KV cache is not a memory concern for Llama-1B.
+ +Showing one layer's K cache (the V cache has the same structure). Each cell is one position; rows are the 8 KV heads.
+ +(1) Cache allocation — once per generate() call:
# llama32_1b_inference.py:369
+k_cache = np.zeros((n_layers, n_kv_heads, max_seq, head_dim), dtype=bfloat16)
+v_cache = np.zeros((n_layers, n_kv_heads, max_seq, head_dim), dtype=bfloat16)(2) Prefill writes to the cache — extracts k_roped and v from each layer's intermediates:
+ +# llama32_1b_inference.py:401 — runs after each layer in the prefill loop
+k_cache[layer_idx, :, :seq_len, :] = (
+ k_roped.astype(bfloat16)
+ .reshape(seq_len, n_kv_heads, head_dim)
+ .transpose(1, 0, 2) # layout: (n_kv_heads, seq_len, head_dim)
+)
+v_cache[layer_idx, :, :seq_len, :] = (
+ v_raw.astype(bfloat16).reshape(seq_len, n_kv_heads, head_dim).transpose(1, 0, 2)
+)(3) Decode appends to the cache and reads from it — inside decode_attention_cpu and run_decode_block:
# llama32_1b_decode.py — paraphrased
+def run_decode_block(x, lw, cache, config, k_cache_layer, v_cache_layer, current_pos, ...):
+ # 1. Compute new k, v from this token (NPU rms_gemv_rope call)
+ out = cache.load_and_run("rms_gemv_rope", ...)
+ new_k_roped = out[12] # shape (kv_dim,) = (512,) flat
+ new_v = out[8] # shape (kv_dim,)
+
+ # 2. Append to cache at current_pos
+ k_cache_layer[:, current_pos] = new_k_roped.reshape and transpose
+ v_cache_layer[:, current_pos] = new_v.reshape and transpose
+
+ # 3. CPU attention reads positions 0..current_pos
+ attn_out = decode_attention_cpu(q_roped, k_cache_layer, v_cache_layer,
+ current_pos, n_heads, n_kv_heads, head_dim)
+
+# Inside decode_attention_cpu:
+seq_len = current_pos + 1
+k_cached = k_cache[:, :seq_len, :] # only positions 0..current_pos
+v_cached = v_cache[:, :seq_len, :]
+# Then standard QKᵀ V softmax against this slice...current_pos = prompt_len (NOT 0). The cache positions 0..prompt_len-1 are populated by the prefill. The first decode step writes the new k, v at position prompt_len and reads positions 0..prompt_len for attention (the new entry plus all the prefill entries).
+This implementation uses fixed seq_len=2048 because NPU kernels are compiled for one specific shape — recompiling for every prompt length would be prohibitive. So we always pad shorter prompts up to 2048. Let's trace exactly how that works.
+ +In llama32_1b_inference.py:731:
+ +def _tokenize_prompt(session, prompt_text):
+ if session.model_variant == "instruct":
+ messages = [{"role": "user", "content": prompt_text}]
+ chat_text = session.tokenizer.apply_chat_template(messages, tokenize=False,
+ add_generation_prompt=True)
+ return session.tokenizer.encode(chat_text)
+ return session.tokenizer.encode(prompt_text)For "What is the capital of France?" with the instruct model, this returns ~30 tokens (the chat template adds system/user role markers).
In llama32_1b_inference.py:754 (run_once):
tokens = _tokenize_prompt(session, prompt_text) # length = real prompt
+prompt_len_actual = len(tokens) # save the real length
+if len(tokens) < session.seq_len:
+ tokens = tokens + [session.tokenizer.eos_token_id] * (session.seq_len - len(tokens))
+# Now len(tokens) == 2048 always.So if the real prompt is 30 tokens long, tokens becomes [real_0, real_1, ..., real_29, EOS, EOS, ..., EOS] with 2018 EOS tokens of padding.
The prefill function doesn't receive prompt_len_actual directly — it gets only the padded token_ids array. It recovers the real length by counting non-EOS tokens (llama32_1b_inference.py:422):
prompt_len = len([t for t in token_ids if t != tokenizer.eos_token_id])
+pred_pos = prompt_len - 1 # index of the last real prompt token<|begin_of_text|>, <|start_header_id|>, etc. — none of which are EOS — so this works in practice. If a prompt legitimately contained EOS, this counting would be wrong.
+The NPU runs the full forward pass over all 2048 positions including the EOS padding. The padding positions produce garbage k, v values. But we only use the logits at pred_pos = prompt_len - 1, which is BEFORE any padding (llama32_1b_inference.py:427):
# Final RMSNorm + LM Head — only on the last real-token row
+last_hidden = np.asarray(x_bf16, dtype=np.float32)[pred_pos:pred_pos + 1]
+last_normed_bf16 = _rms_norm(last_hidden, weights.final_norm).flatten().astype(bfloat16)
+
+# NPU LM Head GEMV (8 partitions) on the single normalized row
+results = decode_cache.load_and_run("lm_head_gemv", ...)
+logits_row = np.concatenate(results, axis=0)[:vocab_size]
+prefill_token = int(np.argmax(logits_row))This is one of the production optimizations: instead of running the LM Head GEMM on all 2048 positions and then taking row pred_pos, we extract just that one row first (CPU RMSNorm in <1 ms) and run a 1×128256 GEMV on the NPU. Saves ~150 ms of pointless compute.
After prefill, the KV cache has positions 0..2047 populated, but only positions 0..prompt_len-1 contain MEANINGFUL k/v (the rest are garbage from EOS padding). Decode starts at current_pos = prompt_len (llama32_1b_inference.py:573):
generated_tokens = [prefill_token]
+current_pos = prompt_len # skip past the garbage padding positions
+x_decode = weights.embed_table[prefill_token].astype(bfloat16)
+
+for token_idx in range(n_tokens):
+ # Run all 16 transformer blocks in decode mode
+ for layer_idx in range(config.n_layers):
+ x = run_decode_block(x, ..., k_cache[layer_idx], v_cache[layer_idx],
+ current_pos, ...)
+ # LM Head GEMV → next token
+ # ...
+ current_pos += 1 # cache grows by 1 per tokenInside decode_attention_cpu, the slicing k_cache[:, :current_pos+1, :] ensures we only attend to real prefill positions + actually-decoded positions. The garbage at indices prompt_len..2047 (left over from prefill processing the EOS padding) is never read — those slots are reused by decode if it generates enough tokens to overwrite them.
For a 30-token prompt padded to 2048, the prefill compute does 2048 / 30 ≈ 68× more work than necessary, because every layer processes 2018 padding positions whose results we throw away. This is a deliberate tradeoff: fixed-shape kernels are vastly easier to compile and faster per-position than dynamic-shape kernels would be on this hardware.
+ +Decode doesn't suffer from this — each decode call only processes ONE token (seq=1), and that token is the real new one.
+ +pred_pos = prompt_len - 1 = 6) is the FIRST generated token. It becomes generated_tokens[0]. Then decode generates tokens 1, 2, 3, ... and writes their k/v at cache positions prompt_len, prompt_len+1, .... The cache positions don't move; the cache just grows in-place into the previously-allocated max_seq array.
+Short answer: No. The hidden state at pred_pos = prompt_len − 1 is bit-identical to what you'd get if you ran with seq=prompt_len instead of seq=2048. (Same bytes, not just same logits.) This is why padding-with-EOS is a sound workaround, not a numerical approximation.
The reason: of the 14 ops in a transformer block (Part A2), only attention crosses positions. All other ops are per-position: each output row depends ONLY on its own input row. So the only path by which a padding position could contaminate pred_pos's output is through attention — and attention is causally masked, so pred_pos never sees positions later than itself. EOS padding tokens are by construction at indices ≥ prompt_len = pred_pos + 1, all of which the causal mask blocks.
Let x[i] denote the hidden state at position i. For each op, the question is: does the output at position pred_pos depend on any x[j] with j ≠ pred_pos?
| Op | Math | Cross-position? | Why / why not |
|---|---|---|---|
| Embedding lookup | +x[i] = embed_table[token_ids[i]] |
+ No | +Per-token table lookup. Position i depends only on token_ids[i]. | +
| RMSNorm | +x[i] · rsqrt(mean(x[i]²)+ε) · w |
+ No | +The mean is over the embedding dimension (2048 elements of one row), NOT over positions. RMSNorm at position i depends only on x[i]. Easy to verify: the norm formula has no sum across positions. | +
| Q/K/V projection | +Q[i] = x[i] @ Wq (etc.) |
+ No | +A matmul (seq, emb) @ (emb, out) is independent matmul per row. Q[i] = x[i] @ Wq. |
+
| RoPE | +rotate Q[i] by angle θ(i) from LUT |
+ No | +RoPE rotates each (position, head) pair by an angle that is a function of position alone. Q_roped[i] depends only on Q[i] and the constant LUT[i]. | +
| Attention | +out[i] = softmax(Q[i] · Kᵀ / √d, mask) · V |
+ Yes — but masked | +The ONLY cross-position op. With the causal mask, out[i] attends to positions 0..i ONLY. Position pred_pos attends to 0..pred_pos — strictly before any padding. Padding positions are at indices pred_pos+1..2047, all blocked. |
+
| O projection | +proj[i] = attn_out[i] @ Wo |
+ No | +Per-row matmul. | +
| Residual add | +res[i] = x[i] + proj[i] |
+ No | +Elementwise per row. | +
| FFN RMSNorm | +same as above | +No | +Per-row. | +
| Gate / Up GEMMs | +per-row matmul | +No | +Per-row. | +
| SwiGLU | +SiLU(gate[i]) * up[i] |
+ No | +Elementwise per row. | +
| Down GEMM | +per-row matmul | +No | +Per-row. | +
| Residual add #2 | +elementwise per row | +No | +Per-row. | +
| Final RMSNorm | +per-row | +No | +Per-row. | +
| LM Head | +logits[i] = x[i] @ W_lm.T |
+ No | +Per-row matmul. (And we only compute row pred_pos — see A7.) |
+
pred_pos, no padding position can leak into pred_pos's output through any pathway.
+The padding positions DO produce garbage output. EOS embeddings get RMSNormed, projected, RoPE-rotated, and run through attention (which can attend to real tokens earlier in the sequence — so the garbage is "garbage with prompt context"). But we never USE that garbage:
+ +pred_pos (see A7), so padding-position logits don't exist.prompt_len..2047 are written with garbage K/V from the padding positions. Decode skips them — it starts at current_pos = prompt_len and only reads cache slices 0..current_pos+1, never touching the garbage region. (Visualized in A4 and A5.)x_in at padding positions: this gets passed to the next transformer block, where it again produces garbage. Wasted compute, but causally walled off from pred_pos.No, because:
+pred_pos sums over only positions 0..pred_pos due to the causal mask. Padding positions don't enter the sum.You can prove the bit-identity empirically: run prefill on a 30-token prompt padded to 2048, then run prefill on the same 30 tokens with seq_len=30 (no padding) — assuming you have kernels compiled for seq=30, which production doesn't but the CPU reference does. Compare x_bf16[pred_pos] from both runs. They should be byte-equal.
This is something you have to script yourself if you ever need to re-prove it (make diagnosis probes the NPU vs HF bf16 per-layer cosine — see VERIFICATION.html — but it does not directly compare seq=30 vs seq=2048 padded).
Short answer: general optimization. Always sufficient for autoregressive single-stream generation, regardless of padding. Even a real seq=2048 prompt with no padding would only need the logits at the last position to generate the next token.
+ +Autoregressive language generation has a one-step lookahead: given hidden states for positions 0..T−1, the next token's distribution depends only on logits[T−1]. The logits at positions 0..T−2 would tell you "if I had sampled here, what would the next token be?" — but you've already committed to the actual tokens at those positions (they're the prompt). You don't re-sample them.
So the LM Head's job during inference is always the same: project ONE hidden state row (the last position's) into vocab space, argmax (or sample), produce ONE next token.
+ +| Use case | Why multi-row? | Used in this implementation? |
|---|---|---|
| Training (computing cross-entropy loss against teacher-forced labels) | +Loss is summed over all positions; need logits everywhere | +No — this is inference-only | +
| Speculative decoding (verify a draft model's K-token speculation) | +Need logits at K positions to score the speculation | +No — single-stream sampling only | +
| Beam search (track top-K candidate sequences) | +Need full distributions at each step for multiple beams | +No — greedy argmax (1 stream) | +
| Dumping logits for analysis / probing | +Researcher wants per-position logits for downstream analysis | +No | +
For the standard autoregressive sampling that this implementation does (greedy or top-k), you only need the last position's logits. This optimization holds whether your prompt fits in 30 tokens or 2048 tokens.
+ +| Approach | Compute | Why |
|---|---|---|
| Naive: full-seq LM Head | +(2048, 2048) @ (2048, 128256) = (2048, 128256) ≈ 1 TFLOP | +Computes 2047 rows you'll never look at | +
| This implementation: single-row GEMV | +(1, 2048) @ (2048, 128256) = (1, 128256) ≈ 0.5 GFLOP | +Only the row you need; ~2000× less compute | +
In wall time, this is the "Saves ~150 ms" optimization mentioned in profile.md. Implemented at llama32_1b_inference.py:425-446: extract the single hidden-state row, do RMSNorm on it (CPU, <1 ms because it's one row of 2048 elements), then call the decode-side lm_head_gemv.elf on that single row. The same ELF is reused for both prefill's last-token projection and per-token decode — they're the same operation (1×128256 GEMV).
Now to your bigger question: what's the difference between this implementation's padding-with-EOS and what a real production inference server does?
+ +Our approach is the simplest possible: compile kernels for one fixed shape (seq=2048), pad shorter prompts with EOS. This is appropriate for a research prototype on novel hardware where building a dynamic-shape compiler is itself a research problem.
+ +Production inference servers (vLLM, TensorRT-LLM, SGLang, llama.cpp, etc.) use much more sophisticated approaches:
+ +| Technique | What it does | This implementation? | Why production needs it |
|---|---|---|---|
| Dynamic-shape kernels | +Same kernel handles any seq length, branching at runtime on shape | +No — fixed seq=2048 | +Avoids waste on short prompts; supports any prompt length up to a limit | +
| Chunked prefill | +Split a long prompt (e.g., 32K tokens) into chunks of fixed size (e.g., 512), process sequentially with attention reading the cache for earlier chunks | +No — single-shot at seq=2048; longer prompts unsupported | +Supports prompts longer than the kernel's max seq length | +
| Continuous batching | +Pack multiple users' requests into one batch; add new requests / remove finished ones every step | +No — single user, single stream | +Maximize GPU/NPU utilization with multiple concurrent users | +
| Paged KV cache | +KV cache split into fixed-size pages (like virtual memory pages); attention gathers them at runtime | +No — contiguous (n_layers, n_kv_heads, max_seq, head_dim) array | +Avoids fragmentation and overcommit when serving many users with variable sequence lengths | +
| Speculative decoding | +Use a small draft model to speculate K tokens, verify in one big-model forward pass | +No — vanilla autoregressive | +~2-3× decode speedup at the cost of ~10-30% extra compute | +
| Quantization (INT8/INT4) | +Compress weights to lower precision, dequantize in kernel | +No — bf16 throughout | +~2-4× speedup, ~2-4× memory reduction | +
| Multi-node tensor/pipeline parallelism | +Shard model across multiple devices | +No — single NPU | +Required for models larger than one device's memory | +
To return to your distinction: these are two completely separate things.
+ +| Optimization | Always applicable? | Why |
|---|---|---|
| Single-row LM Head GEMV at the end of prefill | +Yes, always. Production servers do this too. | +Autoregressive sampling only needs the last row's logits, regardless of how the prompt was processed. | +
| Pad short prompts with EOS to 2048 | +No — specific to fixed-shape kernels. Production usually avoids this. | +It wastes compute (~68× for a 30-token prompt). Only acceptable when dynamic-shape kernels would be even more expensive (e.g., due to compile time, runtime branching cost, or tooling immaturity). | +
So when you read the LM Head GEMV code, don't think "this is a workaround". Think "this is the right thing to do, and it happens to also dodge an extra 2047 wasted rows that the padding would have created if we used the full-seq GEMM here".
+ + +Part A was the model. Now we look at how this codebase realizes those ops on AMD NPU2. The translation is not 1-to-1: the model has 14 ops per layer; production runs them as 3 NPU kernel calls per layer (rms_gemms_rope = ops 1-6, flash_attn = op 7, o_ffn = ops 8-15). That's the "multi-launch merging" optimization at work.
+ +One inference's prefill phase: from the input prompt to the first generated token. The diagram shows which steps run on CPU (gray, host-side numpy) vs which run on NPU (purple, stitched ELFs). FA is its own ELF (pink-purple); the per-layer triple (rms_gemms_rope.elf, flash_attn.elf, o_ffn.elf) is grouped inside the "decoder block × 16" container. KV cache extraction happens on the host after each layer.
+ + + +Read the colors: gray = CPU/host (numpy, embedding lookup, KV cache management, argmax), purple = NPU stitched ELF, pink = NPU FlashAttention (always its own ELF, never stitched — see B3). The dashed purple outline marks the 16-layer loop boundary.
+ +Decode generates ONE token per pass. Per layer it makes 2 NPU calls + 1 CPU step (because attention runs on CPU during decode — see B9 for why). The KV cache is read+appended on each layer.
+ + + +| Phase | NPU calls per layer | NPU calls total | CPU work per layer |
|---|---|---|---|
| Prefill (1 pass, 16 layers) | 3 (rms_gemms_rope + flash_attn + o_ffn) | 48 + 1 (lm_head_gemv) = 49 | KV cache write (numpy slice assign) |
| Decode (1 token, 16 layers) | 2 (rms_gemv_rope + o_gemv_ffn) | 32 + 1 (lm_head_gemv) = 33 | decode_attention_cpu (single-query GQA against KV cache) |
NPU2 (AMD Strix, AIE2P architecture) has a 32-tile compute array arranged as 8 columns × 4 rows. Plus 8 mem-tiles (L2) and shim tiles for DMA. Each compute tile is a VLIW vector core with its own L1 SRAM. Different kernels use different subsets of the 32 tiles depending on parallelism strategy:
+ +| Herd shape | Tiles used | Used for (typical) |
|---|---|---|
[8, 4] | 32 / 32 (full) | Prefill GEMMs (Q/K/V/O/Gate/Up/Down). M-dim split 8 ways × N-dim split 4 ways. |
[8, 1] | 8 / 32 | RMSNorm, RoPE (prefill), SwiGLU, eltwise add, GEMV (decode). Row-parallel across one column of tiles. |
[1, 1] | 1 / 32 | RoPE (decode) — single tile is enough for the tiny single-token rotation. |
Cascade [c_nq, c_ns] | varies | FlashAttention — uses an internal segment + cascade-stages design (4 stages × per-head segments). Hard to give one number; FA stresses the array more than any other single ELF. |
Each kernel's exact tile usage is listed in B2's per-kernel cards. The choice of herd shape is made by the Python builder (passed as herd_x / herd_m / herd_n kwargs) and locked at compile time — it can't change between calls of the same ELF.
llama32_1b_inference.py:mainFrom make run to printed output:
Phase 1: build_session llama32_1b_inference.py:669
+One-time setup: create KernelCache instances, compile (or load cached) all ELFs, load model weights from HuggingFace, build the RoPE LUT, call prepare_runtime.
Phase 2: prepare_runtime llama32_1b_inference.py:129
+Pre-loads ALL weights for ALL 16 layers into per-layer NPU Buffer Objects (BOs), so subsequent inference calls only need to write activations. This is the single biggest cost-amortization in the pipeline (see B7).
+ +Phase 3: run_once / generate llama32_1b_inference.py:742, 523
+Tokenize the prompt → pad to seq_len=2048 (see Part A5) → call run_npu_prefill → enter the decode loop.
Phase 4: decode/print
+For instruct models, apply chat template; emit tokens incrementally via the streaming callback in interactive mode.
+# One-time compile (~3 min)
+make compile
+
+# Run inference
+make run
+make run PROMPT="..."
+
+# With profiling breakdown
+make profile
+
+# Top-k token-level correctness gate vs HF transformers bf16
+make verify
+
+# Per-layer ffn_out cosine vs HF bf16 (informational)
+make diagnosis
+
+# Interactive REPL
+make chatBefore discussing optimizations (multi-launch ELF stitching, BO management), let's see what the basic units are. The codebase has 7 unique compute kernels that together implement every model op from Part A. Each kernel is one of two implementation patterns:
+ +| Pattern | How it works | Used for |
|---|---|---|
| MLIR-only (codegen) | +The Python builder constructs an MLIR module that describes the operation in the linalg / scf / air dialects. aircc + aiecc lower it to AIE-tile instructions through standard linalg-vectorize and AIR placement passes. Peano compiles the resulting per-tile LLVM IR. No hand-written C++. | +RMSNorm, GEMM, eltwise add | +
| MLIR + external C++ kernel | +The MLIR module declares func.func private @kernel_name { link_with = "kernel.o" } and calls it from inside an air.herd. The .o is a hand-written C++ kernel compiled separately by Peano (LLVM-AIE). aiecc links the .o into the per-tile ELFs. |
+ GEMV, RoPE, SwiGLU, FlashAttention | +
External C++ is used when a hand-tuned implementation beats codegen — typically for kernels with non-trivial vectorization patterns, double-buffering, or tile-level fused operations (FA's softmax + MMA fusion is the canonical example).
+ +For external-C++ kernels, the .o file is compiled by Peano in advance (see kernel_builder/external_kernels.py) and placed in the build directory before aircc runs; aiecc finds it via the link_with attribute when packaging per-tile ELFs.
The whole pipeline is invoked by XRTBackend.compile(mlir_module) inside KernelCache.compile_and_cache — see kernel_builder/cache.py:251. (B3 covers stitching multiple kernels into one ELF; this section is just the per-kernel building blocks.)
| Kernel | Pattern | Maps to model op (Part A) | Source builder | External C++ (if any) |
|---|---|---|---|---|
| RMSNorm | MLIR-only | RMSNorm (attn-norm, ffn-norm, final-norm) | weighted_rms_norm/weighted_rms_norm.py | — |
| GEMM | MLIR-only | Q/K/V/O proj, Gate/Up/Down proj (prefill, S=2048) | kernel_builder/gemm_builder.py | — |
| GEMV | MLIR + C++ | Q/K/V/O proj, Gate/Up/Down proj (decode, S=1); LM Head | matrix_vector_multiplication/bf16/matvec.py | mv.cc → mv.o + mv_k8192.o |
| RoPE | MLIR + C++ | RoPE Q, RoPE K | rope_lut/rope_lut.py | kernel_builder/rope_halfsplit.cc → rope.o |
| SwiGLU | MLIR + C++ | SiLU(gate) ⊙ up — fused | kernel_builder/ffn_swiglu/silu_and_mul.py | kernel_builder/ffn_swiglu/silu_and_mul.cc → silu_and_mul.o |
| FlashAttention | MLIR + C++ | Scaled dot-product attention (causal, GQA) | flash_attention/kernel_fusion_based/attn_npu2_seqfirst.py | flash_attention/kernel_fusion_based/attn_npu2.cc → attn.o |
| Eltwise Add | MLIR-only | Residual add #1, Residual add #2 | eltwise_add/eltwise_add.py | — |
External-C++ .o compilation is centralized in kernel_builder/external_kernels.py, which uses Peano (LLVM-AIE, found via $PEANO_INSTALL_DIR) with --target=aie2p-none-unknown-elf -O2 -std=c++20. Each function (compile_silu_and_mul, compile_rope, etc.) checks if the .o already exists and skips if so.
| Source builder | programming_examples/weighted_rms_norm/weighted_rms_norm.py |
| External C++ | None — pure MLIR/codegen |
| Maps to model op | RMSNorm (Part A2 op #1, #10; final norm in Part A3) |
| Production usage | Inside rms_gemms_rope.elf + o_ffn.elf (prefill); rms_gemv_rope.elf + o_gemv_ffn.elf (decode); the final RMSNorm at the end of inference is computed on CPU instead (single row only — see A7) |
| NPU compute tile usage | herd [8, 1] = 8 of 32 tiles. One column of 8 tiles, each tile reducing across one slice of rows. Same shape used in both prefill and decode (the per-row reduction doesn't benefit from row-direction parallelism beyond the column count). |
How it's compiled. The Python builder uses FuncOp.from_py_func + @herd to construct an air.herd that does the per-row reduction (sum-of-squares), then the rsqrt + multiply. There's no external C++ — aircc lowers the linalg/scf/arith ops to AIE-tile vector intrinsics, and Peano then turns the per-tile LLVM IR into AIE2P machine code.
The op: y[i] = x[i] · rsqrt(mean(x[i]², dim=-1) + ε) · γ per row. γ (the learned scale) is a per-feature [H]-shaped weight broadcast across rows. The implementation tiles the row dim across an herd_x-tile-tall herd; each tile reduces and normalizes its rows.
Quirk: the builder produces a bare air.herd (not wrapped in air.launch). When stitched into a multi-launch ELF, the stitching code wraps it in air.launch { air.segment { herd } } via _wrap_ir_in_launch from kernel_builder/stitching.py. (See B5 for why this wrapping is needed.)
| Source builder | programming_examples/llama32_1b/kernel_builder/gemm_builder.py (function _build_gemm_module(m, k, n, ...)) — thin wrapper around the upstream BF16 GEMM |
| Wraps | programming_examples/matrix_multiplication/bf16/run.py (function build_module(m, k, n, tile_m, tile_k_l2, tile_k_l1, tile_n, herd_m, herd_n, np_dtype_in, np_dtype_out, arch, direct_codegen)) — the generic BF16 GEMM module builder shared with the standalone GEMM example |
| External C++ | None — codegen via aircc's linalg.matmul lowering |
| Maps to model ops | Q proj, K proj, V proj, O proj, Gate proj, Up proj, Down proj (Part A2 ops #2-#4, #8, #11-#12, #14) — during prefill only, where S=2048 makes a true matrix-matrix GEMM |
| Production usage | rms_gemms_rope.elf contains 3 GEMMs (Q, K, V); o_ffn.elf contains 4 GEMMs (O, Gate, Up, Down) |
| NPU compute tile usage | herd [8, 4] = 32 of 32 tiles. Production sets herd_m=8, herd_n=4 — the herd's M dim (8) parallelizes output-row tiles and the N dim (4) parallelizes output-col tiles. This is the only kernel that uses the full NPU2 compute array. Configured per-GEMM in rms_gemms_rope_multi.py:200-209 and o_ffn_multi.py:182-202. |
Relationship to the upstream programming_examples GEMM. There is NOT a separate Llama-specific GEMM kernel. gemm_builder.py is a 30-line wrapper that:
build_module from programming_examples/matrix_multiplication/bf16/run.py with bfloat16 input AND output, arch="aie2p" (NPU2), and direct_codegen=True. This produces a base MLIR module containing one air.herd wrapping a tiled linalg.matmul.GEMM_TRANSFORM_IR string in gemm_builder.py) on top of that module. The transform script does additional tiling, herd-vectorization, vector-contract → f32 cast lifting, and several rounds of cast-pair hoisting that move arith.extf / arith.truncf ops out of the innermost loops.Without the transform-IR step, the GEMM compiles but the inner-loop quality is significantly worse (extra bf16↔f32 conversions per MMA iteration). The transform script is what makes the production GEMM competitive with hand-written kernels — but the actual linalg.matmul tiling structure comes from the shared upstream builder, not from the wrapper.
Tile config (prefill default). The wrapper accepts tile_m, tile_k_l2, tile_k_l1, tile_n, herd_m, herd_n. Production uses different configs per GEMM (smaller L2 tiles for the small Q/K/V/O 2048-emb GEMMs, larger for the wider Gate/Up/Down 8192-D_ff GEMMs). All configs come from multi_launch_builder/rms_gemms_rope_multi.py:200-209 and multi_launch_builder/o_ffn_multi.py:182-202.
Why no external C++. The aircc + aiecc pipeline can lower a tiled linalg.matmul with the right transform IR to the same AIE MMA intrinsic that a hand-written kernel would use. There's no measurable win from hand-rolling the matmul C++.
| Source builder | programming_examples/matrix_vector_multiplication/bf16/matvec.py (function build_module(M, K, tile_m, m_input, herd_m, ...)) |
| External C++ | programming_examples/matrix_vector_multiplication/bf16/mv.cc → compiled to mv.o (and mv_k8192.o, see below) |
| Maps to model ops | Q/K/V/O/Gate/Up/Down projections — during decode (S=1 makes it M=1 GEMV); also the LM Head (which is structurally a 1×V GEMV regardless of phase, see A7) |
| Production usage | rms_gemv_rope.elf contains 3 GEMVs (Q, K, V); o_gemv_ffn.elf contains 4 GEMVs (O, Gate, Up, Down); lm_head_gemv.elf is an 8-partition GEMV stitched 8 times |
| NPU compute tile usage | herd [8, 1] = 8 of 32 tiles. Production sets tile_m=8, m_input=4, herd_m=8 — the herd's 8 tiles parallelize the M output dim. With M=1 (S=1 in decode) the GEMV gets ZERO M-direction parallelism within a single tile — the 8 tiles instead each handle a slice of the output rows of the projection. The Down GEMV (K=8192) uses a renamed mv_k8192.o variant with tile_m=2 but the same 8-tile herd shape. |
How it's compiled. The MLIR builder constructs an air.launch wrapping an air.herd whose body calls the C++ kernel @matvec_vectorized_bf16_bf16 (declared private with link_with = "mv.o"). The C++ in mv.cc implements a hand-vectorized y = W @ x using AIE bf16 MMA intrinsics. Peano compiles this to a .o file via kernel_builder/external_kernels.py:compile_mv:
def compile_mv(tile_m=8):
+ src = _PROJ_ROOT / "matrix_vector_multiplication" / "bf16" / "mv.cc"
+ _compile_kernel(src, "mv.o", extra_flags=[f"-DDIM_M_OUTPUT={tile_m}"])The mv_k8192.o trick. The decode o_gemv_ffn.elf needs TWO GEMV variants in one ELF: K=2048 (for O/Gate/Up/normal slots) and K=8192 (for the Down GEMV). MLIR can't have two private functions with the same name and different signatures — so the same mv.cc source is compiled a SECOND time with renamed entry points via -D macros (see kernel_builder/external_kernels.py:155):
def compile_mv_k8192():
+ _compile_kernel(src, "mv_k8192.o", extra_flags=[
+ "-DDIM_M_OUTPUT=2",
+ "-Dmatvec_vectorized_bf16_bf16=dg_matvec_vectorized_bf16_bf16", # renamed
+ "-Dlinalg_fill_bf16=dg_linalg_fill_bf16",
+ ])The renamed function appears in the merged ELF as a separate symbol, side-by-side with the K=2048 version.
+ + +| Source builder | programming_examples/rope_lut/rope_lut.py (decode/per-row); for prefill multi_launch_builder/rms_gemms_rope_multi.py:_build_rope_2d wraps it for 2D inputs |
| External C++ | programming_examples/llama32_1b/kernel_builder/rope_halfsplit.cc → compiled to rope.o |
| Maps to model op | RoPE Q, RoPE K (Part A2 ops #5, #6) |
| Production usage | rms_gemms_rope.elf + rms_gemv_rope.elf (one RoPE for Q-side, one for K-side per ELF) |
| NPU compute tile usage | Prefill: herd [8, 1] = 8 of 32 tiles (rope_herd_x=8, herd_y=1 in rms_gemms_rope_multi.py; the 8 tiles split the seq dim S=2048 across rows). Decode: herd [1, 1] = 1 of 32 tiles (rope_herd_x=1 in rms_gemv_rope_multi.py; only one row to rotate, so single-tile is sufficient and avoids DMA fan-out overhead). |
How it's compiled. The MLIR builder constructs an air.herd that DMA-loads one row of (cos, sin) LUT plus one row of input data into L1, then calls @rope (declared with link_with = "rope.o"). The C++ in rope_halfsplit.cc implements the per-position rotation.
The rope_halfsplit.cc story. Two RoPE conventions exist:
(d[i], d[i + d_h/2]) for rotation. LUT layout: [cos_0, ..., cos_{d_h/2-1}, sin_0, ..., sin_{d_h/2-1}].llama.cpp and the original RoPE paper): pair (d[2i], d[2i+1]). LUT layout: [cos_0, sin_0, cos_1, sin_1, ...].Mixing the two produces wrong outputs. The upstream aie_kernels/aie2p/rope.cc uses the interleaved convention. Llama-3.2-1B needs half-split, so this codebase has its own rope_halfsplit.cc compiled to the same rope.o filename → drop-in replacement, no MLIR changes needed. See kernel_builder/external_kernels.py:119 (compile_rope):
def compile_rope():
+ src = Path(__file__).resolve().parent / "rope_halfsplit.cc" # NOT the upstream rope.cc
+ _compile_kernel(src, "rope.o")The LUT (cos/sin table) is precomputed once per session by generate_rope_lut in llama32_1b_weights.py and passed as a kernel input — not compiled into the kernel.
| Source builder | programming_examples/llama32_1b/kernel_builder/ffn_swiglu/silu_and_mul.py |
| External C++ | programming_examples/llama32_1b/kernel_builder/ffn_swiglu/silu_and_mul.cc → compiled to silu_and_mul.o |
| Maps to model ops | SiLU(gate) + elementwise multiply (Part A2 ops #13 — fused into one kernel) |
| Production usage | o_ffn.elf + o_gemv_ffn.elf (one fused SwiGLU step between gate/up GEMMs and down GEMM) |
| NPU compute tile usage | herd [8, 1] = 8 of 32 tiles (swiglu_herd_x=8, swiglu_herd_y=1). The 8 tiles split the elementwise work across the row dim. SiLU+multiply is memory-bound at this scale — adding more tiles wouldn't help because L2/L1 DMA bandwidth is already saturated. |
How it's compiled. The MLIR builder constructs an air.herd that takes the gate and up tensors as inputs (each [B, S, D_ff]) and produces one output tensor. The herd body calls @silu_and_mul_bf16 (declared with link_with = "silu_and_mul.o"). The C++ implementation does out[i] = SiLU(gate[i]) · up[i] in a vectorized inner loop using AIE bf16 SiLU + multiply intrinsics — fusing the two ops eliminates one full pass over the 8192-wide tensor (vs. doing SiLU and the multiply as two separate kernels).
Compile (with extra include for utils header): see kernel_builder/external_kernels.py:106 (compile_silu_and_mul):
def compile_silu_and_mul():
+ src = _PROJ_ROOT / "llama32_1b" / "kernel_builder" / "ffn_swiglu" / "silu_and_mul.cc"
+ include_dir = _get_aie_include_dir()
+ utils_header = Path(include_dir) / "aie_kernels" / "aie_kernel_utils.h"
+ extra = []
+ if utils_header.exists():
+ extra = ["-include", str(utils_header)]
+ _compile_kernel(src, "silu_and_mul.o", extra_flags=extra)| Source builder | programming_examples/flash_attention/kernel_fusion_based/attn_npu2_seqfirst.py (function build_module(lk, lkp, lq, lqp, dk, dv, num_q_tiles, num_cascade_stages, num_heads, num_kv_heads, causal)) |
| External C++ | programming_examples/flash_attention/kernel_fusion_based/attn_npu2.cc → compiled to attn_npu2.o (also copied to attn.o) |
| Maps to model op | Scaled dot-product attention (Part A2 op #7) with causal mask + GQA |
| Production usage | flash_attn.elf — its OWN ELF, never stitched with rms_gemms_rope or o_ffn (un-mergeable, see B5) |
| NPU compute tile usage | Cascade design — uses ~16-24 tiles depending on config. Production sets num_q_tiles=4, num_cascade_stages=4, num_heads_per_unroll=2. The kernel uses MULTIPLE air.segments (sized [num_heads_per_unroll, 1]) each containing a herd sizes=[c_nq, c_ns]. Effectively the cascade pipelines Q-tile streaming across stages — different from the single-herd pattern of the other 6 kernels. Decode reuses prefill's flash_attn.elf only for full-prefill recomputation (rare); the per-token decode attention runs on CPU instead. |
How it's compiled. Of all 7 kernels, FlashAttention is by far the most complex. The MLIR builder produces a multi-tile cascade of air.herds that stream Q tiles through K/V tiles using air.channels for inter-tile DMA. The actual softmax + MMA fusion is in C++ (attn_npu2.cc), which exposes ~16 functions for the FA tile primitives (Q tile load, K tile load, dot-product, online softmax update, V multiply-accumulate, rescale, etc.).
Many compile-time flags. See kernel_builder/external_kernels.py:130 (compile_attn_npu2):
def compile_attn_npu2(head_dim=64):
+ src = _PROJ_ROOT / "flash_attention" / "kernel_fusion_based" / "attn_npu2.cc"
+ _compile_kernel(src, "attn_npu2.o", extra_flags=[
+ "-DBIT_WIDTH=8",
+ f"-Dlqp={head_dim}", # Q-per-tile
+ f"-Dlkp={head_dim}", # K-per-tile
+ f"-Ddk={head_dim}", # head dim, K side
+ f"-Ddk_full={head_dim}",
+ f"-Ddv={head_dim}", # head dim, V side
+ f"-Ddv_full={head_dim}",
+ "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
+ "-DROUND_CONV_EVEN",
+ ])
+ # Some link_with attrs use "attn.o", so make a copy
+ shutil.copy2("attn_npu2.o", "attn.o")Most of these -D flags are head_dim parameters that the C++ uses to size internal tile buffers at compile time. head_dim=64 for Llama-3.2-1B; the same kernel works for Llama-3.2-3B with head_dim=128.
Why this can't go in a multi-launch ELF. The cascade design uses many air.channels and stresses the air-opt-shim-dma-bds compiler pass quadratically. With 9+ launches (i.e., FA + the rms_gemms_rope launches) in one ELF, this pass takes >10 minutes. So FA stays as its own single-launch ELF and is invoked between rms_gemms_rope and o_ffn from the host (see B5). This is the main reason production has 3 NPU calls per layer instead of 1.
| Source builder | programming_examples/eltwise_add/eltwise_add.py; specialized 2D and 2D→1D variants are defined locally in multi_launch_builder/o_ffn_multi.py (_build_add_2d_to_2d, _build_add_2d_to_1d) |
| External C++ | None — pure MLIR/codegen |
| Maps to model op | Residual #1 (after attention), Residual #2 (after FFN) (Part A2 ops #9, #15) |
| Production usage | Two adds inside o_ffn.elf (one for each residual); two analogous adds inside o_gemv_ffn.elf |
| NPU compute tile usage | herd [8, 1] = 8 of 32 tiles. The 8 tiles split the row dim. Pure DMA-bound: the add itself is one cycle per element, so total time = DDR↔L1 transfer time. More tiles wouldn't help. |
How it's compiled. The simplest kernel: an air.herd with a tiled elementwise loop, lowered by aircc to the AIE add intrinsic. The 2D and 2D→1D variants exist because the residual outputs may be consumed as flat 1D arrays by the next sub-launch (e.g., the final o_ffn output is 1D n_total = seq*emb); the variant just calls memref.collapse_shape internally to handle the type mismatch.
Quirk: like RMSNorm, the simple add builder produces a bare air.herd; multi-launch stitching wraps it via _wrap_ir_in_launch.
Two files coordinate the actual external-C++ compilation:
+ +| File | What it does |
|---|---|
| kernel_builder/external_kernels.py | Per-kernel compile_* functions (one per .o) + a compile_all_external_kernels(head_dim) top-level that runs all 5 (silu_and_mul, rope, attn, mv, mv_k8192). Each uses Peano via $PEANO_INSTALL_DIR/bin/clang++. Skips compilation if the .o already exists. |
| kernel_builder/cache.py:prepare_air_project | Called from compile_and_cache before each ELF compile. Cleans air_project/, calls compile_all_external_kernels, then copies all .o files into air_project/ where aiecc's link_with search path will find them. |
So the flow for compiling one ELF is: prepare_air_project → external C++ .o files exist in air_project/ → backend.compile(mlir_module) runs aircc + aiecc, which links the .os into the per-tile ELFs → output .elf + .insts.bin are copied into cache_dir/.
.o files (GEMV, RoPE, SwiGLU, FlashAttention). A single ELF can contain one or many of these — see B5 for stitching.
+Side-by-side view of how each of the 7 kernels maps onto the NPU2 8×4 compute array:
+ +| Kernel | Phase | Herd shape | Tiles | Why this shape |
|---|---|---|---|---|
| RMSNorm | Both | [8, 1] | 8 | Per-row reduction; 8-tile column splits rows |
| GEMM | Prefill | [8, 4] | 32 | Full 2D output-tile parallelism (M and N) |
| GEMV | Decode | [8, 1] | 8 | M=1 forces output-row-only parallelism |
| RoPE | Prefill | [8, 1] | 8 | S=2048 rows split across 8 tiles |
| RoPE | Decode | [1, 1] | 1 | Only 1 row to rotate; multi-tile would just add fan-out overhead |
| SwiGLU | Both | [8, 1] | 8 | Memory-bound; more tiles wouldn't help |
| Eltwise Add | Both | [8, 1] | 8 | DMA-bound; 1-cycle add |
| FlashAttention | Prefill | cascade [c_nq, c_ns] | ~16-24 | Multi-segment Q-tile cascade pipeline |
Observation: only the prefill GEMM uses the entire 32-tile array. Most kernels use 8 tiles (one column) — they are limited by either the reduction structure (RMSNorm) or by DMA bandwidth (SwiGLU, eltwise add). For decode, the loss of M-direction parallelism (M=1) means there is simply no work for the additional column dim, so even GEMV drops to 8 tiles. Implication: the M=1 decode path leaves 24/32 = 75% of the compute array idle on every dispatch, which is one reason the per-token throughput is dispatch-overhead-bound.
+ + +B2 covered each kernel as a standalone unit — what it computes, how it's compiled, and how many tiles it uses. But you cannot just chain those 7 kernels together and get a working 1.27 s prefill. Several practical problems sit between "I have a working RMSNorm kernel" and "I have a 16-layer transformer running on the NPU":
+ +| Gap | Problem if unsolved | Solution | Section |
|---|---|---|---|
| #1 — Layout matching | +Kernel A's output shape/layout doesn't match what kernel B expects to read. Naive chaining produces wrong values or silently misaligned data. | +CPU pre-transpose of weights, free MLIR reshapes, deliberate physical KV-cache transpose on the host side, mv_k8192 macro-rename trick. | +B4 | +
| #2 — XRT dispatch overhead | +Each xrt.run() call has ~100 µs fixed overhead. With 49 kernels per prefill pass × 16 layers, dispatch alone would dominate runtime. |
+ Stitch multiple air.launchs into one ELF so 6-8 logical kernels run from a single xrt.run() call. Intermediates flow via DDR, host stays out of the loop. |
+ B5 | +
| #3 — Per-call BO management | +Naive flow re-allocates and re-uploads every kernel argument on every call. A 14 MB weight tensor uploaded per kernel call would dominate the ~30 ms-per-call budget. | +Allocate XRT Buffer Objects once, classify each arg as static (write-once), intermediate (no host transfer at all), or output (host-readable). Skip everything that hasn't changed. | +B6 | +
| #4 — Compile time + per-layer state | +Each ELF compile takes ~30-50 s. Recompiling on every script start costs 3+ minutes. Also: 16 layers × 6 ELFs × N weights each → which BO holds which layer's weights? | +KernelCache persists compiled ELFs to disk, caches loaded XRT contexts in process, and maintains per-layer BO sets keyed by bo_key="rms_gemms_rope_L{layer}". |
+ B7 | +
Sections B4-B7 cover each gap one at a time. Once they're all in place, the prefill (B8) and decode (B9) detail sections show the four gaps working together on real per-layer code paths. B10 is the final code map.
+ +KernelCache needs to track per layer (B7). Skipping ahead leaves you with isolated tricks; reading in order shows why each was necessary.
+The 7 building-block kernels were each developed in their own standalone programming_examples demo. Their input/output layouts were chosen for that demo's convenience — not for chaining into a transformer. Several layout mismatches show up the moment you try to feed one kernel's output into another:
HuggingFace stores Llama weights as (out_features, in_features): e.g. wq has shape (2048, 2048) with the FIRST dim being the output. The standalone GEMV kernel, however, expects A[M, K] with M=output, K=input — but reads A contiguously in K-major order (last dim is the contiguous one). HuggingFace storage is output-major. Naive use → reading the wrong elements per MMA, silent garbage output.
Fix: CPU pre-transpose every decode-side weight matrix once, before any timing starts. Implemented in llama32_1b_inference.py:171-197 inside prepare_runtime:
# Pre-transpose all decode GEMV weights (one-time, before timing)
+for lw in weights.layers:
+ lw._wq_t = np.ascontiguousarray(lw.wq.astype(bfloat16).reshape(emb_dim, emb_dim).T)
+ lw._wk_t = np.ascontiguousarray(lw.wk.astype(bfloat16).reshape(emb_dim, kv_dim).T)
+ lw._wv_t = np.ascontiguousarray(lw.wv.astype(bfloat16).reshape(emb_dim, kv_dim).T)
+ lw._wo_t = np.ascontiguousarray(lw.wo.astype(bfloat16).reshape(emb_dim, emb_dim).T)
+ lw._wgate_t = np.ascontiguousarray(lw.w_gate.astype(bfloat16).reshape(emb_dim, hidden_dim).T)
+ lw._wup_t = np.ascontiguousarray(lw.w_up.astype(bfloat16).reshape(emb_dim, hidden_dim).T)
+ lw._wdown_t = np.ascontiguousarray(lw.w_down.astype(bfloat16).reshape(hidden_dim, emb_dim).T)The .T + ascontiguousarray physically reorders the weight matrix bytes in DDR so the GEMV kernel reads them in K-major order naturally. This costs ~50 ms per layer × 16 layers ≈ 800 ms ONCE at startup, then never again — the transposed buffers live on as _wq_t, _wk_t, etc. and get uploaded to NPU BOs during weight preload.
Why CPU and not on the NPU? The NPU DMA engine has stride=1 mandatory for sub-32-bit types (it can't do a strided BF16 DMA). Doing the transpose during DMA-in would require shape rearrangement that the DMA hardware refuses. So the transpose lives in numpy on the CPU.
+ +The same physical KV tensor is touched by three different consumers, each with its own preferred layout:
+ +| Consumer | Wants layout |
|---|---|
| RoPE K kernel output (prefill) | [seq, n_kv_heads, head_dim] — sequence-major |
| FlashAttention input (prefill) | [seq, n_kv_heads, head_dim] — sequence-major (matches RoPE) |
| KV cache storage (host) | [n_kv_heads, max_seq, head_dim] — head-major (so per-head slicing is contiguous) |
| Decode CPU attention (per-token reads) | [n_kv_heads, current_pos+1, head_dim] — needs head-major for fast per-head dot-products |
Solution: the prefill kernels keep the seq-major layout that RoPE produces (so RoPE→FlashAttention has a free zero-cost layout match), and the host transposes once after each layer's prefill output to populate the head-major KV cache. From llama32_1b_inference.py:401-410:
+ +k_cache[layer_idx, :, :seq_len, :] = (
+ intermediates["k_roped"]
+ .astype(bfloat16)
+ .reshape(seq_len, n_kv_heads, head_dim)
+ .transpose(1, 0, 2) # seq-major → head-major
+)
+v_cache[layer_idx, :, :seq_len, :] = (
+ intermediates["v"].astype(bfloat16)
+ .reshape(seq_len, n_kv_heads, head_dim)
+ .transpose(1, 0, 2)
+)This transpose runs on the CPU (~1 ms per layer) for the same DMA-stride reason as Mismatch #1. The bf16 stride=1 hardware limit means you cannot do a layout transpose during NPU DMA-out; the host has to materialize the head-major view itself. (See BF16 DMA stride limitation note in project docs.)
Q/K GEMM emits [seq, n_heads * head_dim] as a flat 2D tensor. RoPE expects [seq, n_heads, head_dim] so it can apply the per-(head, dim/2) rotation. This one is FREE — it's a pure shape view, no data movement. The MLIR builder uses memref.expand_shape on the L2 buffer between the GEMM air.launch and the RoPE air.launch inside the same stitched ELF (no DDR round-trip, no DMA reshape). Same trick at the eltwise-add → next-RMSNorm boundary.
o_ffn.elf's final output (after the second residual add) is shaped [seq, emb] as far as the math cares, but the next layer's rms_gemms_rope.elf wants its input as a flat 1D [seq * emb] buffer (because that's how the leading RMSNorm's L2 tile shape was specified). The eltwise-add kernel gained a _build_add_2d_to_1d variant that calls memref.collapse_shape internally so the producer and consumer agree on a flat 1D buffer. See multi_launch_builder/o_ffn_multi.py.
The decode o_gemv_ffn.elf contains FOUR GEMVs: O, Gate, Up, and Down. Three of them have K=2048 (the embedding dim); the Down GEMV alone has K=8192 (the FFN hidden dim, accumulating back to embedding). MLIR can't have two private functions with the same name and different signatures in one module.
Solution (from kernel_builder/external_kernels.py:155): compile mv.cc a SECOND time with macro renames, producing a separate symbol for the K=8192 variant:
def compile_mv_k8192():
+ _compile_kernel(src, "mv_k8192.o", extra_flags=[
+ "-DDIM_M_OUTPUT=2",
+ "-Dmatvec_vectorized_bf16_bf16=dg_matvec_vectorized_bf16_bf16", # renamed
+ "-Dlinalg_fill_bf16=dg_linalg_fill_bf16",
+ ])Both .o files end up in air_project/ at link time. The MLIR module references each one by its (renamed) symbol, and the linker happily places both into the same ELF.
The problem. Each xrt.run() call has fixed dispatch overhead (kernel-handle lookup, host↔device synchronization) of ~100 µs. With 7 kernels per layer × 16 layers = 112 NPU calls per prefill pass, dispatch alone is ~11 ms — small relative to a 1.2 s prefill, but devastating for decode where each kernel does only hundreds of µs of NPU work. For decode, raw dispatch overhead can rival the actual compute time.
The fix. Combine multiple kernels into one ELF that runs in one xrt.run() call. The host issues one dispatch; intermediates flow between sub-kernels via DDR using NPU DMA, with no host involvement. From the host's view, "rms_gemms_rope" looks like one kernel even though it's really 6 stitched air.launchs back-to-back.
An MLIR module can contain multiple air.launch operations inside a single func.func. Each air.launch wraps an air.segment wrapping air.herd(s) — i.e., one logical kernel. When that combined module is compiled to one ELF and invoked by one xrt.run(), the launches execute sequentially and intermediates flow between them via DDR using NPU DMA — without CPU involvement.
The Python builders in multi_launch_builder/*_multi.py do this stitching. They take individual MLIR modules (from B2's per-kernel builders) as text strings and concatenate the function bodies into one combined func, with SSA values renamed to avoid collisions.
The production code stitches the 7 kernel building blocks from B2 into 6 ELFs:
+ +| ELF | Phase | Stitched kernels | Builder | Compile time |
|---|---|---|---|---|
rms_gemms_rope.elf | Prefill | 6: RMSNorm + Q GEMM + K GEMM + V GEMM + RoPE Q + RoPE K | multi_launch_builder/rms_gemms_rope_multi.py:193 | ~33 s |
flash_attn.elf | Prefill | 1: FlashAttention | flash_attention/.../attn_npu2_seqfirst.py | ~46 s |
o_ffn.elf | Prefill | 8: O GEMM + Add + RMSNorm + Gate GEMM + Up GEMM + SwiGLU + Down GEMM + Add | multi_launch_builder/o_ffn_multi.py:178 | ~50 s |
rms_gemv_rope.elf | Decode | 6: RMSNorm + Q/K/V GEMV + RoPE Q + RoPE K (GEMV variants) | multi_launch_builder/rms_gemv_rope_multi.py:369 | ~3 s |
o_gemv_ffn.elf | Decode | 8: O GEMV + Add + RMSNorm + Gate/Up GEMV + SwiGLU + Down GEMV + Add (GEMV variants) | multi_launch_builder/o_gemv_ffn_multi.py | ~7 s |
lm_head_gemv.elf | Both | 8: identical 8-partition GEMV stitched 8 times | multi_launch_builder/lm_head_gemv_multi.py | ~13 s |
So one prefill layer = 3 NPU calls (rms_gemms_rope + flash_attn + o_ffn) covering 15 sub-launches. Without stitching it would be 15 NPU calls per layer × 16 layers = 240 calls per prefill. With stitching it's 48 calls per prefill (16 × 3).
+ +FA's MLIR uses many air.channels for its cascade-of-tiles design. The air-opt-shim-dma-bds compiler pass scales super-linearly with the number of channels in a module. With 9+ stitched launches in one ELF (i.e., FA + the rms_gemms_rope launches), this pass takes >10 minutes — empirically prohibitive. So the production split is: FA stays as a 1-launch ELF, called between the stitched rms_gemms_rope and o_ffn. That's why one prefill layer is 3 NPU calls, not 1.
All in kernel_builder/stitching.py as text-manipulation utilities. No MLIR Python API for moving operations between modules — every operation belongs to a Context, and you can't lift a region from one func and graft it into another. Text-based stitching sidesteps this.
+ +The algorithm:
+func.func body (just the operations between signature and return).%argN references to the combined function's arg indices (this is what threads the data flow between launches).mlir.ir.Module.parse(...) to validate.rms_gemms_rope is stitched# multi_launch_builder/rms_gemms_rope_multi.py:466-481 (paraphrased)
+bodies, maps_all = [], []
+for ir, prefix, arg_map in [
+ (rms_ir, "r", {0:0, 1:1, 2:2}), # RMSNorm: x_in, norm_w, normed
+ (q_ir, "q", {0:2, 1:3, 2:4}), # Q GEMM: normed (=arg2), wq (=arg3), q (=arg4)
+ (k_ir, "k", {0:2, 1:5, 2:6}), # K GEMM: normed, wk (=arg5), k (=arg6)
+ (v_ir, "v", {0:2, 1:7, 2:8}), # V GEMM: normed, wv (=arg7), v (=arg8)
+ (rope_q_ir, "rq", {0:4, 1:9, 2:11}), # RoPE Q: q (=arg4), lut_q (=arg9), q_roped (=arg11)
+ (rope_k_ir, "rk", {0:6, 1:10, 2:12}), # RoPE K: k (=arg6), lut_k (=arg10), k_roped (=arg12)
+]:
+ body = _extract_between_func_and_return(ir)
+ maps = _extract_affine_maps(ir)
+ body = _rename_all_with_externs(body, prefix, _EXTERN_FUNCS) # prefix all SSA
+ maps = [_rename_all_with_externs(m, prefix, _EXTERN_FUNCS) for m in maps]
+ body = _fix_launch_func_args(body, prefix, arg_map) # remap arg refs
+ bodies.append(body)
+ maps_all.extend(maps)
+
+# Then assemble: module { #maps... func.func @rms_gemms_rope(13 args) { bodies... return } }The arg_map values are what enable data flow: {0:2, 1:3, 2:4} for Q GEMM means "the Q GEMM's slot 0 (its activation input) connects to the combined func's slot 2 (which is the RMSNorm output, normed)". Same DDR buffer, no host hop between RMSNorm and Q GEMM.
kernel_builder/stitching.py| Function | What it does |
|---|---|
_extract_between_func_and_return(mlir) | Returns the body of the public func.func — everything between signature and return. |
_extract_affine_maps(mlir) | Returns the #map0 = ..., #map1 = ... declarations from the module header. |
_extract_private_funcs(mlir) | Returns func.func private declarations (e.g., external C++ kernel decls like @matvec_vectorized_bf16_bf16). |
_rename_all(text, prefix) | Renames every SSA value (%arg0 → %q_arg0), every affine map (#map0 → #q_map0), every symbol (@herd_0 → @q_herd_0) — but preserves external kernel function names. |
_fix_launch_func_args(text, prefix, arg_map) | After rename, fixes air.launch args(...) references to point at the COMBINED func's arg slots, not the per-sub-kernel ones. |
_wrap_ir_in_launch(mlir) | Some sub-builders (RMSNorm, eltwise add) emit a bare air.herd not wrapped in air.launch. This wraps it in air.launch { air.segment { herd } } — required because airrt-to-npu only sees segment_load ops. |
This is the easiest place to get confused, so it's worth being explicit. The "stay on NPU" property of stitched intermediates applies only inside one ELF. As soon as you cross from one xrt.run() to another (e.g., rms_gemms_rope → flash_attn → o_ffn), the intermediates go through the host by default.
| Boundary | How intermediates flow | Cost per transfer | What "production" does |
|---|---|---|---|
| Intra-ELF between sub-launches inside one merged ELF (e.g., RMSNorm → Q GEMM inside rms_gemms_rope) |
+ NPU DMA reads from / writes to the same DDR-resident BO. Host is completely uninvolved during the xrt.run(). |
+ ~µs (NPU-internal DMA, dominated by L2/L1 fan-out) | +Always uses NPU-only flow. Marked via intermediate_indices so KernelCache neither host-writes on entry nor host-reads on exit. |
+
| Inter-ELF between two separate xrt.run() calls (e.g., rms_gemms_rope → flash_attn) |
+ By default: producer's output BO → sync(FROM_DEVICE) → host numpy view → next call's memcpy + sync(TO_DEVICE) into a SEPARATE BO. Two cache-coherent transfers + a memcpy per intermediate. |
+ ~µs/MB at PCIe-equivalent bandwidth; per prefill layer the inter-ELF traffic adds up to ~40 MB round-trip | +Production uses the host-broker pattern even though BO aliasing is technically possible (the alternative has been validated in development). See D2 for why production accepts this and what it would take to remove. | +
Concrete prefill numbers per pass (16 layers × 3 ELF dispatches per layer):
+ +| Where | Per layer | Per pass (16 layers) |
|---|---|---|
Inside rms_gemms_rope (6 launches stitched) | 0 host transport (5 NPU-only handoffs) | 0 |
rms_gemms_rope → flash_attn (Q + K + V, host-broker) | ~12 MB ↓↑ (Q=8 MB, K=2 MB, V=2 MB) | ~192 MB |
flash_attn → o_ffn (attn_out, host-broker) | ~8 MB ↓↑ | ~128 MB |
Inside o_ffn (8 launches stitched) | 0 host transport (7 NPU-only handoffs) | 0 |
| K, V to KV cache (host transpose, B4) | ~4 MB ↓ each, plus CPU transpose | ~64 MB ↓ + ~16 ms CPU |
| Total inter-ELF host↔device traffic per pass | ~640 MB round-trip |
At ~20 GB/s of host↔device bandwidth, ~640 MB ≈ ~32 ms ≈ 3% of the 1.13 s prefill. Decode is much smaller because per-token intermediates are KB-scale: ~10 KB per inter-ELF transfer × 33 NPU calls per token = a few MB, well under measurement noise. So inter-ELF host-broker is a real prefill cost, but tiny in decode.
+ +The problem. A stitched ELF (B5) hides 6-8 sub-launches behind one xrt.run(). But that single call still has to: get every input from host RAM into NPU-accessible DDR, hand the kernel handles to those buffers, run the kernel, and read outputs back. Done naively, every call would re-allocate buffers and re-upload weights — for a 14 MB wq tensor, that's ~5 ms of PCIe traffic per call, or ~80 ms × 16 layers = 1.3 s extra per prefill pass. The kernel finishes in tens of milliseconds; we cannot afford 5+ ms of host overhead per call.
This section explains what happens during ONE xrt.run() at the BO (Buffer Object) level — the unit of memory the NPU can read and write. Once you understand this anatomy, the per-layer BO trick in B7 (KernelCache) is straightforward.
A BO is an XRT abstraction for a chunk of NPU-accessible memory. Physically it lives in DDR — the same RAM the host uses, but with a NPU-readable mapping. Created by xrt.bo(device, size_bytes, ...). Two operations matter:
| Op | Cost | What it does |
|---|---|---|
bo.map() | ~free | Returns a host pointer you can memcpy into. Host writes go to RAM directly. |
bo.sync(TO_DEVICE) | ~µs/MB (cache flush) | Flush host CPU caches so the NPU sees the up-to-date bytes when it DMAs from DDR. |
bo.sync(FROM_DEVICE) | ~µs/MB (cache invalidate) | Invalidate host CPU caches so the host sees the up-to-date bytes the NPU wrote. |
The kernel doesn't get bytes — it gets a list of BOs (one per func.func argument), and the kernel's compiled code uses NPU DMA to stream chunks of those BOs into per-tile L1 / L2 SRAM as it runs.
xrt.run()| Step | What happens | Cost (typical) |
|---|---|---|
| 1. Resolve XRT context | Look up the loaded xclbin for this kernel name; get the device handle and kernel symbol. | ~µs (cached) |
| 2. Resolve BO list | Look up or allocate the BO array for this bo_key. One BO per kernel argument. | ~µs (cached) or ~ms (first allocation) |
| 3. Write inputs | For each non-static, non-intermediate input: memcpy(bo.map(), input_array) + bo.sync(TO_DEVICE). Static slots (weights) and intermediate slots (kernel-overwritten) are SKIPPED on every call after the first. | ~µs/MB per slot actually written |
| 4. Submit kernel | invoker.run(*bos) — XRT enqueues the kernel and the call blocks until completion. | ~100 µs dispatch overhead + actual NPU compute time |
| 5. Read outputs | For each slot in output_indices: bo.sync(FROM_DEVICE) + return a numpy view onto bo.map(). Other slots get a 0-length placeholder. | ~µs/MB per output |
Every load_and_run call (B7) accepts three optional sets that control which slots get host↔device data movement:
| Set | Meaning | Effect |
|---|---|---|
output_indices | Slots the caller wants to read back to host (e.g., q_roped, k_roped). | Triggers sync(FROM_DEVICE) for those slots only. Other slots get a 0-length placeholder in the return tuple. |
static_input_indices | Slots holding weights/LUTs that are pre-loaded once and never change (e.g., wq, norm_w, RoPE LUT). | Skipped by the host write loop on every call after the first. Combined with bo_key, lets per-layer weights persist on device across calls. |
intermediate_indices | Slots the kernel will OVERWRITE — entry contents don't matter (e.g., the normed output of RMSNorm that the next launch reads). | Skipped by the host write loop on every call after the first. Saves a memcpy + sync for buffers the host never needs to read or initialize. |
These sets are what makes per-call cost go from "upload everything" (~ms) to "upload only the new activation" (~µs).
+ +rms_gemms_rope, layer 5, mid-prefill)# Argument layout for rms_gemms_rope (13 slots, see B5/B7 for full list):
+# 0: x_in ← layer activation, CHANGES every call
+# 1: norm_w ← layer 5's RMSNorm weight, STATIC
+# 2: normed ← intermediate (RMSNorm → GEMM)
+# 3: wq ← layer 5's Q weight (~14 MB), STATIC
+# 4: q ← intermediate (GEMM → RoPE)
+# 5: wk ← layer 5's K weight (~3.5 MB), STATIC
+# 6: k ← intermediate
+# 7: wv ← layer 5's V weight (~3.5 MB), STATIC
+# 8: v ← intermediate
+# 9: rope_lut_q ← STATIC (LUT)
+# 10: rope_lut_k ← STATIC
+# 11: q_roped ← intermediate, but caller wants to READ it (output_index)
+# 12: k_roped ← intermediate, but caller wants to READ it (output_index)
+
+cache.load_and_run(
+ "rms_gemms_rope", RGR_BACKEND,
+ x_in_bf16, # slot 0 (only this gets written)
+ lw.attn_norm, np.zeros(...), # slots 1, 2
+ lw.wq, np.zeros(...), # slots 3, 4
+ lw.wk, np.zeros(...), # slots 5, 6
+ lw.wv, np.zeros(...), # slots 7, 8
+ rope_lut_q, rope_lut_k, # slots 9, 10
+ np.zeros(...), np.zeros(...), # slots 11, 12 (output buffers)
+ output_indices=[11, 12],
+ static_input_indices={1, 3, 5, 7, 9, 10},
+ intermediate_indices={2, 4, 6, 8, 11, 12},
+ bo_key=f"rms_gemms_rope_L5", # this layer's BO set
+)Per-call work: ONE memcpy (slot 0, ~8 KB) + ONE sync(TO_DEVICE) + run + TWO sync(FROM_DEVICE) (slots 11, 12). All 21 MB of weights stay resident on the NPU's BOs — the host doesn't touch them. Without static_input_indices + bo_key, the same call would memcpy and sync ~21 MB of weights every single time.
bo_key are the entire vocabulary for that separation. Whoever owns the load_and_run contract (B7) gets to make every call cheap — even the kernel-call burst inside a tight per-token decode loop.
+Each load_and_run call resolves its own BO list via bo_key. Two different kernels (or two calls with different bo_keys) get independent BOs even if they conceptually pass the same intermediate. So:
xrt.run(): the merged ELF's sub-launches all see the SAME BO list, so an intermediate written by sub-launch N is automatically visible to sub-launch N+1 (just two MLIR launches reading/writing the same arg slot). No host involvement.xrt.run() calls: kernel A's BOs and kernel B's BOs are different XRT objects in different _cached_bos entries. To get A's output into B's input you EITHER (1) sync to host and re-upload to B's BO (the default — host-broker), OR (2) explicitly alias B's input BO to point at A's output BO via a manual _share_bo trick.Production uses (1) for cross-kernel-group transfers — see the per-pass cost breakdown in B5 "Intra-ELF vs inter-ELF intermediate flow". Path (2) is the optimization tracked in D2 (Future work).
+ + +KernelCache: compile-once, per-layer BO setsThe problem. Two costs would otherwise dominate every script start AND every kernel call:
+python llama32_1b_inference.py run is unworkable.KernelCache (in kernel_builder/cache.py:183) is the single class that solves both. It's the bridge between the per-call BO anatomy (B6) and the realities of running a 16-layer transformer.
| Layer | What's cached | Lifetime | Key |
|---|---|---|---|
| 1. Disk artifact | Compiled .elf + .insts.bin + kernel symbol name | Persistent (until make clean) | name (e.g. "rms_gemms_rope") |
| 2. XRT context | Loaded XRT device + xclbin + kernel handle | Process lifetime | name |
| 3. Buffer Objects | Allocated xrt.bo objects (one per kernel arg) | Process lifetime | bo_key (defaults to name; overridden per layer) |
Layer 1 saves the 3-minute compile. Layer 2 saves the ~100 ms xclbin reload per kernel call. Layer 3 (combined with static_input_indices from B6) saves the per-call weight upload.
class KernelCache:
+ def __init__(self, cache_dir=None, verbose=False, profiler=None):
+ self.cache_dir = Path(cache_dir) # where .elf files persist on disk
+ self.profiler = profiler or Profiler()
+ self.artifacts = {} # Layer 1: name → XRTCompileArtifact (paths + symbol)
+ self._loaded = {} # Layer 2: name → (backend, invoker) — XRT handles
+ self._cached_bos = {} # Layer 3: bo_key → list[xrt.bo] — per-session BOscompile_and_cache(name, mlir_module, backend_kwargs) — called ONCE per ELF# kernel_builder/cache.py:251 (paraphrased)
+def compile_and_cache(self, name, mlir_module, backend_kwargs, output_binary_name="air"):
+ prepare_air_project() # clear air_project/ + compile .o files
+ backend = XRTBackend(**backend_kwargs)
+ artifact = backend.compile(mlir_module, ...) # aircc → aiecc → .elf (the slow step)
+
+ cached_binary = self.cache_dir / f"{name}{ext}"
+ shutil.copy2(artifact.output_binary, cached_binary)
+
+ self.artifacts[name] = XRTCompileArtifact(str(cached_binary), artifact.kernel, cached_insts)
+ backend.unload()Records name → cached_binary_path in self.artifacts. _save_manifest() writes the dict to cache_dir/manifest.json so a subsequent run with --run-only skips compilation entirely via load_manifest(). This is the difference between a 3-minute startup and a 5-second startup.
load_and_run(name, backend_kwargs, *inputs, ...) — called dozens of times per inferenceThis is the implementation of the per-NPU-call anatomy from B6. Annotated:
+ +# kernel_builder/cache.py:294 (paraphrased — the contract)
+def load_and_run(self, name, backend_kwargs, *inputs,
+ output_indices=None,
+ static_input_indices=None,
+ intermediate_indices=None,
+ bo_key=None):
+
+ # 1. Lookup or load XRT context for this kernel name (Layer 2)
+ if name not in self._loaded:
+ backend = XRTBackend(**backend_kwargs)
+ backend.load(self.artifacts[name])
+ self._loaded[name] = (backend, backend.invoker)
+
+ # 2. Lookup or allocate BO list for this bo_key (Layer 3)
+ bo_key = bo_key or name # default: shared BOs per kernel
+ if bo_key not in self._cached_bos:
+ bos = [allocate_bo(arr.nbytes) for arr in inputs]
+ self._cached_bos[bo_key] = bos
+ first_call = True
+ else:
+ bos = self._cached_bos[bo_key]
+ first_call = False
+
+ # 3. Write inputs (skipping static + intermediate after first call)
+ static = static_input_indices or set()
+ intermediate = intermediate_indices or set()
+ skip = (static | intermediate) if not first_call else set()
+
+ for i, arr in enumerate(inputs):
+ if i in skip:
+ continue # BO already has the right data
+ memcpy(bos[i].map(), arr)
+ bos[i].sync(TO_DEVICE) # host → DDR
+
+ # 4. Run the kernel
+ invoker.run(*bos)
+
+ # 5. Read back only the requested outputs
+ output_indices = output_indices or [len(inputs) - 1]
+ results = []
+ for i, arr in enumerate(inputs):
+ if i in output_indices:
+ bos[i].sync(FROM_DEVICE) # DDR → host
+ results.append(np_view(bos[i].map(), arr.shape, arr.dtype))
+ else:
+ results.append(np.empty(0, dtype=arr.dtype)) # placeholder
+ return tuple(results)len(inputs), not len(output_indices). Slots not in output_indices get an empty placeholder. Callers index by original arg position: out[2], out[14], etc.static_input_indices and intermediate_indices only kick in after the first call for a given bo_key. The first call must write everything (the BOs have garbage). The pre-load pattern in prepare_runtime exists specifically to make the first call happen during init, not during timed inference.bo_key trick — per-layer weight BOsThe single most consequential decision in the whole codebase. In plain language: give each of the 16 transformer layers its own independent set of NPU BOs, pre-load every layer's weights once at startup, then never re-upload weights again during inference.
+ +bo_key defaults to the kernel name (e.g. "rms_gemms_rope") — meaning ALL 16 layers share ONE set of BOs. With 6 weight slots in rms_gemms_rope totaling ~21 MB, the per-layer behavior would be:
That's pure host overhead with zero NPU benefit. For decode, the per-token version of the same problem dominates the entire decode loop.
+ +bo_keyOverride bo_key to f"rms_gemms_rope_L{layer_idx}" so each layer gets its own slot in self._cached_bos. After the one-time preload, _cached_bos looks like this:
# Conceptual view of the cache state after preload
+self._cached_bos = {
+ "rms_gemms_rope_L0": [bo_x, bo_norm0, bo_normed, bo_wq0, bo_q, ...], # Layer 0's weights pre-uploaded
+ "rms_gemms_rope_L1": [bo_x, bo_norm1, bo_normed, bo_wq1, bo_q, ...], # Layer 1's weights pre-uploaded
+ "rms_gemms_rope_L2": [bo_x, bo_norm2, bo_normed, bo_wq2, bo_q, ...], # ...
+ ...
+ "rms_gemms_rope_L15": [bo_x, bo_norm15, bo_normed, bo_wq15, bo_q, ...],
+ "o_ffn_L0": [...], # Same pattern for the other prefill ELF
+ ...
+}16 layers × independent BO sets, each holding its own layer's weights resident on the NPU. Now the per-call code:
+ +# preload_prefill_weights — runs ONCE before timing starts
+for layer_idx in range(16):
+ cache.load_and_run(
+ "rms_gemms_rope", RGR_BACKEND,
+ np.zeros(...), # slot 0: x_in placeholder
+ weights.layers[layer_idx].attn_norm, # slot 1
+ np.zeros(...), # slot 2
+ weights.layers[layer_idx].wq, # slot 3 (~14 MB)
+ ... # slots 4-12
+ bo_key=f"rms_gemms_rope_L{layer_idx}", # UNIQUE per layer
+ )
+# After this loop: 16 separate BO sets are cached, each with its layer's weights uploaded.
+
+# During TIMED inference, exact same call shape but with the real activation in slot 0:
+for layer_idx in range(16):
+ out = cache.load_and_run(
+ "rms_gemms_rope", RGR_BACKEND,
+ x_bf16, # slot 0: actual activation
+ ... # slots 1-12 (just placeholders, BOs already have weights)
+ static_input_indices={1, 3, 5, 7, 9, 10}, # skip weight write
+ intermediate_indices={2, 4, 6, 8, 11, 12},
+ bo_key=f"rms_gemms_rope_L{layer_idx}", # picks layer's pre-loaded BOs
+ )Now the timed call uploads ONLY the activation (slot 0, ~8 KB), even though there are 13 args. The 12 weight/intermediate slots are skipped because (static | intermediate) covers them and the BO list lookup hit the cached entry for that layer's bo_key. Internal measurements indicate this single optimization is the dominant per-token speedup contributor in decode.
Two mechanisms work together: bo_key decides which set of BOs to look up; static_input_indices decides which slots in that set don't need to be re-written. Either alone wouldn't work — without per-layer keys, every layer overwrites every other layer's weights; without the static-skip flag, KernelCache would dutifully re-memcpy every weight slot every call even though the contents are already correct.
This is fundamentally a trade memory for speed design. Concrete numbers:
+ +| Cost | Default (shared bo_key) | Per-layer bo_key |
|---|---|---|
| NPU-resident BO memory | ~120 MB (one set per ELF × 6 ELFs) | ~1.0 GB (16 layers × 6 ELFs) |
| Host→device upload per prefill pass | ~336 MB (16 × 21 MB rewrites) | ~128 KB (just activations) |
| One-time preload cost | 0 | ~200-300 ms (once at startup) |
~1 GB of pinned BO memory is acceptable for a 1.24 B-parameter model on a system with 16+ GB of RAM. If memory were tight, you could fall back to shared bo_key and accept the per-call upload cost — the contract would still work, just slower.
Yes — NPU2 (Strix) is a unified-memory architecture, so the NPU and CPU share the same physical DDR. So why is there still a memcpy + memory duplication?
+ +Because "shared DDR" doesn't mean "shared allocation". A normal numpy array and an XRT BO live in the same DDR but in different memory regions with different attributes:
+ +| Buffer kind | Allocator | Attributes | Who can read it? |
|---|---|---|---|
| numpy weight array | Python / glibc malloc | Pageable, virtual, CPU-cached | CPU only |
| XRT Buffer Object | xrt.bo(device, size) | Physically contiguous, pinned (non-pageable), specific cache attributes, mapped into BOTH CPU and NPU virtual address spaces | CPU and NPU |
The NPU's DMA engine can ONLY access physically-contiguous, pinned memory — it can't read a random pageable numpy buffer (which is virtually contiguous but physically scattered, and may be swapped out at any moment). So a BO is a special chunk of DDR, requested separately and held alive for the BO's lifetime.
+ +That means the data flow is genuinely:
+wq)memcpy(bo.map(), weight_array) → physical byte copy into the BO's pinned region (~3 ms for 14 MB)bo.sync(TO_DEVICE) → flushes CPU L1/L2/L3 caches so the NPU's DMA reads the up-to-date DDR contents (NOT a copy — pure cache management)bo.sync(FROM_DEVICE) → invalidates CPU caches so a subsequent host read sees what the NPU wroteSo yes — even with shared DDR, the production codebase keeps two physical copies of each weight (the numpy array + the BO), and the preload step really does memcpy them. ~1 GB extra memory + ~200-300 ms one-time preload is the price.
+ +Could it be zero-copy? In principle yes — you could allocate the BO first and then construct a numpy view via np.frombuffer(bo.map(), ...), so the safetensors loader writes directly into the pinned region. The codebase doesn't do this for two reasons:
.reshape().T.ascontiguousarray() always materializes a fresh buffer, so the transposed result has to be copied into the BO regardless of how the original was allocated.So the codebase trades the simplicity of standard numpy for a small one-time memory + memcpy cost. "Unified memory" eliminates cross-PCIe DMA (which discrete GPUs suffer); it doesn't eliminate the pinned-vs-pageable distinction or the cache-coherency flush.
+ +load_and_run) implementing the B6 anatomy with the index-set contract, and one trick (bo_key=f"name_L{layer_idx}") that turns "16 layers × ~50 MB of weights to upload per call" into "0 weight uploads per call after preload". The trade is ~1 GB of pinned BO memory for ~hundreds of ms saved per inference. Without this class, the codebase wouldn't be 1.27 s prefill — it would be tens of seconds.
+x_in (seq, 2048); writes q_roped (seq, 2048), k_roped (seq, 512), v (seq, 512). Realizes Part A2 ops 1-6.cache.load_and_run("rms_gemms_rope", ...)q_roped, k_roped, v; writes attn_out (seq, 2048). Also extracts k_cache, v_cache for decode. Realizes Part A2 op 7.cache.load_and_run("flash_attn", ...)attn_out, x_residual; writes the layer output. Realizes Part A2 ops 8-15.cache.load_and_run("o_ffn", ...)After all 16 layers: CPU RMSNorm on the last token's hidden state (Part A5), then lm_head_gemv.elf (8 partitions, 1 NPU call) → argmax → first generated token.
Tile usage: rms_gemms_rope's GEMMs use the full [8,4] = 32-tile array; its RMSNorm + RoPE use [8,1] = 8 tiles. flash_attn uses a multi-segment cascade ~16-24 tiles. o_ffn's GEMMs use [8,4] = 32 tiles; its add/RMSNorm/SwiGLU use [8,1] = 8 tiles. See B2.8 tile-mapping summary for the full table.
run_npu_prefill# llama32_1b_inference.py:341 — main prefill entry
+def run_npu_prefill(token_ids, weights, config, prefill_cache, decode_cache,
+ rope_lut_bf16, max_seq, tokenizer, ...):
+ seq_len = len(token_ids) # 2048
+
+ # Pre-allocate KV cache (16 layers × 8 KV heads × 2048 × 64), see Part A4
+ k_cache = np.zeros((config.n_layers, n_kv_heads, max_seq, head_dim), dtype=bfloat16)
+ v_cache = np.zeros((config.n_layers, n_kv_heads, max_seq, head_dim), dtype=bfloat16)
+
+ # Token embedding (host-side numpy lookup)
+ x_bf16 = weights.embed_table[token_ids].astype(bfloat16)
+
+ # --- TIMED SECTION START ---
+ for layer_idx in range(config.n_layers): # 16 layers
+ x_bf16, intermediates = run_transformer_block(
+ x_bf16, weights.layers[layer_idx], rope_lut_bf16,
+ config, prefill_cache, layer_idx=layer_idx, ...
+ )
+ # Extract KV cache from this layer's intermediates (see Part A4)
+ k_cache[layer_idx, :, :seq_len, :] = intermediates["k_roped"]...
+ v_cache[layer_idx, :, :seq_len, :] = intermediates["v"]...
+
+ # Find last real token (see Part A5 padding)
+ prompt_len = len([t for t in token_ids if t != tokenizer.eos_token_id])
+ pred_pos = prompt_len - 1
+
+ # Final RMSNorm + LM Head — only the last real-token row
+ last_normed = _rms_norm(x_bf16[pred_pos:pred_pos+1], weights.final_norm)
+
+ # NPU LM Head GEMV — reuse decode-cache 8-partition GEMV ELF
+ results = decode_cache.load_and_run("lm_head_gemv", LM_GEMV_BACKEND, ...)
+ logits_row = np.concatenate(results, axis=0)[:vocab_size]
+ prefill_token = int(np.argmax(logits_row))
+
+ return prefill_token, k_cache, v_cache, prompt_lenBefore any timing starts, preload_prefill_weights writes ALL 16 layers' weights into per-layer NPU BOs:
# llama32_1b_prefill.py — preload_prefill_weights (paraphrased)
+def preload_prefill_weights(weights, config, cache, seq_len, rope_lut):
+ for layer_idx in range(config.n_layers): # 16 layers
+ lw = weights.layers[layer_idx]
+ cache.load_and_run(
+ "rms_gemms_rope", RMS_GEMMS_ROPE_BACKEND,
+ np.zeros((seq_len, emb_dim), dtype=bfloat16), # slot 0: x_in (placeholder)
+ lw.attn_norm.astype(bfloat16), # slot 1: norm_w (STATIC)
+ np.zeros((seq_len, emb_dim), dtype=bfloat16), # slot 2: normed (intermediate)
+ lw.wq.astype(bfloat16), # slot 3: wq (STATIC)
+ # ... 9 more args (intermediates + weights + LUTs)
+ output_indices=[11, 12], # read q_roped, k_roped back
+ static_input_indices={1, 3, 5, 7, 9, 10}, # weights/LUTs: written once
+ intermediate_indices={2, 4, 6, 8, 11, 12}, # overwritten by kernel
+ bo_key=f"rms_gemms_rope_L{layer_idx}", # per-layer BO set
+ )
+ # Same pattern for o_ffn ELF — 16 different BO sets, one per layerbo_key trick (this is what "per-layer weight BOs" means): KernelCache caches BO objects keyed by bo_key. By using f"rms_gemms_rope_L{layer_idx}", each layer gets its OWN set of NPU BOs. The weights for layer 5 stay in layer 5's BOs and are never overwritten by layer 6. During inference, the timed call uses the same bo_key, so the per-layer weights are already on device — only the x_in activation needs to be host-uploaded.
+Decode works on one token at a time. Per token, per layer, it makes 3 calls (2 NPU + 1 CPU):
+ +x_in (2048,); writes single-token q_roped (2048,), k_roped (512,), v (512,).cache.load_and_run("rms_gemv_rope", ...)x_decode.cache.load_and_run("o_gemv_ffn", ...)After all 16 layers (per token): CPU RMSNorm on the resulting hidden state, then lm_head_gemv.elf → argmax → next token.
Tile usage: EVERY decode kernel uses ≤ 8 tiles (one column of the 8×4 array): the GEMVs are [8,1], RMSNorm + SwiGLU + add are [8,1], and RoPE drops to [1,1] (only one row to rotate). The decode path leaves at least 24/32 = 75% of the compute array idle on every NPU dispatch — one reason decode is dispatch-overhead-bound (the large per-token speedup we achieved comes from removing dispatch overhead, not from doing more compute).
# llama32_1b_inference.py:585 — the decode loop inside generate()
+for token_idx in range(n_tokens):
+ t_token_start = time.perf_counter()
+
+ x = x_decode.copy() # single-token activation (emb_dim,)
+ for layer_idx in range(config.n_layers): # 16 layers
+ x = run_decode_block(
+ x, weights.layers[layer_idx], decode_cache, config,
+ k_cache[layer_idx], v_cache[layer_idx], # growing each iter
+ current_pos, rope_lut_bf16,
+ )
+
+ # Final RMSNorm (CPU, <1ms for 2048 elements)
+ x_normed = rms_norm(x.astype(np.float32).reshape(1, emb_dim),
+ weights.final_norm.astype(np.float32))
+
+ # LM Head — NPU 8-partition GEMV (single XRT call, 8 launches in one ELF)
+ x_lm = x_normed.flatten().astype(bfloat16)
+ lm_inputs = [x_lm] # slot 0: shared input
+ for p in range(_LM_N_PARTITIONS): # 8 partitions
+ lm_inputs.append(weights._lm_weight_parts_gemv[p]) # weight
+ lm_inputs.append(np.zeros(_LM_N_PART, dtype=bfloat16)) # output buffer
+
+ lm_results = decode_cache.load_and_run(
+ "lm_head_gemv", LM_GEMV_BACKEND, *lm_inputs,
+ output_indices=[2 + 2*p for p in range(8)], # 8 outputs
+ static_input_indices={1 + 2*p for p in range(8)}, # weights static
+ intermediate_indices={2 + 2*p for p in range(8)}, # skip output writes
+ )
+
+ # Concatenate 8 partition outputs into one logits array, argmax
+ logits = _assemble_logits(lm_results, vocab_size)
+ next_token = int(np.argmax(logits[0]))
+ generated_tokens.append(next_token)
+ x_decode = weights.embed_table[next_token].astype(bfloat16)
+ current_pos += 1
+
+ if next_token in (tokenizer.eos_token_id, 128009): # <|eot_id|>
+ breakprofile.md as a known limitation; an NPU decode FA was added for the larger Llama-3B variant (head_dim=128) but isn't used here.
+Reference section: a top-down map of every file involved in the production runtime, useful for grepping or for finding the right entry point.
+ +| File | Lines | Purpose |
|---|---|---|
llama32_1b_inference.py | 975 | Main entry point. Unified prefill + decode pipeline. main() at the bottom. |
llama32_1b_prefill.py | 514 | Standalone prefill (with profiler report). compile_all_kernels, run_transformer_block, preload_prefill_weights. |
llama32_1b_decode.py | 286 | Standalone decode. compile_decode_kernels, run_decode_block, decode_attention_cpu. |
llama32_1b_weights.py | 522 | HuggingFace safetensors loader. LlamaConfig, LayerWeights, LlamaWeights, load_weights, synthetic_weights, generate_rope_lut. |
llama32_1b_cpu_helpers.py | ~90 | Small NumPy helpers shared by production + verify: rms_norm (LM-head GEMV final norm), attention_reference (prefill cpu_attn=True fallback), softmax (used by attention_reference). The file used to host a full F32 forward pass + standalone --verify CLI; both became redundant once the verify subsystem started comparing directly against HF transformers bf16. |
verify/ | — | End-to-end verification subsystem. verify_runner.py orchestrates the top-k token gate (make verify) and the diagnosis lens (make diagnosis). See VERIFICATION.html. |
Makefile | 112 | Convenience targets: compile, run, profile, chat, verify, diagnosis, clean. |
| File | Lines | Purpose |
|---|---|---|
cache.py | 453 | The KernelCache class. Manages compile, cache, load, run, and BO reuse for all kernels. See B7. |
stitching.py | 206 | Text-based MLIR stitching utilities for assembling multi-launch ELFs. See B5. |
gemm_builder.py | 137 | Wraps the upstream matrix_multiplication/bf16/run.py:build_module + applies an additional MLIR transform IR script for prefill GEMMs. See B2.2. |
external_kernels.py | 180 | Compiles all C++ .o kernel files via Peano (rope, silu_and_mul, mv, mv_k8192, attn). |
backend_presets.py | 65 | All *_BACKEND kwarg dicts (RGR_BACKEND, OGF_BACKEND, etc.) — XRTBackend init params per kernel. |
rope_halfsplit.cc | ~100 | Custom RoPE C++ kernel matching HuggingFace's half-split convention. |
| File | Phase | Launches | Builds |
|---|---|---|---|
rms_gemms_rope_multi.py | Prefill | 6 | RMSNorm + Q/K/V GEMM + RoPE Q + RoPE K (Part A2 ops 1-6) |
o_ffn_multi.py | Prefill | 8 | O GEMM + Add + RMSNorm + Gate/Up GEMM + SiLU×mul + Down GEMM + Add (Part A2 ops 8-15) |
rms_gemv_rope_multi.py | Decode | 6 | RMSNorm(1D) + Q/K/V GEMV + RoPE Q + RoPE K — single-token version |
o_gemv_ffn_multi.py | Decode | 8 | GEMV variants of o_ffn — single-token version |
lm_head_gemv_multi.py | Both | 8 | 8-partition vocab GEMV (16384 outputs each) |
| Path | Purpose |
|---|---|
standalone_kernels/K1..K10/ | Individual chunk-level kernels for debug; not used by production runtime. |
ffn_swiglu/silu_and_mul.cc | Custom SwiGLU C++ kernel. |
docs/ | Documentation: profile.md, explain.md, usage.md, plus HTML walkthroughs in docs/detail/. |
| Model concept | NPU realization | File:Function |
|---|---|---|
| One transformer block (14 ops) | 3 NPU calls per layer (rms_gemms_rope + flash_attn + o_ffn) | llama32_1b_prefill.py:run_transformer_block |
| 14 ops within a block | Stitched into 6+1+8 = 15 sub-launches across 3 ELFs (B5) | The multi_launch_builder/*_multi.py files |
| Token embedding lookup | numpy fancy-indexing on host | llama32_1b_inference.py:373 (embed_table[token_ids]) |
| Final RMSNorm | Host CPU (1 row only — only the prediction row matters) | llama32_1b_inference.py:425-430 |
| LM Head | NPU 8-partition GEMV (1 ELF, 8 launches in 1 xrt.run) | multi_launch_builder/lm_head_gemv_multi.py |
| K cache write (prefill, with transpose) | numpy slice assign on host (B4 layout mismatch #2) | llama32_1b_inference.py:401 |
| K cache write (decode) | numpy slice assign on host inside run_decode_block | llama32_1b_decode.py |
| Decode attention | CPU (numpy) — single-query GQA against the cache slice | llama32_1b_decode.py:96 decode_attention_cpu |
| Prefill attention | NPU FlashAttention causal GQA (its own ELF, see B5) | flash_attention/kernel_fusion_based/attn_npu2_seqfirst.py |
| Decode GEMV pre-transposed weights | One-time CPU pre-transpose at startup (B4 layout mismatch #1) | llama32_1b_inference.py:171-197 |
The verification subsystem lives in its own subdirectory (verify/) and is documented end-to-end in VERIFICATION.html. This part is a one-page pointer; treat the companion doc as the source of truth.
+ +Two entry points, both routed through the parent Makefile and both comparing against HuggingFace transformers in bf16 (same dtype as the NPU — fair fight):
| Target | What it does | Pass/fail? |
|---|---|---|
make verify [MODEL=base|instruct] | 2 prompts × 32 greedy-decoded tokens (CI gate; use make verify-full for the full 8-prompt sweep). At each step both runners' chosen tokens must appear in the OTHER side's top-5 (k=5). Mirrors vLLM's check_logprobs_close. ~2 min (verify-full: ~6 min). | Yes. Exits 1 on any FAIL. |
make diagnosis [MODEL=...] [PROMPT="..."] | Single prompt, prefill only. Per-layer ffn_out cosine + max_abs (NPU vs HF bf16) for all 16 layers. ~3 min. | Informational only. Read the table by hand to localize a regression flagged by verify. |
The verify NPU runner (verify/runners/npu_runner.py) is a thin adapter — it imports and invokes the same prepare_runtime, run_npu_prefill, and run_npu_decode_step functions that make run calls. Any change to the production prefill/decode path is automatically tracked by make verify; there is no parallel maintenance.
bf16 ULP noise routinely flips per-step top-1 between two mathematically equivalent implementations, so a corr > 0.99-style threshold either trips on noise or sits so loose that real regressions slip through. Discrete top-k inclusion is the escape: bf16 noise can flip top-1 but rarely displaces a token from the top-5, so the gate distinguishes "drift" from "implementation bug" cleanly. See VERIFICATION.html §3 for the full argument.
The LIT test run_npu2_verify.lit runs make verify MODEL=instruct on the NPU2 self-hosted runner and FileCheck-asserts [verify] PASS. REQUIRES: ryzen_ai_npu2, peano, hf_token — local runs without an HF token skip cleanly.
A running list of optimizations and design changes that the current production codebase does NOT do, but that we have identified as worth pursuing — typically because they unlock a new capability (larger models, lower latency) or remove a known scalability bottleneck. Each entry captures the motivation, current behavior, proposed change, and rough impact estimate, so a future contributor can pick one up without re-deriving the context.
+ +Format: impact tag (how much it matters), effort tag (rough engineering size), status tag (idea / scoped / in-progress). This section grows over time as new ideas emerge.
+ + +The current preload pipeline keeps two or three physical copies of each weight tensor in DDR (see B7 "Subtle point: aren't CPU and NPU sharing the same DDR?"):
+wq)_wq_t created by the GEMV pre-transpose step (B4 layout mismatch #1, ~14 MB)For Llama-3.2-1B (~2.5 GB of bf16 weights), the per-layer BO trick (~1 GB resident) plus duplicated numpy/transposed copies puts total memory at ~5-6 GB. This is fine on a 16-32 GB host, but it does NOT scale:
+ +| Model | BF16 weights | Estimated total RAM with current scheme (rough) |
|---|---|---|
| Llama-3.2-1B (current) | ~2.5 GB | ~5-6 GB ✓ fits |
| Llama-3.2-3B | ~6.4 GB | ~13-15 GB (tight on 16 GB host) |
| Llama-3.1-8B | ~16 GB | ~32-40 GB (won't fit on most consumer NPU2 systems) |
| Llama-3.3-70B | ~140 GB | — (impossible without zero-copy) |
Memory will become the bottleneck once we move beyond 1-3 B-parameter models. Solving this is a prerequisite for larger model deployment, not a nice-to-have.
+ +From preload_prefill_weights via cache.load_and_run with static_input_indices:
# Three physical copies in DDR for each weight tensor:
+weights.layers[5].wq # 1) HuggingFace numpy, ~14 MB pageable
+lw._wq_t = np.ascontiguousarray( # 2) transposed numpy, ~14 MB pageable
+ lw.wq.astype(bfloat16)
+ .reshape(emb_dim, emb_dim).T
+)
+memcpy(bo.map(), lw._wq_t) # 3) XRT BO, ~14 MB pinned
+bo.sync(TO_DEVICE)Use np.frombuffer(bo.map(), ...) to make the BO the only physical storage; numpy is just a view onto it:
# Allocate the destination BO first
+bo = xrt.bo(device, weight_size_bytes)
+
+# Construct a numpy view that points INTO the BO's pinned region
+weight_view = np.frombuffer(
+ bo.map(), dtype=bfloat16, count=weight_n_elements
+).reshape(out_dim, in_dim)
+
+# safetensors loader writes directly into the BO via the numpy view
+load_safetensors_layer_into(weight_view, layer_idx, "wq")
+bo.sync(TO_DEVICE)
+# NO memcpy. NO second copy. The BO IS the weight storage..T.ascontiguousarray()). For zero-copy to work end-to-end, the transposed result must land directly in the destination BO too. Either:
+ verify/runners/npu_runner.py calls prepare_runtime + run_npu_prefill + run_npu_decode_step with the production LlamaWeights object — the same one this BO-aliasing scheme would mutate. If a weight tensor switches from a numpy array to a bf16 BO view mid-call, both verify (HF-bf16 reference, dtype-agnostic) and diagnosis (per-layer ffn_out cosine) need to keep producing the same numbers. Audit the Hf-comparison path before flipping the storage.bo.map() but the bo Python object is GC'd, the view becomes a dangling pointer. Need explicit owner-tracking (e.g. attach the BO as an attribute of the numpy view, or maintain a parallel _bo_keepalive list).weights.lm_head is sliced into 8 partitions for the LM Head GEMV. If the source is a BO view, all 8 partition views must coexist without anyone freeing the underlying BO.| Saves | Amount |
|---|---|
| One-time preload memcpy time | ~200-300 ms (currently amortized; not in critical path) |
| Pageable RAM (numpy original) | ~2.5 GB for 1B model, scales with model size |
| Pageable RAM (transposed copy) | ~1.3 GB extra (decode-side weights only — prefill GEMM uses original layout) |
| Total RAM saving for 1B | ~3.8 GB → roughly halves total memory footprint |
| Unlocks | Llama-8B+ on consumer NPU2 hardware that today can't fit those models |
wq), implement the BO-allocate-then-numpy-view path, confirm bit-exact outputs vs. current path on the verify gate.Background discussion: the trade-off and the pinned-vs-pageable subtlety are documented in B7. The reason "shared DDR" doesn't make this problem go away on its own is also there.
+xrt.run() callsAs documented in B5 "Intra-ELF vs inter-ELF intermediate flow", production currently routes intermediates between separate ELFs (e.g. rms_gemms_rope → flash_attn → o_ffn) through the host: producer output is sync'd to host, then memcpy'd + sync'd back into the consumer's input BO. This adds up to ~640 MB host↔device round-trip per prefill pass — about 3% of the 1.13 s prefill wall time. Decode is unaffected (intermediates are KB-scale).
Multi-launch ELF stitching (B5 / Gap #2) eliminates this for sub-launches inside one ELF, but FlashAttention is un-mergeable into the surrounding kernel-groups (compiler pass complexity), so prefill stays as 3 separate ELFs per layer with host-broker round-trips between them. Cross-ELF BO aliasing is the technique that recovers that 3% without merging the ELFs.
+ +From cells/multi_layer.py / production prefill loop:
+ +for L in range(16):
+ rg_out = run_rms_gemms_rope(cache, layer_in, layer_idx=L)
+ # rg_out["q_roped"] is a numpy view onto host RAM — sync(FROM_DEVICE) just happened
+
+ q_roped_2d = rg_out["q_roped"].reshape(seq, emb) # free metadata reshape
+ k_roped_2d = rg_out["k_roped"].reshape(seq, kv)
+ v_2d = rg_out["v"].reshape(seq, kv)
+
+ fa_out = run_flash_attn(cache, q_roped_2d, k_roped_2d, v_2d, layer_idx=L)
+ # ↑ entering FA: memcpy host numpy → FA's BO + sync(TO_DEVICE)
+ # Same data that just left rms_gemms_rope's output BO is now duplicated in FA's input BOUse the same _share_bo helper already validated in development:
# During preload, after both ELFs have allocated their BOs:
+_share_bo(cache,
+ f"rms_gemms_rope_L{L}", slot=11, # producer's q_roped output BO
+ f"flash_attn_L{L}", slot=0, # consumer's Q input BO — now points at same DDR
+)
+_share_bo(cache, f"rms_gemms_rope_L{L}", 12, f"flash_attn_L{L}", 1) # K
+_share_bo(cache, f"rms_gemms_rope_L{L}", 8, f"flash_attn_L{L}", 2) # V
+_share_bo(cache, f"flash_attn_L{L}", 3, f"o_ffn_L{L}", 0) # attn_out
+
+# During timed inference, mark these slots intermediate so KernelCache skips host I/O:
+fa_out = cache.load_and_run("flash_attn", FA_BACKEND, ...,
+ intermediate_indices={0, 1, 2, 3}, # Q, K, V (in), attn_out (out)
+ # NO output_indices for attn_out — it stays on device for o_ffn
+)Not all inter-ELF transfers can be 100% eliminated, because the host still needs SOME of them for non-NPU work:
+ +| Transfer | Can fully alias? | Reason |
|---|---|---|
| Q (rms_gemms_rope → FA) | ✅ Yes | Host never touches Q during prefill |
| K (rms_gemms_rope → FA) | ⚠️ Partial | FA reads it, AND host needs to sync(FROM_DEVICE) + transpose to write KV cache (B4 mismatch #2). Save the host→FA write only |
| V (rms_gemms_rope → FA) | ⚠️ Partial | Same as K |
| attn_out (FA → o_ffn) | ✅ Yes | Host never touches attn_out |
| o_ffn output → next layer's rms_gemms_rope's x_in | ✅ Yes | Pure layer-to-layer activation pass |
Best-case saving: drop ~640 MB / pass to ~150 MB / pass (KV cache extraction still needs the device→host read). Wall-time saving: from ~3% to ~0.7% — recovering ~25 ms of the prefill.
+ +_share_bo wiring call during preload. For 16 layers × 4-5 cross-ELF edges, that's ~70 wiring lines that must stay synchronized with the kernel-group load_and_run argument layouts. If a layout changes, every aliasing line has to be audited.rms_gemms_rope emits 1D flat arrays (q_roped[seq*emb]); FA expects 2D (seq, emb). Today the host does the metadata-only reshape between them. With aliasing the host is no longer in the loop — the shape conversion has to happen on the MLIR side via memref.expand_shape at the FA entry, which means modifying FA's kernel signature or wrapping its launch.sync(FROM_DEVICE) must happen at the right moment — after the producer has finished writing but before/during FA reading. Currently the host-broker pattern enforces this naturally; with aliasing it needs explicit ordering.air.channels and many internal sub-buffers. Aliasing its input BOs needs to verify that FA doesn't internally reuse those slots in a way that would corrupt the producer's data mid-execution.| Saves | Amount |
|---|---|
| Inter-ELF host↔device round-trip per prefill pass | ~640 MB → ~150 MB (factor 4× reduction) |
| Wall time per prefill pass | ~25 ms (~2.3% of 1.13 s) |
| Wall time per decode token | < 1 ms (negligible — intermediates are KB-scale in decode) |
| Doesn't change anything for | Decode performance, model size scaling, code complexity tradeoffs |
attn_out (FA → o_ffn). It has no host consumer, so it's a clean win.make verify (top-k token gate) and inspect make diagnosis for unexpected per-layer drift.Background: this pattern has been validated in development WITHIN one kernel-group (between separate xrt.run()s of the un-merged baseline). The same _share_bo mechanism would extend to ACROSS kernel-groups in production.
HF_TOKEN so make verify actually runs in CIThe whole point of refactoring NpuRunner into a thin adapter over the production prefill/decode functions (VERIFICATION.html) is that any change to production code is automatically tracked by make verify — no parallel maintenance. But that guarantee only pays off if CI actually runs make verify on every PR. Today it does not.
REQUIRES: ryzen_ai_npu2, peano, hf_token.hf_token lit feature only when the HF_TOKEN env var is present (so local runs without a token skip cleanly instead of failing).ninja check-programming-examples-peano but does NOT inject HF_TOKEN into the job's env. As a result, lit doesn't enable the hf_token feature, and run_npu2_verify.lit is skipped on every CI run — no failure, but no actual verify either.HF_TOKEN at the job (or just the lit-test step) level:
+env:
+ HF_TOKEN: ${{ secrets.HF_TOKEN }}HF_TOKEN with a read token for meta-llama/Llama-3.2-1B-Instruct (and the base model if running the MODEL=base variant in CI). Required on the fork that runs CI; if upstream wants the verify gate too, the same secret needs to be configured there.~/.cache/huggingface/ in the workflow to avoid re-downloading the 2.5 GB checkpoint on every run. Self-hosted runners typically persist this directory naturally, so this is only needed for ephemeral runners.Every PR runs the 8-prompt × 32-token top-k inclusion gate against HF transformers bf16, end to end through the production prefill + decode kernels. ~4 min added to the existing Ryzen AI CI step. Without it, any regression in run_npu_prefill, run_npu_decode_step, the multi-launch kernel builders, or the external kernels (rope.o, silu_and_mul.o, attn_npu2.o, mv.o, mv_k8192.o) can land if its symptom is “tokens drift outside top-5” rather than a structural breakage caught by other tests.
Tiny. Adding the env var is one line; missing the secret in the env just keeps the current skip-behavior (the test fails cleanly with “REQUIRES: hf_token” not satisfied, but does not break the rest of CI).
+xrt.bo(device, size_bytes). Has .map() (returns a host pointer for memcpy) and .sync(direction) (cache flush + barrier). One BO per kernel argument. "Allocating a BO" is cheap; "syncing a BO" is what costs time.bo_key. When preload_prefill_weights calls load_and_run(..., bo_key="rms_gemms_rope_L5") with layer 5's wq tensor in slot 3, KernelCache allocates a fresh BO list for that key and writes the weights. Later, when inference does the same call with the same bo_key, KernelCache finds the cached BOs (already on device with the right weights), and static_input_indices={3, ...} tells it to skip writing slot 3 from host. 16 layers × 2 kernels × ~6 weight slots ≈ ~200 cached weight BOs holding ~1 GB of weights resident on device.static_input_indices)bo_key, these slots are skipped by the host write loop in load_and_run. The BO already has the right data from the preload call.intermediate_indices)load_and_run skips writing zeros to these slots (saves a memcpy + sync). For a multi-launch ELF, intermediate slots include both internal handoff buffers (like normed) and the final output (until the host reads it back via output_indices).xrt.run() calls where call N's output is call N+1's input, you can manually alias call N's output BO into call N+1's input BO (via the _share_bo helper), so the data goes from device to device without a host round-trip. Useful for isolating "BO sharing" from "ELF merging" as separate optimizations during analysis..elf binary that contains multiple air.launch operations stitched into a single func.func. Invoked by ONE xrt.run() call. The launches execute sequentially within the single XRT submission, with intermediates flowing through DDR (NPU DMA reads/writes) without CPU involvement. Saves XRT dispatch overhead and host orchestration cost.air.launch operation. The 6 sub-launches in rms_gemms_rope.elf are the 6 logical kernels (RMSNorm, Q GEMM, K GEMM, V GEMM, RoPE Q, RoPE K) — each was originally a separate air.launch in its own MLIR module before stitching.air.herd @h tile(%tx, %ty) in (%sx=8, %sy=4) means an 8×4 grid of tiles. Inside an air.launch, each herd is mapped to physical AIE tiles by the air-place-herds compiler pass.air.segment represents a partition of the NPU array. The wrapping air.launch { air.segment { air.herd { ... } } } is the canonical AIR program structure. Required so that airrt-to-npu emits airrt.segment_load ops.aircc runs the AIR-dialect passes (dependency analysis, broadcast detection, herd placement, AIR→AIE lowering). aiecc runs the AIE-dialect passes (vectorization, routing, generates per-tile ELFs, packages into the final .elf + .insts.bin).rope.cc, silu_and_mul.cc, mv.cc) into per-tile object files that get linked into the AIE ELF.generate_rope_lut in llama32_1b_weights.py builds an array of shape (max_seq, head_dim) = (2048, 64) in bf16. The first half is cos, second half is sin (concatenated, not interleaved — matches the half-split RoPE convention).kv_h = h // group_size.SwiGLU(gate, up) = SiLU(gate) * up elementwise. Two GEMMs (gate, up) feed it; one GEMM (down) follows. Compared to GELU, requires 1 extra GEMM but learns better.RMSNorm(x, w) = x · rsqrt(mean(x²) + ε) · w. Like LayerNorm but without the mean-subtraction and without a bias parameter. Cheaper and works equally well for transformers.pred_pos = prompt_len - 1 are used. See Part A5.| If you want to understand… | Read these in this order |
|---|---|
| The model itself (math, no NPU) | +1. Part A2 of this guide 2. Optionally: the original Llama paper for context |
+
| The whole pipeline end-to-end | +1. Makefile (entry points)2. llama32_1b_inference.py — start with main() at the bottom, then build_session, run_once, generate, run_npu_prefill3. llama32_1b_decode.py:run_decode_block |
+
| How weights are loaded and pre-staged | +1. llama32_1b_weights.py — load_weights()2. llama32_1b_inference.py:prepare_runtime (line 129)3. llama32_1b_inference.py:_preload_decode_weights (line 219)4. llama32_1b_prefill.py:preload_prefill_weights |
+
| How a single ELF gets compiled | +1. multi_launch_builder/rms_gemms_rope_multi.py:build_rms_gemms_rope_module (line 193) — the highest-level builder2. kernel_builder/stitching.py — text manipulation helpers3. kernel_builder/cache.py:compile_and_cache (line 251)4. kernel_builder/external_kernels.py — C++ .o compilation |
+
| How an ELF gets invoked at runtime | +1. kernel_builder/cache.py:load_and_run (line 294) — the central dispatch function2. Any caller in llama32_1b_inference.py or llama32_1b_decode.py3. kernel_builder/backend_presets.py — the backend kwargs dicts |
+
| How multi-launch merging works | +1. kernel_builder/stitching.py in full2. multi_launch_builder/rms_gemms_rope_multi.py lines 466-481 (the stitch loop)3. docs/explain.md for the design rationale |
+
| Why decode uses CPU attention | +1. llama32_1b_decode.py:decode_attention_cpu (line 96)2. docs/profile.md "Decode Breakdown" section |
+
| Performance breakdown / where time goes | +1. docs/profile.md top-to-bottom — has all the numbers2. kernel_builder/cache.py:Profiler class (line 54)3. Run make profile to see live numbers |
+
| How to add a new kernel-group | +1. Look at any multi_launch_builder/*_multi.py as a template2. Need a build_module entry point + sub-builder calls + a stitch loop3. Add a backend preset to kernel_builder/backend_presets.py4. Add compile + load_and_run wiring in llama32_1b_inference.py |
+
| If you grep for… | Meaningful hits in… |
|---|---|
load_and_run | cache.py (def), llama32_1b_inference.py + llama32_1b_decode.py + llama32_1b_prefill.py (callers) |
bo_key | cache.py (cache impl), and every preload/run call in inference scripts |
static_input_indices | Same as bo_key + load_and_run |
compile_and_cache | cache.py (def), llama32_1b_prefill.py:compile_all_kernels, llama32_1b_decode.py:compile_decode_kernels |
build_module | Each multi_launch_builder/*_multi.py file's main entry point |
_wrap_ir_in_launch | stitching.py (def), used by builders that wrap bare herds |
RGR_BACKEND / OGF_BACKEND / LM_GEMV_BACKEND | backend_presets.py (def), and at every call site |
output_indices | The contract document for what the caller wants back from each kernel |
k_cache / v_cache | llama32_1b_inference.py (allocation + prefill writes) and llama32_1b_decode.py:decode_attention_cpu (reads + appends) |
pred_pos | llama32_1b_inference.py:run_npu_prefill — the "find last real prompt token" logic from Part A5 |
Per-step wall-time attribution of the production prefill + decode pipeline, end-to-end. Diagrams mirror the dataflow in IMPLEMENTATION_GUIDE.html Part B1; numbers are reproduced from a single make profile run on NPU2 (AMD Strix), seq_len=2048, MODEL=instruct.
make profile reportsmake profile runs the same code path as make run — the production prefill + decode functions, end to end, real HuggingFace weights — and just enables the otherwise-disabled Profiler instance that cache.load_and_run already records into. There is no profile-only code path; any change to the production functions is automatically reflected in the profile.
The report (printed at the end of the run) opens with an architecture-aware dataflow summary (matches this page’s SVG order) and then dumps generic detail tables per phase (prefill / decode):
+ +| Section | What it tells you |
|---|---|
| END-TO-END DATAFLOW (at the top) | Architecture-aware walkthrough: tokenize → eos_pad → embed → 16×(rms_gemms_rope + flash_attn + o_ffn + kv_cache_extract) → final_norm → lm_head_gemv. Each row tagged CPU/NPU/— with measured ms. Same ordering as the SVGs in Part A / Part B below. Also prints the one-time Preprocessing (prepare_runtime) wall as a reminder. |
| Wall-Time Attribution | How the total wall budget splits across NPU XRT calls, CPU host ops, and the layer-loop envelope (sanity check; remainder is python scheduling). |
| Per-Layer Execution | One row per layer for prefill; aggregated avg/min/max across tokens for decode. |
| NPU XRT Call Breakdown | Each multi-launch ELF’s wall time per invocation, plus call count. The granularity is one XRT run = one merged ELF (sub-launches inside the ELF stay opaque, since that’s how production dispatches them). |
| CPU Op Breakdown | Each tracked CPU host operation (tokenize, eos_pad, embed lookup, KV-cache extract, final RMSNorm, decode CPU attention). |
| Fine-Grained NPU Breakdown | Each XRT call further split into BO Write / NPU Run / BO Read (concept explained in Part C). |
| Per-Token Wall Trend (decode only) | Per-token layer-loop wall for token 1 / middle / last + min/max/avg + first→last drift. Lets you see whether per-token latency grows with KV-cache length (decode CPU attention is O(current_pos)). With a 2048-token prompt and 30 decode tokens the drift is typically <1%. |
Snapshot from the report (single run, instruct model, 30 decode tokens):
+ +| Metric | Wall | Notes |
|---|---|---|
| TTFT (time-to-first-token, prefill end-to-end) | ~1.28 s | tokenize + EOS-pad + embed + 16×layer + final RMSNorm + LM head. Matches the vLLM / TGI / TRT-LLM TTFT metric (user-facing latency from request submit to first output token). 95% NPU-bound. Tokenize varies by prompt length; ~10 ms typical. |
| TPOT (per output token, steady-state decode) | ~92 ms (10.8 tok/s) | 16 layers × 4.95 ms each + 13.6 ms LM head + ~0.1 ms host wrappers. Slope vs token index is <1% over 30 tokens (KV cache grows by ~1.5% on a 2048-token prompt). |
Preprocessing (one-time, prepare_runtime) | ~7.6 s | Compile external kernels + pre-load weights into per-layer BOs. Happens once per process and is NOT included in TTFT. |
One inference’s prefill phase: prompt → first generated token. Each box shows the step, where it runs, and the measured wall time. The 16 layers are identical; one iteration is shown in the “decoder block” container.
+ + + + +| ELF | Launches | avg / call | BO Write | NPU Run | BO Read | BO MB written |
|---|---|---|---|---|---|---|
rms_gemms_rope | 6 stitched | 7.3 ms | 0.5 ms | 6.5 ms | 0.1 ms | 8.0 MB |
flash_attn (separate ELF) | 1 | 21.6 ms | 1.3 ms | 20.1 ms | 0.1 ms | 20.0 MB |
o_ffn | 8 stitched | 41.0 ms | 1.0 ms | 39.8 ms | 0.1 ms | 16.0 MB |
lm_head_gemv (prefill end) | 8 stitched | 13.6 ms | ~0 ms | 13.5 ms | ~0 ms | ~0 MB |
| Op | Count | avg | Total |
|---|---|---|---|
tokenize | 1 | ~10 ms | ~10 ms |
eos_pad | 1 | ~0 ms | ~0 ms |
embed_lookup | 1 | 5.8 ms | 5.8 ms |
kv_cache_extract | 16 | 1.1 ms | 17.6 ms |
final_rms_norm | 1 | 3.1 ms | 3.1 ms |
| Total CPU | 20 | — | ~37 ms |
Wall-time attribution check: NPU XRT 1119 ms (16 layer-invocations × 3 kernels + 1 LM head = 49 calls) + CPU host ~37 ms = ~1156 ms accounted, vs. TTFT ~1280 ms → ~125 ms unattributed python/numpy scheduling, mostly inside the layer loop.
+ + +Per-token decode step: takes the last produced token, returns the next. Diagram and numbers cover one token; the loop repeats until EOT. Each kernel reflects an avg over 30 decode tokens, 16 layers.
+ + + + +| ELF | Launches | avg / call | BO Write | NPU Run | BO Read |
|---|---|---|---|---|---|
rms_gemv_rope | 6 stitched | 0.9 ms | 0.02 ms | 0.83 ms | 0.01 ms |
o_gemv_ffn | 8 stitched | 3.7 ms | 0.02 ms | 3.64 ms | 0.01 ms |
lm_head_gemv | 8 stitched | 13.6 ms | 0.01 ms | 13.50 ms | 0.03 ms |
| Op | Count / token | avg | Total / token |
|---|---|---|---|
decode_attention_cpu | 16 | 0.28 ms | 4.5 ms |
embed_lookup | 1 | ~0 ms | ~0 ms |
final_rms_norm | 1 | 0.07 ms | 0.07 ms |
| Total CPU / token | 18 | — | ~4.6 ms |
Wall-time check: NPU XRT per token = 16 × (0.9 + 3.7) + 13.6 = 87.2 ms · CPU = 4.6 ms · sum 91.8 ms ≈ observed 92 ms wall. Decode is overwhelmingly NPU-bound; the LM head GEMV alone is ~15% of the per-token cost.
+ +Observation: across decode, BO Write is <1% — this is the payoff for pre-loading all weights into per-layer BOs (and marking them static_input_indices) during prepare_runtime. Without that, each layer would re-write its 116 MB of weights per token.
Each cache.load_and_run("kernel", backend, arg0, ..., argN) invocation is split into three timed segments:
t_write_msFor each input/intermediate argument that needs new bytes, the host does memcpy(numpy_data → BO.map()). Args marked static_input_indices (e.g. layer weights) skip this step on every call after prepare_runtime, so steady-state t_write_ms mainly reflects the dynamic inputs that change call-to-call (the input activation, RoPE LUT row, KV-cache slice, …).
What this measures in practice: host-to-DDR memcpy bandwidth for the dynamic inputs only. If you see this rise, either an argument lost its static_input_indices mark, or a normally-small dynamic input grew (e.g. a bigger seq_len).
t_kernel_msWall time of xrt.run.start() + xrt.run.wait(). This is the NPU actually executing the multi-launch ELF: DDR → L2/L1 DMAs, AIE-tile compute, and L1/L2 → DDR DMAs of outputs. Host does nothing here except spin-wait the completion signal.
What this measures: real NPU hardware execution time for the ELF. All the multi-launch’s stitched sub-launches (e.g. RMSNorm + Q + K + V + RoPE_Q + RoPE_K inside rms_gemms_rope.elf) run sequentially on-device and are not separately resolved here — that’s by design, because production never dispatches them separately.
t_read_msFor each output argument, the host constructs a numpy view over the BO’s mapped memory using np.frombuffer(BO.map(), …). This is zero-copy — no memcpy — and consistently <0.1 ms. If t_read_ms ever climbs into the ms range, that signals an accidental copy was introduced (e.g. an .astype() on a large output).
| Phase | BO Write | NPU Run | BO Read |
|---|---|---|---|
| Prefill (one full pass) | ~46 ms (4%) | ~1062 ms (95%) | ~5 ms (0%) |
| Decode (per token) | ~0.6 ms (1%) | ~86 ms (98%) | ~0.3 ms (0%) |
Both phases are dominated by NPU Run — the host’s job is mostly to feed the right BOs and wait. Decode is even closer to pure-NPU because the per-token dynamic inputs are tiny (a single activation row vs. an entire sequence’s worth).
+ + +cd programming_examples/llama32_1b
+
+# One-time kernel compilation (~3-4 min, cached)
+make compile
+
+# Full profiling report (single run, instruct model)
+make profile N_TOKENS=30 PROMPT="Explain photosynthesis in detail."
+
+# Or with the base checkpoint
+make profile MODEL=base N_TOKENS=30 PROMPT="Once upon a time"
+The report is printed to stdout at the end of the run. To save a copy:
+ +make profile 2>&1 | tee profile_$(date +%Y%m%d-%H%M%S).logNumbers will jitter ±3-5% between runs (NPU power state, OS scheduling, etc); the breakdown structure is stable. make verify is the orthogonal gate that ensures the production code path producing these numbers is still numerically correct.
+ Companion: profile.md (textual perf summary, optimization history, vs IRON comparison) ·
+ IMPLEMENTATION_GUIDE.html B1 (same dataflow, no timing — shows just the structural picture).
+
Two ways to look at the production NPU2 inference pipeline, both comparing against HuggingFace transformers in bf16. Companion to IMPLEMENTATION_GUIDE.html Part C.
+ +make verify [MODEL=instruct|base] — the industry-standard correctness gate. 2 prompts × 32 greedy tokens (fast CI gate; make verify-full runs the full 8-prompt sweep), top-5 set inclusion vs HuggingFace transformers bf16 on the NPU end-to-end production path (NPU FlashAttention on, no CPU attention fallback). Lite-mode runners — no inside probing. ~2 minutes / run (verify-full: ~6 minutes). Default MODEL=instruct matches what production stacks deploy.
+make diagnosis [MODEL=...] — the inside-probing lens. Single prompt's prefill, per-layer ffn_out cosine + max_abs (NPU vs HF bf16) for all 16 layers. Same end-to-end NPU production path as verify (NPU FlashAttention on). Informational only — diagnosis never fails the run. The verify gate is the correctness signal; this table is what you read by hand when verify flags an issue and you need to localize. ~2 minutes / run.
+verify answers "would this model deploy" using the exact criterion industry uses to qualify a BF16 LLM for production — discrete top-k judgment that is robust to bf16 ULP noise. diagnosis gives localization: a continuous-cosine table per layer that tells you where the NPU implementation drifts most from HF. The verify gate gates; the diagnosis lens informs.
+make verify MODEL=instruct: 8/8 PASS, ~3m41smake verify MODEL=base: 8/8 PASS, ~3m39smake diagnosis MODEL=instruct (NPU FA on): cos_p5 in [0.926, 0.993], U-shape with single L1-L2 dip and L10 peak.make diagnosis MODEL=base (NPU FA on): cos_p5 in [0.929, 0.992], double-dip shape (L1-L3 and L12-L14). Same-checkpoint dependence on prompt + fine-tune is what diagnosis surfaces; both pass verify regardless. See Part B.make verify — the correctness gatecheck_logprobs_close)make verify runs 2 prompts (fast CI gate); make verify-full runs the full 8.NPU runs the full production path (GEMV + RMSNorm + RoPE + FlashAttention + LM-head GEMV). Discrete top-k inclusion is robust to bf16 ULP noise: noise routinely flips per-step top-1 between mathematically equivalent implementations but rarely displaces a token from the top-5.
+ +| # | Base (verify/prompts/base.txt) | Instruct (verify/prompts/instruct.txt) |
|---|---|---|
| 0 | GPU stands for | Introduce me what is GPU |
| 1 | The capital of France is | Briefly describe the major milestones in the development of artificial intelligence from 1950 to 2020. |
| 2 | Artificial intelligence is a branch of computer science that | Compare and contrast artificial intelligence with human intelligence in terms of processing information. |
| 3 | A neural network consists of | Describe the basic components of a neural network and how it can be trained. |
| 4 | Once upon a time, there was a robot who dreamed about | Write a short story about a robot that dreams for the first time. |
| 5 | The COVID-19 pandemic, which began in late 2019, | Analyze the impact of the COVID-19 pandemic on global economic structures and future business models. |
| 6 | The Mona Lisa was painted by | Explain the cultural significance of the Mona Lisa painting, and how its perception might vary in Western versus Eastern societies. |
| 7 | The French translation of "The early bird catches the worm" is | Translate the following English sentence into Japanese, French, and Swahili: 'The early bird catches the worm.' |
Topics deliberately mirror each other so base-vs-instruct comparisons read naturally row-by-row. Base prompts are intentionally incomplete sentences (the base model continues raw text rather than answering instructions). Instruct prompts are imperative requests (7 verbatim from vllm/tests/prompts/example.txt + 1 swapped for project relevance).
For each prompt we display the first divergence step (0-based; step 0 is the prefill prediction, step 1 is the first decode token); each side's chosen token at that step (decoded text, quoted so leading whitespace stays visible) plus its 1-based rank in the OTHER runner's top-5; and the agreed prefix — the actual generated text both runners produced identically before splitting.
+ +| # | Prompt | Diverge | NPU choice (rank in HF) | HF choice (rank in NPU) | Agreed prefix |
|---|---|---|---|---|---|
| 0 | GPU stands for | 7 | ✓ " special" (#2) | ✓ " specialized" (#2) | " Graphics Processing Unit. It is a" |
| 1 | The capital of France is | 1 | ✓ "," (#2) | ✓ "." (#2) | " Paris" |
| 2 | Artificial intelligence is… | 7 | ✓ "," (#2) | ✓ "." (#2) | " deals with the creation of intelligent machines" |
| 3 | A neural network consists of | 3 | ✓ " nodes" (#2) | ✓ " interconnected" (#3) | " a set of" |
| 4 | Once upon a time, there was a robot… | 7 | ✓ " little" (#2) | ✓ " robot" (#2) | " being a human. He was a" |
| 5 | The COVID-19 pandemic… | 9 | ✓ "," (#2) | ✓ "." (#2) | " has had a significant impact on the global economy" |
| 6 | The Mona Lisa was painted by | 7 | ✓ " and" (#2) | ✓ "." (#3) | " Leonardo da Vinci in 1503" |
| 7 | The French translation… | 6 | ✓ " prend" (#3) | ✓ " g" (#2) | " "Le premier oisif" |
| # | Prompt | Diverge | NPU choice (rank in HF) | HF choice (rank in NPU) | Agreed prefix |
|---|---|---|---|---|---|
| 0 | Introduce me what is GPU | 0 | ✓ " acceleration" (#2) | ✓ " (" (#2) | (no prefix) |
| 1 | Briefly describe… | 0 | ✓ " Some" (#4) | ✓ " Key" (#3) | (no prefix) |
| 2 | Compare and contrast… | 8 | ✓ " (" (#4) | ✓ " are" (#2) | " Artificial intelligence (AI) and human intelligence" |
| 3 | Describe the basic components… | 20 | ✓ " multiple" (#2) | ✓ " three" (#2) | " \n\n## Step 1: Define the basic components of a neural network\nA neural network consists of" |
| 4 | Write a short story… | 11 | ✓ " model" (#3) | ✓ " android" (#2) | " It's a robot named Zeta, a highly advanced" |
| 5 | Analyze the impact of COVID… | — | ✓ (all 32 match) | ✓ (all 32 match) | (no divergence within sample) |
| 6 | Explain the cultural significance… | 29 | ✓ " Created" (#4) | ✓ " It" (#2) | " \n\nThe Mona Lisa, painted by Leonardo da Vinci in the early 16th century, is one of the most famous paintings in the world." |
| 7 | Translate the following… | 26 | ✓ " Here" (#2) | ✓ "The" (#2) | " This is a common English idiom that means…" |
Both checkpoints PASS the gate. Most divergences are #2/#2 swaps (both runners agreed on the same two top candidates; bf16 noise picked which ranked first); a few are #3/#4. None hit out-of-top-5. On Instruct, prompts 3, 6, 7 reach 20-29 tokens of agreement before splitting, and prompt 5 had zero divergence in the 32-token sample.
make diagnosis — the inside lensSingle prompt's prefill on NPU + HF bf16, then computes per-position cosine + element-wise abs error for each layer's ffn_out (the block output). For layers 0…n_layers−2, both sides expose the raw layer output. For layer n_layers−1, both sides expose the post-final-RMSNorm hidden state — HF surfaces this as hidden_states[n_layers] (post-norm by HF v5.3 convention); NPU produces the equivalent via the same final_norm step it does inside its production LM-head GEMV path. So both L15 cells correspond to "the value the LM-head sees".
Diagnosis is informational only. No threshold, no pass/fail, no exit code based on the cosine. Verify is the correctness signal; the diagnosis table tells you where the NPU implementation drifts most from HF (which layer, by how much), which is what you want when triaging a real verify failure or weighing a kernel-side optimization.
+ +Same prompt, same NPU end-to-end path, both checkpoints. Run side-by-side so the per-layer precision shape can be compared directly.
+ +meta-llama/Llama-3.2-1B-Instruct)| Layer | cos_p5 | cos_min | cos_median | max_abs |
|---|---|---|---|---|
| 0 | 0.993269 | 0.993257 | 0.993733 | 0.75 |
| 1 | 0.926400 | 0.908160 | 0.990950 | 22 |
| 2 | 0.927211 | 0.908539 | 0.988378 | 22 |
| 3 | 0.940698 | 0.927680 | 0.988209 | 24 |
| 4 | 0.951836 | 0.940504 | 0.987463 | 26 |
| 5 | 0.959359 | 0.950193 | 0.988150 | 28 |
| 6 | 0.965235 | 0.958839 | 0.988398 | 30 |
| 7 | 0.969200 | 0.964980 | 0.988053 | 30 |
| 8 | 0.975010 | 0.973589 | 0.989355 | 32 |
| 9 | 0.981512 | 0.980698 | 0.990487 | 34 |
| 10 | 0.983873 | 0.983115 | 0.990943 | 36 |
| 11 | 0.981148 | 0.978896 | 0.990446 | 36 |
| 12 | 0.976977 | 0.973395 | 0.990023 | 38 |
| 13 | 0.975324 | 0.970957 | 0.989895 | 42 |
| 14 | 0.971639 | 0.966981 | 0.990319 | 44 |
| 15 | 0.970669 | 0.966320 | 0.987503 | 10.83 |
meta-llama/Llama-3.2-1B)| Layer | cos_p5 | cos_min | cos_median | max_abs |
|---|---|---|---|---|
| 0 | 0.991912 | 0.991241 | 0.994038 | 1.75 |
| 1 | 0.966095 | 0.959596 | 0.989646 | 7 |
| 2 | 0.960257 | 0.952361 | 0.988373 | 6 |
| 3 | 0.958956 | 0.950566 | 0.986123 | 7 |
| 4 | 0.970088 | 0.965457 | 0.985988 | 8 |
| 5 | 0.972773 | 0.969458 | 0.985526 | 9 |
| 6 | 0.974773 | 0.973999 | 0.983875 | 10 |
| 7 | 0.971905 | 0.968814 | 0.982661 | 10 |
| 8 | 0.955578 | 0.949168 | 0.987208 | 11 |
| 9 | 0.960433 | 0.959102 | 0.989534 | 12 |
| 10 | 0.965993 | 0.965948 | 0.990815 | 13 |
| 11 | 0.954954 | 0.949146 | 0.990970 | 13 |
| 12 | 0.941147 | 0.929415 | 0.989791 | 15 |
| 13 | 0.936710 | 0.923149 | 0.988866 | 16 |
| 14 | 0.929362 | 0.912219 | 0.987908 | 17 |
| 15 | 0.939495 | 0.924292 | 0.990349 | 4.013 |
max_abs sits in the 6-17 range; Instruct sits in 22-44. Instruction tuning amplifies certain pathways; the bigger absolute deltas are not a precision problem (cosine is direction-only).max_abs (~10 for Instruct, ~4 for base) is much smaller than mid-stack because final_norm rescales the hidden state to unit-variance-ish magnitude.Three things have to hold for make verify to be a meaningful correctness signal: the version we test must be the version that ships, the reference we compare against must be trustworthy, and the comparison criterion must be sound for bf16. We address each below.
NpuRunner directly imports and invokes the production functions — no reimplementation:
from llama32_1b_inference import prepare_runtime
+from llama32_1b_prefill import run_transformer_block as run_prefill_block
+from llama32_1b_decode import compile_decode_kernels, run_decode_blockNpuRunner.__init__ compiles the same kernels production compiles and runs the same prepare_runtime setup. NpuRunner.prefill calls run_prefill_block for each of the 16 layers, then runs the production 8-partition LM-head GEMV. NpuRunner.decode_step calls run_decode_block. If NpuRunner produces the right tokens, llama32_1b_inference.py produces the right tokens — by construction.
| Criterion | Choice |
|---|---|
| Canonical | transformers.AutoModelForCausalLM is the reference implementation that Meta + HuggingFace + the open-source LLM ecosystem maintain. Every bf16 LLM deployment (vLLM, llama.cpp, TRT-LLM, …) is qualified against this codebase. |
| Same dtype | Loaded as torch_dtype=torch.bfloat16, matching NPU production. Both sides hit the same bf16 round-off characteristics; the comparison is not testing a dtype gap. |
| Same weights | Both runners load meta-llama/Llama-3.2-1B[-Instruct] from the same HF cache. Identical bytes on disk. |
HfRunner is ~110 lines that delegate to self.model(input_ids, use_cache=True). No transformer-block reimplementation, no custom kernel — the simpler the reference, the harder it is for the reference to be wrong.
Continuous metrics (cosine, KL) on bf16 logits are fragile: bf16 ULP noise routinely flips per-step top-1 between two mathematically equivalent implementations. Discrete top-k inclusion is robust — bf16 noise can flip top-1 but rarely displaces a token from the top-5. compute_topk_set_check in comparators.py mirrors vLLM's tests/models/utils.py::check_logprobs_close; k=5 and n_tokens=32 are vLLM's defaults for the standard model gate.
make verify run, end to endverify/prompts/{instruct,base}.txt (selected by MODEL). make verify uses the first 4 (fast CI gate); make verify-full uses all 8.chosen[i] + topk[i] (top-5 IDs) per step.chosen[i] + topk[i] capture.compute_topk_set_check(npu_chosen, npu_topk, hf_chosen, hf_topk, k=5) walks both sequences in lockstep:
+ OK or FAIL; stop.Report.has_failure() returns True iff any record is FAIL.verify_topk_token_*.{json,md}; exit 1 on FAIL else exit 0 (PASS).Catches (every step exercises the entire production stack):
+Can miss:
+| File | Responsibility |
|---|---|
Makefile (parent) | verify / diagnosis / clean targets. MODEL=base|instruct, PROMPT=… for diagnosis. |
verify/verify_runner.py | Orchestrator. Builds NPU + HF runners, loops prompts, calls the comparator, writes the report, exits 1 on FAIL. |
verify/comparators.py | topk_token_ids (top-k with argmax-consistent tie-break), compute_topk_set_check (top-k token-level inclusion, mirrors vLLM's check_logprobs_close), plus diagnosis-only helpers (per_position_cosine, error_metrics, compare_pair). |
verify/report.py | Report accumulator + JSON / markdown dumpers. has_failure() returns True iff any npu_vs_hf record is FAIL. |
verify/runners/npu_runner.py | Imports + invokes the production prefill / decode / LM-head functions. |
verify/runners/hf_runner.py | Loads AutoModelForCausalLM in torch.bfloat16; delegates to model(input_ids, use_cache=True). |
verify/runners/_records.py | PrefillRecord / DecodeStepRecord dataclasses shared by both runners. |
verify/prompts/instruct.txt | 8 instruction-style prompts (MODEL=instruct); 7 from vllm/tests/prompts/example.txt + 1 GPU-related swap. |
verify/prompts/base.txt | 8 continuation-style prompts (MODEL=base); incomplete sentences matched to base behavior. |
Production-side touch points: llama32_1b_prefill.py::run_transformer_block populates ffn_out in the intermediates dict it already returns; diagnosis (which re-runs prefill layer-by-layer) reads it. Verify never reads any per-layer intermediates — it only consumes the final logits + chosen tokens.
cd programming_examples/llama32_1b
+
+make verify MODEL=instruct # ~3m41s — top-k token-level inclusion gate, NPU vs HF bf16 (NPU FA on)
+make verify MODEL=base # ~3m39s — base checkpoint, continuation prompts
+
+make diagnosis MODEL=instruct # ~2m55s — per-layer ffn_out cosine table (NPU FA on)
+make diagnosis MODEL=base # same lens, base checkpoint
+Reports land in verify/reports/{verify_topk_token_,diagnosis_}YYYYMMDD-HHMMSS.{json,md} (gitignored). The chosen MODEL, model_name, and (for verify) prompts_file are recorded in the report config so the file is unambiguous.
Companion: IMPLEMENTATION_GUIDE.html Part C (the original CI smoke that this subsystem extends).
+ + + + + diff --git a/programming_examples/llama32_1b/docs/explain.md b/programming_examples/llama32_1b/docs/explain.md index 58a399c81..0622885d8 100644 --- a/programming_examples/llama32_1b/docs/explain.md +++ b/programming_examples/llama32_1b/docs/explain.md @@ -249,8 +249,9 @@ The kernel exports the same `@rope` function name and signature as upstream, so no MLIR or multi-launch builder changes are needed. It is compiled to `rope.o` in `external_kernels.py:compile_rope()`. -The CPU reference (`llama32_1b_reference.py:apply_rope()`) uses the same half-split -convention, ensuring NPU and CPU produce identical results. +The NPU output is then gated against HuggingFace transformers in bf16 +(`make verify` — see [`VERIFICATION.html`](detail/VERIFICATION.html)), +which exercises the same half-split RoPE convention end-to-end. --- diff --git a/programming_examples/llama32_1b/docs/issues.md b/programming_examples/llama32_1b/docs/issues.md deleted file mode 100644 index 5a050a009..000000000 --- a/programming_examples/llama32_1b/docs/issues.md +++ /dev/null @@ -1,155 +0,0 @@ -# Known Issues and Future Work - -## 1. BF16 Precision Divergence (mostly resolved) - -**Previous symptom**: NPU decode generated repetitive or incorrect text compared -to HuggingFace. The instruct model was completely broken (emitting wrong control -tokens). - -**Root cause found and fixed**: The RoPE (Rotary Position Embedding) implementation -used the wrong rotation convention. Our pipeline used **interleaved** rotation -(pairing adjacent dimensions `d[2i], d[2i+1]`) while HuggingFace Llama uses -**half-split** rotation (pairing `d[i], d[i+32]`). This produced semantically -wrong Q/K rotations in every layer. - -**Fix**: Created `rope_halfsplit.cc` (custom NPU kernel matching HF convention) -and updated the LUT layout from interleaved `[cos,sin,cos,sin,...]` to -concatenated `[cos,...,sin,...]`. CPU reference updated to match. See -`docs/explain.md` for details. - -**Current status**: Both base and instruct models produce correct output: -- CPU reference vs HuggingFace: correlation 0.9997 (was 0.616 before fix) -- Instruct model generates correct Q&A responses -- Base model output quality improved (less repetitive) - -**Remaining**: Minor BF16 numerical differences still exist between NPU and CPU -(expected — different hardware precision paths). Some prompts may produce slightly -different token choices than HuggingFace, but the overall output quality is correct. - ---- - -## 2. Fixed Sequence Length (seq_len=2048) - -**Symptom**: All prompts are padded to 2048 tokens regardless of actual length. -A 6-token prompt processes 2042 EOS padding tokens, wasting ~99% of prefill compute. - -**Current behavior**: -``` -"The capital of France is" → 6 real + 2042 padding = 2048 tokens -"Hello" → 2 real + 2046 padding = 2048 tokens -``` - -Prefill takes the same wall time (~1.27s on NPU2) regardless of prompt length — -all 2048 positions are computed even when most are padding. - -**Why**: All NPU kernels are compiled with fixed dimensions: -- GEMM launch grids: M=2048 -- FlashAttention: lq=2048, lk=2048 -- Buffer Object sizes: (2048, 2048) matrices -- RoPE LUT: 2048 positions - -Changing seq_len requires recompiling all kernels (~4 min). - -**Impact**: -- Wasted prefill compute for short prompts -- May amplify BF16 precision loss (more unnecessary computation, see Issue #1) -- Cannot process prompts longer than 2048 tokens - -**Potential solutions**: - -1. **Bucket compilation**: Pre-compile kernels for multiple seq_len buckets - (e.g., 64, 256, 512, 1024, 2048). Route each prompt to the smallest bucket - that fits. Increases disk usage but dramatically reduces prefill time for - short prompts. - -2. **Dynamic seq_len**: Modify kernel builders to support runtime-parameterized - sequence length. Requires changes to GEMM launch grid computation, FlashAttention - tiling, and BO allocation. Significant engineering effort. - -3. **Chunked prefill**: Process the prompt in fixed-size chunks (e.g., 256 tokens - at a time), accumulating KV cache. Reuses one set of kernels compiled for - chunk_size. Requires incremental attention (append to KV cache each chunk). - ---- - -## 3. No Sampling (Greedy Decode Only) - -**Symptom**: Generated text tends to be repetitive, especially for base models. - -**Current behavior**: The decode loop uses `argmax` (greedy decoding) — always picks -the single highest-probability token. - -```python -next_token = int(np.argmax(logits[0])) # greedy, no randomness -``` - -**Impact**: Greedy decoding is deterministic but prone to repetition loops in base -models (like LLAMA-3.2-1B). The model gets stuck repeating high-probability patterns. - -**Fix**: Add temperature scaling + top-k sampling: -```python -# Temperature scaling -logits = logits / temperature - -# Top-k filtering -top_k_indices = np.argsort(logits)[-top_k:] -mask = np.full_like(logits, -np.inf) -mask[top_k_indices] = logits[top_k_indices] -logits = mask - -# Softmax + sample -probs = np.exp(logits - logits.max()) / np.sum(np.exp(logits - logits.max())) -next_token = np.random.choice(len(probs), p=probs) -``` - -IRON uses `temperature=0.7, top_k=50` which produces diverse, coherent text. -This is a straightforward Python-side change (no kernel modifications needed). - ---- - -## 4. Base Model vs Instruct Model - -**Current**: We use `meta-llama/Llama-3.2-1B` — the **base** (pre-training) model. -It is a text completion model, not a chatbot. It does not follow instructions or -answer questions. - -**Impact**: Prompts like "Which is larger: 9.11 or 9.9?" produce text completions, -not answers. The model treats the question as text to continue, not a query to answer. - -**Fix**: Switch to `meta-llama/Llama-3.2-1B-Instruct` for instruction-following. -This requires no kernel changes (same architecture, same weights shape). Only the -weight loading path changes. - ---- - -## 5. CPU Attention for Decode (Grows with Context) - -**Current**: Decode attention runs on CPU. At short contexts (pos < 100), this is -fast (~0.3ms/layer). At longer contexts, it grows linearly: - -| Context length | CPU attention/layer | x16 layers | -|---------------|--------------------|-----------| -| 50 tokens | 0.2ms | 3ms | -| 512 tokens | 0.9ms | 14ms | -| 2048 tokens | 1.8ms | 29ms | -| 4096 tokens | 3.6ms | 58ms | - -**Impact**: For long conversations (multi-turn chat with context > 2000 tokens), -CPU attention becomes a significant bottleneck (adding 30-60ms/token to the -92ms baseline). - -**Fix**: Implement an NPU decode attention kernel (single-query GQA with on-device -KV cache). This is different from the prefill FlashAttention (which has lq=2048). -The decode kernel would have lq=1 and attend to an on-device KV cache. - ---- - -## Priority - -| Issue | Impact | Effort | Priority | -|-------|--------|--------|----------| -| #3 Add sampling | Fixes repetition, matches IRON | Low | **High** | -| #2 Variable seq_len | 10-100x prefill speedup for short prompts | High | **High** | -| #1 BF16 divergence | Cosmetic (correct numerically) | N/A (inherent) | Low | -| #4 Instruct model | Better user experience | Low | Medium | -| #5 NPU decode attention | Long context performance | High | Medium | diff --git a/programming_examples/llama32_1b/docs/profile.md b/programming_examples/llama32_1b/docs/profile.md index ce281b550..2e5081238 100644 --- a/programming_examples/llama32_1b/docs/profile.md +++ b/programming_examples/llama32_1b/docs/profile.md @@ -6,16 +6,28 @@ | Phase | AIR (NPU2) | IRON | Speedup | |-------|------------|------|---------| -| **Prefill** (seq_len=2048) | **1.27s wall** | 2.744s | **2.17x** | -| **Decode** (steady-state) | **92ms/token (10.8 tok/s)** | 370ms/token (2.7 tok/s) | **4.0x** | - -- **Wall time**: End-to-end from embedding to LM Head argmax (includes minimal - Python host overhead — KV-cache extraction, embedding lookup, numpy views) +| **Prefill / TTFT** (seq_len=2048) | **1.27s wall** | 2.744s | **2.17x** | +| **Decode / TPOT** (steady-state) | **92ms/token (10.8 tok/s)** | 370ms/token (2.7 tok/s) | **4.0x** | + +- **TTFT** (time-to-first-token): end-to-end from `make run` invocation to + first decoded token — includes tokenize + EOS-pad + embed + 16 layers + + final RMSNorm + LM head GEMV. Matches the vLLM / TGI / TRT-LLM TTFT + definition. With tokenize added back in, current measured TTFT is + ~1.28 s (the 1.27 s row above is the NPU-only fraction used + in the IRON comparison, since IRON does not bundle the tokenizer). +- **TPOT** (time-per-output-token): steady-state per-token decode latency + (excludes prefill / first-token cost). Drift across 30 decode tokens is + <1% — see `Per-Token Wall Trend` in `make profile` output. - **IRON baseline**: measured against the IRON reference at commit [`2b62dc7`](https://github.com/amd/IRON/commit/2b62dc77ecc72f0fa8fb3381b05579ab84778d27) of `amd/IRON`, same NPU2 hardware (Strix), same LLAMA-3.2-1B BF16 model, same `seq_len=2048`. +For the visual end-to-end dataflow with per-step measured timing and the +BO Write / NPU Run / BO Read concept walkthrough, see +[`PROFILE.html`](detail/PROFILE.html). This file is the textual reference +(per-kernel tables, optimization history, vs IRON comparison). + **Recent optimizations** (vs. an earlier 1.54s wall headline): 1. Last-token-only LM Head: drop full-sequence NPU rmsnorm + 8-partition GEMM in prefill; do CPU rmsnorm on the 1×emb_dim last row (<1 ms) and reuse the @@ -88,13 +100,15 @@ Key differences favoring AIR: ## Prefill Breakdown (seq_len=2048, 16 layers) -### Wall Time Breakdown: 1.27s +### Wall Time Breakdown: 1.27s (NPU-only) / ~1.28s TTFT | Component | Time | Notes | |-----------|------|-------| -| **Kernel time** (sum of `load_and_run`) | ~1.16s | BO Write + NPU Run + BO Read (49 kernel calls: 16×3 transformer + 1 lm_head_gemv) | -| **Python host overhead** | ~0.11s | KV cache extraction, embedding lookup, CPU rmsnorm, numpy views | -| **Total wall time** | **1.27s** | | +| **NPU XRT calls** (sum of `load_and_run`) | ~1.12s | BO Write + NPU Run + BO Read across 49 calls: 16×3 transformer + 1 lm_head_gemv | +| **CPU host ops** (profiled) | ~37ms | tokenize + eos_pad + embed_lookup + 16×kv_cache_extract + final_rms_norm | +| **Python / numpy scheduling** | ~125ms | Per-layer dict access, numpy view setup, loop overhead (`layer-loop wall − inside-layer NPU − inside-layer CPU`) | +| **Total TTFT** (incl. tokenize) | **~1.28s** | matches `make run` Time-to-First-Token line | +| Total wall (NPU-only fraction, vs IRON) | ~1.27s | excludes tokenize; the row used in the IRON comparison | Overhead reduced from 0.67s → 0.24s by: - Suppressing print I/O in non-profile mode (4 prints × 16 layers) @@ -104,29 +118,41 @@ Overhead reduced from 0.67s → 0.24s by: - Skipping intermediate dict storage when not verifying - Removing redundant `.astype(bfloat16)` on already-bf16 kernel results -### Per-Kernel Timing +### Per-Kernel Timing (NPU XRT calls only) -| Kernel | Launches | Per-call | x Calls | Total | % | +| Kernel | Launches | Per-call | x Calls | Total | % of NPU | |--------|----------|----------|---------|-------|---| -| **o_ffn** | 8 | 41ms | 16 | **656ms** | **51%** | -| **flash_attn** | 1 | 22ms | 16 | **352ms** | **27%** | -| **lm_head** | 8 | 171ms | 1 | **171ms** | **13%** | -| **rms_gemms_rope** | 6 | 8ms | 16 | **128ms** | **10%** | -| rmsnorm | 1 | 3ms | 1 | 3ms | <1% | +| **o_ffn** | 8 (stitched) | 41.0ms | 16 | **656ms** | **59%** | +| **flash_attn** | 1 (separate ELF) | 21.6ms | 16 | **346ms** | **31%** | +| **rms_gemms_rope** | 6 (stitched) | 7.3ms | 16 | **117ms** | **10%** | +| **lm_head_gemv** | 8 partitions (stitched) | 13.6ms | 1 | **14ms** | **1%** | + +Per-CPU-op: -### Host vs NPU Breakdown (kernel time only) +| CPU op | Per-call | x Calls | Total | +|--------|----------|---------|-------| +| tokenize | ~10 ms | 1 | ~10 ms | +| eos_pad | <0.1 ms | 1 | <0.1 ms | +| embed_lookup | 5.8 ms | 1 | 5.8 ms | +| kv_cache_extract | 1.1 ms | 16 | 17.6 ms | +| final_rms_norm | 3.1 ms | 1 | 3.1 ms | + +### Host vs NPU Breakdown (XRT calls only — `cache.load_and_run` internals) | | BO Write | NPU Run | BO Read | Total | |---|----------|---------|---------|-------| -| **Sum** | 48ms | 1237ms | 9ms | 1294ms | -| **%** | **4%** | **96%** | **1%** | 100% | +| **Sum** | 46ms | 1062ms | 5ms | 1113ms | +| **%** | **4%** | **95%** | **0%** | 100% | + +(BO Read is zero-copy view construction — see PROFILE.html Part C for what +these three segments actually measure.) ### Per-Layer Data Flow ``` Layer input: x_bf16 (2048x2048, 8MB) -┌─ KERNEL 1: rms_gemms_rope (8ms/layer) ─────────────────────────┐ +┌─ KERNEL 1: rms_gemms_rope (7.3ms/layer) ───────────────────────┐ │ │ │ WRITE: x_in (8MB) ← activation, changes/layer │ │ SKIP: norm_w, wq, wk, wv ← STATIC (per-layer BO) │ @@ -142,7 +168,7 @@ Layer input: x_bf16 (2048x2048, 8MB) │ READ: v (2MB), q_roped (8MB), k_roped (2MB) │ └────────────────────────────┬────────────────────────────────────┘ ▼ -┌─ KERNEL 2: flash_attn (22ms/layer) ────────────────────────────┐ +┌─ KERNEL 2: flash_attn (21.6ms/layer) ──────────────────────────┐ │ │ │ WRITE: q_roped (8MB), k_roped (2MB), v (2MB) │ │ SKIP: attn_out ← INTERMEDIATE │ @@ -173,8 +199,11 @@ Layer input: x_bf16 (2048x2048, 8MB) └─────────────────────────────────────────────────────────────────┘ × 16 layers, then: - rmsnorm (3ms): Final layer normalization - lm_head (171ms): 8-partition GEMM → vocab logits → argmax → first token + final_rms_norm (CPU, 3.1ms): RMSNorm on single prediction-position row + lm_head_gemv (NPU, 13.6ms): 8-partition GEMV → vocab logits → argmax → first token + (reuses the decode-side 8-partition ELF; see + A7 in IMPLEMENTATION_GUIDE.html for why + full-seq GEMM was dropped in favor of single-row GEMV) ``` --- diff --git a/programming_examples/llama32_1b/docs/usage.md b/programming_examples/llama32_1b/docs/usage.md index 990e2a823..e2e4caa55 100644 --- a/programming_examples/llama32_1b/docs/usage.md +++ b/programming_examples/llama32_1b/docs/usage.md @@ -102,41 +102,68 @@ What happens internally: ### `make profile` -Same as `make run` but prints per-token timing and kernel breakdown. +Same as `make run` but enables the otherwise-disabled `Profiler` so the +end-to-end inference path is broken down into per-XRT-call and per-CPU-op +wall times. Production code path is identical to `make run`. ```bash make profile -make profile N_TOKENS=10 +make profile N_TOKENS=30 PROMPT="Explain photosynthesis in detail." ``` -Example output (with `N_TOKENS=10`): -``` -NPU prefill done in 1.27s. First token: 12366 - -Decoding 10 tokens (token 1 to 10)... - Token 1: id=13, time=92ms - Token 2: id=1102, time=91ms - ... - Token 10: id=578, time=92ms +After the model output, the report prints (per phase: prefill / decode): + +1. **END-TO-END DATAFLOW** — architecture-aware summary in dataflow order + (tokenize → eos_pad → embed → 16×(rms_gemms_rope + flash_attn + o_ffn + + kv_cache_extract) → final_norm → lm_head_gemv → per-query total). + Mirrors the SVGs in [`PROFILE.html`](detail/PROFILE.html). +2. **Wall-Time Attribution** — totals: NPU XRT vs CPU host ops vs layer-loop. +3. **Per-Layer Execution** — one row per prefill layer; aggregated avg/min/max + per layer across tokens for decode. +4. **NPU XRT Call Breakdown** — each multi-launch ELF, wall time per call. +5. **CPU Op Breakdown** — each tracked CPU host op (embed, kv_cache_extract, + final_rms_norm, tokenize, eos_pad, decode_attention_cpu). +6. **Fine-Grained NPU Breakdown** — each XRT call split into + `BO Write` / `NPU Run` / `BO Read` (concept explained in PROFILE.html + Part C). +7. **Per-Token Wall Trend** (decode only) — token 1 / middle / last wall + + first→last drift %, so you can spot any KV-cache-growth-driven slowdown. + +For reproduction commands + visual dataflow + concept walkthrough see +[`PROFILE.html`](detail/PROFILE.html). + +### `make verify` (and `make verify-full`) + +Top-k token-level inclusion gate against HuggingFace transformers in **bf16** +(same dtype as NPU). Greedy-decodes a pre-selected prompt set × 32 tokens; at +each step, both runners' chosen tokens must appear in the OTHER side's top-5. +Pass/fail signal for end-to-end production correctness. Mirrors vLLM's +`check_logprobs_close` method. -Generated 10 tokens in 0.92s -Tokens/second: 10.87 -Time/token: 92ms +```bash +make verify # 2 prompts (fast CI gate, ~2 min) +make verify-full # full 8-prompt sweep (~6 min) +make verify MODEL=base # base checkpoint, continuation prompts ``` -### `make verify` +`make verify` runs the first 2 prompts from the model's prompt file and is the +default CI gate. `make verify-full` runs every prompt in the file (currently 8) +for exhaustive local validation. Token count and `k` are fixed by the gate +(32 / 5) — not user-tunable. + +### `make diagnosis` -Runs inference and compares every intermediate result against a CPU F32 reference. -Useful for validating correctness after kernel changes. +Per-layer `ffn_out` cosine + max_abs error vs HF bf16 for a single prompt. +Informational only (never fails the run); reach for it when `make verify` +flags a regression and you need to localize which layer drifted. ```bash -make verify N_TOKENS=10 +make diagnosis # uses default PROMPT +make diagnosis PROMPT="The capital of France is" ``` -Checks: -- Per-layer KV cache correlation (NPU vs CPU) -- Logits correlation at prediction position -- Top-1 token match +See [VERIFICATION.html](detail/VERIFICATION.html) for the full design rationale, +gate criteria, and report layout. ### `make clean` @@ -175,7 +202,7 @@ llama32_1b/ ├── llama32_1b_prefill.py ← Prefill-only pipeline ├── llama32_1b_decode.py ← Decode-only pipeline ├── llama32_1b_weights.py ← Weight loading from safetensors -├── llama32_1b_reference.py ← CPU F32 reference +├── llama32_1b_cpu_helpers.py ← Small NumPy helpers: rms_norm, attention_reference, softmax │ ├── kernel_builder/ ← Shared kernel infrastructure │ ├── stitching.py ← MLIR text stitching for multi-launch ELFs @@ -212,5 +239,7 @@ llama32_1b/ **Slow first token**: The NPU enters power-save after ~10s idle. The warmup pass handles this automatically. If running manually, ensure `prepare_runtime()` is called. -**Wrong results**: Run `make verify` to compare against CPU reference. Check that -`.o` files are fresh (`make clean` then `make compile`). +**Wrong results**: Run `make verify` to gate against HuggingFace transformers +bf16 (top-k token inclusion). If verify fails, run `make diagnosis` to +localize which layer drifted. Check that `.o` files are fresh +(`make clean` then `make compile`). diff --git a/programming_examples/llama32_1b/kernel_builder/cache.py b/programming_examples/llama32_1b/kernel_builder/cache.py index d35dca937..2203b53af 100644 --- a/programming_examples/llama32_1b/kernel_builder/cache.py +++ b/programming_examples/llama32_1b/kernel_builder/cache.py @@ -45,7 +45,6 @@ def prepare_air_project(): "attn_npu2.o", "mv.o", "mv_k8192.o", - "attn_decode_npu2.o", ]: src = Path(obj_name) if src.exists(): @@ -58,7 +57,8 @@ class Profiler: def __init__(self, enabled=False): self.enabled = enabled self.compile_times = {} # name -> seconds - self.kernel_times = {} # name -> list of seconds + self.kernel_times = {} # NPU XRT call: name -> list of seconds + self.cpu_times = {} # CPU op: name -> list of seconds self.layer_times = [] # list of (layer_idx, seconds) self.kernel_breakdowns = ( {} @@ -72,6 +72,15 @@ def record_kernel(self, name, duration): if self.enabled: self.kernel_times.setdefault(name, []).append(duration) + def record_cpu(self, name, duration): + """Record a CPU host-side operation's wall time. Use for things like + embed lookup, KV-cache extract, CPU attention fallback, final RMSNorm + — anything that is not an `xrt.run()` but consumes inference wall + time. Reported in a separate section from NPU XRT calls so the two + are easy to compare.""" + if self.enabled: + self.cpu_times.setdefault(name, []).append(duration) + def record_breakdown( self, name, write_ms, kernel_ms, read_ms, n_written, bytes_written, n_readback ): @@ -89,12 +98,45 @@ def record_breakdown( def start_layer(self): if self.enabled: - return time.time() + return time.perf_counter() return None def end_layer(self, layer_idx, t0): if self.enabled and t0 is not None: - self.layer_times.append((layer_idx, time.time() - t0)) + self.layer_times.append((layer_idx, time.perf_counter() - t0)) + + def time_cpu(self, name): + """Context manager: `with prof.time_cpu("embed_lookup"): ...` + Records the elapsed wall time as a CPU op named `name`. Safe to + use whether enabled or disabled (zero overhead when disabled).""" + prof = self + + class _Ctx: + def __enter__(self_inner): + self_inner.t0 = time.perf_counter() if prof.enabled else None + return self_inner + + def __exit__(self_inner, *exc): + if self_inner.t0 is not None: + prof.record_cpu(name, time.perf_counter() - self_inner.t0) + return False + + return _Ctx() + + def per_token_walls_ms(self, n_layers): + """Sum every consecutive `n_layers` layer-time entries into one + per-token wall (in ms). Returns [] if not enabled or no data. + Used by the dataflow summary to expose decode slowdown trends.""" + if not self.enabled or not self.layer_times: + return [] + if len(self.layer_times) % n_layers != 0: + # Shouldn't happen in a clean run; bail rather than mis-bucket. + return [] + out = [] + for tok_start in range(0, len(self.layer_times), n_layers): + chunk = self.layer_times[tok_start : tok_start + n_layers] + out.append(sum(t for _, t in chunk) * 1000.0) + return out def report(self): if not self.enabled: @@ -104,6 +146,36 @@ def report(self): print("PROFILING REPORT") print(f"{'='*60}") + # Top-level phase summary: total wall time attributed to NPU XRT + # calls vs CPU host ops vs the layer envelope. Sums won't add up + # exactly (layer envelope is the wall budget; NPU + CPU are the + # accounted-for parts inside it; remainder is python scheduling / + # numpy view setup / loop overhead). Useful as a sanity check. + if self.kernel_times or self.cpu_times or self.layer_times: + npu_total_ms = sum(t * 1000 for v in self.kernel_times.values() for t in v) + cpu_total_ms = sum(t * 1000 for v in self.cpu_times.values() for t in v) + layer_total_ms = sum(t * 1000 for _, t in self.layer_times) + npu_count = sum(len(v) for v in self.kernel_times.values()) + cpu_count = sum(len(v) for v in self.cpu_times.values()) + print(f"\n--- Wall-Time Attribution ---") + if npu_count: + print( + f" NPU XRT calls {npu_total_ms:9.2f}ms ({npu_count} calls)" + ) + if cpu_count: + print( + f" CPU host ops {cpu_total_ms:9.2f}ms ({cpu_count} calls)" + ) + if self.layer_times: + accounted = npu_total_ms + cpu_total_ms + # CPU ops happen both inside and outside the layer envelope; + # so layer_total_ms is the inside-layer wall budget, and the + # remainder vs (NPU+CPU) inside layers is python overhead. + print( + f" Layer-loop wall {layer_total_ms:9.2f}ms " + f"({len(self.layer_times)} layer-invocations)" + ) + if self.compile_times: print(f"\n--- Compilation Phase ---") total_compile = 0 @@ -115,34 +187,71 @@ def report(self): ) if self.layer_times: - print(f"\n--- Per-Layer Execution ---") + # Group by layer_idx. Prefill: each idx appears once -> one row per + # layer. Decode: each idx appears once per token -> aggregate with + # avg / min / max / count. + from collections import defaultdict + + grouped = defaultdict(list) for idx, t in self.layer_times: - print(f" Layer {idx:3d}: {t:8.2f}s") - total_layers = sum(t for _, t in self.layer_times) - print(f" {'Total prefill':40s} {total_layers:8.2f}s") + grouped[idx].append(t * 1000.0) # ms + multi_invocation = any(len(v) > 1 for v in grouped.values()) + print(f"\n--- Per-Layer Execution ---") + if multi_invocation: + for idx in sorted(grouped): + ts = grouped[idx] + print( + f" Layer {idx:3d}: avg={sum(ts)/len(ts):7.2f}ms " + f"min={min(ts):7.2f}ms max={max(ts):7.2f}ms (x{len(ts)})" + ) + else: + for idx in sorted(grouped): + print(f" Layer {idx:3d}: {grouped[idx][0]:7.2f}ms") + total_ms = sum(t * 1000.0 for _, t in self.layer_times) + print(f" {'Total layer-time':40s} {total_ms:8.2f}ms") if self.kernel_times: - print(f"\n--- Kernel Breakdown (avg per invocation) ---") + print(f"\n--- NPU XRT Call Breakdown (avg per invocation) ---") total_avg = 0 for name, times in sorted(self.kernel_times.items()): - avg = sum(times) / len(times) - total_avg += avg * len(times) - mn = min(times) - mx = max(times) - count = len(times) + times_ms = [t * 1000.0 for t in times] + avg = sum(times_ms) / len(times_ms) + total_avg += avg * len(times_ms) + count = len(times_ms) print( - f" {name:40s} avg={avg:6.3f}s " - f"min={mn:6.3f}s max={mx:6.3f}s (x{count})" + f" {name:40s} avg={avg:7.2f}ms " + f"min={min(times_ms):7.2f}ms max={max(times_ms):7.2f}ms (x{count})" ) if self.layer_times: n_layers = len(self.layer_times) - print(f" {'Total kernel time':40s} {total_avg:8.2f}s") + print(f" {'Total kernel time':40s} {total_avg:8.2f}ms") print( - f" {'Avg per layer (kernel time)':40s} {total_avg/n_layers:8.2f}s" + f" {'Avg per layer (kernel time)':40s} {total_avg/n_layers:8.2f}ms" ) + if self.cpu_times: + print(f"\n--- CPU Op Breakdown (avg per invocation) ---") + total_cpu_ms = 0 + for name, times in sorted(self.cpu_times.items()): + times_ms = [t * 1000.0 for t in times] + avg = sum(times_ms) / len(times_ms) + total_cpu_ms += avg * len(times_ms) + count = len(times_ms) + print( + f" {name:40s} avg={avg:7.2f}ms " + f"min={min(times_ms):7.2f}ms max={max(times_ms):7.2f}ms (x{count})" + ) + print(f" {'Total CPU op time':40s} {total_cpu_ms:8.2f}ms") + if self.kernel_breakdowns: - print(f"\n--- Fine-Grained Breakdown (avg per invocation) ---") + print(f"\n--- Fine-Grained NPU Breakdown (avg per invocation) ---") + print( + f" Three-segment timing of each XRT call:\n" + f" BO Write = host→DDR memcpy of dynamic inputs (weights\n" + f" pre-loaded once via static_input_indices)\n" + f" NPU Run = xrt.run.start() + wait() — actual NPU exec\n" + f" BO Read = numpy view construction (zero-copy, ~0)" + ) print( f" {'Kernel':20s} {'BO Write':>10s} {'NPU Run':>10s} {'BO Read':>10s} {'Total':>10s} {'Written':>8s} {'Read':>6s}" ) @@ -316,8 +425,16 @@ def load_and_run( output_indices: Optional list of buffer indices to read back from device. If None, only the last buffer is read back (default). Use for multi-output kernels (e.g. attn_gemms: [2, 4, 6]). + static_input_indices: Optional set of buffer indices that are static + (e.g. weights, LUTs). On the first call for a given bo_key the BO is + written; on subsequent calls the host->device sync is skipped because + the kernel reads from the already-resident BO. intermediate_indices: Optional set of buffer indices that are intermediate (overwritten by kernel). Skips host->device sync. + bo_key: Optional cache key for BO reuse. Calls sharing a bo_key reuse + the same xrt.bo objects, which combined with static_input_indices + enables write-once-read-many for weights. Default uses the kernel + name (one BO set shared across all calls to that kernel). Returns: Tuple of numpy arrays (all kernel outputs) diff --git a/programming_examples/llama32_1b/kernel_builder/external_kernels.py b/programming_examples/llama32_1b/kernel_builder/external_kernels.py index 02287e390..3613658fc 100644 --- a/programming_examples/llama32_1b/kernel_builder/external_kernels.py +++ b/programming_examples/llama32_1b/kernel_builder/external_kernels.py @@ -12,7 +12,6 @@ """ import os -import shutil import subprocess from pathlib import Path @@ -27,28 +26,30 @@ def _get_peano_clang(): def _get_aie_include_dir(): """Find the AIE API include directory (for aie_api/aie.hpp).""" - # Primary: locate via aie-opt on PATH. Matches the convention used by - # every other Makefile in this repo (AIEOPT_DIR = $(dir $(which aie-opt))/..) - # and works for both local source builds and CI's mlir_aie wheel install. - aie_opt = shutil.which("aie-opt") - if aie_opt: - p = Path(aie_opt).resolve().parent.parent / "include" - if (p / "aie_api" / "aie.hpp").exists(): - return str(p) - # Fallback: explicit local dev install path. - p = ( + # Try mlir-aie install path relative to this file (main-repo layout) + candidates = [ Path(__file__).resolve().parent.parent.parent.parent / "my_install" / "mlir-aie" / "install" - / "include" - ) - if (p / "aie_api" / "aie.hpp").exists(): - return str(p) - raise RuntimeError( - "Cannot find aie_api/aie.hpp include directory " - "(no aie-opt on PATH and no my_install/mlir-aie/install)" - ) + / "include", + ] + # Also honour MLIR_AIE_INSTALL_DIR env var (set by env_setup.sh; works + # in git worktrees where the relative path above resolves to the worktree + # root rather than the main repo root). + mlir_aie_dir = os.environ.get("MLIR_AIE_INSTALL_DIR", "") + if mlir_aie_dir: + candidates.append(Path(mlir_aie_dir) / "include") + for p in candidates: + if (p / "aie_api" / "aie.hpp").exists(): + return str(p) + # Fallback: search from PEANO_INSTALL_DIR + peano_dir = os.environ.get("PEANO_INSTALL_DIR", "") + if peano_dir: + p = Path(peano_dir).parent.parent / "include" + if (p / "aie_api" / "aie.hpp").exists(): + return str(p) + raise RuntimeError("Cannot find aie_api/aie.hpp include directory") _PEANO_FLAGS = [ @@ -171,20 +172,6 @@ def compile_mv(tile_m=8): _compile_kernel(src, "mv.o", extra_flags=[f"-DDIM_M_OUTPUT={tile_m}"]) -def compile_attn_decode_npu2(head_dim=64): - """Compile attn_decode_npu2.o (RoPE helpers for the fused decode kernel).""" - src = _PROJ_ROOT / "attention_decode" / "attn_decode_npu2.cc" - _compile_kernel( - src, - "attn_decode_npu2.o", - extra_flags=[ - f"-DDIM_N={head_dim}", - f"-DHEAD_SIZE={head_dim}", - "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", - ], - ) - - def compile_all_external_kernels(head_dim=64): """Compile all external C++ kernels from source. @@ -195,6 +182,5 @@ def compile_all_external_kernels(head_dim=64): compile_silu_and_mul() compile_rope() compile_attn_npu2(head_dim=head_dim) - compile_attn_decode_npu2(head_dim=head_dim) compile_mv() compile_mv_k8192() diff --git a/programming_examples/llama32_1b/llama32_1b_cpu_helpers.py b/programming_examples/llama32_1b/llama32_1b_cpu_helpers.py new file mode 100644 index 000000000..72a854e96 --- /dev/null +++ b/programming_examples/llama32_1b/llama32_1b_cpu_helpers.py @@ -0,0 +1,88 @@ +# Copyright (C) 2026, Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Small NumPy CPU helpers shared by production prefill/decode + verify. + +This file used to be a full F32 CPU forward-pass implementation of the model +(plus a standalone `--verify` CLI that compared the F32 forward against HF +transformers F32). With the verify subsystem rewritten to compare directly +against HF transformers in bf16 (see verify/), that whole F32 reference +chain became redundant. What is kept here is the small set of NumPy helpers +that production still imports: + + - rms_norm : LM-head GEMV final-norm (inference.py prefill end, + and every decode step). + - attention_reference: prefill cpu_attn=True fallback (full GQA attention + in F32 on host; used when the NPU FlashAttention + kernel is unavailable for the configured head_dim). + - softmax : kept because attention_reference uses it; not + imported anywhere else. +""" + +import numpy as np + + +def rms_norm(x, weight, eps=1e-5): + """RMS normalization: x / sqrt(mean(x^2) + eps) * weight. + + Args: + x: (M, N) input array in F32. + weight: (N,) learned scale parameter. + eps: Small constant for numerical stability. + + Returns: + (M, N) normalized and scaled array in F32. + """ + x = np.asarray(x, dtype=np.float32) + weight = np.asarray(weight, dtype=np.float32) + rms = np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + eps) + return (x / rms) * weight + + +def softmax(x, axis=-1): + """Numerically stable softmax (used by attention_reference).""" + x = np.asarray(x, dtype=np.float32) + x_max = np.max(x, axis=axis, keepdims=True) + exp_x = np.exp(x - x_max) + return exp_x / np.sum(exp_x, axis=axis, keepdims=True) + + +def attention_reference(q, k, v, n_heads, n_kv_heads): + """Multi-head attention with Grouped Query Attention (GQA), causal mask. + + Args: + q: (seq_len, n_heads * head_dim) -- already projected and RoPE'd. + k: (seq_len, n_kv_heads * head_dim) -- already projected and RoPE'd. + v: (seq_len, n_kv_heads * head_dim) -- already projected. + n_heads: Number of query heads. + n_kv_heads: Number of key/value heads (for GQA). + + Returns: + (seq_len, n_heads * head_dim) attention output (F32). + """ + q = np.asarray(q, dtype=np.float32) + k = np.asarray(k, dtype=np.float32) + v = np.asarray(v, dtype=np.float32) + + seq_len = q.shape[0] + head_dim = q.shape[1] // n_heads + group_size = n_heads // n_kv_heads + + # Reshape to per-head views: (seq, n_*_heads, head_dim) -> (n_*_heads, seq, head_dim) + q = q.reshape(seq_len, n_heads, head_dim).transpose(1, 0, 2) + k = k.reshape(seq_len, n_kv_heads, head_dim).transpose(1, 0, 2) + v = v.reshape(seq_len, n_kv_heads, head_dim).transpose(1, 0, 2) + + scale = 1.0 / np.sqrt(head_dim) + causal_mask = np.triu(np.full((seq_len, seq_len), -np.inf, dtype=np.float32), k=1) + + out_heads = np.empty((n_heads, seq_len, head_dim), dtype=np.float32) + for h in range(n_heads): + kv_idx = h // group_size + scores = q[h] @ k[kv_idx].T * scale + scores = scores + causal_mask + probs = softmax(scores, axis=-1) + out_heads[h] = probs @ v[kv_idx] + + # (n_heads, seq, head_dim) -> (seq, n_heads * head_dim) + return out_heads.transpose(1, 0, 2).reshape(seq_len, n_heads * head_dim) diff --git a/programming_examples/llama32_1b/llama32_1b_decode.py b/programming_examples/llama32_1b/llama32_1b_decode.py index ccb80cdee..37de7d75c 100644 --- a/programming_examples/llama32_1b/llama32_1b_decode.py +++ b/programming_examples/llama32_1b/llama32_1b_decode.py @@ -157,7 +157,7 @@ def run_decode_block( rope_lut_bf16: (max_seq, head_dim) RoPE LUT Returns: - output: (emb_dim,) — block output + output: (emb_dim,) — block output. """ emb_dim = config.emb_dim n_heads = config.n_heads @@ -232,15 +232,19 @@ def _run(name, backend, *inputs, static_indices=None, **kwargs): v_cache_layer[:, current_pos, :] = v.reshape(n_kv_heads, head_dim) # --- CPU Attention --- - attn_out = decode_attention_cpu( - q_roped.flatten(), - k_cache_layer, - v_cache_layer, - current_pos, - n_heads, - n_kv_heads, - head_dim, - ) + # Single-query attention against the growing K/V cache. CPU-side because + # at head_dim=64 the NPU FA kernel's per-call overhead dominates the + # single-query workload. + with cache.profiler.time_cpu("decode_attention_cpu"): + attn_out = decode_attention_cpu( + q_roped.flatten(), + k_cache_layer, + v_cache_layer, + current_pos, + n_heads, + n_kv_heads, + head_dim, + ) # --- Call 2: o_gemv_ffn (8 launches, 15 args) --- # O GEMV + Add + RMSNorm + Gate/Up GEMV + SiLU*mul + Down GEMV + Add @@ -281,6 +285,4 @@ def _run(name, backend, *inputs, static_indices=None, **kwargs): static_indices={0, 7, 9, 12}, intermediate_indices={2, 4, 6, 8, 10, 11, 13, 14}, ) - output = results[14].astype(bfloat16) - - return output + return results[14].astype(bfloat16) diff --git a/programming_examples/llama32_1b/llama32_1b_inference.py b/programming_examples/llama32_1b/llama32_1b_inference.py index 18c9de206..a4b768a43 100644 --- a/programming_examples/llama32_1b/llama32_1b_inference.py +++ b/programming_examples/llama32_1b/llama32_1b_inference.py @@ -17,7 +17,6 @@ # Run inference with cached kernels: python3 ../llama32_1b_inference.py --run-only --n-tokens 10 --profile python3 ../llama32_1b_inference.py --run-only --n-tokens 100 --profile - python3 ../llama32_1b_inference.py --run-only --n-tokens 5 --verify python3 ../llama32_1b_inference.py --run-only --n-tokens 20 --prompt "Once upon a time" """ @@ -37,10 +36,9 @@ from llama32_1b_weights import ( LlamaConfig, load_weights, - synthetic_weights, generate_rope_lut, ) -from kernel_builder.cache import KernelCache +from kernel_builder.cache import KernelCache, Profiler from kernel_builder.external_kernels import compile_all_external_kernels from kernel_builder.backend_presets import ( LM_GEMV_BACKEND, @@ -82,21 +80,6 @@ def _delta_text(tokenizer: Any, ids: list[int], state: _StreamState) -> str: return delta -class _SyntheticTokenizer: - """Stub tokenizer used with --synthetic-weights (no HuggingFace dependency). - - The synthetic path skips real tokenization entirely (token IDs come from a - deterministic numpy array); this stub satisfies the few attribute lookups - the pipeline still does — eos_token_id (decode-loop stop) and decode() - (verify/profile prints). - """ - - eos_token_id = -1 # never matches real token ids; decode loop runs full N - - def decode(self, ids, skip_special_tokens=False): # noqa: ARG002 - return f"