Llama-3.2-1B on AMD NPU2 — Implementation Guide

A model-first walkthrough: understand what Llama-3.2-1B inference IS, then how this codebase runs it on AMD NPU2 hardware.

Part A — The Model (no NPU yet)

A1. Llama-3.2-1B at a glance

Hyperparameters (defined in LlamaConfig at llama32_1b_weights.py:36)

Total parameter accounting (~1.24 B)

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.

A2. The transformer block — math and shapes

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 convention (used in every shape annotation below)

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

A2.1 — Attention sub-block

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

Per-kernel explanations (attention sub-block)

A2.2 — FFN sub-block (SwiGLU)

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.

Per-kernel explanations (FFN sub-block)

A2.3 — Block-level annotations

Compute-heavy ops (FLOPs ranking, prefill at S=2048)

The three FFN GEMMs dominate FLOPs because D_ff is 4× larger than H. Per-block prefill FLOPs:

• Gate proj: 2 · S · H · D_ff ≈ 2 · 2048 · 2048 · 8192 = 69 GFLOP
• Up proj: same as gate ≈ 69 GFLOP
• Down proj: 2 · S · D_ff · H ≈ 69 GFLOP
• Q proj: 2 · S · H · H ≈ 17 GFLOP
• K proj, V proj: each ≈ 4 GFLOP (smaller because of GQA)
• O proj: 17 GFLOP
• Attention compute: 4 · S² · H ≈ 34 GFLOP (dominated by S² scaling — biggest if S grew)

The 3 FFN projections together = 207 GFLOP per layer ≈ 60% of per-layer compute. × 16 layers × 1.27 s prefill ≈ 2.6 TFLOP/s achieved on the NPU.

Memory-bound ops (bandwidth-limited at small S)

RMSNorm and the elementwise SwiGLU multiply have low arithmetic intensity (~1 FLOP/byte). Attention's softmax + the sub-multiplies inside FlashAttention also become memory-bound when S is small or d_h is small. In decode (S=1), everything except the GEMVs is memory-bound — this is why the per-token decode time is dominated by weight bandwidth, not FLOPs.

Fusable kernel boundaries

Common fusions seen in this and other implementations:

• SiLU + elementwise multiply → one kernel (silu_and_mul.cc). Saved per-pass over the 8192-wide tensor.
• Gate proj + Up proj → one big GEMM with output dim 2·D_ff (some implementations; ours doesn't currently).
• FlashAttention fuses transpose + QK^T + mask + softmax + SV into one tiled kernel (this is exactly what makes "FA" different from naive attention).
• RMSNorm + next GEMM can be fused with epilogue tricks; our impl does NOT fuse this (norm is its own sub-launch). Trade-off vs the multi-launch ELF approach.

The marginal contribution of our specific multi-launch grouping has been validated in internal measurements.

Convention gotchas (where this implementation differs from "vanilla" Llama)

• RoPE half-split vs interleaved. HuggingFace Llama (and our impl, via 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 mask is implicit in FlashAttention. Our FA kernel takes causal=True and never materializes a mask matrix; it just skips attending to t > s in the inner loop.
• RMSNorm has no bias. Unlike LayerNorm. Just x · rsqrt(mean(x²) + ε) · γ. ε is a small constant (1e-5 typically) for numerical stability.
• No dropout at inference. (Only relevant at training.)

GQA effects on KV cache size

With G = 4 (each KV head shared by 4 Q heads), the KV cache is 4× smaller than it would be without GQA. For Llama-3.2-1B at max_seq=2048:
KV cache size = 2 · L · N_kv · max_seq · d_h · 2 bytes = 2 · 16 · 8 · 2048 · 64 · 2 = ~32 MB
Without GQA (N_kv = N_h = 32), this would be ~128 MB. The savings matter much more for larger models / longer sequences.

Weight sharing

Llama-3.2-1B uses untied embeddings — the LM head 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.

A2.4 — Mapping back to our codebase

The 14 ops above map to the production NPU kernels as follows:

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.

One transformer block as math (paraphrased)

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 output

A3. Full forward pass — what one inference call does

Top-level pipeline

The diagram below shows the whole inference call as 6 stages. The decoder block is collapsed (×L) — its internals are diagrammed in A2.

Per-stage explanations (top-level pipeline)

The two operating modes (model-level)

The forward pass above works for ANY input length. But there are two common usage patterns:

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.

A4. KV cache — what it is, why we need it, how it grows

The problem

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 masked

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

The solution: cache K and V

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.

Memory layout in our codebase

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)

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.

Visual: how the K/V cache grows

Showing one layer's K cache (the V cache has the same structure). Each cell is one position; rows are the 8 KV heads.

State after prefill (prompt_len = 7 tokens, max_seq = 20 in this toy example):

State after 4 decode steps (current_pos = 11):

The key code points

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

A5. Padding to fixed seq_len + finding the real prompt

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.

Step 1 — Tokenization (host, CPU)

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

Step 2 — Padding to seq_len

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.

Step 3 — Recovering the real prompt length inside prefill

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

Step 4 — Prefill processes ALL 2048 positions but only reads pred_pos's logits

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.

Step 5 — KV cache for decode uses prompt_len, not seq_len

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 token

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

Cost of padding

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.

Visual summary of the prompt+padding+decode lifecycle

A6. Does padding affect the math at real positions?

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.

Per-op analysis: which ops cross positions?

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?

What about the padding positions themselves?

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:

• LM Head logits: only computed at pred_pos (see A7), so padding-position logits don't exist.
• KV cache for decode: the cache slots at indices 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.)
• Layer N+1's 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.

