Graph-based Language Models

[hunhoffe]

This PR moves IRON from per-operator scripts and hand-assembled sequences to declared operators composed into traced graphs. A model is a Python class whose body() calls operators on handles. It is traced once per input shape and compiled once, and per-call values (a position, a row count) are patched in without a rebuild. A cost-model tuner sits on top, and two applications run end to end on Phoenix (npu1) and Strix/Krackan (npu2): Llama 3.2 1B and
EmbeddingGemma 2.

The work builds directly on what devel validated: the operator arrays and tilings, the kernels (now upstream in mlir-aie's aie_kernels/ and aie.iron.kernels), llama_npu.py's fused-ELF decode, the FastFlowLM ports and the stream-dse SwiGLU. This PR ports those proven designs into one declared form and builds the graph and tuning layers on it.

Size. 882 commits, 380 files, +55k / −20k. The rough Python split at the tip: iron/operators 18.5k lines, iron/tests 18.3k, iron/common/graph 8.5k, iron/common/declare 3.3k, iron/common/image 3.1k, iron/lm 2.9k and iron/applications 2.4k.

Suggested reading order. AGENTS.md (Architecture, Graphs, Design rationale), then iron/common/declare/, iron/operators/relu.py →clamp.py → gemv.py, then iron/common/graph/compiled.py, iron/tests/common/graphs.py, iron/lm/causal.py, and iron/common/graph/narrowing.py (its module docstring is the cost model).

---

1. Operators: one declared class

Each operator is a single class: its fields, its operands and the values its cores read, plus array(), which builds ObjectFIFOs and Workers around an mlir-aie kernel. The library derives the runtime sequence from that declaration, and an operator overrides sequence(rt) only when it needs a different order (GEMM and MHA do).

The smallest operator names its kernel and its test cases, and nothing else:

# iron/operators/relu.py
class ReLU(UnaryElementwise):
    """AIE-accelerated ReLU activation function."""

    test = Testing(Sweep(), draw=dict(centered=("x",)))  # both signs

    def kernel(self):
        return eltwise.relu_sized(self.tile_size)

A fuller declaration (iron/operators/gemv.py, abridged):

class GEMV(Operator):
    test = Testing(_cases, draw=dict(normal=("A", "B")))

    M: int = Param()
    K: int = Param()
    num_batches: int = Param(default=1)
    num_aie_columns: int = Auto()
    num_channels: int = Auto()
    tile_size_input: int = Auto(domain=Divisors(of=lambda op, dev: op.M, cap=...))
    ...

The vocabulary (iron/common/declare/):

• Param(): what a host shape names. Param(..., array=True) puts a field in the array tier.
• Auto(): a tunable. A call site or a profile may give it; otherwise resolve(dev) fills it from the device and the extents. Auto(domain=...) declares what the tuner searches (section 8).
• In / Out: operands declared by shape. tile= makes each operand its own stream into the array, per= gives its lane count, and when=flag makes it optional (RMSNorm(x, weight=w)).
• Value(derive=...): trip counts and other scalars the cores read, derived from the extents.
• Extent(field): an extent a graph can bound per call (section 4).
• Testing(cases, tolerance=, draw=): the shapes the operator is checked at on a device. Case(..., extensive=, bench=, lower=) and Sweep(...) describe the cases.

An inline-kernel "hello world" lives in iron/tests/toolchain/inline_kernel.py. It shows the whole contract a kernel carries (roles, scalar bindings and reference):

class VectorAdd(BinaryElementwise):
    def kernel(self):
        tiles = [self.a.tile, self.b.tile, self.y.tile, np.int32]
        contract = KernelContract(
            roles=(In, In, Out, Param),
            parameter_bindings=((3, self.tile_size),),
            reference=lambda a, b: a + b,
        )
        return ExternalFunction("vadd", source_string=VADD, arg_types=tiles, contract=contract)

Invariants

• One array serves every extent. array() sees only the array tier, and reading an extent there raises. Trip counts are Values that the sequence writes. iron/tests/toolchain/array_identity.py compiles operators at two extents and compares the core ELFs byte for byte.
• Tunable precedence: the call site's value, then the graph's Profile, then resolve(dev), then the Auto() default. Identity (array_key(), design_key()) is taken after resolution, so two calls that resolve alike share a build.
• IRON hosts no kernels. Kernels come from mlir-aie's factories (aie.iron.kernels). Each factory's ExternalFunction carries the reference and tolerance (.contract), and the catalog test judges an operator against them unless Testing(tolerance=) overrides.
• numpy only. References are numpy and ml_dtypes.bfloat16. Torch is no longer a requirement.
• No build context. An operator takes the current device and nothing else. On a host without an NPU, bind one with aie_utils.set_current_device(from_name("npu2", n_cols=8)) (the test tree's npu2 / npu1 fixtures do this).

2. The lazy catalog

import iron and iron.operators load names on first use (PEP 562), so importing one operator imports only its module:

from iron.operators import GEMM  # imports iron.operators.gemm, nothing else
from iron.operators.flm import GEMM as FLMGEMM, DequantBFP  # flm is its own catalog

iron/tests/infrastructure/lazy_imports.py checks this in a fresh interpreter.
dir() and type checkers still see every name (TYPE_CHECKING imports).

Main catalog, iron/operators/__init__.py: AXPY, Clamp, Copy, DepthwiseConv1d, Dequant, ElementwiseAdd / RowwiseAdd, ElementwiseMax, ElementwiseMul / RowwiseMul, Emit, GELU, GEMM, GEMV, LayerNorm, LeakyReLU,
Limbs, Log, Magnitude, Merge, MHA, PatchPositions / Resample / ResampleTaps, ReLU, RMSNorm, Repeat, RoPE, Sample, Sigmoid, SiLU, Softmax, Tanh and Transpose. Copy replaces MemCopy and StridedCopy. The SwiGLU composites are now graphs in iron.lm.blocks.

3. Graphs: trace once, compile once, call many times

A graph is a subclass of iron.Graph:

• body()'s positional parameters are its inputs.
• Its keyword-only parameters, annotated Scratchpad[T], Carried[T] or DispatchTime[T], are per-call values.
• Its return values are its outputs.
• Weights and iron.state(...) buffers are whatever the instance holds, named by attribute path (self.layers[3].q → layers.3.q).

From iron/tests/common/graphs.py:

class Rotate(iron.Graph):
    """RoPE over the first `n` positions, `n` given per call."""

    def body(self, x, angles, *, n: Scratchpad[np.int32]):
        return RoPE(x[:n], angles[:n])


class Project(iron.Graph):
    """The first `n` rows of `x` projected by `w`."""

    def body(self, x, w, *, n: Scratchpad[np.int32]):
        return GEMM(x[:n], w, b_col_maj=True)

compile(**shapes) traces and builds one version per input signature. Calling the graph runs the matching version, and compiles it on the first call if needed. From iron/tests/operators/gemm_bounded_rows.py:

project = Project().compile(x=(M, K), w=(N, K), boundaries=boundaries)
for n in BOUNDS:                                   # (512, 1, 300, 64, 65, 512, 17, 128, 3)
    got = np.asarray(project(x, w, n=n)).reshape(M, N)

Graph.compile(*, dev=None, boundaries=None, image=None, verbose=False, record="memory", feeds=None, coresident=None, fold=False, **shapes):

• dev= binds a device for this compile only. compile() links the image and stops there; the runtime that loads it is created on the first call, so a host with the toolchain and no NPU can compile ahead of time.
• verbose=True prints the packaging decision and how each per-call value lowers (section 5).
• graph.reference(*tensors, **values) runs the same body() through each operator's reference() on host arrays, with states modelled.
• A graph called inside another graph's body traces into the caller: its steps, weights and states become the caller's, and one compile covers both. iron.lm.blocks.SwiGLU, LogMel, Conv2d and ImagePatches are such reusable graphs.

Invariants

• All versions of a graph share one scratch arena, so a weight is on the device once and a state one version writes is where the next one reads it. The arena is split into banks of at most 3.5 GiB, one buffer object each.
• Tracing refuses what cannot be traced. Branching on a per-call value raises TypeError (Value.__bool__), and so do *args/**kwargs in body, an output that is a graph weight or input, and an input defaulting to anything but None.

4. Per-call values and bounds: Scratchpad and x[:n]

There are three kinds of per-call values (iron.common.declare.value):

Kind	What it is	Full ELF	xclbin
Scratchpad[T]	patched at call time	written to the ctrl scratchpad	a dispatch-time scalar: an offset use regenerates the stream, a core-read use is an RTP write
Carried[T]	computed by the graph for its own next call (section 6)	as Scratchpad, and the device can loop	host writes it back
DispatchTime[T]	regenerates the sequence per call	forces xclbin	sizes, strides and offsets regenerated

Arithmetic on a value builds an Affine (pos + 1, chunk * C), which can
index or slice a handle, state or weight:

• x[:n] bounds an extent per call. The bound carries through reshape and transpose and reaches every operator that consumes the view.
• self.kv[:, pos] / cache[step.position] is a per-call offset.
• table[ids] with ids a graph input or a handle the device made is a gather on the NPU (Copy(self.embedding[ids[:rows]])).

Invariants and limits

• Under a bound, GEMM limits its compute (row blocks, and K per call, on both packagings), and MHA's K and V move only the blocks up to the bound. Attention over a cache therefore costs the context, not the cache length.
• A per-call size is mlir-aie's runtime transfer length (fill/drain(length_parameter=, length_unit=)) in whole 16-byte units, so a bounded tile is at least 16 bytes. A transfer whose size a call bounds must fit one DMA descriptor, because the patch lands in its length field. A pattern past the descriptor limits (BdLimits) is split by the compiler, with the per-call offset patched into every piece.
• Every per-call index is checked against its view before the call reaches the device.
• On a full ELF the preamble syncs the per-call words before opening the cores' barriers. a0a7d024f fixed this order: before it, from the second call on, each core read the previous call's words. gemm_bounded_rows.py and mha_bounded_rows.py drive one compile across alternating bounds to catch any regression.

5. Packaging: one full ELF, or a chain of xclbins

iron/common/image/packaging.py derives the image from the traced graph and
the device:

• Full ELF (iron.ELF), one fused sequence and one dispatch per call. This is the default on AIE2P (npu2).
• xclbin chain (iron.XCLBIN, boundaries=iron.each_step), one dispatch per step. Any of these forces it: a DispatchTime value, an npu1 device (full_elf(dev) is dev.arch is AIEArch.AIE2p), or boundaries=iron.each_step. A new xclbin chain starts every 32 designs.

Limits

• A fused sequence inside one xclbin (several steps in one dispatch) has no proven construction yet and is refused with NotImplementedError. So npu1 runs every step as its own dispatch.
• A shipped-image operator (image=Xclbin(...)) is refused in a graph; run it on its own with OperatorImage(op).

Standalone use, outside a graph (iron/operators/flm/gemm/README.md):

from iron.common.image import OperatorImage
from iron.operators.flm import GEMM
from iron.operators.flm.gemm.design import Epilogue

op = GEMM(M=1024, K=1536, N=6144, epilogue=Epilogue.SILU)
OperatorImage(op)(A, op.pack_B(B), C_out)

Build outputs live in mlir-aie's JIT cache (~/.npu/cache/<hash>/, or
NPU_CACHE_HOME), keyed on content. Nothing is written to the working
directory.

6. Carried values and the device loop

A Carried[T] value is one the graph computes for its own next call, returned last via iron.carry(...). On a full ELF the device can start the next call itself (CarriedLoop). From iron/tests/common/graphs.py:

class Walk(iron.Graph):
    def __init__(self, successor):
        self.successor = iron.weight(successor)

    def body(self, *, node: Carried[np.int32], steps: Carried[np.int32]):
        nxt = Copy(self.successor[node], dtype=np.int32)
        return iron.carry(node=nxt, steps=steps + 1)

CausalLM uses this to sample on the NPU: Sample draws the next token from the logits and a per-position draw row, and carries token and position + 1. On npu2 the host only restarts the loop and reads the tokens
back at the end. On npu1 (xclbin chain) the host reads the drawn token and dispatches the next step; the draw still happens on the NPU.

7. The operators language models lean on

• MHA (mha.py) has causal or bidirectional attention, a sliding window, scale, d above 64, and kv_len with a DMA-bounded kv_valid for attention over a cache. Decode is MHA with one query: a KV group's heads are packed into a block, and a group that does not divide the block is covered by two overlapping runs. The softmax running state and the O accumulator are float32. On AIE2P the softmax uses mlir-aie's accurate exp2 (867b88de0). B_q, num_pipelines, pv_cores and placement are searched by the tuner. MHA is placed on npu1's 4 columns as well as npu2's 8.
• GEMM (gemm.py) bounds rows and K per call, zeroes the tail of the last K tile, and can finish its C tiles with a folded elementwise step. It searches tile_m/tile_k/tile_n and placement. An M that no tile_m splits is refused at construction.
• flm.GEMM (flm/gemm/) is the FastFlowLM port with a fused f32→bf16 epilogue (activation, clamp) whose M, K and N are runtime parameters, so one xclbin serves every shape. It reads B as stored (b_col_maj) on AIE2 and AIE2P, and iron.lm.blocks.project() uses it for prompt-chunk projections on both devices. flm.Shipped is the downloaded binary, kept for comparison only.
• GEMV (gemv.py) supports num_batches, repeat= (one matrix serving consecutive batches, for GQA scores without a materialized copy), two lanes per column (num_channels=2), and finish / prepare links (a SiLU, a residual add or an RMSNorm folded into its cores).
• Copy (copy.py) is the single data mover: views on both sides, multi-channel splits on the run both sides share, per-call offsets, and gathers by host or device ids (GatherWords, up to four pairs of shim columns).
• Sample, Emit and Merge do on-device sampling and carry the next call's words.
• Front-end operators for EmbeddingGemma 2's towers: Limbs / Magnitude / Log (float32 log-mel steps), Resample / ResampleTaps / PatchPositions (image resize into patches), DepthwiseConv1d, ElementwiseMax, Clamp (bounds as runtime parameters), and RowwiseAdd / RowwiseMul.

8. The tuner

The tuner chooses each design's tunables, which steps fold together and which designs share a device configuration. It decides from measured costs, never from bit-unequal guesses: a setting is a candidate only if its output is bit-identical to the default's, or within the gate of both.

8.1 Domains: what is searched

A tunable declares its search space with Auto(domain=...) (iron/common/declare/domain.py). Each domain is evaluated on the resolved operator, so the default is always among its values:

Domain	Values
Width()	the field's own value and every power of two up to the device's columns, widest first (the default for a per= tunable)
Divisors(of=, step=, cap=, span=)	the resolved value times powers of two that divide of, are multiples of step and at most cap
Choices(options)	a fixed set, the resolved value first
Placement({...})	named tile placements (8.2)

# iron/operators/mha.py
B_q: int = Auto(array=True, domain=BLOCK_SIZES)
num_pipelines: int = Auto(1, array=True, domain=Width())
placement: str = Auto(array=True, domain=PLACEMENT)
pv_cores: int = Auto(array=True, repr=False, domain=Choices((1, 2)))

A tunable that a call or profile gives explicitly is held, not searched. Auto(derived=True) fields are always derived by resolve() and are never tuned.

8.2 Placement as a tunable

Where an array's tiles go is a choice too, given as a name for a Pins holding cores, memtiles and shims at a Level (TILE, COLUMN or FREE for the placer's choice):

# iron/operators/mha.py
PLACEMENT = Placement(
    {
        "memtiles": Pins(cores=Level.COLUMN, shims=Level.FREE),
        "tiles": Pins(),
        "columns": Pins(cores=Level.COLUMN, memtiles=Level.FREE, shims=Level.FREE),
    }
)

The hand-found placement is the default. Before a setting is measured, the tuner asks the placer and router whether it fits (narrowing.fitting), and the verdicts persist under NPU_CACHE_HOME/iron/fits.

8.3 Measuring: probes, cost tables and the CostCache

• probe.py times each design alone (t_step, alone), the configure cost between pairs of designs (a calibration triangle of operator classes), loads, and packs of designs measured as one device. Timing runs only while no other process holds the NPU and while the interrupted cores are held out of deep idle, and a figure whose calls waited on a CPU is refused.
• CostTable (narrowing.py) is one graph's working set, as JSON, for one device and one packaging: costs_<device>.json for the device's own packaging, costs_<device>_<dispatch>.json for the other. Entries are keyed by design identity, so a design that has changed since measurement is missing from the table and left as the profile gives it. Each table records a single power mode.
• CostCache (costcache.py) holds every design ever measured on this NPU, one JSON per entry under NPU_CACHE_HOME/iron/costs/<platform>/<power mode>/. It is keyed on the build recipe, the per-call values and the input contents, so a design shared by two graphs is measured once. Code is not part of the key: after editing a reference or a tolerance, run with --remeasure.
• A design whose work follows a per-call value (attention over the context) is measured at a ladder of contexts and priced at their weighted mean (StepCost.points).

# one operator at every setting it tunes to
python -m iron.common.graph.tune GEMV M=2048 K=8192
python -m iron.common.graph.tune Softmax rows=32 cols=2048 --value valid_cols=1500

# a model's table, by the plan its Runner gives
python -m iron.lm.tune llama --checkpoint /path/to/model.safetensors
python -m iron.lm.tune llama --checkpoint /path/to/model.safetensors --dispatch separate
python -m iron.lm.tune embedding_gemma2

Run these with XRT sourced and the NPU otherwise idle. No cost table is checked in; you measure one on the NPU you use it on.

8.4 The fold pass

iron/common/graph/fold.py finds three kinds of fold:

• a step folded into its input's producer (Fold, e.g. SiLU into the gate GEMV's epilogue, a residual add into a GEMV), including chains of folds;
• a step folded into the matvecs that read its output (Prologue, e.g. RMSNorm into the following GEMVs);
• a copy into a state folded into its producer's drain.

compile(fold=True) takes every admissible fold. Under the tuner, each fold is priced against the graph without it and taken only where the modelled time drops.

Limit: a fold that would make a core read more streams than its two input channels is refused (Operator._finish_at). For example, EmbeddingGemma 2's residual add does not fold into a weighted RMSNorm.

8.5 Joint narrowing

JointNarrowing(table) is passed as coresident= to compile(). It chooses widths, folds and packs together to minimise the modelled time of a full-ELF runlist:

D0 + Σ steps t_step + Σ device entries (base + Σ members load) + R if the entries are odd

For an xclbin chain the model is F + Σ steps c + Σ switches L(incoming), and each design's setting is chosen independently. NarrowingKinds(widths=, folds=, packs=) turns individual decisions off for ablations. A model takes a table on its command line:

python -m iron.applications.llama.model model.safetensors tokenizer.model \
    --cost-table iron/applications/llama/costs_npu2.json

iron/tests/infrastructure/narrowing.py and iron/tests/toolchain/narrowing.py cover the search and the cost model, and probe.check_model compares the model against a measured run of the graph it tuned.

9. Profiles: tunables as data

A graph's profile is a Profile, or a directory holding one <device>.json per device. Entries are keyed by operator class and shape, and the most specific match fills any tunable a call leaves open:

{"entries": [
  {"operator": "GEMV", "M": 2048, "K": 2048, "num_batches": 1, "tile_size_input": 4, "tile_size_output": 256},
  {"operator": "MHA", "num_heads": 32, "num_KV_heads": 8, "d": 64, "seq_pad": 1, "num_pipelines": 4},

(iron/applications/llama/profiles/npu2.json.) Profile.load(path), profile.save(path) and with profile: are the API.
test_llama_names_only_the_tunables_that_matter checks that every keyword the Llama graph still passes is one the profile could not have supplied.

10. iron.lm and the applications

iron/lm/ holds what every model shares, and each model is an application package under iron/applications/. Nothing in the library imports a model.

• base.py:
  • Config, the model's shape.
  • Checkpoint, a mapped .safetensors file or a sharded index. tree(fill=) returns the weights by the checkpoint's own names: "read", "random" (drawn at a seed, in the checkpoint's shapes) or "unread" (header only, for tracing and tuning).
  • Oracle, a float32 host forward pass, widened one layer at a time.
  • Model and Runner/main, the command line, plus --random-weights SEED for running without real weights or a tokenizer.
• causal.py: CausalLM, a decoder as one graph. It has two versions, a decode step and a prefill_chunk-row prompt chunk, each compiled once. max_seq_len sizes only the caches and the RoPE table.
• encoder.py: Encoder, compiled once at max_tokens rows, with the length bound per call.
• blocks.py: project() (GEMV for one row, flm.GEMM for more, an optional bias via RowwiseAdd), swiglu / SwiGLU, geglu, sandwich, attention, LogMel, Conv1d / Conv2d, MaxPool2d, ConcatChannels, ImagePatches, and Block / BlockOracle.
• generation.py, testing.py, tune.py: sampling, the accuracy and determinism checks, a model's device-test helpers, and the tune CLI.

The whole decoder body (iron/lm/causal.py):

def body(self, ids=None, pixel_values=None, *, token: Carried[np.int32],
         position: Carried[np.int32], chunk: Scratchpad[np.int32], rows: Scratchpad[np.int32]):
    c = self.config
    L, C, A = c.max_seq_len, c.prefill_chunk, c.rope_width
    prompt = ids is not None
    if prompt:
        x = Copy(self.embedding[ids[:rows]])
        angles = Copy(self.rope.reshape(L // C, C, A)[chunk, :rows])
    else:
        x = Copy(self.embedding[token]).reshape(1, c.emb_dim)
        angles = Copy(self.rope[position]).reshape(1, A)
    x = self.embed(x, pixel_values, rows if prompt else None)
    step = Step(prompt, angles, chunk, rows, position)
    for i, weights in enumerate(self.layers):
        x = self.layer(step, i, weights, x)
    if prompt:
        x = Copy(x[rows - 1]).reshape(1, c.emb_dim)  # the last row's logits alone
    logits = self.head(x)
    _, drawn = Sample(logits.reshape(c.vocab_size), self.draws, self.drawn,
                      row=position * ROW_WORDS, at=position)
    return logits, iron.carry(token=drawn, position=position + 1)

The caches are (max_seq_len, n_kv_groups, head_dim), so each call's write is contiguous and no descriptor steps by the context. test_llama_does_not_grow_with_the_context checks that no activation or array follows max_seq_len.

Llama 3.2 1B (iron/applications/llama/)

Llama(Block, CausalLM) adds only its head. The layer is Block.layer, and the oracle is LlamaOracle(BlockOracle, Oracle):

from iron.applications.llama.model import Runner

runner = Runner("model.safetensors", "tokenizer.model")
model = runner.npu()                    # compiles and loads both versions; weights uploaded
logits = model.logits(runner.encode("The capital of France is"))

python -m iron.applications.llama.model /path/to/model.safetensors /path/to/tokenizer.model \
    --prompt-len 2048 --num-tokens 40     # --check-accuracy, --check-determinism N, --each-step, --cost-table

• Llama 3 RoPE frequency scaling is applied, and contexts go up to 32768 by default (--max-seq-len, a multiple of 2048). The compile is the same at any context.
• Prompts run in fixed chunks against the caches. A chat turn reruns only the chunks that changed, and its logits match a from-scratch run bit for bit (check_chat_turn).
• Limit: a prompt chunk is one dispatch and its attention grows with the context. Past about 48k tokens a chunk exceeds amdxdna's default watchdog (tdr_timeout_ms), so very long prompts need it raised (AGENTS.md, "Not built, or limited").
• XRT maps device buffers locked, so ulimit -l must cover one arena bank.

EmbeddingGemma 2 (iron/applications/embedding_gemma2/)

The text encoder (EmbeddingGemma, an Encoder) is compiled once at 2048 rows, with the prompt length bound per call. The audio and vision towers (audio.py, vision.py) are graphs that Multimodal calls in its body, so
text, a clip and an image encode as one graph. The NPU runs the token-embedding gather, log-mel features, image resize to patches, patch positions and the multimodal merge. Each part has a float32 oracle, checked against Hugging Face computed live when torch and transformers are installed (otherwise those tests skip), on COCO photos and LibriSpeech clips fetched by sha256.

class EmbeddingGemma(Encoder):
    def body(self, ids, *, n: Scratchpad[np.int32]):
        return self.encoder(Copy(self.embedding[ids]), n)

python -m iron.applications.embedding_gemma2.model /path/to/embeddinggemma-2 \
    --query "What causes the northern lights?" \
    --document "Charged particles from the sun." "Photosynthesis in plants."

gemma4_flm (the FastFlowLM engine example from devel) is unchanged in substance; its test moves to tests/applications/gemma4_flm/.

11. Phoenix (npu1) and npu2

	npu1 (AIE2, Phoenix / Hawk Point)	npu2 (AIE2P, Strix / Strix Halo / Krackan)
columns	4	8
graph image	xclbin chain, one dispatch per step	one full ELF (or --each-step)
per-call Scratchpad	dispatch-time scalar / RTP write	ctrl scratchpad word
decode loop	host restarts each step on the token the NPU drew	device loop (CarriedLoop)
Llama prompt	flm.GEMM projections, MHA on 4 columns, a chunk at a time	the same, on 8 columns
EmbeddingGemma 2	text, audio and vision towers and the multimodal versions	all, plus TestEachStep

The decode GEMVs run two lanes per column on Phoenix. --each-step builds the npu1 form on npu2, which is how TestEachStep checks it there. CI covers both devices: phoenix-{small,extensive,test-examples}.yml and krackan-{small,extensive,examples}.yml.

12. Seeing a graph

TracedGraph and CompiledGraph draw themselves: graph.trace(**shapes).html(), version.dot(), and version.picture(table=...) to annotate steps with priced times. The pictures show each step's design, the steps folded into it, moved tunables and packs, the scratch arena over the runlist, and where a step sits on the array (iron/common/graph/visualize.py).

13. Tests, and how to run them

Layout:

Where	What
iron/tests/operators/catalog.py	every operator's Testing cases against reference() on a device
iron/tests/operators/*.py	operator behaviour the catalog does not reach (bounded rows and K, gathers, long-context MHA) and rejected_shapes.py (device-free refusals)
iron/tests/infrastructure/	graphs, sequences, the allocator, carried values, narrowing and probes
iron/tests/toolchain/	device-free lowering gates (needs aiecc and Peano): every operator lowers; fused graphs build to full ELFs and xclbins; array_identity.py, inline_kernel.py
iron/tests/lm/	the iron.lm.blocks graphs (SwiGLU, LogMel, Conv1d/2d, MaxPool2d, ConcatChannels, ImagePatches)
iron/tests/common/	shared graphs and helpers, plus device-free graph and Llama-reference tests
iron/operators/**/test.py	operators with a device test of their own (flm, swiglu_prefill_stream)
tests/applications/<name>/test.py	Llama, EmbeddingGemma 2, Gemma 4 on FastFlowLM

Markers and options (pytest.ini, root conftest.py):

• @pytest.mark.supported_devices("npu1", "npu2") marks a device test. It is the only mark a device test needs: on a host without an NPU runtime it skips with the reason, and the runtime is released after it.
• extensive is excluded from the default suite; bench marks the cases CI charts over time.
• --random-weights SEED, --max-seq-len N and --cost-table PATH apply to the language-model tests. --csv-output sets the metrics CSV path (default tests_latest.csv).
• npu1 / npu2 fixtures bind a device for device-free resolve and compile.

# what CI runs per device (operators suite), and the extensive suite
pytest iron/operators/ iron/tests/ -m "not extensive"
pytest iron/operators/ iron/tests/

# one operator's catalog cases; in parallel
pytest iron/tests/operators/catalog.py -k AXPY
pytest iron/operators/ iron/tests/ -n auto -m "not extensive"

# device-free gates (no NPU needed; aiecc and Peano on PATH)
pytest iron/tests/toolchain/

# applications (weights under $IRON_EXAMPLE_WEIGHTS_DIR, default /srv)
pytest tests/applications/
pytest tests/applications/llama/ --random-weights 0

14. Measured numbers, as recorded in the commits

These are point measurements taken at the named commits. Later commits change some of them; for example, the accurate exp2 costs decode throughput, as listed below.

What	Result	Host	Commit
Llama 3.2 1B decode, profile vs tuned (random weights, 32768 caches, 1024-token prompt, 100 steps)	14.80 → 16.50 tok/s (+11.5%); model predicted 57.5 ms/step, measured 60.60; tuned and untuned greedy logits bitwise equal	Strix Halo, xcoradaie218	1398c5f0c
Llama 3.2 1B decode on Phoenix, profile vs tuned (16-token prompt, 20 tokens)	361.5 → 337.8 ms/token; KL to oracle unchanged	Phoenix, xcoradaie205	88f23339c
Llama 3.2 1B, xclbin chain on npu2, profile vs tuned	3.58–3.61 → 4.76–4.83 tok/s; KL 5.47e-5, top-1 100/100	Strix Halo, xcoradaie218	34ff047b1
GEMV + residual add, folded vs unfolded	M=2048 K=2048: 430.6 → 367.3 µs (−14.7%); outputs bit-identical	Strix Halo, xcoradaie218	61d210b05
Gather 5120 rows of a 262144×512 bf16 table, 1 vs 4 feeds	2002 → 903 µs, bit-exact vs np.take	Strix Halo, xcoradaie220	727a96180
MHA float32 softmax state and O accumulator (accuracy fix)	+4.9% to +9.1% time; mean abs error 0.038 → 0.0068 at 2048 rows over 32768 keys	Strix Halo, xcoradaie220	11a66db95
Accurate exp2 in MHA on AIE2P (accuracy fix)	Llama 1B max KL 0.056 → 0.013 at 1024 tokens; decode 14.97 → 13.99 tok/s (−6.6%)	npu2 (commit names "220")	867b88de0
EmbeddingGemma 2 audio tower, 2 s / 7.3 s / 30 s clip	129.3 / 195.7 / 397.9 ms; mean cosine ≥ 0.997 vs oracle	Strix Halo, xcoradaie220	86edca193

15. Known limits (also in AGENTS.md)

• A fused sequence inside one xclbin is refused, so npu1 runs one dispatch per step and has no device loop.
• A prompt chunk's attention grows with the context. Past about 48k tokens a chunk exceeds amdxdna's default watchdog.
• A fold is refused when its cores would read more streams than a core's two input channels.
• A per-call-bounded transfer must fit one DMA descriptor.
• Cost-table entries are keyed by design identity, not code. After changing a reference or a tolerance, remeasure.
• Open asks in mlir-aie: a builder for aiex.configure / aiex.run, and an accessor for L1 banking.

Note, this PR relies on: Xilinx/mlir-aie#3884

PR Merge Checklist

1. The PR is rebased on the latest devel commit and pointing to devel.
2. Your PR has been reviewed and approved.
3. All checks are passing.

[github-actions]

CI performance trends

e0ae8f3 (2026-10-16T18:18:36Z)

Phoenix - Operators

Operators added: flm/gemm

Operators dropped: axpy, dequant, elementwise_add, elementwise_mul, flm/dequant, gelu, gemm, gemv, layer_norm, leaky_relu, mem_copy, mha, relu, repeat, rms_norm, rope, sigmoid, silu, softmax, strided_copy, swiglu_decode, swiglu_prefill, tanh, transpose

Phoenix - Applications

Operators dropped: gemma4_flm/gemm_dequant

⚠️ No results from: Krackan - Operators (missing), Krackan - Applications (missing). This comment covers the remaining suites only.

[hunhoffe]

Early notes

[jgmelber]

LGTM