Subtle case: do dropout, layer norm running stats, etc. matter?

No, because:

• Dropout is not used at inference time.
• RMSNorm has no running statistics (unlike BatchNorm — RMSNorm is purely per-row at inference; no batch statistics to corrupt).
• FlashAttention's softmax normalizes per-row (per-query-position) — the denominator at row pred_pos sums over only positions 0..pred_pos due to the causal mask. Padding positions don't enter the sum.

How to verify this claim

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

A7. Single-row LM Head GEMV — workaround or general optimization?

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.

Why this is true

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.

Where multi-row LM Head WOULD be needed

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.

The math savings

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

Padding workaround vs production-grade variable-length support

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:

What our implementation IS vs IS NOT

The "single-row LM Head" optimization is general; the "padding-to-2048" optimization is specific

To return to your distinction: these are two completely separate things.

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 B — How we run it on the NPU

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.

B1. End-to-end runtime flow

Implementation overview — prefill

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.

Implementation overview — decode (per token)

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.

NPU calls per pass — concrete count

NPU2 tile array — context

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:

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.

The 4 phases of llama32_1b_inference.py:main

From make run to printed output:

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

B2. The kernel building blocks

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

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

The compile pipeline (one ELF, regardless of pattern)

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

The 7 kernels — quick index

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.

B2.1 — RMSNorm

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

B2.2 — GEMM (matrix-matrix multiply, prefill)

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:

1. Calls the upstream 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.
2. Applies an extra transform IR script (the ~100-line 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++.

B2.3 — GEMV (matrix-vector multiply, decode)

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.

B2.4 — RoPE (Rotary Position Embedding)

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:

• Half-split (used by HuggingFace Llama and our impl): pair (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}].
• Interleaved (used by 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.

B2.5 — SwiGLU (silu_and_mul, fused activation)

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)

B2.6 — FlashAttention

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.

B2.7 — Eltwise Add (residual)

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.

B2.8 — Compile-time helpers and orchestration

Two files coordinate the actual external-C++ compilation:

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

Tile-mapping summary

Side-by-side view of how each of the 7 kernels maps onto the NPU2 8×4 compute array:

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.

B3. From standalone kernels to end-to-end inference — the four gaps

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

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.

B4. Gap #1 — Layout matching between kernels

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:

Mismatch #1 — Weight matrix orientation (GEMV)

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.

Mismatch #2 — KV cache layout (prefill ↔ FlashAttention ↔ decode)

The same physical KV tensor is touched by three different consumers, each with its own preferred layout:

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

Mismatch #3 — GEMM output flat shape vs. RoPE multi-head input

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.

Mismatch #4 — FFN flat output for the next layer

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.

Mismatch #5 — Two GEMV variants in one ELF (K=2048 and K=8192)

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.

B5. Gap #2 — Multi-launch ELF stitching

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.

The mechanism

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 6 production ELFs (stitched products)

The production code stitches the 7 kernel building blocks from B2 into 6 ELFs:

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

Why FlashAttention is its own ELF (un-mergeable)

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.

How stitching works (text-based)

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:

1. Build each sub-kernel as its own complete MLIR module (using B2's per-kernel builders).
2. Extract each module's func.func body (just the operations between signature and return).
3. Rename all SSA values, affine maps, and symbols with a unique prefix to avoid collisions.
4. Remap the original %argN references to the combined function's arg indices (this is what threads the data flow between launches).
5. Concatenate all bodies into one combined func, surrounded by combined affine map declarations and external function decls.
6. Parse the resulting text with mlir.ir.Module.parse(...) to validate.

Concrete example: how 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.

Stitching helpers in kernel_builder/stitching.py

Intra-ELF vs inter-ELF intermediate flow — what the production design actually does

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.

Concrete prefill numbers per pass (16 layers × 3 ELF dispatches per layer):

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.

B6. Gap #3 — Anatomy of one NPU call (BOs and host↔device data flow)

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.

What is a Buffer Object (BO)?

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:

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.

The five steps of one xrt.run()

The three index sets — the per-call control knobs

Every load_and_run call (B7) accepts three optional sets that control which slots get host↔device data movement:

These sets are what makes per-call cost go from "upload everything" (~ms) to "upload only the new activation" (~µs).

What ONE prefill kernel call actually does (concrete: 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.

One important scope note: BOs are per-call, not shared across calls

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:

• Inside one 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.
• Across two 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).

B7. Gap #4 — KernelCache: compile-once, per-layer BO sets

The problem. Two costs would otherwise dominate every script start AND every kernel call:

1. Compile time. Compiling all 6 production ELFs takes ~3 minutes (B5 table). Recompiling on every python llama32_1b_inference.py run is unworkable.
2. BO management state. 16 layers × 6 ELFs × ~6 weight slots ≈ ~600 weight BOs holding ~1 GB of pre-uploaded weights need to stay alive and be addressable. Naively re-allocating per call would also dominate.

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.

Three layers of caching

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 signature and state

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 BOs

The two methods

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

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

The bo_key trick — per-layer weight BOs

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

Why the default is too slow

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:

• Layer 0: write layer-0 weights into BOs (~21 MB host→DDR), run kernel
• Layer 1: BOs now hold layer-0 weights → must overwrite with layer-1 (~21 MB again), run
• ... 16 layers total: ~336 MB of weight upload per prefill pass, just to feed the GEMMs

That's pure host overhead with zero NPU benefit. For decode, the per-token version of the same problem dominates the entire decode loop.

The trick: encode layer index in bo_key

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

Trade-off: memory for speed

This is fundamentally a trade memory for speed design. Concrete numbers:

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

Subtle point: aren't CPU and NPU sharing the same DDR?

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:

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:

1. Weight loaded by HuggingFace → numpy array in pageable RAM (one copy, ~14 MB for wq)
2. Preload calls memcpy(bo.map(), weight_array) → physical byte copy into the BO's pinned region (~3 ms for 14 MB)
3. 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)
4. NPU runs; reads the BO via DMA; writes outputs back
5. For outputs: bo.sync(FROM_DEVICE) → invalidates CPU caches so a subsequent host read sees what the NPU wrote

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

• The CPU-side weight pre-transpose (B4 mismatch #1) creates new arrays anyway — .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.
• Engineering cost vs. payoff — making the weight loader BO-aware would require a custom allocator path through HuggingFace + safetensors, significant complexity for ~200-300 ms savings on a one-time startup cost that's not in the inference critical path.

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.

B8. Prefill in NPU detail — putting all four gaps together

Per-layer kernel sequence — 3 NPU calls

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.

Code walk: 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_len

How weights flow into the kernel: prefill preload

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

B9. Decode in NPU detail — putting it all together for per-token generation

Per-token, per-layer kernel sequence

Decode works on one token at a time. Per token, per layer, it makes 3 calls (2 NPU + 1 CPU):

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

Code walk: the decode loop

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

B10. Code map — where everything lives

Reference section: a top-down map of every file involved in the production runtime, useful for grepping or for finding the right entry point.

Top-level Python files programming_examples/llama32_1b/

Shared infrastructure kernel_builder/

Multi-launch builders multi_launch_builder/

Other directories

How model concepts (Part A) map to NPU code (Part B)

Part C — Verification

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.

What runs

Two entry points, both routed through the parent Makefile and both comparing against HuggingFace transformers in bf16 (same dtype as the NPU — fair fight):

How it stays in sync with production

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.

Why discrete top-k inclusion (and not continuous correlation)

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.

CI

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.

Part D — Future work

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.

D1. Zero-copy weight loading — eliminate CPU↔BO duplication

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

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

D2. Cross-ELF BO aliasing — eliminate inter-ELF host round-trips

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 BO

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

D3. CI: wire up HF_TOKEN so make verify actually runs in CI

Part E — Reference

E1. Glossary — terms defined in one place

Buffer Object (BO)

An XRT abstraction for a chunk of NPU-accessible memory (in DDR — the same physical RAM the host sees, but with NPU access permissions). Created by 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.

Per-layer weight BO

A BO that holds the weight tensor for a SPECIFIC layer of the transformer. The trick: KernelCache caches BOs keyed by 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 (static_input_indices)

The set of arg slot indices that hold weights/LUTs (data that doesn't change between calls). On any call after the first for a given 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 (intermediate_indices)

The set of arg slot indices that hold buffers the kernel will OVERWRITE — it doesn't matter what's in them on entry. The host doesn't need to initialize them; 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).

Shared intermediate BO

NOT a feature of production code (production uses multi-launch merging instead). A development-only pattern: if you have two SEPARATE 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.

Multi-launch ELF

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

Sub-launch

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

Herd

An AIR dialect concept: a 2D array of NPU compute tiles all running the same kernel code in parallel. E.g., 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.

Segment

An AIR dialect concept above the herd: 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 / aiecc

Two MLIR-AIR compiler drivers. 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).

Peano

The AMD fork of LLVM that targets the AIE2P ISA. Used to compile C++ kernels (rope.cc, silu_and_mul.cc, mv.cc) into per-tile object files that get linked into the AIE ELF.

RoPE LUT

Pre-computed cosine/sine table for Rotary Position Embedding. 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).

GQA (Grouped Query Attention)

Llama-3.2-1B has 32 Q heads but only 8 KV heads. Each KV head is shared by 4 Q heads. Reduces KV cache size 4× without much quality loss. Implemented in both NPU FA and CPU attention by indexing kv_h = h // group_size.

SwiGLU

The FFN activation used by Llama: 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

Root-Mean-Square layer normalization: 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.

KV cache

Per-layer cache of K and V tensors at every token position seen so far. During decode, attention reads the entire cache (positions 0..current_pos) but only computes one new K and V (for the new token). Without it, decode would be O(N) per token instead of O(1). See Part A4.

Prefill / Decode

Two operating modes of LLM inference. Prefill: process the whole prompt at once (seq=N), populate KV cache. Decode: process one new token (seq=1), append to KV cache, get next token. Repeated decode generates text. See Part A3.

Padding (in this implementation)

NPU kernels are compiled for fixed shapes. Llama-1B's prefill kernels expect seq=2048. Shorter prompts get padded with EOS tokens up to 2048; the prefill processes all 2048 positions but only the logits at pred_pos = prompt_len - 1 are used. See Part A5.

E2. Reading guide — where to start for specific questions

Quick-reference: which file does what when you grep