diff --git a/README.md b/README.md index befeb9b21..b48b15ead 100755 --- a/README.md +++ b/README.md @@ -97,7 +97,9 @@ If starting from `Ubuntu 24.04` you may need to update the Linux kernel to 6.11+ |---|---|---| | Gemma 4 E2B | 131072 | 182 | | Llama 3.2 3B | 131072 | 101 | + | Granite 4.2 3B | 131072 | 86 | | Llama 3.2 1B | 131072 | 26 | + | Granite 4.0-H 1B | 131072 | 15 | | EmbeddingGemma 2, each tower and the encoder | its largest input | under 0.7 | On NPU1 a command is one operator, the longest an attention layer of the @@ -151,11 +153,17 @@ Each model is a package under [iron/applications/](./iron/applications/) over [i | Model | Application | Checkpoint | |---|---|---| | Llama 3.2 1B, 3B | [llama](./iron/applications/llama/) | `llama3.2-1b`, `llama3.2-3b` | +| Granite 4.2 3B | [granite](./iron/applications/granite/) | `granite-4.2-3b` | +| Granite 4.0-H 1B | [granite_hybrid](./iron/applications/granite_hybrid/) | `granite-4.0-h-1b` | +| Qwen3 1.7B, 0.6B | [qwen3](./iron/applications/qwen3/) | `qwen3-1.7b`, `qwen3-0.6b` | +| Qwen3-Reranker 0.6B (`score`) | [qwen3_reranker](./iron/applications/qwen3_reranker/) | `qwen3-reranker-0.6b` | +| Qwen3-Embedding 0.6B (`embed`) | [qwen3_embedding](./iron/applications/qwen3_embedding/) | `qwen3-embedding-0.6b` | | Gemma 4 E2B text decoder | [gemma4](./iron/applications/gemma4/) | `gemma-4-e2b-it` | +| Qwen3-30B-A3B (reference only) | [qwen3_moe](./iron/applications/qwen3_moe/) | `qwen3-30b-a3b` | | EmbeddingGemma 2 | [embedding_gemma2](./iron/applications/embedding_gemma2/) | `embeddinggemma-2` | | Gemma 4 E2B in FastFlowLM | [gemma4_engine](./iron/operators/flm/gemma4_engine/) | `$FLM_MODEL_PATH` (the tests default to `fastflowlm`) | -What the text decoders (Llama and Gemma 4 above) share: +What the text decoders (Llama, Granite, Qwen3 and Gemma 4 above) share: - **Setup**: `requirements_examples.txt`, as [installing](#installation-linux) says. XRT maps every device buffer locked, and a model's weights and caches are banks of at most 3.5 GiB (`iron.common.image.allocator.BANK`), so `ulimit -l` must cover one bank. diff --git a/iron/applications/granite/README.md b/iron/applications/granite/README.md new file mode 100644 index 000000000..739db34db --- /dev/null +++ b/iron/applications/granite/README.md @@ -0,0 +1,41 @@ + + +# Granite 4.2 3B on the NPU + +IBM's Granite 4.2 3B, `Granite` in `model.py`, whose docstring says what is +its own: Llama's layer with four scalars from its `config.json`, an untied +head and a sharded checkpoint. How to run, test and tune it is the README's [Models](../../../README.md#models); the shared library is [AGENTS.md](../../../AGENTS.md#language-models). + +## Weights and tokenizer + +From Hugging Face +([`ibm-granite/granite-4.2-3b`](https://huggingface.co/ibm-granite/granite-4.2-3b), +Apache-2.0, ungated), over plain https, with no account: + +```bash +D=/path/to/granite-4.2-3b +REV=e459acceac81e5fe67c07d9cfc72329a332e7eb1 # the revision every KL figure here is against +for f in config.json generation_config.json tokenizer_config.json \ + special_tokens_map.json tokenizer.json model.safetensors.index.json \ + model-00001-of-00002.safetensors model-00002-of-00002.safetensors; do + curl -fL -o "$D/$f" \ + "https://huggingface.co/ibm-granite/granite-4.2-3b/resolve/$REV/$f" +done +``` + +The command line takes `$D`. Its native context, 131072 positions, holds +10 GiB of caches beside the weights' 6.8 GiB, the default where the host +has room for both (`Model.fit`). + +## Testing + +`tests/applications/granite/` also runs the each-step form +(`TestEachStep`, `--no-full-elf`), a two-chunk prompt and a chat turn. + +NPU1 runs `TestEachStep` from an empty `profiles/npu1.json` (every +operator's device defaults, as smollm2's, phi3's and qwen2's are), and +all of it passes there: generation, accuracy and determinism. Not yet +run on NPU1: `tests/applications/long_context`'s deep checks. diff --git a/iron/applications/granite/__init__.py b/iron/applications/granite/__init__.py new file mode 100644 index 000000000..7f1f2325a --- /dev/null +++ b/iron/applications/granite/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.2 3B end to end on the NPU; see ``iron/applications/granite/README.md``.""" + +from iron.lm.cli import exports + +__all__ = ["Granite"] +__getattr__ = exports(__name__, __all__) diff --git a/iron/applications/granite/model.py b/iron/applications/granite/model.py new file mode 100644 index 000000000..0658d2975 --- /dev/null +++ b/iron/applications/granite/model.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.2 3B over the shared decoder and layer (``iron.lm``): its +embedding and head on the NPU (``Granite``) and its shape, read from its +checkpoint's ``config.json``. Run with ``python -m +iron.applications.granite.model``. + +Granite is Llama's layer with four scalars from its config: attention +scaled by ``attention_multiplier``, each block's residual by +``residual_multiplier`` (``GraniteLayer``), the embedding rows by +``embedding_multiplier`` and the logits by ``logits_scaling``. Its head is +untied (``lm_head``); its checkpoint is sharded. + +Tunables are keyed by operator shape in ``profiles/.json``, at +Granite 4.2 3B's: a decode step's and a 2048-row prompt chunk's. +""" + +import dataclasses +import math +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +from iron import Precision, lm, param +from iron.lm import CausalLM, DecoderConfig, DecoderLayer + + +@dataclasses.dataclass(frozen=True) +class GraniteConfig(DecoderConfig): + """``DecoderConfig``, Granite's own four scalars, and every + projection's ``precision``. + """ + + attention_multiplier: float = dataclasses.field(kw_only=True) + residual_multiplier: float = dataclasses.field(default=1.0, kw_only=True) + embedding_multiplier: float = dataclasses.field(default=1.0, kw_only=True) + logits_scaling: float = dataclasses.field(default=1.0, kw_only=True) + # 4.2-3B's prompt over two chunks needs bf16 macs on its projections, as + # Llama 3B's does (decision 40). + precision: Precision = lm.not_json(Precision.HIGH, kw_only=True) + + +class GraniteLayer(DecoderLayer): + """Llama's layer, its attention logits scaled by + ``attention_multiplier`` and each block's output by + ``residual_multiplier``. + """ + + def __init__(self, c: GraniteConfig, i: int): + super().__init__(c, i) + self.self_attn.scaling = c.attention_multiplier + self.residual_multiplier = c.residual_multiplier + + +class Granite(CausalLM): + """Granite on ``config``'s shape and a checkpoint's ``weights``: the + untied output head is ``lm_head``, every matrix read ``(out, in)`` as + the checkpoint ships it. From its sharded checkpoint and + ``tokenizer.json``. + """ + + config = GraniteConfig + profile = Path(__file__).with_name("profiles") + layer_class = GraniteLayer + + def __init__(self, config: GraniteConfig, weights): + super().__init__(config, weights) + # The embedding's multiplier as a row every row is multiplied by. + m = config.embedding_multiplier + self.embed_scale = param(np.full(config.hidden_size, m, bfloat16)) + if config.logits_scaling != 1: + # The logits' divisor folded into the final norm's loaded weight. + norm = self.model.norm + norm.weight = param(norm.weight / bfloat16(config.logits_scaling)) + + def lookup(self, ids): + x = super().lookup(ids) + if self.config.embedding_multiplier == 1: + return x + E = self.config.hidden_size + x = x.reshape(math.prod(x.shape) // E, E) + return x * self.embed_scale + + +if __name__ == "__main__": + lm.main(Granite) diff --git a/iron/applications/granite/profiles/npu1.json b/iron/applications/granite/profiles/npu1.json new file mode 100644 index 000000000..a8509caeb --- /dev/null +++ b/iron/applications/granite/profiles/npu1.json @@ -0,0 +1,2 @@ +{"entries": [ +]} diff --git a/iron/applications/granite/profiles/npu2.json b/iron/applications/granite/profiles/npu2.json new file mode 100644 index 000000000..4abe590e6 --- /dev/null +++ b/iron/applications/granite/profiles/npu2.json @@ -0,0 +1,19 @@ +{"entries": [ + {"operator": "RMSNorm", "weighted": true, "rows": 1, "tile_size": 2560, "num_aie_columns": 1}, + {"operator": "GEMV", "M": 2560, "K": 2560, "num_batches": 1, "tile_size_input": 4, "tile_size_output": 256}, + {"operator": "GEMV", "M": 512, "K": 2560, "num_batches": 1, "tile_size_input": 4}, + {"operator": "MHA", "num_heads": 40, "num_KV_heads": 8, "d": 64, "seq_pad": 1, "num_pipelines": 4}, + {"operator": "GEMV", "M": 2560, "K": 64, "num_batches": 40, "tile_size_input": 4, "tile_size_output": 256}, + {"operator": "ElementwiseMul", "size": 102400, "tile_size": 4096}, + {"operator": "ElementwiseAdd", "size": 2560, "tile_size": 256}, + {"operator": "GEMV", "M": 8192, "K": 2560, "num_batches": 1, "tile_size_input": 4, "tile_size_output": 1024}, + {"operator": "SiLU", "size": 8192, "tile_size": 1024}, + {"operator": "ElementwiseMul", "size": 8192, "tile_size": 1024}, + {"operator": "GEMV", "M": 2560, "K": 8192, "num_batches": 1, "tile_size_input": 1, "tile_size_output": 256}, + {"operator": "GEMV", "M": 100352, "K": 2560, "num_batches": 1, "tile_size_input": 4, "tile_size_output": 32}, + {"operator": "RMSNorm", "weighted": true, "rows": 2048, "tile_size": 2048, "num_aie_columns": 8}, + {"operator": "MHA", "num_heads": 40, "num_KV_heads": 8, "d": 64, "seq_pad": 2048, "num_pipelines": 8}, + {"operator": "ElementwiseAdd", "size": 5242880, "tile_size": 2048}, + {"operator": "SiLU", "size": 16777216, "tile_size": 4096}, + {"operator": "ElementwiseMul", "size": 16777216, "tile_size": 4096} +]} diff --git a/iron/applications/granite_hybrid/README.md b/iron/applications/granite_hybrid/README.md new file mode 100644 index 000000000..391ddd9fd --- /dev/null +++ b/iron/applications/granite_hybrid/README.md @@ -0,0 +1,66 @@ + + +# Granite 4.0-H + +IBM's Granite 4.0-H, `GraniteHybrid` in `model.py`, whose docstring says how +it differs from Granite: Mamba2 mixers in most layers' place of attention, +and a shared MLP. This package ships the dense 1B; a checkpoint with +experts is refused. The shared library is +[AGENTS.md](../../../AGENTS.md#language-models). + +The dense 1B runs on the NPU from its real checkpoint: every layer, the +Mamba2 scan (`SelectiveScan`) and causal convolution included, one compile +serving every prompt length and its native 131072 positions, the scan's +state checkpointed at each 2048-position chunk boundary so a chat turn +reruns only its last chunk. `tests/applications/granite_hybrid/test.py` +checks it against Hugging Face's float32 model, and `reference/` the graph +at a dense toy, device-free. + +## Weights and tokenizer + +From Hugging Face (Apache-2.0, ungated), over plain https, with no account: + +```bash +D=/path/to/granite-4.0-h-1b +REV=d18cca4c121edb87d022116d281ce212c9136f57 # the revision every KL figure here is against +for f in config.json generation_config.json tokenizer_config.json \ + special_tokens_map.json tokenizer.json chat_template.jinja model.safetensors; do + curl -fL -o "$D/$f" \ + "https://huggingface.co/ibm-granite/granite-4.0-h-1b/resolve/$REV/$f" +done +``` + +The weights are 2.9 GB. At its native context the Mamba2 states' +checkpoints (float32, 36 layers by 65 chunk boundaries, 3.4 GiB) and the 4 +attention layers' caches (1 GiB) make the model 7.2 GiB on the device. + +## Precision + +Its projections run at `Precision.HIGH` (Granite's config's): bf16 macs, +each K tile's partial sums rounded once. On the checks below, against +MEDIUM (bfp16 macs): + +| Precision | KL over 40 steps (mean / p90 / max) | 2 chunks, 3 steps (mean / max) | TTFT, 1024 chars | TPS | +|---|---|---|---|---| +| HIGH | 0.0029 / 0.0051 / 0.0105 | 0.0098 / 0.025 | 3.7 s | 6.7 | +| MEDIUM | 0.0040 / 0.0066 / 0.0186 | 0.0063 / 0.0099 | 3.35 s | 6.7 | + +HIGH holds the 40 steps closer, at 10% of the first token's time; a +decode step is a GEMV, which sums in float32 whatever the precision. +HIGHEST (float32 partial sums) is not measured yet. + +## Accuracy and speed + +On Strix Halo (npu2, xcoradaie220), one compile at the native context, +against the float32 golden recorded beside the checkpoint (`golden.npz`): +KL over 40 teacher-forced steps after 292 tokens, mean 0.0024, p90 0.0048, +max 0.0063; over a prompt of 3261 tokens, two chunks, and two decode steps +after it, max 0.0049. A chat turn rerun from the states checkpointed at the +chunk boundary is bit for bit the run from scratch, and repeated runs are +bitwise identical. The first token takes 2.9 s after a 13-character +prompt (a prompt chunk's scan streams all 2048 rows whatever the prompt's +length), 3.6 s after 292 tokens and 8.0 s after 2048, and decode runs at +6.7 tokens a second; none of it is tuned. diff --git a/iron/applications/granite_hybrid/__init__.py b/iron/applications/granite_hybrid/__init__.py new file mode 100644 index 000000000..f9789360f --- /dev/null +++ b/iron/applications/granite_hybrid/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.0-H; see ``iron/applications/granite_hybrid/README.md``.""" + +from iron.lm.cli import exports + +__all__ = ["GraniteHybrid"] +__getattr__ = exports(__name__, __all__) diff --git a/iron/applications/granite_hybrid/model.py b/iron/applications/granite_hybrid/model.py new file mode 100644 index 000000000..ee3152198 --- /dev/null +++ b/iron/applications/granite_hybrid/model.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.0-H on the shared decoder (``iron.lm``), its shape read from +its checkpoint's ``config.json``. This package ships the dense sizes +only; a checkpoint with experts is refused (``num_local_experts``). + +Granite 4.0-H is Granite's layer (its four scalars, ``GraniteConfig``) +with most layers' attention replaced by a Mamba2 mixer (``Mamba``) and +its feed-forward a shared gated MLP (``SharedMLP``). The attention +layers rotate nothing unless ``position_embedding_type`` is ``"rope"``. +A Mamba2 mixer's convolution and state-space scan (``SelectiveScan``) +carry state between calls as ``Checkpoints``, like LFM2's convolution. +""" + +import dataclasses +from pathlib import Path + +import numpy as np + +import iron +import iron.functional as F +from iron.applications.granite.model import Granite, GraniteConfig, GraniteLayer +from iron.lm import ( + CausalConv1d, + Checkpoints, + Linear, + RMSNorm, + main, + residual_add, +) +from iron.operators.copy import Copy +from iron.operators.selective_scan import SelectiveScan + + +@dataclasses.dataclass(frozen=True) +class GraniteHybridConfig(GraniteConfig): + """``GraniteConfig``, which layers are attention, the Mamba2 mixer's + shape and the shared MLP's width. A checkpoint with experts + (``num_local_experts``) is refused; this package ships the dense + sizes only. + """ + + # transformers' default; __post_init__ sets None but under "rope". + rope_theta: float | None = dataclasses.field(default=10000.0, kw_only=True) + layer_types: tuple[str, ...] = dataclasses.field( + kw_only=True, metadata={"convert": tuple} + ) + position_embedding_type: str | None = dataclasses.field(default=None, kw_only=True) + shared_intermediate_size: int = dataclasses.field(kw_only=True) + num_local_experts: int = dataclasses.field(default=0, kw_only=True) + num_experts_per_tok: int = dataclasses.field(default=0, kw_only=True) + mamba_n_heads: int = dataclasses.field(kw_only=True) + mamba_d_head: int = dataclasses.field(kw_only=True) + mamba_d_state: int = dataclasses.field(kw_only=True) + mamba_n_groups: int = dataclasses.field(default=1, kw_only=True) + mamba_d_conv: int = dataclasses.field(default=4, kw_only=True) + + def __post_init__(self): + super().__post_init__() + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError( + f"GraniteHybridConfig: {len(self.layer_types)} layer_types for " + f"{self.num_hidden_layers} layers" + ) + if self.num_local_experts: + raise ValueError( + f"GraniteHybridConfig: num_local_experts {self.num_local_experts}, " + "this package ships the dense sizes only" + ) + if self.position_embedding_type != "rope": + object.__setattr__(self, "rope_theta", None) + + @property + def conv_channels(self) -> int: + """A Mamba2 mixer's convolved channels: ``x``, then ``B`` and ``C``.""" + return self.mamba_n_heads * self.mamba_d_head + 2 * ( + self.mamba_n_groups * self.mamba_d_state + ) + + def attends(self, i: int) -> bool: + """Whether layer ``i`` is attention (``"attention"``, or transformers' + newer ``"full_attention"``), not Mamba2. + """ + return self.layer_types[i] in ("attention", "full_attention") + + +class Mamba(iron.Graph): + """A Mamba2 mixer on ``c``'s shape: ``in_proj`` to the gate, ``x``, + ``B``, ``C`` and ``dt``; ``conv1d`` over ``x``, ``B`` and ``C``, then + silu; the scan; ``norm`` of its output times silu of the gate; and + ``out_proj``. + """ + + def __init__(self, c: GraniteHybridConfig): + E, H, K = c.hidden_size, c.mamba_n_heads, c.mamba_d_conv + self.inner = I = H * c.mamba_d_head + self.channels = c.conv_channels + self.groups = c.mamba_n_groups + self.in_proj = Linear(E, I + self.channels + H, False, precision=c.precision) + self.conv1d = CausalConv1d(c, self.channels, K, bias=True) + self.dt_bias, self.A_log, self.D = (iron.param(H) for _ in range(3)) + self.norm = RMSNorm(I, c.eps) + self.out_proj = Linear(I, E, False, precision=c.precision) + state = (H, c.mamba_d_head, c.mamba_d_state) + self.ssm_state = Checkpoints.of( + state, c.max_seq_len, c.prefill_chunk, np.float32 + ) + + def load_state_dict(self, tree, *, path: str = "") -> None: + super().load_state_dict(tree, path=path) + # Each head's (dt_bias, A, D) in float32. + A = -np.exp(np.asarray(self.A_log, np.float32)) + self.head = iron.param( + np.stack([self.dt_bias, A, self.D], 1).astype(np.float32) + ) + + def body(self, x, call): + I, C = self.inner, self.channels + w, p = self.in_proj.weight, self.in_proj.precision + # Three projections, not one split three ways, as LFM2's conv. + gate, xbc, dt = ( + F.linear(x, part, precision=p) for part in (w[:I], w[I : I + C], w[I + C :]) + ) + xbc = F.silu(self.conv1d(xbc, call)) + read, write = self.ssm_state.at(call) + y, final = SelectiveScan.step( + xbc, + dt, + self.head, + Copy.step(read, dtype=np.float32), + groups=self.groups, + ) + Copy.step(final, write, dtype=np.float32) + return self.out_proj(self.norm(y * F.silu(gate))) + + +class SharedMLP(iron.Graph): + """The gated MLP every layer runs: ``input_linear``, the gate's rows + over the up's, then ``output_linear``. + """ + + def __init__(self, c: GraniteHybridConfig): + E, S = c.hidden_size, c.shared_intermediate_size + self.width = S + self.input_linear = Linear(E, 2 * S, False, precision=c.precision) + self.output_linear = Linear(S, E, False, precision=c.precision) + + def body(self, x): + w, S, p = self.input_linear.weight, self.width, self.input_linear.precision + gate = F.linear(x, w[:S], precision=p) + up = F.linear(x, w[S:], precision=p) + return self.output_linear(F.silu(gate) * up) + + +class GraniteHybridLayer(GraniteLayer): + """Granite's layer, ``mamba`` in attention's place where the layer is + Mamba2, its feed-forward ``shared_mlp``. + """ + + def __init__(self, c: GraniteHybridConfig, i: int): + super().__init__(c, i) + del self.mlp + self.shared_mlp = SharedMLP(c) + self.mamba = None + if not c.attends(i): + del self.self_attn, self.cache + self.mamba = Mamba(c) + + def mix(self, h, call): + return super().mix(h, call) if self.mamba is None else self.mamba(h, call) + + def body(self, x, call): + x = residual_add( + x, self.mix(self.input_layernorm(x), call), self.residual_multiplier + ) + h = self.post_attention_layernorm(x) + y = self.shared_mlp(h) + return residual_add(x, y, self.residual_multiplier) + + +class GraniteHybrid(Granite): + """Granite 4.0-H on ``config``'s shape and a checkpoint's ``weights``, + from its checkpoint and ``tokenizer.json``. + """ + + config = GraniteHybridConfig + profile = Path(__file__).with_name("profiles") + layer_class = GraniteHybridLayer + + +if __name__ == "__main__": + main(GraniteHybrid) diff --git a/iron/applications/granite_hybrid/profiles/npu1.json b/iron/applications/granite_hybrid/profiles/npu1.json new file mode 100644 index 000000000..a8509caeb --- /dev/null +++ b/iron/applications/granite_hybrid/profiles/npu1.json @@ -0,0 +1,2 @@ +{"entries": [ +]} diff --git a/iron/applications/granite_hybrid/profiles/npu2.json b/iron/applications/granite_hybrid/profiles/npu2.json new file mode 100644 index 000000000..e12883c8e --- /dev/null +++ b/iron/applications/granite_hybrid/profiles/npu2.json @@ -0,0 +1,4 @@ +{"entries": [ + {"operator": "MHA", "num_heads": 12, "num_KV_heads": 4, "d": 128, "seq_pad": 1, "num_pipelines": 4}, + {"operator": "MHA", "num_heads": 12, "num_KV_heads": 4, "d": 128, "seq_pad": 2048, "num_pipelines": 8} +]} diff --git a/iron/applications/qwen3/README.md b/iron/applications/qwen3/README.md new file mode 100644 index 000000000..1fe1735da --- /dev/null +++ b/iron/applications/qwen3/README.md @@ -0,0 +1,39 @@ + + +# Qwen3 1.7B on the NPU + +Qwen3 1.7B and 0.6B, `Qwen3` in `model.py`, whose docstring says how it +differs from Llama: a per-head RMSNorm on q and k before RoPE, and an +unread `lm_head.weight` beside its tied head. How to run, test and tune it is the README's [Models](../../../README.md#models); the shared library is [AGENTS.md](../../../AGENTS.md#language-models). + +## Weights and tokenizer + +From Hugging Face ([`Qwen/Qwen3-1.7B`](https://huggingface.co/Qwen/Qwen3-1.7B), Apache-2.0), +in one directory: `model-00001-of-00002.safetensors`, +`model-00002-of-00002.safetensors`, `model.safetensors.index.json`, +`tokenizer.json` and `config.json`. Qwen3 0.6B is the same model at another shape: its +`model.safetensors`, `tokenizer.json` and `config.json` from +[`Qwen/Qwen3-0.6B`](https://huggingface.co/Qwen/Qwen3-0.6B). + +## NPU1's profile + +Qwen3's `head_dim` is 128, twice Llama 3.2 1B's: MHA's block-size search +picks a smaller block, which at `num_pipelines=1` pushes a prompt chunk's +K/V fill past the DMA descriptor's 64-iteration cap. `profiles/npu1.json` +pins `num_pipelines=4` for the prompt-chunk shape to stay under it. + +Open on NPU1: 1.7B's deepest prompt chunk at its native 40960 positions +differs from the graph's reference by KL 0.148 (`long_context`'s +`test_prompt_deep_in_the_cache`, bound 0.1). It passes at 34816 positions +(4.8e-6), 32768 (2.0e-5) and 16384 (2.1e-6), its deepest decode step +passes at 40960, and the prompt's MHA alone, given its inputs, is as +accurate at 40959 as at 2047. + +## Testing + +`tests/applications/qwen3/` runs 1.7B's checkpoint and 0.6B's, each also +in the each-step form, over a two-chunk prompt and a chat turn, and checks +that the unread `lm_head.weight` is the tied embedding. diff --git a/iron/applications/qwen3/__init__.py b/iron/applications/qwen3/__init__.py new file mode 100644 index 000000000..e678a0600 --- /dev/null +++ b/iron/applications/qwen3/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3 end to end on the NPU; see ``iron/applications/qwen3/README.md``.""" + +from iron.lm.cli import exports + +__all__ = ["Qwen3"] +__getattr__ = exports(__name__, __all__) diff --git a/iron/applications/qwen3/model.py b/iron/applications/qwen3/model.py new file mode 100644 index 000000000..e95130ea9 --- /dev/null +++ b/iron/applications/qwen3/model.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3 1.7B on the shared decoder (``iron.lm``), its shape read from its +checkpoint's ``config.json``. Run it with ``python -m +iron.applications.qwen3.model``. + +Qwen3 is Llama's layer plus a per-head RMSNorm on q and k before RoPE, +``self_attn.q_norm`` and ``k_norm`` (``Qwen3Layer``). Its output +head is the input embedding, tied; its checkpoint also ships +``lm_head.weight``, byte-identical to ``model.embed_tokens.weight``, which +is never read. + +The operators' tunables are the graph's profile, ``profiles/.json``, +keyed by operator shape at Qwen3 1.7B's shape: a decode step's and a +2048-row prompt chunk's, at any ``max_seq_len``, which sizes the caches and +the RoPE table alone. +""" + +import dataclasses +from pathlib import Path + +from iron import Precision, lm +from iron.lm import Attention, CausalLM, DecoderConfig, DecoderLayer + + +@dataclasses.dataclass(frozen=True) +class Qwen3Config(DecoderConfig): + """``DecoderConfig`` and its projections' ``precision``.""" + + precision: Precision = lm.not_json(Precision.MEDIUM, kw_only=True) + + +class Qwen3Layer(DecoderLayer): + """Llama's layer, its attention norming q and k per head.""" + + def __init__(self, c: DecoderConfig, i: int): + super().__init__(c, i, Attention(c, i, norms="qk", o_bias=c.attention_bias)) + + +class Qwen3(CausalLM): + """The Qwen3 decoder, with the shape ``config`` and the checkpoint's + ``weights``. ``Qwen3.load`` builds one from a checkpoint and its + ``tokenizer.json``. + """ + + config = Qwen3Config + profile = Path(__file__).with_name("profiles") + layer_class = Qwen3Layer + + +if __name__ == "__main__": + lm.main(Qwen3) diff --git a/iron/applications/qwen3/profiles/npu1.json b/iron/applications/qwen3/profiles/npu1.json new file mode 100644 index 000000000..1d7e3dda9 --- /dev/null +++ b/iron/applications/qwen3/profiles/npu1.json @@ -0,0 +1,17 @@ +{"entries": [ + {"operator": "RMSNorm", "weighted": true, "rows": 1, "tile_size": 2048, "num_aie_columns": 1}, + {"operator": "RMSNorm", "weighted": true, "rows": 2048, "tile_size": 2048, "num_aie_columns": 4}, + {"operator": "GEMV", "M": 2048, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 64}, + {"operator": "GEMV", "M": 1024, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 64}, + {"operator": "MHA", "num_heads": 16, "num_KV_heads": 8, "d": 128, "seq_pad": 1, "num_pipelines": 2}, + {"operator": "MHA", "num_heads": 16, "num_KV_heads": 8, "d": 128, "seq_pad": 2048, "num_pipelines": 4}, + {"operator": "ElementwiseAdd", "size": 2048, "tile_size": 512}, + {"operator": "GEMV", "M": 6144, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 1024}, + {"operator": "SiLU", "size": 6144, "tile_size": 2048}, + {"operator": "ElementwiseMul", "size": 6144, "tile_size": 2048}, + {"operator": "GEMV", "M": 2048, "K": 6144, "num_batches": 1, "num_channels": 2, "tile_size_input": 1, "tile_size_output": 256}, + {"operator": "GEMV", "M": 151936, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 32}, + {"operator": "ElementwiseAdd", "size": 4194304, "tile_size": 2048}, + {"operator": "SiLU", "size": 12582912, "tile_size": 4096}, + {"operator": "ElementwiseMul", "size": 12582912, "tile_size": 4096} +]} diff --git a/iron/applications/qwen3/profiles/npu2.json b/iron/applications/qwen3/profiles/npu2.json new file mode 100644 index 000000000..1d7e3dda9 --- /dev/null +++ b/iron/applications/qwen3/profiles/npu2.json @@ -0,0 +1,17 @@ +{"entries": [ + {"operator": "RMSNorm", "weighted": true, "rows": 1, "tile_size": 2048, "num_aie_columns": 1}, + {"operator": "RMSNorm", "weighted": true, "rows": 2048, "tile_size": 2048, "num_aie_columns": 4}, + {"operator": "GEMV", "M": 2048, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 64}, + {"operator": "GEMV", "M": 1024, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 64}, + {"operator": "MHA", "num_heads": 16, "num_KV_heads": 8, "d": 128, "seq_pad": 1, "num_pipelines": 2}, + {"operator": "MHA", "num_heads": 16, "num_KV_heads": 8, "d": 128, "seq_pad": 2048, "num_pipelines": 4}, + {"operator": "ElementwiseAdd", "size": 2048, "tile_size": 512}, + {"operator": "GEMV", "M": 6144, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 1024}, + {"operator": "SiLU", "size": 6144, "tile_size": 2048}, + {"operator": "ElementwiseMul", "size": 6144, "tile_size": 2048}, + {"operator": "GEMV", "M": 2048, "K": 6144, "num_batches": 1, "num_channels": 2, "tile_size_input": 1, "tile_size_output": 256}, + {"operator": "GEMV", "M": 151936, "K": 2048, "num_batches": 1, "num_channels": 2, "tile_size_input": 4, "tile_size_output": 32}, + {"operator": "ElementwiseAdd", "size": 4194304, "tile_size": 2048}, + {"operator": "SiLU", "size": 12582912, "tile_size": 4096}, + {"operator": "ElementwiseMul", "size": 12582912, "tile_size": 4096} +]} diff --git a/iron/applications/qwen3_embedding/__init__.py b/iron/applications/qwen3_embedding/__init__.py new file mode 100644 index 000000000..b2e65ea17 --- /dev/null +++ b/iron/applications/qwen3_embedding/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Embedding on the NPU; see ``model.py``.""" + +from iron.lm.cli import exports + +__all__ = ["Qwen3Embedding"] +__getattr__ = exports(__name__, __all__) diff --git a/iron/applications/qwen3_embedding/hugging_face.py b/iron/applications/qwen3_embedding/hugging_face.py new file mode 100644 index 000000000..ae53a320b --- /dev/null +++ b/iron/applications/qwen3_embedding/hugging_face.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face's own Qwen3-Embedding in float32, what its golden is +recorded from. It needs torch and transformers, which are not IRON's +requirements, only to record. +""" + +import numpy as np +import torch + +from iron.lm.golden import HuggingFace + + +class LastHidden(HuggingFace): + """transformers' Qwen3 model without its head: the last token's final + hidden state, unit length. + """ + + model_class = "AutoModel" + + def logits(self, tokens) -> np.ndarray: + with torch.no_grad(): + out = self.model(torch.tensor([list(tokens)]), use_cache=False) + last = out.last_hidden_state[0, -1].numpy() + return last / np.linalg.norm(last) diff --git a/iron/applications/qwen3_embedding/model.py b/iron/applications/qwen3_embedding/model.py new file mode 100644 index 000000000..dfcf669a3 --- /dev/null +++ b/iron/applications/qwen3_embedding/model.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Embedding on the NPU: Qwen3's decoder without its head, over +``CausalLM``'s prompt chunks and caches, so one compile serves every +length up to its native context. A text's embedding is its last token's +final hidden state, unit length; the tokenizer's end of text closes it, +and a query goes behind its task's instruction. +""" + +import numpy as np + +import iron +from iron.applications.qwen3.model import Qwen3 +from iron.lm import subtree + + +class Qwen3Embedding(Qwen3): + """Qwen3-Embedding on ``config``'s shape from a checkpoint's + ``weights``, Qwen3's ``model`` modules at its root. A call returns the + last row's final hidden state; nothing is drawn. + """ + + # The end-of-text token each input ends in. + end_of_text = 151643 + + def __init__(self, config, weights): + super().__init__(config, subtree(weights, model=weights)) + + @property + def hugging_face(self): + """transformers' Qwen3 model, imported only where a golden is recorded.""" + from .hugging_face import LastHidden + + return LastHidden + + def head(self, x): + return self.model.norm(x) + + def draw(self, logits, call): + return logits, iron.carry(token=call.token, position=call.position + 1) + + def embed(self, text: str, task: str | None = None) -> np.ndarray: + """The unit-length embedding of ``text``, a query behind ``task``'s + instruction where given, else a document, float32. + """ + if task is not None: + text = f"Instruct: {task}\nQuery:{text}" + tokens = [*self.tokenizer.encoding.encode(text), self.end_of_text] + y = self.logits(tokens).astype(np.float32) + return y / np.linalg.norm(y) diff --git a/iron/applications/qwen3_moe/README.md b/iron/applications/qwen3_moe/README.md new file mode 100644 index 000000000..3fcf9979b --- /dev/null +++ b/iron/applications/qwen3_moe/README.md @@ -0,0 +1,22 @@ + + +# Qwen3-30B-A3B on the NPU + +Qwen3-30B-A3B, `Qwen3Moe` in `model.py`, whose docstring says how it +differs from Qwen3: a mixture of 128 experts, 8 a token, in place of the +feed-forward. The shared library is [AGENTS.md](../../../AGENTS.md#language-models). + +The experts' feed-forwards are `Experts`, whose first kernel streams every +expert's weights for each row, too slow to generate with: its test is +the device-free `tests/applications/qwen3_moe/reference/`. Its weights are +61 GB in bf16, every expert held, and stacked into one copy as they load. + +## Weights and tokenizer + +From Hugging Face +([`Qwen/Qwen3-30B-A3B`](https://huggingface.co/Qwen/Qwen3-30B-A3B), +Apache-2.0), in one directory: `model.safetensors.index.json` and its +shards, `tokenizer.json` and `config.json`. diff --git a/iron/applications/qwen3_moe/__init__.py b/iron/applications/qwen3_moe/__init__.py new file mode 100644 index 000000000..b86cfb10a --- /dev/null +++ b/iron/applications/qwen3_moe/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-MoE on the shared decoder; see ``model.py``.""" + +from iron.lm.cli import exports + +__all__ = ["Qwen3Moe"] +__getattr__ = exports(__name__, __all__) diff --git a/iron/applications/qwen3_moe/model.py b/iron/applications/qwen3_moe/model.py new file mode 100644 index 000000000..174101da6 --- /dev/null +++ b/iron/applications/qwen3_moe/model.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-30B-A3B, a small mixture of experts, on the shared decoder +(``iron.lm``), its shape read from its checkpoint's ``config.json``. + +Qwen3-MoE is Qwen3's layer with its feed-forward a mixture of experts +(``SparseMoE``): a router picks each token's ``num_experts_per_tok`` +experts and weighs them by a softmax over their logits alone, which is +``TopKRouter``'s. Its output head is untied. + +The experts stay in the mapped checkpoint, 58 GB of them at 30B. Between +the router and the experts the host copies the ones a call chose into +states every layer shares (``SparseMoE.stage``), so the device holds only +those and the page cache whichever the host read lately. The staged +experts run as one gated feed-forward, each row's routing weights +scaling its experts' activations and zero at the others'. +""" + +import dataclasses +from pathlib import Path + +import numpy as np + +import iron +import iron.functional as F +from iron import Precision, lm +from iron.applications.qwen3.model import Qwen3Layer +from iron.common.graph.handle import State +from iron.lm import CausalLM, DecoderConfig, GatedMLP, Linear, activation +from iron.operators.router import TopKRouter + + +@dataclasses.dataclass(frozen=True) +class Qwen3MoeConfig(DecoderConfig): + """``DecoderConfig`` and its experts: every layer's feed-forward a + mixture, its chosen experts' weights normalized, as Qwen3-MoE's + checkpoints have them. ``Qwen3MoeConfig`` refuses any other. + """ + + precision: Precision = lm.not_json(Precision.MEDIUM, kw_only=True) + num_experts: int = dataclasses.field(kw_only=True) + num_experts_per_tok: int = dataclasses.field(kw_only=True) + moe_intermediate_size: int = dataclasses.field(kw_only=True) + norm_topk_prob: bool = dataclasses.field(default=False, kw_only=True) + decoder_sparse_step: int = dataclasses.field(default=1, kw_only=True) + mlp_only_layers: tuple[int, ...] = dataclasses.field( + default=(), kw_only=True, metadata={"convert": tuple} + ) + tie_word_embeddings: bool = False + + def __post_init__(self): + super().__post_init__() + if not self.norm_topk_prob: + raise ValueError( + "Qwen3MoeConfig: TopKRouter normalizes the chosen experts' " + "weights; no operator weighs them unnormalized" + ) + if self.decoder_sparse_step != 1 or self.mlp_only_layers: + raise ValueError( + f"Qwen3MoeConfig: decoder_sparse_step {self.decoder_sparse_step} " + f"and mlp_only_layers {self.mlp_only_layers} make some layers " + "dense; here every layer's feed-forward is a mixture" + ) + + +@dataclasses.dataclass(frozen=True) +class Staged: + """The experts a call of ``rows`` rows is staged in, each ``width`` wide + in a slot of its own: ``gate`` and ``up`` stacked as one projection's + rows, ``down`` as one's columns, and ``routed``, each row's routing + weight for each slot over that slot's columns. + """ + + gate: State + up: State + down: State + routed: State + + @classmethod + def of(cls, rows: int, slots: int, width: int, hidden: int) -> "Staged": + n = slots * width + return cls( + iron.state((n, hidden)), + iron.state((n, hidden)), + iron.state((hidden, n)), + iron.state((rows, n)), + ) + + +class SparseMoE(iron.Graph): + """A mixture of experts on ``c``'s shape: ``gate``, the router, and + ``experts``, each a gated feed-forward ``moe_intermediate_size`` wide, + left in the checkpoint. ``Qwen3Moe`` gives every layer's ``staged``, + the experts a call stages, by its rows. + """ + + staged: dict[int, Staged] + + def __init__(self, c: Qwen3MoeConfig): + self.gate = Linear(c.hidden_size, c.num_experts, False, precision=c.precision) + expert = dataclasses.replace(c, intermediate_size=c.moe_intermediate_size) + self.experts = [GatedMLP(expert) for _ in range(c.num_experts)] + self.top_k = c.num_experts_per_tok + self.width = c.moe_intermediate_size + self.act = activation(c.hidden_act) + + def body(self, x): + n, experts = x.shape[0], self.gate.weight.shape[0] + logits = F.linear(x, self.gate.weight, precision=self.gate.precision) + logits = logits.reshape(n, experts) + ids, weights = TopKRouter.step(logits, top_k=self.top_k) + s = self.staged[n] + iron.on_host(self.stage, ids, weights, writes=(s.gate, s.up, s.down, s.routed)) + # A prompt chunk's staged experts are 98304 wide at 30B, past what + # flm.ProjectionGEMM's per-call transfers fit in one descriptor. + p = Precision.HIGHEST + h = self.act(F.linear(x, s.gate, precision=p)) * F.linear(x, s.up, precision=p) + routed = s.routed[: x.valid[0]] if x.valid else s.routed + return F.linear(h * routed.reshape(h.shape), s.down, precision=p) + + def stage(self, ids, weights, gate, up, down, routed) -> None: + """Copy the experts ``ids`` chose into ``gate``, ``up`` and ``down``, a + slot each, and each of ``ids``' rows' ``weights`` into ``routed`` at + its experts' slots, zero at the others. A slot none chose keeps the + expert it held, which no row weighs. + """ + chosen, slots = np.unique(ids, return_inverse=True) + I, n = self.width, len(ids) + for slot, e in enumerate(chosen): + expert, at = self.experts[e], slice(slot * I, (slot + 1) * I) + gate[at] = expert.gate_proj.weight + up[at] = expert.up_proj.weight + down[:, at] = expert.down_proj.weight + routed[:n] = 0 + rows = np.arange(n)[:, None] + routed[:n].reshape(n, -1, I)[rows, slots.reshape(ids.shape)] = weights[ + :, ::2, None + ] + + +class Qwen3MoeLayer(Qwen3Layer): + """Qwen3's layer, its feed-forward ``SparseMoE``.""" + + def __init__(self, c: Qwen3MoeConfig, i: int): + super().__init__(c, i) + self.mlp = SparseMoE(c) + + +class Qwen3Moe(CausalLM): + """Qwen3-MoE on ``config``'s shape and a checkpoint's ``weights``, from + its checkpoint and ``tokenizer.json``. + """ + + config = Qwen3MoeConfig + layer_class = Qwen3MoeLayer + profile = Path(__file__).with_name("profiles") + + def __init__(self, config: Qwen3MoeConfig, weights): + super().__init__(config, weights) + c = config + H, I, K = c.hidden_size, c.moe_intermediate_size, c.num_experts_per_tok + staged = { + rows: Staged.of(rows, min(c.num_experts, rows * K), I, H) + for rows in (1, c.prefill_chunk) + } + for layer in self.layers: + layer.mlp.staged = staged diff --git a/iron/applications/qwen3_moe/profiles/npu2.json b/iron/applications/qwen3_moe/profiles/npu2.json new file mode 100644 index 000000000..e61fd85bc --- /dev/null +++ b/iron/applications/qwen3_moe/profiles/npu2.json @@ -0,0 +1,4 @@ +{"entries": [ + {"operator": "MHA", "num_heads": 32, "num_KV_heads": 4, "d": 128, "seq_pad": 1, "num_pipelines": 4}, + {"operator": "MHA", "num_heads": 32, "num_KV_heads": 4, "d": 128, "seq_pad": 2048, "num_pipelines": 4} +]} diff --git a/iron/applications/qwen3_reranker/__init__.py b/iron/applications/qwen3_reranker/__init__.py new file mode 100644 index 000000000..7cb841a4e --- /dev/null +++ b/iron/applications/qwen3_reranker/__init__.py @@ -0,0 +1,9 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Reranker on the NPU; see ``model.py``.""" + +from iron.lm.cli import exports + +__all__ = ["Qwen3Reranker"] +__getattr__ = exports(__name__, __all__) diff --git a/iron/applications/qwen3_reranker/model.py b/iron/applications/qwen3_reranker/model.py new file mode 100644 index 000000000..9924a1562 --- /dev/null +++ b/iron/applications/qwen3_reranker/model.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Reranker on the NPU: Qwen3 asked, in the prompt its model card +gives, whether a document meets a query and an instruction; its score is +the probability of "yes" over "no" as the answer's first token. +""" + +import math + +from iron.applications.qwen3.model import Qwen3 + +PROMPT = ( + "<|im_start|>system\nJudge whether the Document meets the requirements " + "based on the Query and the Instruct provided. Note that the answer can " + 'only be "yes" or "no".<|im_end|>\n<|im_start|>user\n' + ": {instruction}\n: {query}\n: {document}" + "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n" +) +INSTRUCTION = ( + "Given a web search query, retrieve relevant passages that answer the query" +) + + +class Qwen3Reranker(Qwen3): + """Qwen3 scoring a document against a query (``score``).""" + + def score(self, query: str, document: str, instruction=INSTRUCTION) -> float: + """The probability that ``document`` meets ``query`` and + ``instruction``. + """ + encode = self.tokenizer.encoding.encode + text = PROMPT.format(instruction=instruction, query=query, document=document) + logits = self.logits(encode(text)) + [yes], [no] = encode("yes"), encode("no") + return 1 / (1 + math.exp(logits[no] - logits[yes])) diff --git a/iron/functional.py b/iron/functional.py index 151552e6f..c494e8ab8 100644 --- a/iron/functional.py +++ b/iron/functional.py @@ -8,6 +8,7 @@ them. A handle's own ``+ - * /`` are steps too, and ``@`` of one row. """ +import dataclasses import math import ml_dtypes @@ -16,6 +17,7 @@ from iron.common import Precision from iron.common.graph.handle import Handle, State +from iron.operators.copy import Copy from iron.operators.flm.gemm.design import Epilogue from iron.operators.flm.gemm.op import ProjectionGEMM from iron.operators.gelu import GELU @@ -121,3 +123,52 @@ def gelu(x: Handle, approximate: str = "none") -> Handle: def rope(x: Handle, angles: Handle) -> Handle: """``x`` ``(rows, head_dim)`` turned by RoPE, a row of ``angles`` per row.""" return RoPE.step(x, angles) + + +@dataclasses.dataclass(frozen=True) +class ConvState: + """What a short causal conv carries between calls: ``read``, the + ``K - 1`` rows before a call's first, and ``write``, where the ``K - 1`` + ending it go; and the scratch states a call stages them in, ``window`` + ``(2, chunk, channels)`` and ``tail`` ``(K - 1, channels)``. + """ + + read: Handle + write: Handle + window: State + tail: State + + +def causal_conv1d(x: Handle, taps, state: ConvState) -> Handle: + """``x`` ``(rows, channels)`` through a depth-wise causal filter of + ``K`` taps, each channel its own column of ``taps`` ``(K, channels)``: + output row ``t`` is ``sum_k taps[k] * x[t - K + 1 + k]``. The ``K - 1`` + rows before ``x``'s first row are read from ``state.read``, and the + ``K - 1`` rows ending this call (counting those read) are written to + ``state.write`` for the next call. When ``x``'s rows are bounded per + call (``x[:n]``), the call is a prompt chunk of at most ``chunk`` rows; + otherwise ``x`` is one row, a decode step. + """ + read, write, window, tail = state.read, state.write, state.window, state.tail + K, C = taps.shape[0], window.shape[1] + # Block 0 ends with the rows before the call; block 1 is the call's. + flat = window.reshape(2 * C, x.shape[1]) + first = C - (K - 1) + flat[first:C] = read + rows = x.valid.get(0) + if rows is None: + flat[C] = x.reshape(x.shape[1]) + y = x * taps[K - 1] + for k in range(K - 1): + y = y + flat[first + k : first + k + 1] * taps[k] + Copy.step(flat[first + 1 : C + 1], write) + return y + window[1, :rows] = x + y = flat[first : first + C] * taps[0] + for k in range(1, K): + y = y + flat[first + k : first + k + C] * taps[k] + y = y[:rows] + for j in range(K - 1): + tail[j] = flat[first + rows + j] + Copy.step(tail, write) + return y diff --git a/iron/lm/__init__.py b/iron/lm/__init__.py index f2ee58078..aab1da4cc 100644 --- a/iron/lm/__init__.py +++ b/iron/lm/__init__.py @@ -33,7 +33,7 @@ residual_add, ) from iron.common.feature_map import Conv2d -from .causal import Call, CausalLM, Decoder +from .causal import Call, CausalConv1d, CausalLM, Checkpoints, Decoder from .encoder import Encoder, Tower from .golden import Golden, HuggingFace from .generation import Sampler, accuracy, determinism, generate_on_host @@ -43,8 +43,10 @@ __all__ = [ "Attention", "Call", + "CausalConv1d", "CausalLM", "Checkpoint", + "Checkpoints", "Config", "Conv2d", "Decoder", diff --git a/iron/lm/causal.py b/iron/lm/causal.py index bbbe165e9..cedc336e1 100644 --- a/iron/lm/causal.py +++ b/iron/lm/causal.py @@ -16,11 +16,14 @@ import numpy as np from aie.iron.kernels.sample import ROW_WORDS +from ml_dtypes import bfloat16 import iron +import iron.functional as F from iron import Carried, Scratchpad from iron.common.graph import CarriedLoop, CompiledGraph, Handle from iron.common.graph.compiled import walk +from iron.common.graph.handle import State from iron.operators.copy import Copy from iron.operators.sample import Sample @@ -37,8 +40,9 @@ @dataclasses.dataclass(kw_only=True) class Call: """One call's values, built once in ``body`` and read by every module that - keeps memory from call to call (``KVCache``, ``Rope``): a prompt chunk's - ``rows`` from ``start``, or the decode step's one at ``position``. + keeps memory from call to call (``KVCache``, ``Checkpoints``, ``Rope``): a + prompt chunk's ``rows`` from ``start``, or the decode step's one at + ``position``. """ prompt: bool @@ -63,6 +67,95 @@ class Call: embeddings: Any = None +@dataclasses.dataclass(frozen=True) +class Checkpoints: + """A state a decoder carries from position to position, a slot per + prompt-chunk boundary: slot ``k`` holds it after the positions up to + ``k * chunk`` run so far, slot 0 (never written) before the sequence. + """ + + slots: State + chunk: int + + @classmethod + def of( + cls, shape: tuple[int, ...], max_seq_len: int, chunk: int, dtype=bfloat16 + ) -> "Checkpoints": + """A state of ``shape``, a slot for every boundary of ``chunk``-position + prompt chunks up to ``max_seq_len`` positions. + + Raises: + ValueError: ``chunk`` is not a power of two, which a decode step's + slot (``position // chunk``) needs. + """ + if chunk < 1 or chunk & (chunk - 1): + raise ValueError( + f"Checkpoints: a {chunk}-position chunk is no power of two" + ) + slots = (max_seq_len - 1) // chunk + 2 + return cls(iron.state((slots, *shape), dtype), chunk) + + def at(self, call: Call) -> tuple[Handle, Handle]: + """The slot ``call`` reads, the state before its first position, and + the slot it writes, the state after its last, in a graph's body. A + call ends within the chunk it starts in (``CausalLM`` splits them). + """ + C = self.chunk + return self.slots[(call.start - 1) // C + 1], self.slots[call.position // C + 1] + + def restart(self, held: int, ran: int) -> int: + """Where a call may start to change the positions from ``held`` on, + after ``ran`` positions ran: there where the last call ended, else at + the boundary before it. + """ + return held if held == ran else held // self.chunk * self.chunk + + def slot(self, position: int) -> int: + """The slot holding the state after the first ``position`` positions.""" + return 0 if position == 0 else (position - 1) // self.chunk + 1 + + +class CausalConv1d(iron.Graph): + """A depth-wise causal convolution of ``kernel`` taps over a call's + rows, each of ``channels`` its own taps (``weight``, torch's + ``(channels, 1, kernel)``), plus ``bias`` where it holds one. The + ``kernel - 1`` rows before a call are its ``state``, a slot per + prompt-chunk boundary; ``CausalLM`` gives every one of a model the + scratch ``F.causal_conv1d`` stages a call in (``window``, ``tail``). + + Raises: + ValueError: ``prefill_chunk`` is smaller than ``kernel - 1``, the + history the staging window must hold. + """ + + window: State + tail: State + + def __init__(self, c, channels: int, kernel: int, bias: bool = False): + if c.prefill_chunk < kernel - 1: + raise ValueError( + f"CausalConv1d: prefill_chunk {c.prefill_chunk} is smaller than " + f"kernel - 1 ({kernel - 1})" + ) + self.weight = iron.param(channels, 1, kernel) + self.bias = iron.param(channels) if bias else None + self.state = Checkpoints.of( + (kernel - 1, channels), c.max_seq_len, c.prefill_chunk + ) + + def load_state_dict(self, tree, *, path: str = "") -> None: + super().load_state_dict(tree, path=path) + # A row per tap, each the row RowwiseMul takes. + C, _, K = self.weight.shape + self.taps = iron.param(np.ascontiguousarray(self.weight.reshape(C, K).T)) + + def body(self, x, call): + read, write = self.state.at(call) + state = F.ConvState(read, write, self.window, self.tail) + y = F.causal_conv1d(x, self.taps, state) + return y if self.bias is None else y + self.bias + + class Decoder(iron.Graph): """A decoder's modules on ``c``'s shape, as a checkpoint's ``model`` names them: ``embed_tokens``, ``layers``, layer ``i`` a ``layer(c, i)``, @@ -118,15 +211,28 @@ def __init__(self, config: DecoderConfig, weights): head = np.asarray(self.lm_head.weight) self.lm_head.weight = iron.param(np.pad(head, ((0, 1), (0, 0)))) self.layers = self.model.layers + # The scratch every causal conv of a shape stages its calls in. + scratch = {} + for conv in {id(p): p for _, _, p, _ in walk(self)}.values(): + if isinstance(conv, CausalConv1d): + (K, C), chunk = conv.taps.shape, config.prefill_chunk + if (K, C) not in scratch: + scratch[K, C] = iron.state((2, chunk, C)), iron.state((K - 1, C)) + conv.window, conv.tail = scratch[K, C] @property - def recurrent(self) -> list[KVCache]: - """Each ``KVCache`` keeping a ring, which restarts only at some - positions (``restart``). + def recurrent(self) -> list[Checkpoints | KVCache]: + """Each state a layer carries from position to position, which + restarts only at some (``restart``): ``Checkpoints``, and a + ``KVCache`` keeping a ring. """ found = {} for _, _, p, _ in walk(self): - if isinstance(p, KVCache) and p.ring is not None: + if ( + isinstance(p, Checkpoints) + or isinstance(p, KVCache) + and p.ring is not None + ): found.setdefault(id(p), p) return list(found.values()) @@ -335,6 +441,11 @@ def device_loop(self) -> bool: raise RuntimeError(f"{type(self).__name__}: load_versions() first") return self.decode_step.emit is not None + @property + def draws_at_once(self) -> bool: + """A device loop's: ``generate`` reads its draws once the last is drawn.""" + return self.device_loop + def release(self) -> None: """Free the device loops' runs, then the graph's memory.""" self._loops.clear() diff --git a/iron/lm/model.py b/iron/lm/model.py index 23729f549..596da0207 100644 --- a/iron/lm/model.py +++ b/iron/lm/model.py @@ -225,7 +225,8 @@ def token(self, id: int) -> str | None: @dataclasses.dataclass(frozen=True) class Stats: """A generation's timings: the seconds to the first token, and the - tokens per second after it (NaN for one token). + tokens per second drawn after it (NaN for one token), each token drawn + counted, those past the first ``eos`` too where they were drawn anyway. """ ttft: float @@ -260,6 +261,13 @@ def __init__(self, config: Any, weights: Any) -> None: ... def _draw(self, *args, **kwargs) -> Iterator[int]: ... + @property + def draws_at_once(self) -> bool: + """Whether, after a non-EOS first token, ``generate`` draws the rest + before the second arrives: even those past ``eos`` cost no more to read. + """ + return False + @classmethod def load( cls, @@ -500,12 +508,17 @@ def _stream(self, prompt, max_tokens, sampler, inputs, each: bool) -> Iterator: sampler = Sampler(temperature=0) if sampler is None else sampler # The text of drawn[held:] is yet to be shown; drawn[seen:held] is # decoded with it, since a token's text can depend on the one before. - drawn, seen, held = [], 0, 0 + drawn, seen, held, count = [], 0, 0, 0 + # Past eos, tokens drawn at once still count in the time they took. + at_once = not each and self.draws_at_once began = first = last = time.perf_counter() for token in self._draw(tokens, max_tokens, sampler, each, **inputs): last = time.perf_counter() - if not drawn: + count += 1 + if count == 1: first = last + if drawn and drawn[-1] in self.eos: + continue drawn.append(token) if not text: yield token @@ -517,13 +530,13 @@ def _stream(self, prompt, max_tokens, sampler, inputs, each: bool) -> Iterator: if piece: yield piece seen, held = held, len(drawn) - if token in self.eos: + if token in self.eos and (count == 1 or not at_once): break if text and held < len(drawn): shown = tokenizer.decode(drawn[seen:held]) if piece := tokenizer.decode(drawn[seen:])[len(shown) :]: yield piece - later = len(drawn) - 1 + later = count - 1 rate = later / (last - first) if later else math.nan self.stats = Stats(first - began, rate) @@ -605,6 +618,10 @@ def logits(self, tokens) -> np.ndarray: ... def reference(self, *tensors, dev=None, **values) -> Any: ... + @property + def draws_at_once(self) -> bool: + return False + def _draw(self, tokens, max_tokens: int, sampler: Sampler, each=True): history = [int(t) for t in tokens] for _ in range(max_tokens): diff --git a/iron/lm/multimodal.py b/iron/lm/multimodal.py index 94829bc4c..fcc7b5757 100644 --- a/iron/lm/multimodal.py +++ b/iron/lm/multimodal.py @@ -247,6 +247,10 @@ def logits(self, tokens, raw: dict[int, tuple]) -> np.ndarray: assert isinstance(self.text, CausalLM), "a causal text model's" return self.text.logits(tokens) + @property + def draws_at_once(self) -> bool: + return self.text.draws_at_once + def _draw( self, tokens, diff --git a/iron/lm/session.py b/iron/lm/session.py index 394c4a149..147ae8e96 100644 --- a/iron/lm/session.py +++ b/iron/lm/session.py @@ -108,10 +108,11 @@ def draw( def forget(self, position: int) -> None: """Hold nothing the caches have from ``position`` on: the next call - runs from there, as a change before it (a placeholder's soft tokens) - would leave them stale. + runs from there or an earlier restart boundary, as a change before + it (a placeholder's soft tokens) would leave them stale. """ - self._seen = self._seen[:position] + held = self._seen[:position].size + self._seen = self._seen[: self.resume(held)] def resume(self, held: int) -> int: """The position calls start from to run the tokens from ``held`` on: diff --git a/iron/lm/testing.py b/iron/lm/testing.py index 5ab7c0853..99854a00f 100644 --- a/iron/lm/testing.py +++ b/iron/lm/testing.py @@ -22,6 +22,7 @@ from iron.common.graph.compiled import walk from .blocks import KVCache +from .causal import Checkpoints from .generation import ( SEED, Sampler, @@ -297,10 +298,14 @@ def check_chat_turn(model, chars: int, turn: int, num_tokens: int = 2, *, record assert not steps, f"the turn's logits differ from scratch's at steps {steps}" -def check_deep_decode(model, position: int, bound: float, chars: int = 1024, *, record): +def check_deep_decode( + model, position: int, bound: float, chars: int = 1024, *, cosine=False, record +): """A decode step at ``position``, deep in the caches, against the graph's own reference (``Graph.reference``), for a context transformers - would take too long over, within ``bound`` of its KL. + would take too long over, within ``bound`` of its KL, or with + ``cosine`` (a model whose row is a hidden state, not logits) at a + cosine of at least ``bound``. """ tokens = prompt(model, chars, 1) token = greedy(model.logits(tokens)) @@ -310,17 +315,16 @@ def check_deep_decode(model, position: int, bound: float, chars: int = 1024, *, got, _ = model(**values) got = got.numpy() expected, _ = model.reference(**values) - kl, top1 = divergence(expected, got) - record("DeepKL", kl) - record("DeepTop1", top1) - assert kl <= bound, f"at {position}: KL {kl} > {bound}" + _hold_deep(expected, got, bound, cosine, "Deep", f"at {position}", record) -def check_deep_prompt(model, rows: int, bound: float, chars: int = 1024, *, record): +def check_deep_prompt( + model, rows: int, bound: float, chars: int = 1024, *, cosine=False, record +): """The deepest prompt chunk (holding position ``max_seq_len - 1``), at most ``rows`` rows of it, against the graph's own reference within - ``bound`` of its KL, every earlier position filled as - ``check_deep_decode`` fills them. + ``bound`` of its KL (or ``cosine``, as ``check_deep_decode`` holds it), + every earlier position filled as ``check_deep_decode`` fills them. """ tokens = prompt(model, chars, 1) model.logits(tokens) @@ -343,17 +347,41 @@ def check_deep_prompt(model, rows: int, bound: float, chars: int = 1024, *, reco record("DeepPromptSeconds", time.perf_counter() - start) got = got.numpy() expected, _ = model.reference(ids, **values) + _hold_deep( + expected, got, bound, cosine, "DeepPrompt", f"the chunk at {begin}", record + ) + + +def _hold_deep( + expected, got, bound: float, cosine: bool, name: str, where: str, record +): + """``got``'s row against the reference's ``expected``: KL within + ``bound``, or with ``cosine`` a cosine of at least ``bound``; recorded + under ``name``. + """ + if cosine: + e = np.asarray(expected, np.float64).reshape(-1) + g = np.asarray(got, np.float64).reshape(-1)[: e.size] + value = float(e @ g / (np.linalg.norm(e) * np.linalg.norm(g))) + record(f"{name}Cosine", value) + assert value >= bound, f"{where}: cosine {value} < {bound}" + return kl, top1 = divergence(expected, got) - record("DeepPromptKL", kl) - record("DeepPromptTop1", top1) - assert kl <= bound, f"the chunk at {begin}: KL {kl} > {bound}" + record(f"{name}KL", kl) + record(f"{name}Top1", top1) + assert kl <= bound, f"{where}: KL {kl} > {bound}" @contextmanager def _fill_caches(model, n: int, position: int): """The caches' first ``n`` rows, a prompt's, repeated up to ``position`` on the device and as the graph reference's state, for the body of the - ``with``; the prompt's own rows stay as it wrote them. + ``with``; the prompt's own rows stay as it wrote them. A ring holds the + ``ring`` positions before ``position`` so filled; past ``ring``, a + prompt's positions in it are as it left them. A state checkpointed at + each chunk boundary (``Checkpoints``) holds, in the slot a call at + ``position`` reads, the state the prompt left; the reference reads a + copy of its slots. """ version = next(iter(model.versions.values())) caches = [(x, p) for _, x, p, _ in walk(model) if isinstance(p, KVCache)] @@ -370,9 +398,17 @@ def _fill_caches(model, n: int, position: int): rows[q % p.ring] = rows[p.ring + q % p.ring] = held[q % n] version.push(cache) cache.host = rows + checkpoints = [c for c in model.recurrent if isinstance(c, Checkpoints)] + for c in checkpoints: + slots = version.read(c.slots) + slots[c.slot(position)] = slots[c.slot(n)] + version.push(c.slots) + # A step may read and write one slot: the reference, run after the + # device's step, reads it as it was, not as the device left it. + c.slots.host = slots.copy() try: yield finally: - for state in (x for x, _ in caches): + for state in [*(x for x, _ in caches), *(c.slots for c in checkpoints)]: version.push(state) state.host = None diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index 813625a50..4ca4ea1f6 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -11,6 +11,7 @@ from .elementwise_add import ElementwiseAdd, RowwiseAdd from .elementwise_mul import ElementwiseMul, RowwiseMul from .emit import Emit +from .experts import Experts from .gelu import GELU from .gemm import GEMM from .gemv import GEMV @@ -24,7 +25,9 @@ from .resample.op import PatchPositions, Resample, ResampleTaps from .rms_norm import RMSNorm from .rope import RoPE +from .router import TopKRouter from .sample import Sample +from .selective_scan import SelectiveScan from .sigmoid import Sigmoid from .silu import SiLU from .softmax import Softmax diff --git a/iron/operators/experts.py b/iron/operators/experts.py new file mode 100644 index 000000000..38a47e001 --- /dev/null +++ b/iron/operators/experts.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experts: a mixture of experts' feed-forwards, each row through the +experts ``TopKRouter`` chose for it. +""" + +import numpy as np +from aie.dialects.aie import T +from aie.extras.dialects import arith +from aie.helpers.taplib import TensorAccessPattern +from aie.iron import In, ObjectFifo, Out, TaskGroup, Worker +from aie.iron.controlflow import range_ +from aie.iron.kernels import moe +from aie.utils.verify import Tolerance +from aie.utils.compile.jit.params import param, per_call +from ml_dtypes import bfloat16 + +from iron.common import Built, Operator +from iron.common.testing import Case, Testing + +#: The rows and columns of a weight tile, so that at Qwen3's shapes a +#: shim BD repeats an expert's row blocks at most 64 times. +TILE = 32 + + +def _routed(op: "Experts") -> dict: + """Each row's ``top_k`` distinct experts, and weights at a layer's scale.""" + rng = np.random.default_rng(7) + ids = [rng.choice(op.experts, op.top_k, replace=False) for _ in range(op.rows)] + gate_up = rng.standard_normal(op.gate_up.shape) / np.sqrt(op.hidden) + down = rng.standard_normal(op.down.shape) / np.sqrt(op.width) + return dict( + given={ + op.ids: np.array(ids, np.int32), + op.gate_up: gate_up.astype(bfloat16), + op.down: down.astype(bfloat16), + }, + normal=(op.x, op.weights), + ) + + +class Experts(Operator): + """Row ``r`` of ``x``: the sum over ``j < top_k`` of ``weights[r, 2 * j]`` + times expert ``e = ids[r, j]``'s gated feed-forward, ``down[e] @ + (silu(g) * u)`` where ``g`` and ``u`` are the halves of ``gate_up[e] @ + x``, in float32, rounded once. ``ids`` and ``weights`` are + ``TopKRouter``'s; ``gate_up`` and ``down`` stack the experts as + transformers holds them. + + One core, scalar, reading two streams: a row's ``x``, ``weights`` and + ``ids``, which a memtile joins, and every expert's ``gate_up`` then + ``down`` in ``TILE`` x ``TILE`` tiles, streamed for each row. The + kernel skips the tiles of experts the row did not choose. + """ + + test = Testing( + [ + Case(dict(rows=2, hidden=256, width=64, experts=4, top_k=2)), + Case(dict(rows=8, hidden=512, width=256, experts=8, top_k=2), bench=True), + ], + draw=_routed, + ) + + rows: int = param(array=True) + hidden: int = param() + width: int = param(array=True) + experts: int = param(array=True) + top_k: int = param() + weights_width: int = param(derive=lambda op: 2 * op.top_k, repr=False) + gate_up_width: int = param(derive=lambda op: 2 * op.width, repr=False) + tile: int = param(derive=lambda op: TILE * TILE, repr=False) + + x = In(rows, hidden, obj_shape=(hidden,)) + ids = In(rows, top_k, dtype=np.int32, obj_shape=(top_k,)) + weights = In(rows, weights_width, obj_shape=(weights_width,)) + gate_up = In(experts, gate_up_width, hidden, obj_shape=(tile,)) + down = In(experts, hidden, width, obj_shape=(tile,)) + y = Out(rows, hidden, obj_shape=(hidden,)) + # Every row is still computed under a bound. + valid = per_call(bounds=rows) + + def kernel(self): + return moe.experts( + hidden=self.hidden, + width=self.width, + num_experts=self.experts, + top_k=self.top_k, + rows=TILE, + cols=TILE, + ) + + def build(self) -> Built: + H, K = self.hidden, self.top_k + steps = 3 * self.width * H // self.tile + row = ObjectFifo(np.ndarray[(H + 4 * K,), np.dtype[bfloat16]], depth=1) + parts = row.prod().join( + [0, H, H + 2 * K], + obj_types=[o.obj_type for o in (self.x, self.weights, self.ids)], + ) + for o, f in zip((self.x, self.weights, self.ids), parts): + o.bind(f.prod()) + tiles = ObjectFifo(self.gate_up.obj_type) + self.gate_up.bind(tiles.prod()) + self.down.bind(tiles.prod()) + of_y = ObjectFifo(self.y.obj_type) + self.y.bind(of_y.cons()) + + def core_fn(of_row, of_tile, of_y, experts): + for _ in range_(self.rows): + row, y = of_row.acquire(1), of_y.acquire(1) + for step in range_(self.experts * steps): + tile = of_tile.acquire(1) + experts(row, tile, y, arith.index_cast(step, to=T.i32())) + of_tile.release(1) + of_row.release(1) + of_y.release(1) + + worker = Worker(core_fn, [row.cons(), tiles.cons(), of_y.prod(), self.kernel()]) + return Built([worker]) + + def sequence(self, rt): + """A row at a time: its inputs, every expert's tiles, its output.""" + H, I = self.hidden, self.width + + def tiles(e, rows, cols): + """Expert ``e``'s ``rows`` x ``cols`` matrix, tile by tile.""" + return TensorAccessPattern( + (self.experts, rows, cols), + e * rows * cols, + [rows // TILE, cols // TILE, TILE, TILE], + [TILE * cols, TILE, cols, 1], + ) + + for r in range(self.rows): + tg = TaskGroup() + for operand in (self.x, self.weights, self.ids): + operand.fill(operand[r : r + 1], group=tg) + for e in range(self.experts): + expert = TaskGroup() + self.gate_up.fill(tiles(e, 2 * I, H), group=expert) + self.down.fill(tiles(e, H, I), group=expert, wait=True) + expert.finish() + self.y.drain(self.y[r : r + 1], group=tg, wait=True) + tg.finish() + + def tolerance(self): + return Tolerance.relative( + 8e-3, range_frac=1e-3, note="f32 throughout, rounded once: a bf16 ulp" + ) + + def ops(self) -> int: + return 6 * self.rows * self.top_k * self.width * self.hidden + + def reference(self, x, ids, weights, gate_up, down): + x, weights = (np.asarray(a, np.float32) for a in (x, weights)) + y = np.zeros(x.shape, np.float32) + for r, chosen in enumerate(np.asarray(ids)): + for j, e in enumerate(chosen): + g, u = np.split(gate_up[e].astype(np.float32) @ x[r], 2) + h = g / (1 + np.exp(-g)) * u + y[r] += weights[r, 2 * j] * (down[e].astype(np.float32) @ h) + return y.astype(bfloat16) diff --git a/iron/operators/router.py b/iron/operators/router.py new file mode 100644 index 000000000..5827be222 --- /dev/null +++ b/iron/operators/router.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""TopKRouter: each row's ``top_k`` experts and their weights, on the device. + +Bit for bit as ``aie.iron.kernels.moe.router_top_k_ref``: the exact top +``top_k`` of a row of bf16 router logits, the lowest indices winning a tie, +and a softmax over the chosen logits. A select core reduces each row to a +``sample_select`` summary and a router core turns it into the ids and the +weights, a pair of cores for each lane of rows. +""" + +import numpy as np +from aie.iron import Buffer, In, ObjectFifo, Out, Worker, ceildiv +from aie.iron.controlflow import range_ +from aie.iron.kernels import moe, sample +from aie.utils.verify import Tolerance +from aie.utils.compile.jit.params import param, per_call, tunable, widths +from ml_dtypes import bfloat16 + +from iron.common import Built, Operator, Rtp, Unresolvable +from iron.common.testing import Case, Testing + + +def _logits(op: "TopKRouter") -> dict: + """Rows of logits, each with ties straddling its ``top_k``-th.""" + rng = np.random.default_rng(13) + logits = rng.normal(0, 3, (op.rows, op.experts)).astype(bfloat16) + for row in logits: + boundary = np.sort(row.astype(np.float32))[::-1][op.top_k - 1] + row[rng.choice(op.experts, min(4, op.experts), replace=False)] = boundary + return dict(given={op.logits: logits}) + + +class TopKRouter(Operator): + """Each of ``rows`` rows of ``experts`` router logits to its ``top_k`` + expert ids and weights: ``weights`` holds a row's ``i``-th weight at + ``2 * i``, a zero after each. + """ + + test = Testing( + [ + Case(dict(experts=32, top_k=8), id="one_row"), + Case(dict(experts=64, top_k=6, rows=16, pairs=4), id="rows", lower=True), + Case(dict(experts=32, top_k=8, rows=256, pairs=8), id="chunk", bench=True), + ], + draw=_logits, + ) + + experts: int = param() + top_k: int = param() + rows: int = param(default=1) + weights_width: int = param(derive=lambda op: 2 * op.top_k, repr=False) + pairs: int = tunable(domain=widths) + + logits = In(rows, experts, obj_shape=(experts,), per=(pairs,)) + ids = Out(rows, top_k, dtype=np.int32, obj_shape=(top_k,), per=(pairs,)) + weights = Out(rows, weights_width, obj_shape=(weights_width,), per=(pairs,)) + valid = per_call(bounds=rows) + count = per_call(derive=lambda op: ceildiv(op.valid, op.pairs)) + + def validate(self) -> None: + if not 1 <= self.top_k <= self.experts: + raise ValueError( + f"TopKRouter: top_k ({self.top_k}) must be from 1 to experts " + f"({self.experts})" + ) + + def resolve(self, dev): + pairs = self.pairs or max( + p for p in range(1, dev.cols + 1) if self.rows % p == 0 + ) + if self.rows % pairs or pairs > dev.cols: + raise Unresolvable( + f"TopKRouter: {self.rows} rows over {pairs} pairs of cores, on a " + f"device of {dev.cols} columns" + ) + return self.replace(pairs=pairs) + + def compatible(self) -> None: + config = self.layout() + config.select() + config.route() + + def select(self): + return sample.sample_select( + slice_size=self.experts, chunk=self.experts, k_max=self.top_k + ) + + def route(self): + return moe.router_top_k(slice_size=self.experts, top_k=self.top_k) + + def build(self) -> Built: + select, route = self.select(), self.route() + words = sample.summary_words(self.experts, self.top_k) + summary_ty = np.ndarray[(words,), np.dtype[np.int32]] + + def select_body(of_x, of_sum, kernel, row, state, rtp): + n = rtp.read()[self.count] + for _ in range_(n): + x, summary = of_x.acquire(1), of_sum.acquire(1) + kernel(x, row, state, summary) + of_x.release(1) + of_sum.release(1) + + def route_body(of_sum, of_ids, of_weights, kernel, rtp): + n = rtp.read()[self.count] + for _ in range_(n): + summary = of_sum.acquire(1) + ids, weights = of_ids.acquire(1), of_weights.acquire(1) + kernel(summary, ids, weights) + of_sum.release(1) + of_ids.release(1) + of_weights.release(1) + + workers, rtps, lanes = [], [], ([], [], []) + for _ in range(self.pairs): + of_x = ObjectFifo(self.logits.obj_type, depth=self.logits.depth) + of_ids = ObjectFifo(self.ids.obj_type, depth=self.ids.depth) + of_weights = ObjectFifo(self.weights.obj_type, depth=self.weights.depth) + of_sum = ObjectFifo(summary_ty, depth=2) + for lane, end in zip( + lanes, (of_x.prod(), of_ids.cons(), of_weights.cons()) + ): + lane.append(end) + # The select kernel reads only the row's top-k: always this one's. + row = Buffer( + np.ndarray[(sample.ROW_WORDS,), np.dtype[np.int32]], + initial_value=sample.draw_row(1.0, self.top_k, 0), + ) + state = Buffer( + np.ndarray[(sample.SELECT_STATE_WORDS,), np.dtype[np.int32]], + initial_value=np.zeros(sample.SELECT_STATE_WORDS, np.int32), + ) + pair = [Rtp([self.count]), Rtp([self.count])] + rtps += pair + workers += [ + Worker( + select_body, + [of_x.cons(), of_sum.prod(), select, row, state, pair[0]], + ), + Worker( + route_body, + [ + of_sum.cons(), + of_ids.prod(), + of_weights.prod(), + route, + pair[1], + ], + ), + ] + for operand, ends in zip((self.logits, self.ids, self.weights), lanes): + operand.bind(ends) + return Built(workers, rtps) + + def ops(self) -> int: + return 0 # a selection and a small softmax: its figure is latency + + def reference(self, logits): + """``(ids, weights)``, ``router_top_k_ref`` row by row.""" + logits = np.asarray(logits).reshape(self.rows, self.experts) + ids, weights = zip(*(moe.router_top_k_ref(row, self.top_k) for row in logits)) + return np.stack(ids), np.stack(weights) + + def tolerance(self) -> Tolerance: + return Tolerance.exact() diff --git a/iron/operators/selective_scan.py b/iron/operators/selective_scan.py new file mode 100644 index 000000000..54931b7ae --- /dev/null +++ b/iron/operators/selective_scan.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""SelectiveScan: Mamba2's state space recurrence over a run of rows.""" + +from typing import ClassVar + +import numpy as np +from aie.dialects import vector +from aie.extras import types +from aie.extras.dialects import arith +from aie.helpers.taplib import TensorAccessPattern +from aie.iron import Buffer, In, ObjectFifo, Out, TaskGroup, Worker, kernels +from aie.iron.controlflow import range_ +from aie.utils.verify import Tolerance +from aie.utils.compile.jit.context import compile_context +from aie.utils.compile.utils import resolve_target_arch +from aie.utils.compile.jit.params import param, per_call, tunable, widths +from ml_dtypes import bfloat16 + +from iron.common import Built, Operator, Rtp +from iron.common.testing import Case, Testing +from iron.operators.state_slices import SLICE, depths, resolve_heads, resolve_rows + + +def _slices(op) -> int: + """A core's slices: its heads' ``SLICE``-row runs of state.""" + return op.heads // op.columns * op.head_dim // SLICE + + +def _trained(op: "SelectiveScan") -> dict: + """Heads as a trained Mamba2's are: ``A < 0``, so the state decays.""" + rng = np.random.default_rng(3) + H = op.heads + head = np.stack( + [rng.uniform(-4, 0, H), -rng.uniform(0.5, 16, H), rng.uniform(0, 2, H)], 1 + ) + xbc = rng.standard_normal((op.rows, op.width)) * 0.5 + return dict( + given={op.xbc: xbc.astype(bfloat16), op.head: head.astype(np.float32)}, + normal=(op.dt, op.state), + ) + + +class SelectiveScan(Operator): + """Each row ``t`` of ``xbc``, split into ``x`` ``(heads, head_dim)`` + and ``groups`` rows each of ``B`` and ``C`` ``(state_size,)``, updates + each head's state ``(head_dim, state_size)`` and reads it out: + + dt = softplus(dt[t] + dt_bias) + state = exp(dt * A) * state + dt * outer(x, B) + y[t] = state @ C + D * x + + head ``h`` reading group ``h // (heads // groups)``'s ``B`` and ``C`` + and its ``(dt_bias, A, D)`` from ``head``. ``state`` is the state + before the first row; ``final`` the state after the last valid one. + Float32 throughout, ``y`` rounded once. + + Each of ``columns`` cores takes whole heads, ``SLICE`` rows of a head's + state at a time, and gathers a block of rows' ``x``, ``B``, ``C`` and + ``dt`` for the slice from whole rows of ``xbc``, which one memtile + joins to each row's ``dt`` and every core reads; another joins each + slice's state to its head's ``head``, so a core reads two streams. + The kernel keeps the state from block to block. Scalar. + """ + + test = Testing( + [ + Case(dict(rows=5, heads=2, head_dim=16, state_size=16, columns=1)), + Case(dict(rows=8, heads=4, head_dim=32, state_size=32, groups=2)), + # Granite 4.0-H's decode: a one-row block over a 128-wide state, + # its B and C off the 64-byte line an even block keeps them on. + Case( + dict(rows=1, heads=2, head_dim=16, state_size=128, columns=1), + id="decode", + ), + Case( + dict(rows=3, heads=2, head_dim=16, state_size=64, columns=1), + id="odd_rows", + extensive=True, + ), + Case( + dict(rows=64, heads=16, head_dim=64, state_size=64, groups=2), + bench=True, + ), + # Too many shim BDs at once without reclaiming them. + Case( + dict(rows=2048, heads=8, head_dim=64, state_size=16), + id="reclaim", + extensive=True, + ), + ], + draw=_trained, + ) + + # A long prefill holds more shim BDs than a tile has. + aiecc_flags: ClassVar[tuple[str, ...]] = ("--reclaim-runtime-bds",) + + rows: int = param(array=True) + heads: int = param(array=True) + head_dim: int = param(array=True) + state_size: int = param(array=True) + groups: int = param(default=1, array=True) + columns: int = tunable(domain=widths) + inner: int = param(derive=lambda op: op.heads * op.head_dim, repr=False, array=True) + width: int = param( + derive=lambda op: op.inner + 2 * op.groups * op.state_size, repr=False + ) + # The rows of y one kernel call writes: whole SLICE-row runs that fit L1. + block: int = tunable(repr=False) + slice_state: int = param(derive=lambda op: SLICE * op.state_size, repr=False) + + xbc = In(rows, width, obj_shape=(width,)) + dt = In(rows, heads, obj_shape=(heads,)) + head = In(heads, 3, dtype=np.float32, obj_shape=(3,), per=(columns,)) + state = In( + heads, + head_dim, + state_size, + dtype=np.float32, + obj_shape=(slice_state,), + per=(columns,), + depth=1, + ) + y = Out(rows, inner, obj_shape=(block,), per=(columns,)) + final = Out( + heads, + head_dim, + state_size, + dtype=np.float32, + obj_shape=(slice_state,), + per=(columns,), + depth=1, + ) + valid = per_call(bounds=rows) + count = per_call(derive=lambda op: op.valid) + + def validate(self) -> None: + if self.head_dim % SLICE or self.heads % self.groups or self.heads % 2: + raise ValueError( + f"SelectiveScan: head_dim ({self.head_dim}) must be whole " + f"{SLICE}-row slices, and heads ({self.heads}) whole groups " + f"({self.groups}) and even, so a row's dt after its xbc is " + "whole 4-byte words" + ) + + def resolve(self, dev): + return resolve_rows(resolve_heads(self, dev, self.heads), dev) + + def compatible(self) -> None: + if self.block <= 0 or self.block % SLICE: + raise ValueError( + f"SelectiveScan: block ({self.block}) must be a positive " + f"multiple of {SLICE}" + ) + if self.rows % (self.block // SLICE): + raise ValueError( + f"SelectiveScan: block ({self.block}) must hold a row count " + f"dividing rows ({self.rows})" + ) + + def slice_bytes(self) -> tuple[int, int]: + """Bytes of a slice's (dt_bias, A, D) then state in, and its final out.""" + return 4 * (3 + self.slice_state), 4 * self.slice_state + + def held(self, dev, rows: int) -> int: + """Return the bytes a core of ``dev`` uses at ``rows`` rows per call. + + That is the input row, the block buffer, y, the kernel's data and + stack (its own copy of the state included), and the device's default + stack, which covers the core's own frame. One slice's state and final + state are not counted. + """ + # Only measured, never built, so it stays out of the build's registry. + with compile_context(_iron_kernel_registry=set()): + kernel = self.kernel(rows) + block = 2 * int(np.prod(kernel.arg_shape(0))) + frame = kernel.frame_bytes(resolve_target_arch(dev)) + frame += dev.default_core_stack_bytes + return 2 * (self.width + self.heads) + block + 4 * rows * SLICE + frame + + def kernel(self, rows: int): + return kernels.ssm.selective_scan( + rows=rows, + head_dim=SLICE, + state=self.state_size, + blocks=self.rows // rows, + ) + + def build(self) -> Built: + H, P, N, I = self.heads, self.head_dim, self.state_size, self.inner + T, per_group, heads = self.block // SLICE, H // self.groups, H // self.columns + kernel = self.kernel(T) + f32 = lambda n: np.ndarray[(n,), np.dtype[np.float32]] + bf16 = lambda n: np.ndarray[(n,), np.dtype[bfloat16]] + deep_in, deep_out = depths(self) + # A core takes two streams, a memtile joining each: a row's xbc + # then its dt, one stream every core reads, and a slice's (dt_bias, + # A, D) then its state. + row = ObjectFifo(bf16(self.width + H), depth=1) + parts = row.prod().join( + [0, self.width], obj_types=[self.xbc.obj_type, self.dt.obj_type] + ) + for o, f in zip((self.xbc, self.dt), parts): + o.bind([f.prod()]) + starts, ends = [], ([], []) + for _ in range(self.columns): + start = ObjectFifo(f32(3 + SLICE * N), depth=deep_in) + part = start.prod().join( + [0, 3], obj_types=[self.head.obj_type, self.state.obj_type] + ) + starts.append(start) + for lane, f in zip(ends, part): + lane.append(f.prod()) + for o, lane in zip((self.head, self.state), ends): + o.bind(lane) + of_ys = [ObjectFifo(self.y.obj_type) for _ in range(self.columns)] + of_finals = [ + ObjectFifo(self.final.obj_type, depth=deep_out) for _ in range(self.columns) + ] + self.y.bind([f.cons() for f in of_ys]) + self.final.bind([f.cons() for f in of_finals]) + rtps = [Rtp([self.count]) for _ in range(self.columns)] + + def core_fn(of_row, of_slice, of_y, of_final, scan, block, rtp, first): + valid = rtp.read()[self.count] + i32 = valid.type + zero, size = arith.constant(0, i32), arith.constant(T, i32) + bc, dt = T * SLICE, T * (SLICE + 2 * N) + for q in range_(_slices(self) // heads): + for r in range_(heads): + # Slice q of head r of the core's heads in turn, q and r + # each their own loop counter rather than a flattened + # index's div/mod (AIE2's backend cannot legalize the + # bitreverse Peano's mid-end folds that div/mod into). + h = r + first + p = q * SLICE + h * P + g = (h // per_group) * N + I + start, final = of_slice.acquire(1), of_final.acquire(1) + for j in range_(self.rows // T): + # Gather the block: x, then B and C, then dt, a + # vector of SLICE at a time; xbc first keeps each + # 32-byte aligned. + for t in range_(T): + row = of_row.acquire(1) + block[dt + t] = row[self.width + h] + lane = types.vector(SLICE, types.bf16()) + runs = [(t * SLICE, p)] + [ + ( + bc + t * 2 * N + c * N + n, + g + c * self.groups * N + n, + ) + for c in range(2) + for n in range(0, N, SLICE) + ] + for to, at in runs: + vector.store( + vector.load(lane, row, [at], alignment=32), + block.op, + [to], + alignment=32, + ) + of_row.release(1) + b = arith.index_cast(j, to=i32) + left = valid - b * size + live = arith.minsi(arith.maxsi(left, zero), size) + y = of_y.acquire(1) + scan(block, start, final, y, live, b) + of_y.release(1) + of_slice.release(1) + of_final.release(1) + + workers = [ + Worker( + core_fn, + [ + row.cons(), + starts[k].cons(), + of_ys[k].prod(), + of_finals[k].prod(), + kernel, + Buffer(kernel.arg_types()[0]), + rtps[k], + k * heads, + ], + ) + for k in range(self.columns) + ] + return Built(workers, rtps) + + def sequence(self, rt): + """Every row of ``xbc`` and ``dt`` once a slice, for every core; per + core, its heads' parameters and states, and its slices of ``y`` and + ``final``. + """ + R, H, W, I = self.rows, self.heads, self.width, self.inner + P, SN = self.head_dim, SLICE * self.state_size + S, heads = _slices(self), H // self.columns + per, states, run = S // heads, H * P * self.state_size, S * SN + # A core's heads' states, a slice of each head in turn. + slices = lambda k: TensorAccessPattern( + (states,), k * run, [per, heads, SN], [SN, P * self.state_size, 1] + ) + tg = TaskGroup() + self.xbc.fill(TensorAccessPattern((R, W), 0, [S, R, W], [0, W, 1]), group=tg) + self.dt.fill(TensorAccessPattern((R, H), 0, [S, R, H], [0, H, 1]), group=tg) + for k in range(self.columns): + self.head.fill( + TensorAccessPattern((H, 3), 3 * k * heads, [per, heads, 3], [0, 3, 1]), + lane=k, + group=tg, + ) + self.state.fill(slices(k), lane=k, group=tg) + for k in range(self.columns): + self.y.drain( + TensorAccessPattern( + (R, I), k * heads * P, [per, heads, R, SLICE], [SLICE, P, I, 1] + ), + lane=k, + group=tg, + wait=True, + ) + self.final.drain(slices(k), lane=k, group=tg, wait=True) + tg.finish() + + def tolerance(self): + return Tolerance.relative( + 8e-3, range_frac=1e-3, note="f32 throughout, y rounded once: a bf16 ulp" + ) + + def reference(self, xbc, dt, head, state): + H, P, N, G = self.heads, self.head_dim, self.state_size, self.groups + xbc, dt = np.asarray(xbc, np.float32), np.asarray(dt, np.float32) + dt_bias, A, D = np.asarray(head, np.float32).T + state = np.array(state, np.float32) + y = np.zeros((self.rows, self.inner), np.float32) + for t in range(self.valid): + x = xbc[t, : self.inner].reshape(H, P) + B, C = ( + np.repeat(m.reshape(G, N), H // G, axis=0) + for m in np.split(xbc[t, self.inner :], 2) + ) + step = np.logaddexp(0, dt[t] + dt_bias)[:, None, None] + state = np.exp(step * A[:, None, None]) * state + state += step * x[:, :, None] * B[:, None, :] + y[t] = (np.einsum("hpn,hn->hp", state, C) + D[:, None] * x).reshape(-1) + return y.astype(bfloat16), state diff --git a/iron/operators/state_slices.py b/iron/operators/state_slices.py new file mode 100644 index 000000000..d2a874acd --- /dev/null +++ b/iron/operators/state_slices.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""What SelectiveScan and GatedDeltaRule share: a head's state taken +``SLICE`` at a time, a block of rows a kernel call, and the cores resolved +over whole heads. +""" + +from iron.common import Unresolvable + +#: The rows (SelectiveScan) or columns (GatedDeltaRule) of a head's state a +#: kernel call carries. +SLICE = 16 + + +def block_rows(op) -> int: + """The rows a kernel call takes: the most up to 32 that divide + ``op.rows``, a bound on the block buffer's share of a core's memory. + """ + return max(d for d in range(1, 33) if op.rows % d == 0) + + +def resolve_rows(op, dev): + """Return ``op`` with its ``block`` set. + + ``block`` is the most rows, up to ``block_rows``, that divide + ``op.rows`` and leave room in a core for ``op.held`` plus one slice's + state and final state. + + Raises: + Unresolvable: Not even one row fits. + """ + if op.block is not None: + return op + state, final = op.slice_bytes() + mem = dev.core_memory_bytes + rows = next( + ( + t + for t in range(block_rows(op), 0, -1) + if op.rows % t == 0 and op.held(dev, t) + state + final <= mem + ), + None, + ) + if rows is None: + raise Unresolvable( + f"{type(op).__name__}: no block of rows fits a core's {mem} bytes" + ) + return op.replace(block=rows * SLICE) + + +def depths(op) -> tuple[int, int]: + """How many slices' states and finals ``op``'s core holds at once: two + of each, the next streaming during the kernel, where its memory fits. + """ + held = op.held(op.dev, op.block // SLICE) + state, final = op.slice_bytes() + mem = op.dev.core_memory_bytes + deep_in = 2 if held + 2 * state + final <= mem else 1 + deep_out = 2 if held + deep_in * state + 2 * final <= mem else 1 + return deep_in, deep_out + + +def resolve_heads(op, dev, heads: int): + """``op`` on as many cores as divide its ``heads`` and the shim's input + channels feed. + + Raises: + Unresolvable: No count divides the heads. + """ + columns = op.resolve_columns(dev, op.columns, fits=lambda c: heads % c == 0) + if heads % columns: + raise Unresolvable(f"{type(op).__name__}: {heads} heads over {columns} cores") + return op.replace(columns=columns) diff --git a/iron/tests/common/test_config.py b/iron/tests/common/test_config.py index c222ea666..f9e30c93d 100644 --- a/iron/tests/common/test_config.py +++ b/iron/tests/common/test_config.py @@ -16,6 +16,7 @@ from iron.applications.embedding_gemma2.model import Config as EmbeddingGemmaConfig from iron.applications.gemma4.audio import AudioConfig from iron.applications.gemma4.vision import VisionConfig +from iron.applications.granite.model import GraniteConfig from iron.applications.llama.model import Llama3RopeScaling, LlamaConfig from iron.lm import DecoderConfig from iron.tests.support.llama_model import REAL_1B @@ -52,6 +53,7 @@ def test_from_checkpoint_takes_either_the_file_or_its_directory(): "cls", [ LlamaConfig, + GraniteConfig, EmbeddingGemmaConfig, AudioConfig, VisionConfig, diff --git a/iron/tests/lm/test_generation.py b/iron/tests/lm/test_generation.py new file mode 100644 index 000000000..00995dd2f --- /dev/null +++ b/iron/tests/lm/test_generation.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Generation's EOS exits and rate accounting over a device loop, without an NPU.""" + +import math +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest + +from iron.lm.model import Model +from iron.lm.session import device_draws + + +@pytest.mark.parametrize("each", [False, True], ids=["generate", "stream"]) +@pytest.mark.parametrize( + "eos", [10, 12, None], ids=["first-eos", "later-eos", "no-eos"] +) +@pytest.mark.parametrize("max_tokens", [1, 6]) +def test_device_generation_stops_only_before_undrawn_tokens( + monkeypatch, each, eos, max_tokens +): + clock = [0.0] + monkeypatch.setattr("iron.lm.model.time.perf_counter", lambda: clock[0]) + + def start(*args, **kwargs): + clock[0] += 1.0 + + def steps(n): + for _ in range(n): + clock[0] += 2.0 + yield None + + loop = SimpleNamespace(start=Mock(side_effect=start), steps=Mock(side_effect=steps)) + draws = np.arange(10, 16, dtype=np.int32) + + class DeviceModel(Model): + draws_at_once = True + + def _draw(self, tokens, max_tokens, sampler, each, **inputs): + chunks = [((tokens,), {"position": len(tokens) - 1})] + yield from device_draws(self, chunks, len(tokens), max_tokens, each) + + model = DeviceModel() + model.eos = set() if eos is None else {eos} + model.prompt_chunk = object() + model.drawn = object() + model.decode_step = SimpleNamespace(read=lambda _: draws) + model.loop = Mock(return_value=loop) + tokens = ( + list(model.stream([1], max_tokens)) if each else model.generate([1], max_tokens) + ) + end = max_tokens if eos is None else min(max_tokens, eos - 10 + 1) + assert tokens == list(draws[:end]) + ran = end - 1 if each or eos == 10 else max_tokens - 1 + assert clock[0] == 1.0 + 2.0 * ran + if eos == 10: + loop.steps.assert_not_called() + else: + loop.steps.assert_called_once_with(max_tokens - 1) + assert model.stats.ttft == 1.0 + if ran: + assert model.stats.tokens_per_s == 0.5 + else: + assert math.isnan(model.stats.tokens_per_s) diff --git a/iron/tests/lm/test_modules.py b/iron/tests/lm/test_modules.py index 92b3d238f..2f4e9eca0 100644 --- a/iron/tests/lm/test_modules.py +++ b/iron/tests/lm/test_modules.py @@ -16,6 +16,7 @@ from iron import Precision from iron.lm import Attention, Embedding, LayerNorm, Linear, RMSNorm from iron.lm.blocks import KVCache, Rope +from iron.lm.causal import CausalConv1d from iron.lm.checkpoint import nest, leaves from iron.lm.golden import HuggingFace from iron.lm.testing import assert_close, check_generation @@ -28,6 +29,21 @@ pytestmark = pytest.mark.usefixtures("npu2") +@pytest.mark.parametrize("chunk,kernel", [(1, 3), (2, 4), (2, 5)]) +def test_causal_conv_refuses_a_chunk_shorter_than_its_history(chunk, kernel): + config = dataclasses.replace(SMALL, prefill_chunk=chunk) + with pytest.raises(ValueError, match="prefill_chunk .* kernel - 1"): + CausalConv1d(config, channels=64, kernel=kernel) + + +@pytest.mark.parametrize("chunk,kernel", [(1, 2), (2, 3), (4, 4), (4, 5), (8, 4)]) +def test_causal_conv_accepts_a_chunk_that_holds_its_history(chunk, kernel): + config = dataclasses.replace(SMALL, prefill_chunk=chunk) + conv = CausalConv1d(config, channels=64, kernel=kernel) + assert conv.state.chunk == chunk + assert conv.state.slots.shape[1:] == (kernel - 1, 64) + + class Block(iron.Graph): def __init__(self): self.embed = Embedding(512, 256) diff --git a/iron/tests/operators/test_core_streams.py b/iron/tests/operators/test_core_streams.py new file mode 100644 index 000000000..e572d5b5a --- /dev/null +++ b/iron/tests/operators/test_core_streams.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Experts' array, built device-free: every operand bound, and no core +reading more streams than its two input DMA channels. +""" + +import inspect + +import pytest +from aie.helpers.util import flatten_fn_args +from aie.iron import Worker, kernels +from aie.iron.dataflow.objectfifo import ObjectFifoHandle + +from iron.operators.experts import Experts + +# The kernel and the parameters its build() passes it, which an older +# mlir-aie's lacks (blocks=, from 7807032). Kept in step with that call by +# hand: a parameter added there belongs here too. +KERNELS = { + Experts: ( + "moe", + "experts", + {"hidden", "width", "num_experts", "top_k", "rows", "cols"}, + ), +} + +CASES = [(cls, case) for cls in (Experts,) for case in cls.test.cases] + + +@pytest.mark.parametrize( + "cls,case", CASES, ids=[f"{c.__name__}-{i}" for i, (c, _) in enumerate(CASES)] +) +def test_each_core_reads_at_most_two_streams(cls, case, npu2): + module, name, params = KERNELS[cls] + kernel = getattr(getattr(kernels, module, None), name, None) + if kernel is None or params - inspect.signature(kernel).parameters.keys(): + pytest.skip(f"mlir-aie has no kernels.{module}.{name} taking {sorted(params)}") + view = cls(**case.kwargs).resolved(npu2).layout() + built = view.build() + for operand in (*view.inputs, *view.outputs): + assert operand.handles + for worker in (o for o in built.objects if isinstance(o, Worker)): + reads = [ + a + for a in flatten_fn_args(worker.fn_args) + if isinstance(a, ObjectFifoHandle) and not a._is_prod + ] + assert len(reads) <= 2 diff --git a/iron/tests/operators/test_rejected_shapes.py b/iron/tests/operators/test_rejected_shapes.py index b34bd27f0..83bb73c24 100644 --- a/iron/tests/operators/test_rejected_shapes.py +++ b/iron/tests/operators/test_rejected_shapes.py @@ -41,12 +41,31 @@ from iron.operators.resample.op import PatchPositions from iron.operators.rms_norm import RMSNorm from iron.operators.rope import RoPE +from iron.operators.router import TopKRouter from iron.operators.sample import Sample +from iron.operators.selective_scan import SelectiveScan from iron.operators.softmax import Softmax pytestmark = pytest.mark.usefixtures("npu2") +@pytest.mark.parametrize("block", [0, -16, 1, 17, 32, 96]) +@pytest.mark.parametrize("columns", [None, 1]) +def test_selective_scan_rejects_invalid_blocks(block, columns): + with pytest.raises(ValueError, match="block"): + SelectiveScan( + rows=5, heads=2, head_dim=16, state_size=16, columns=columns, block=block + ).resolved() + + +@pytest.mark.parametrize("rows,block", [(5, 16), (5, 80), (8, 32), (8, 128)]) +def test_selective_scan_accepts_blocks_dividing_rows(rows, block): + op = SelectiveScan( + rows=rows, heads=2, head_dim=16, state_size=16, columns=1, block=block + ) + assert op.resolved().block == block + + def test_transfer_size_not_dividing_the_per_channel_share_is_rejected(): """A BD shorter than the ObjectFifo object hangs the device. @@ -569,6 +588,17 @@ def test_flm_decode_layer_rejects_unknown_configurations(kwargs, match): DecodeLayer(**kwargs) +def test_top_k_router_refuses_more_experts_than_it_routes_between(): + with pytest.raises(ValueError, match="top_k \\(9\\) must be from 1"): + TopKRouter(experts=8, top_k=9) + + +@pytest.mark.parametrize("rows, pairs", [(6, 4), (16, 9)]) +def test_top_k_router_refuses_rows_its_pairs_do_not_split(rows, pairs, npu2): + with pytest.raises(Unresolvable, match="pairs of cores"): + TopKRouter(experts=32, top_k=8, rows=rows, pairs=pairs).resolved(npu2) + + @pytest.mark.parametrize( "make, match", [ diff --git a/iron/tests/support/qwen3_moe.py b/iron/tests/support/qwen3_moe.py new file mode 100644 index 000000000..0942983db --- /dev/null +++ b/iron/tests/support/qwen3_moe.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-MoE at a toy shape, and its checkpoint's tensors, for its reference +and device tests. +""" + +from iron.applications.qwen3_moe.model import Qwen3MoeConfig + +from .llama_model import shapes + +#: Qwen3-MoE's proportions at toy size: GQA, q and k normed, eight experts +#: of which a token takes two, an untied head. +TOY = Qwen3MoeConfig( + vocab_size=1024, + hidden_size=512, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=128, + intermediate_size=1024, + rope_theta=1000000.0, + max_seq_len=256, + prefill_chunk=256, + num_experts=8, + num_experts_per_tok=2, + moe_intermediate_size=256, + norm_topk_prob=True, +) + + +def names(c: Qwen3MoeConfig) -> dict[str, tuple[int, ...]]: + """Qwen3-MoE's tensors: Llama's, q and k normed, each layer's MLP its + router and experts, and an untied head. + """ + out = shapes(c) + out["lm_head.weight"] = (c.vocab_size, c.hidden_size) + E, M, D = c.hidden_size, c.moe_intermediate_size, c.head_dim + assert D is not None, "DecoderConfig derives head_dim" + for i in range(c.num_hidden_layers): + layer = f"model.layers.{i}" + out |= { + f"{layer}.self_attn.q_norm.weight": (D,), + f"{layer}.self_attn.k_norm.weight": (D,), + f"{layer}.mlp.gate.weight": (c.num_experts, E), + } + for p in ("gate_proj", "up_proj", "down_proj"): + del out[f"{layer}.mlp.{p}.weight"] + for j in range(c.num_experts): + expert = f"{layer}.mlp.experts.{j}" + out |= { + f"{expert}.gate_proj.weight": (M, E), + f"{expert}.up_proj.weight": (M, E), + f"{expert}.down_proj.weight": (E, M), + } + return out diff --git a/tests/applications/decoders/test.py b/tests/applications/decoders/test.py index cf30ede60..8bc8f2d3c 100644 --- a/tests/applications/decoders/test.py +++ b/tests/applications/decoders/test.py @@ -16,6 +16,7 @@ import pytest from iron.applications.gemma4 import Gemma4 +from iron.applications.qwen3_moe import Qwen3Moe from iron.lm.testing import ( check_accuracy, check_chat_turn, @@ -25,10 +26,19 @@ weights_dir, ) +INDEX = "model.safetensors.index.json" + # Each decoder's class, checkpoint file, and the lengths of the long and the # short prompt it generates from. DECODERS = [ pytest.param((Gemma4, "gemma-4-e2b-it/model.safetensors", (256, 13)), id="gemma4"), + # 30B, its experts paged from the checkpoint: nightly only, a toy one per PR + # (tests/applications/qwen3_moe). + pytest.param( + (Qwen3Moe, f"qwen3-30b-a3b/{INDEX}", (1024, 13)), + id="qwen3_moe", + marks=pytest.mark.extensive, + ), ] pytestmark = pytest.mark.supported_devices("npu2") @@ -40,7 +50,10 @@ def decoder(request): runtime is released before the next decoder loads. """ cls, weights, prompts = request.param - load = load_from_flags(request, cls, weights_dir(weights)) + # CI holds every checkpoint here but the 30B's. + load = load_from_flags( + request, cls, weights_dir(weights), staged=cls is not Qwen3Moe + ) model = load() yield model, prompts model.release() @@ -49,6 +62,10 @@ def decoder(request): # Gemma 4 E2B on npu2 measured mean 0.0055, P90 0.0086, max 0.100 over 40 steps. MAX_KL = {"Mean": 0.03, "P90": 0.08, "Max": 0.2} +# bf16 transformers itself reaches p90 0.065 on Qwen3-MoE's steps: near-tie +# experts. Pre-#222's own figure; its golden is still missing here, so no +# run has re-measured it against the current numeric fixes. +MOE_KL = {**MAX_KL, "P90": 0.1} @pytest.mark.bench @@ -67,11 +84,14 @@ def test_generation(decoder, prompt, num_tokens, record_property): @pytest.mark.extensive @pytest.mark.bench def test_accuracy(decoder, record_property): - check_accuracy(decoder[0], MAX_KL, record=record_property) + bounds = MOE_KL if isinstance(decoder[0], Qwen3Moe) else MAX_KL + check_accuracy(decoder[0], bounds, record=record_property) @pytest.mark.extensive def test_determinism(decoder, record_property): + if isinstance(decoder[0], Qwen3Moe): + pytest.skip("15 min at 30B; test_host_steps holds a host step's calls exact") check_determinism(decoder[0], record=record_property) @@ -88,7 +108,8 @@ def test_accuracy_across_chunks(decoder, record_property): """ model = decoder[0] assert len(model.prompt(LONG)) > model.config.prefill_chunk - check_accuracy(model, MAX_KL, 3, LONG, record=record_property) + bounds = MOE_KL if isinstance(model, Qwen3Moe) else MAX_KL + check_accuracy(model, bounds, 3, LONG, record=record_property) @pytest.mark.extensive diff --git a/tests/applications/granite/reference/goldens/toy.npz b/tests/applications/granite/reference/goldens/toy.npz new file mode 100644 index 000000000..8f0460002 Binary files /dev/null and b/tests/applications/granite/reference/goldens/toy.npz differ diff --git a/tests/applications/granite/reference/test.py b/tests/applications/granite/reference/test.py new file mode 100644 index 000000000..df02452f1 --- /dev/null +++ b/tests/applications/granite/reference/test.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite's graph against Hugging Face's Granite in float32 over the same +random weights at a toy shape, its logits recorded in ``goldens/``, +device-free: the shared decoder's tests (``Decoder``) with Granite's four +scalars, none at its default, and an untied head. +""" + +from iron.applications.granite.model import Granite, GraniteConfig +from iron.tests.support.llama_model import shapes +from iron.tests.support.llama_reference import Decoder + +#: Granite's proportions at toy size, its scalars Granite 3's own but the +#: attention's, which at 1/64 of a 64-wide head would leave the toy's +#: attention near uniform. +TOY = GraniteConfig( + vocab_size=1024, + hidden_size=512, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=2, + intermediate_size=1024, + rope_theta=10000.0, + max_seq_len=256, + prefill_chunk=256, + tie_word_embeddings=False, + attention_multiplier=0.0625, + residual_multiplier=0.22, + embedding_multiplier=12.0, + logits_scaling=8.0, +) + + +class TestGranite(Decoder): + model, config, model_type, name = Granite, TOY, "granite", "toy" + seed = 0 + + @classmethod + def names(cls): + c = cls.config + return shapes(c) | {"lm_head.weight": (c.vocab_size, c.hidden_size)} diff --git a/tests/applications/granite/test.py b/tests/applications/granite/test.py new file mode 100644 index 000000000..d57625d26 --- /dev/null +++ b/tests/applications/granite/test.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.2 3B on the NPU, from the real checkpoint: speed, accuracy +against the float32 reference, determinism, the device loop, and past one +chunk: a prompt of two and a chat turn (the end of the caches is +``long_context``'s). The model is compiled and loaded once for the module +and every test calls it in-process. + +``--random-weights SEED`` runs it on drawn weights and random prompts, on +a host without the checkpoint, and ``--max-seq-len`` at another context. +""" + +import aie.utils as aie_utils +import pytest + +from iron.applications.granite import Granite +from iron.lm.testing import ( + check_accuracy, + check_chat_turn, + check_determinism, + check_device_sampling, + check_generation, + load_from_flags, + weights_dir, +) + +WEIGHTS = weights_dir("granite-4.2-3b") / "model.safetensors.index.json" + +pytestmark = pytest.mark.supported_devices("npu2") + + +@pytest.fixture(scope="module") +def load(request): + return load_from_flags(request, Granite, WEIGHTS) + + +@pytest.fixture(scope="module") +def model(load): + """The model, compiled and loaded once as a full ELF: the arena's banks + keep its ~10 GB of weights and caches within aiebu's 32-bit relocation + addend. + """ + loaded = load() + assert loaded.device_loop + yield loaded + loaded.release() + + +# KL(fp32 CPU || NPU) over 40 steps measured mean 0.0031, p90 0.0062, max 0.0152. +MAX_KL = {"Mean": 0.01, "P90": 0.02, "Max": 0.05} + +GENERATION = pytest.mark.parametrize( + "prompt_len,num_tokens", + [ + pytest.param( + p, + n, + # The shortest per PR; the rest nightly. + marks=[ + pytest.mark.bench, + *([] if (p, n) == (13, 1) else [pytest.mark.extensive]), + ], + ) + for p in (1024, 13) + for n in (40, 1) + ], + ids=[f"granite_4_2_3b_prompt_{p}_tokens_{n}" for p in (1024, 13) for n in (40, 1)], +) + + +@pytest.mark.extensive +@pytest.mark.supported_devices("npu1", "npu2") +class TestEachStep: + """Every step its own xclbin dispatch, the host starting each on the + token the last one drew. Its model is built and dropped with the class, + first: after any other test it would be held beside the module's + full-ELF one. + """ + + @pytest.fixture(scope="class") + def model(self, load): + model = load(full_elf=False) + assert not model.device_loop + yield model + model.release() + if aie_utils.DefaultNPURuntime is not None: + aie_utils.DefaultNPURuntime.cleanup() + + @GENERATION + def test_granite_4_2_3b_each_step( + self, model, prompt_len, num_tokens, record_property + ): + check_generation(model, prompt_len, num_tokens, record=record_property) + + def test_granite_4_2_3b_each_step_accuracy(self, model, record_property): + check_accuracy(model, MAX_KL, 20, 512, record=record_property) + + def test_granite_4_2_3b_each_step_determinism(self, model, record_property): + check_determinism(model, 4, 3, 128, record=record_property) + + +@GENERATION +def test_granite_4_2_3b(model, prompt_len, num_tokens, record_property): + check_generation(model, prompt_len, num_tokens, record=record_property) + + +# The device draws every token and starts every decode step itself; from the +# same seed its text is the host loop's. The figures are the device loop's. +@pytest.mark.extensive +@pytest.mark.bench +def test_granite_4_2_3b_device_loop(model, record_property): + check_device_sampling(model, 1024, 100, record=record_property) + + +@pytest.mark.extensive +@pytest.mark.bench +def test_granite_4_2_3b_accuracy(model, record_property): + check_accuracy(model, MAX_KL, record=record_property) + + +@pytest.mark.extensive +def test_granite_4_2_3b_determinism(model, record_property): + check_determinism(model, record=record_property) + + +# 12000 characters of prompt.txt are a full chunk and most of a second for +# Granite's tokenizer too: it attends over the first's caches. +LONG = 12000 + + +@pytest.mark.extensive +def test_granite_4_2_3b_accuracy_across_chunks(model, record_property): + """The prompt's two chunks and two decode steps after them, against the + float32 reference over the whole prompt (each step a forward over the + chunks on the host, so only a few). + """ + assert len(model.prompt(LONG)) > model.config.prefill_chunk + check_accuracy(model, MAX_KL, 3, LONG, record=record_property) + + +@pytest.mark.extensive +def test_granite_4_2_3b_chat_turn(model, record_property): + """A turn of 1000 tokens reruns the second chunk alone.""" + check_chat_turn(model, LONG, 1000, record=record_property) diff --git a/tests/applications/granite_hybrid/reference/goldens/dense.npz b/tests/applications/granite_hybrid/reference/goldens/dense.npz new file mode 100644 index 000000000..36de553f5 Binary files /dev/null and b/tests/applications/granite_hybrid/reference/goldens/dense.npz differ diff --git a/tests/applications/granite_hybrid/reference/test.py b/tests/applications/granite_hybrid/reference/test.py new file mode 100644 index 000000000..31023fdcf --- /dev/null +++ b/tests/applications/granite_hybrid/reference/test.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.0-H 1B's graph against Hugging Face's GraniteMoeHybrid in +float32 over the same random weights at a toy shape, its logits recorded +in ``goldens/``, device-free: the shared decoder's tests (``Decoder``), +whose prompts span several of the toy's 64-row chunks, so the Mamba2 +layers' states carry across chunk boundaries and into decode steps. This +package ships the dense sizes only, so the toy has no experts. +""" + +import dataclasses +import json +import subprocess +import sys + +import numpy as np +import pytest + +from iron.applications.granite_hybrid.model import GraniteHybrid, GraniteHybridConfig +from iron.tests.support.llama_reference import Decoder + +#: Granite 4.0-H 1B's proportions at toy size: three Mamba2 layers of two +#: groups and one attention layer without RoPE, a shared MLP and Granite's +#: scalars; its head untied, as Granite's toy, so the logits are not the +#: embedding of the last token. +TOY = GraniteHybridConfig( + vocab_size=1024, + hidden_size=512, + num_hidden_layers=4, + num_attention_heads=8, + num_key_value_heads=2, + intermediate_size=256, + max_seq_len=256, + prefill_chunk=64, + attention_multiplier=0.0625, + residual_multiplier=0.22, + embedding_multiplier=12.0, + logits_scaling=8.0, + tie_word_embeddings=False, + layer_types=("mamba", "attention", "mamba", "mamba"), + position_embedding_type="nope", + shared_intermediate_size=512, + num_local_experts=0, + num_experts_per_tok=2, + mamba_n_heads=16, + mamba_d_head=64, + mamba_d_state=64, + mamba_n_groups=2, +) + + +@pytest.fixture(scope="module") +def golden(golden): + """The toy's golden, its weights first redrawn in place where the + shared draw would leave a layer inert: the time steps and decays as + Mamba2 initializes them, so a state lasts many positions; and the + convolution's bias ten times the draw's, at its taps' scale. + """ + + def recorded(name, build, config, weights): + rng = np.random.default_rng(0) + for layer in weights.model.layers: + if "mamba" in vars(layer): + m, H = layer.mamba, config.mamba_n_heads + m.A_log[:] = np.log(rng.uniform(1, 16, H)) + dt = np.exp(rng.uniform(np.log(1e-3), np.log(1e-1), H)) + m.dt_bias[:] = dt + np.log(-np.expm1(-dt)) + m.conv1d.bias[:] *= 10 + return golden(name, build, config, weights) + + return recorded + + +class TestGraniteHybridDense(Decoder): + """Granite 4.0-H 1B's kind: no experts, the shared MLP alone.""" + + model, config, model_type, name = GraniteHybrid, TOY, "granitemoehybrid", "dense" + recorded = dict(prefill_chunk=256, mamba_n_heads=64) + json = {"mamba_expand": 2} + # No router to tie; at seed 32 bf16 leaves float32 at step 5. + seed = 2 + + @classmethod + def names(cls): + c = cls.config + E, D = c.hidden_size, c.head_dim + S = c.shared_intermediate_size + H, N = c.mamba_n_heads, c.mamba_d_state + I = H * c.mamba_d_head + channels = I + 2 * c.mamba_n_groups * N + out = { + "model.embed_tokens.weight": (c.vocab_size, E), + "model.norm.weight": (E,), + "lm_head.weight": (c.vocab_size, E), + } + for i in range(c.num_hidden_layers): + p = f"model.layers.{i}" + out |= { + f"{p}.input_layernorm.weight": (E,), + f"{p}.post_attention_layernorm.weight": (E,), + f"{p}.shared_mlp.input_linear.weight": (2 * S, E), + f"{p}.shared_mlp.output_linear.weight": (E, S), + } + if c.attends(i): + Q, KV = c.num_attention_heads * D, c.num_key_value_heads * D + out |= { + f"{p}.self_attn.q_proj.weight": (Q, E), + f"{p}.self_attn.k_proj.weight": (KV, E), + f"{p}.self_attn.v_proj.weight": (KV, E), + f"{p}.self_attn.o_proj.weight": (E, Q), + } + else: + m = f"{p}.mamba" + out |= { + f"{m}.in_proj.weight": (I + channels + H, E), + f"{m}.conv1d.weight": (channels, 1, c.mamba_d_conv), + f"{m}.conv1d.bias": (channels,), + f"{m}.dt_bias": (H,), + f"{m}.A_log": (H,), + f"{m}.D": (H,), + f"{m}.norm.weight": (I,), + f"{m}.out_proj.weight": (E, I), + } + return out + + @pytest.mark.parametrize("position,start", [(0, 0), (64, 64), (70, 64), (100, 100)]) + def test_forget_restarts_from_a_valid_recurrent_state(self, cpu, position, start): + """A non-boundary forget must not decode from the later state.""" + model = self.host(cpu.config, cpu.weights) + tokens = np.random.default_rng(3).integers(0, cpu.config.vocab_size, 101) + model.logits(tokens[:100]) + model.forget(position) + assert model.session._seen.size == start + resumed = model.logits(tokens[: position + 1]) + model.forget(0) + if start: + model.logits(tokens[:start]) + scratch = model.logits(tokens[: position + 1]) + np.testing.assert_array_equal(resumed, scratch) + + +def test_model_cli_help(): + result = subprocess.run( + [sys.executable, "-m", "iron.applications.granite_hybrid.model", "--help"], + capture_output=True, + text=True, + check=True, + timeout=30, + ) + assert "GraniteHybrid on the NPU" in result.stdout + assert "--random-weights" in result.stdout + assert "path" in result.stdout + + +def test_a_config_without_rope_reads_back_from_its_json(tmp_path): + (tmp_path / "config.json").write_text(json.dumps(TOY.to_json())) + assert GraniteHybridConfig.from_checkpoint(tmp_path).to_json() == TOY.to_json() + + +def test_a_rope_config_without_rope_theta_takes_transformers_default(tmp_path): + config = TOY.to_json() | {"position_embedding_type": "rope"} + (tmp_path / "config.json").write_text(json.dumps(config)) + assert GraniteHybridConfig.from_checkpoint(tmp_path).rope_theta == 10000.0 + + +def test_layer_types_short_of_the_layers_are_refused(): + with pytest.raises(ValueError, match="3 layer_types for 4 layers"): + dataclasses.replace(TOY, layer_types=TOY.layer_types[:3]) + + +def test_a_config_with_experts_is_refused(): + with pytest.raises(ValueError, match="num_local_experts 4"): + dataclasses.replace(TOY, num_local_experts=4) diff --git a/tests/applications/granite_hybrid/test.py b/tests/applications/granite_hybrid/test.py new file mode 100644 index 000000000..a152ae233 --- /dev/null +++ b/tests/applications/granite_hybrid/test.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Granite 4.0-H 1B on the NPU, from the real checkpoint: speed, accuracy +against the float32 reference, determinism, and past one chunk: a prompt of +two and a chat turn, which restarts the Mamba2 states from a chunk +boundary's checkpoint. The model is compiled and loaded once for the +module and every test calls it in-process. + +``--random-weights SEED`` runs it on drawn weights and random prompts, on +a host without the checkpoint, and ``--max-seq-len`` at another context. +""" + +import aie.utils as aie_utils +import pytest + +from iron.applications.granite_hybrid import GraniteHybrid +from iron.lm.generation import Sampler +from iron.lm.testing import ( + check_accuracy, + check_chat_turn, + check_determinism, + check_generation, + load_from_flags, + weights_dir, +) + +WEIGHTS = weights_dir("granite-4.0-h-1b") / "model.safetensors" + +pytestmark = pytest.mark.supported_devices("npu2") + + +@pytest.fixture(scope="module") +def load(request): + return load_from_flags(request, GraniteHybrid, WEIGHTS) + + +@pytest.fixture(scope="module") +def model(load): + loaded = load() + yield loaded + loaded.release() + + +# KL(fp32 CPU || NPU) measured mean 0.0024, p90 0.0048, max 0.0063 over 40 +# steps after 292 tokens; over two chunks and two decode steps mean 0.0027, +# max 0.0049. +MAX_KL = {"Mean": 0.006, "P90": 0.012, "Max": 0.015} + + +GENERATION = pytest.mark.parametrize( + "prompt_len,num_tokens", + [ + pytest.param( + p, + n, + # The shortest per PR; the rest nightly. + marks=[ + pytest.mark.bench, + *([] if (p, n) == (13, 1) else [pytest.mark.extensive]), + ], + ) + for p in (1024, 13) + for n in (40, 1) + ], + ids=[ + f"granite_4_0_h_1b_prompt_{p}_tokens_{n}" for p in (1024, 13) for n in (40, 1) + ], +) + + +@pytest.mark.extensive +@pytest.mark.supported_devices("npu1", "npu2") +class TestEachStep: + """Every step its own xclbin dispatch, the host starting each on the + token the last one drew. Its model is built and dropped with the class, + first: after any other test it would be held beside the module's + full-ELF one. + """ + + @pytest.fixture(scope="class") + def model(self, load): + model = load(full_elf=False) + assert not model.device_loop + yield model + model.release() + if aie_utils.DefaultNPURuntime is not None: + aie_utils.DefaultNPURuntime.cleanup() + + @GENERATION + def test_granite_4_0_h_1b_each_step( + self, model, prompt_len, num_tokens, record_property + ): + check_generation(model, prompt_len, num_tokens, record=record_property) + + def test_granite_4_0_h_1b_each_step_accuracy(self, model, record_property): + check_accuracy(model, MAX_KL, 20, 512, record=record_property) + + def test_granite_4_0_h_1b_each_step_determinism(self, model, record_property): + check_determinism(model, 4, 3, 128, record=record_property) + + +@GENERATION +def test_granite_4_0_h_1b(model, prompt_len, num_tokens, record_property): + check_generation(model, prompt_len, num_tokens, record=record_property) + + +@pytest.mark.extensive +def test_granite_4_0_h_1b_rate_counts_the_draws_past_eos(model, record_property): + """After 13 characters the model draws eos 8th. generate's device loop + draws all 40 tokens before reading them, so its rate counts all 40; + stream reads each as it is drawn and counts those through eos. The two + rates agree. + """ + if model.random_weights is not None: + pytest.skip("Early EOS is specific to the trained checkpoint") + tokens = model.prompt(13) + model.forget(0) + drawn = model.generate(tokens, 40, Sampler()) + assert len(drawn) < 40 and drawn[-1] in model.eos + generated = model.stats.tokens_per_s + model.forget(0) + assert list(model.stream(tokens, 40, Sampler())) == drawn + streamed = model.stats.tokens_per_s + record_property("GenerateTPS", generated) + record_property("StreamTPS", streamed) + assert abs(generated / streamed - 1) < 0.25, f"{generated} vs {streamed} tokens/s" + + +@pytest.mark.extensive +@pytest.mark.bench +def test_granite_4_0_h_1b_accuracy(model, record_property): + check_accuracy(model, MAX_KL, record=record_property) + + +@pytest.mark.extensive +def test_granite_4_0_h_1b_determinism(model, record_property): + check_determinism(model, record=record_property) + + +# 12000 characters of prompt.txt are a full chunk and most of a second, +# which scans on from the first's states. +LONG = 12000 + + +@pytest.mark.extensive +def test_granite_4_0_h_1b_accuracy_across_chunks(model, record_property): + """The prompt's two chunks and two decode steps after them, against the + float32 reference over the whole prompt. + """ + assert len(model.prompt(LONG)) > model.config.prefill_chunk + check_accuracy(model, MAX_KL, 3, LONG, record=record_property) + + +@pytest.mark.extensive +def test_granite_4_0_h_1b_chat_turn(model, record_property): + """A turn of 1000 tokens reruns the second chunk alone, from the states + the first left at its boundary. + """ + check_chat_turn(model, LONG, 1000, record=record_property) diff --git a/tests/applications/long_context/test.py b/tests/applications/long_context/test.py index 47c036981..1a1362c6d 100644 --- a/tests/applications/long_context/test.py +++ b/tests/applications/long_context/test.py @@ -21,7 +21,13 @@ import iron.applications from iron.applications.gemma4 import Gemma4 +from iron.applications.granite import Granite +from iron.applications.granite_hybrid import GraniteHybrid from iron.applications.llama import Llama +from iron.applications.qwen3 import Qwen3 +from iron.applications.qwen3_embedding import Qwen3Embedding +from iron.applications.qwen3_moe import Qwen3Moe +from iron.applications.qwen3_reranker import Qwen3Reranker from iron.lm import CausalLM, Checkpoint from iron.lm.testing import ( check_deep_decode, @@ -36,6 +42,9 @@ # decode measured 0.002 to 0.015 just past the prompt, 0.009 to 0.033 at # 16k and 0.030 to 0.038 at 32k: it grows with the keys the step sums over. DEEP_KL = 0.1 +# Qwen3-Embedding's row is its final hidden state, held by its cosine to the +# reference, not its KL. +COSINE = 0.97 @dataclasses.dataclass(frozen=True) @@ -43,7 +52,9 @@ class Application: """A decoder application, a checkpoint its test runs and the forms it is held to here: the full ELF on NPU2 where ``full_elf``, and an xclbin a step on the ``each_step`` devices. CI holds every checkpoint, so one - missing there fails rather than skips. + missing there fails rather than skips, but where not ``staged``. A model + whose row is a hidden state, not logits, is held deep by its cosine, at + least ``deep_cosine``. """ model: type[CausalLM] @@ -53,6 +64,8 @@ class Application: full_elf: bool = True each_step: tuple[str, ...] = () deep_kl: float = DEEP_KL + deep_cosine: float | None = None + staged: bool = True @property def path(self) -> Path: @@ -72,6 +85,13 @@ def path(self) -> Path: "llama_3.2_3b": Application( Llama, "llama3.2-3b", INDEX, "tokenizer.json", each_step=("npu1", "npu2") ), + # Granite's decode measured 0.0010 to 0.0022 on npu2, 4k to the whole + # context; npu1 still unmeasured. Each size fills the cache with its + # own content up to max_seq_len - 1, not a continuation of the last, + # so sizes aren't directly comparable. + "granite_4.2_3b": Application( + Granite, "granite-4.2-3b", INDEX, "tokenizer.json", each_step=("npu1", "npu2") + ), # Gemma 4 E2B's decode measured 0.0002 to 0.072 on npu2 and 0.0009 to # 0.0256 on npu1, each 4k to the whole context: each size fills the # cache with its own content up to max_seq_len - 1, not a continuation @@ -83,6 +103,56 @@ def path(self) -> Path: "tokenizer.json", each_step=("npu1", "npu2"), ), + # Granite 4.0-H 1B's decode measured 0.0007 to 0.0069 on npu2, 4k to + # the whole context; its deepest prompt chunk 4.1e-5 to 0.0052; npu1 + # still unmeasured. + "granite_4.0_h_1b": Application( + GraniteHybrid, + "granite-4.0-h-1b", + "model.safetensors", + "tokenizer.json", + each_step=("npu1", "npu2"), + ), + # Qwen3's decode measured 0.0020 (0.6B, native 32768) to 0.0385 (1.7B, + # native 32768), its deepest prompt chunk 0.0021 to 0.0565, on npu2. + "qwen3_0.6b": Application( + Qwen3, "qwen3-0.6b", "model.safetensors", "tokenizer.json", staged=False + ), + "qwen3_1.7b": Application( + Qwen3, "qwen3-1.7b", INDEX, "tokenizer.json", each_step=("npu1", "npu2") + ), + # Qwen3-Embedding's row is its final hidden state, held by its cosine: + # measured 0.98611 at its decode step at 32767 and 0.99859 over its + # deepest prompt chunk, the bound twice the decode's distance from 1. + "qwen3_embedding_0.6b": Application( + Qwen3Embedding, + "qwen3-embedding-0.6b", + "model.safetensors", + "tokenizer.json", + deep_cosine=COSINE, + ), + # Qwen3-Reranker's decode measured 0.0212 at 32767, its deepest prompt + # chunk 0.0038, well under DEEP_KL: it uses Qwen3's own head, not a + # cosine-held hidden state. + "qwen3_reranker_0.6b": Application( + Qwen3Reranker, + "qwen3-reranker-0.6b", + "model.safetensors", + "tokenizer.json", + staged=False, + ), + # 30B, its experts paged from the checkpoint; staged as decoders stages + # it. Its decode measured 0.069 at 40959 (top-1 swapped at a near tie, + # pre-#222's numeric fixes, not yet re-verified), its deepest prompt + # chunk 1.5e-5: twice the decode's, rounded up. + "qwen3_moe_30b_a3b": Application( + Qwen3Moe, + "qwen3-30b-a3b", + INDEX, + "tokenizer.json", + staged=False, + deep_kl=0.2, + ), } @@ -109,7 +179,7 @@ def test_nothing_grows_with_the_context(application, version): row per position: its size is affine in the context. """ app = APPLICATIONS[application] - require(app.path, app.path.with_name("config.json")) + require(app.path, app.path.with_name("config.json"), staged=app.staged) config = app.model.default_config(app.path) weights = Checkpoint(app.path).tree("unread") C = config.prefill_chunk @@ -177,14 +247,18 @@ def loaded(request): """ form, context = request.param app = APPLICATIONS[form.removesuffix("_each_step")] - load = load_from_flags(request, app.model, app.path, app.tokenizer) + load = load_from_flags( + request, app.model, app.path, app.tokenizer, staged=app.staged + ) fitted = app.model.fit(app.model.default_config(app.path), app.path).max_seq_len if context is not None and context >= fitted: pytest.skip(f"{app.directory} holds {fitted} positions on this host") each_step = form.endswith("_each_step") sized = {} if context is None else {"max_seq_len": context} model = load(full_elf=False if each_step else None, **sized) - assert model.device_loop is not each_step + # The packaging, not device_loop: a full ELF with a host step (Qwen3-MoE's + # paged experts) is one but does not loop on the device. + assert model.decode_step.plan.full_elf is not each_step yield app, model model.release() if aie_utils.DefaultNPURuntime is not None: @@ -196,7 +270,12 @@ def test_decode_deep_in_the_cache(loaded, record_property): """The deepest step the caches hold, at ``max_seq_len - 1``.""" app, model = loaded position = model.config.max_seq_len - 1 - check_deep_decode(model, position, app.deep_kl, record=record_property) + if app.deep_cosine is None: + check_deep_decode(model, position, app.deep_kl, record=record_property) + else: + check_deep_decode( + model, position, app.deep_cosine, cosine=True, record=record_property + ) @pytest.mark.extensive @@ -206,7 +285,12 @@ def test_prompt_deep_in_the_cache(loaded, record_property): """ app, model = loaded rows = model.config.prefill_chunk - check_deep_prompt(model, rows, app.deep_kl, record=record_property) + if app.deep_cosine is None: + check_deep_prompt(model, rows, app.deep_kl, record=record_property) + else: + check_deep_prompt( + model, rows, app.deep_cosine, cosine=True, record=record_property + ) @pytest.mark.extensive diff --git a/tests/applications/qwen3/reference/goldens/toy.npz b/tests/applications/qwen3/reference/goldens/toy.npz new file mode 100644 index 000000000..cb939abf0 Binary files /dev/null and b/tests/applications/qwen3/reference/goldens/toy.npz differ diff --git a/tests/applications/qwen3/reference/test.py b/tests/applications/qwen3/reference/test.py new file mode 100644 index 000000000..f931c0540 --- /dev/null +++ b/tests/applications/qwen3/reference/test.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3's graph against Hugging Face's Qwen3 in float32 over the same +random weights at a toy shape, its logits recorded in ``goldens/``, +device-free: the shared decoder's tests (``Decoder``) with q and k normed +per head. +""" + +from iron.applications.qwen3.model import Qwen3, Qwen3Config +from iron.tests.support.llama_model import shapes +from iron.tests.support.llama_reference import Decoder + +#: Qwen3's proportions at toy size: GQA, a head_dim of its own, a tied head. +TOY = Qwen3Config( + vocab_size=1024, + hidden_size=512, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=2, + head_dim=128, + intermediate_size=1024, + rope_theta=1000000.0, + max_seq_len=256, + prefill_chunk=256, +) + + +class TestQwen3(Decoder): + model, config, model_type, name = Qwen3, TOY, "qwen3", "toy" + seed = 3 + + @classmethod + def names(cls): + c = cls.config + out = shapes(c) + for i in range(c.num_hidden_layers): + attention = f"model.layers.{i}.self_attn" + out |= { + f"{attention}.q_norm.weight": (c.head_dim,), + f"{attention}.k_norm.weight": (c.head_dim,), + } + return out diff --git a/tests/applications/qwen3/test.py b/tests/applications/qwen3/test.py new file mode 100644 index 000000000..1a1dcd9e6 --- /dev/null +++ b/tests/applications/qwen3/test.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3 on the NPU, 1.7B and 0.6B, from the real checkpoints: speed, +accuracy against the float32 reference, determinism, and past one chunk: a +prompt of two and a chat turn (the end of the caches is ``long_context``'s). +Every test runs on each checkpoint, which its id names; each one's model is +compiled and loaded once and every test calls it in-process. + +``--random-weights SEED`` runs it on weights drawn in the checkpoint's +shapes (only its header is read) and random prompts, on a host without the +tokenizer, and ``--max-seq-len`` at another context. +""" + +import aie.utils as aie_utils +import numpy as np +import pytest +from ml_dtypes import bfloat16 + +from iron.applications.qwen3 import Qwen3 +from iron.lm import Checkpoint +from iron.lm.testing import ( + check_accuracy, + check_chat_turn, + check_determinism, + check_device_sampling, + check_generation, + load_from_flags, + require, + weights_dir, +) + +CHECKPOINTS = { + "qwen3_1.7b": weights_dir("qwen3-1.7b") / "model.safetensors.index.json", + "qwen3_0.6b": weights_dir("qwen3-0.6b") / "model.safetensors", +} + +pytestmark = pytest.mark.supported_devices("npu2") + + +@pytest.fixture(scope="module", params=list(CHECKPOINTS), ids=str) +def load(request): + weights = CHECKPOINTS[request.param] + return load_from_flags(request, Qwen3, weights) + + +@pytest.fixture(scope="module") +def model(load): + """The model, compiled and loaded once; the runtime is released after its + checkpoint's last test. + """ + yield load() + if aie_utils.DefaultNPURuntime is not None: + aie_utils.DefaultNPURuntime.cleanup() + + +# KL(fp32 CPU || NPU) over 20 steps of 1.7B measured mean 0.011, p90 0.033, max 0.055. +MAX_KL = {"Mean": 0.03, "P90": 0.08, "Max": 0.2} + + +GENERATION = pytest.mark.parametrize( + "prompt_len,num_tokens", + [ + pytest.param( + p, + n, + # The shortest per PR; the rest nightly. + marks=[ + pytest.mark.bench, + *([] if (p, n) == (13, 1) else [pytest.mark.extensive]), + ], + ) + for p in (1024, 13) + for n in (40, 1) + ], + ids=[f"prompt_{p}_tokens_{n}" for p in (1024, 13) for n in (40, 1)], +) + + +@pytest.mark.extensive +@pytest.mark.supported_devices("npu1", "npu2") +class TestEachStep: + """The form NPU1 runs: each step its own dispatch of one xclbin, the + host drawing each token. Its model is built and dropped with the class, + first: after any other test it would be held beside the module's + full-ELF one. + """ + + @pytest.fixture(scope="class") + def model(self, load): + model = load(full_elf=False) + assert not model.device_loop + yield model + if aie_utils.DefaultNPURuntime is not None: + aie_utils.DefaultNPURuntime.cleanup() + + def test_qwen3_each_step_accuracy(self, model, record_property): + check_accuracy(model, MAX_KL, 20, 256, record=record_property) + + def test_qwen3_each_step_determinism(self, model, record_property): + check_determinism(model, 4, 3, 128, record=record_property) + + +@GENERATION +def test_qwen3(model, prompt_len, num_tokens, record_property): + check_generation(model, prompt_len, num_tokens, record=record_property) + + +# The device draws every token and starts every decode step itself; from the +# same seed its text is the host loop's. The figures are the device loop's. +@pytest.mark.extensive +@pytest.mark.bench +def test_qwen3_device_loop(model, record_property): + check_device_sampling(model, 1024, 100, record=record_property) + + +@pytest.mark.extensive +@pytest.mark.bench +def test_qwen3_accuracy(model, record_property): + check_accuracy(model, MAX_KL, record=record_property) + + +@pytest.mark.extensive +def test_qwen3_determinism(model, record_property): + check_determinism(model, record=record_property) + + +# 12000 characters of prompt.txt are 3262 tokens (confirmed against Qwen3's +# own tokenizer): a full chunk and most of a second, which attends over the +# first's caches. +LONG = 12000 + + +@pytest.mark.extensive +def test_qwen3_accuracy_across_chunks(model, record_property): + """The prompt's two chunks and two decode steps after them, against the + float32 reference over the whole prompt (each step a forward over 3k+ + tokens on the host, so only a few). + """ + assert len(model.prompt(LONG)) > model.config.prefill_chunk + check_accuracy(model, MAX_KL, 3, LONG, record=record_property) + + +@pytest.mark.extensive +def test_qwen3_chat_turn(model, record_property): + """A turn of 1000 tokens reruns the second chunk alone.""" + check_chat_turn(model, LONG, 1000, record=record_property) + + +@pytest.mark.parametrize("name", list(CHECKPOINTS)) +def test_qwen3_lm_head_is_the_tied_embedding(name): + """The real checkpoint's unread ``lm_head.weight`` is byte-identical to + ``model.embed_tokens.weight``. + """ + weights = CHECKPOINTS[name] + require(weights, weights.with_name("config.json")) + tree = Checkpoint(weights).tree() + assert tree.lm_head.weight.dtype == bfloat16 + assert np.array_equal( + tree.lm_head.weight.view(np.uint16), + tree.model.embed_tokens.weight.view(np.uint16), + ) diff --git a/tests/applications/qwen3_embedding/reference/goldens/toy.npz b/tests/applications/qwen3_embedding/reference/goldens/toy.npz new file mode 100644 index 000000000..2ec6542f1 Binary files /dev/null and b/tests/applications/qwen3_embedding/reference/goldens/toy.npz differ diff --git a/tests/applications/qwen3_embedding/reference/test.py b/tests/applications/qwen3_embedding/reference/test.py new file mode 100644 index 000000000..b791f0201 --- /dev/null +++ b/tests/applications/qwen3_embedding/reference/test.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Embedding's graph against Hugging Face's Qwen3 model in float32 +over Qwen3's toy, its last hidden state recorded in ``goldens/``, +device-free: a query behind its instruction, closed by end of text. +""" + +import numpy as np +import pytest + +from iron.applications.qwen3_embedding.model import Qwen3Embedding +from iron.lm import OnHost +from iron.lm.golden import HuggingFace +from iron.tests.support.llama_model import checkpoint +from tests.applications.qwen3.reference import test as qwen3 +from tests.applications.qwen3_reranker.reference.test import Words + +#: The end of text, within the toy's vocab. +EOS = 5 + + +class LastHidden(HuggingFace): + """transformers' Qwen3 model without its head: the last token's final + hidden state. + """ + + model_class = "AutoModel" + + def logits(self, tokens) -> np.ndarray: + import torch + + with torch.no_grad(): + out = self.model(torch.tensor([list(tokens)]), use_cache=False) + return out.last_hidden_state[0, -1].numpy() + + +class Host(OnHost, Qwen3Embedding): + """The embedder, its calls the graph's reference.""" + + profile = None + end_of_text = EOS + + +@pytest.mark.parametrize("task", [None, "Given a web search query"]) +def test_an_embedding_is_the_last_hidden_state_unit_length(npu2, golden, task): + names = {k.removeprefix("model."): v for k, v in qwen3.TestQwen3.names().items()} + weights = checkpoint(qwen3.TOY, names).tree("random") + reference = golden( + "toy", lambda: LastHidden.build("qwen3", qwen3.TOY, weights), qwen3.TOY, weights + ) + model = Host(qwen3.TOY, weights) + model.tokenizer = type("Tokenizer", (), {"encoding": Words()})() + text = "what lights the north sky at night" + got = model.embed(text, task) + query = text if task is None else f"Instruct: {task}\nQuery:{text}" + want = reference.logits([*Words().encode(query), EOS]) + want = want / np.linalg.norm(want) + assert np.isclose(np.linalg.norm(got), 1) + assert got @ want > 0.999, got @ want diff --git a/tests/applications/qwen3_embedding/test.py b/tests/applications/qwen3_embedding/test.py new file mode 100644 index 000000000..43675a972 --- /dev/null +++ b/tests/applications/qwen3_embedding/test.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Embedding 0.6B on the NPU against Hugging Face's own model in +float32, from one compile: the unit-length last hidden state of a query +behind its instruction, of a document, and of documents of four prompt +chunks and of the native 32768-token context's last, recorded in +``golden.npz`` beside the checkpoint (``pytest --record-goldens``). +""" + +import aie.utils as aie_utils +import numpy as np +import pytest + +from iron.applications.qwen3_embedding import Qwen3Embedding +from iron.lm.model import prompt_text +from iron.lm.testing import require, weights_dir + +DIRECTORY = weights_dir("qwen3-embedding-0.6b") +TASK = "Given a web search query, retrieve relevant passages that answer the query" + +# Each text, the task it is a query of (None: a document), and the least +# cosine its embedding reaches against transformers'. +TEXTS = { + "query": ("What is the capital of China?", TASK, 0.99), + "document": ( + "The capital of China is Beijing, its seat of government.", + None, + 0.99, + ), + # 7625 tokens and the end of text: four prompt chunks. + "four_chunks": (prompt_text(28000), None, 0.99), + # 32682 tokens and the end of text: the native context's last chunk. + "native_context": (prompt_text(120000), None, 0.99), +} + + +@pytest.fixture(scope="module") +def model(): + require(DIRECTORY / "model.safetensors", DIRECTORY / "tokenizer.json") + yield Qwen3Embedding.load(DIRECTORY) + aie_utils.cleanup_npu_runtime() + + +@pytest.mark.supported_devices("npu1", "npu2") +@pytest.mark.parametrize( + "name", + [ + "query", + "document", + pytest.param("four_chunks", marks=pytest.mark.extensive), + pytest.param("native_context", marks=pytest.mark.extensive), + ], +) +def test_qwen3_embedding_0_6b_accuracy(model, name, record_property): + text, task, min_cosine = TEXTS[name] + query = text if task is None else f"Instruct: {task}\nQuery:{text}" + tokens = [*model.tokenizer.encoding.encode(query), model.end_of_text] + assert len(tokens) <= model.config.max_seq_len + got = model.embed(text, task) + want = model.cpu_reference().logits(tokens) + cosine = float(got @ want) + record_property("Tokens", len(tokens)) + record_property("Cosine", cosine) + np.testing.assert_allclose(np.linalg.norm(got), 1, rtol=1e-3) + assert cosine >= min_cosine, cosine diff --git a/tests/applications/qwen3_moe/reference/goldens/toy.npz b/tests/applications/qwen3_moe/reference/goldens/toy.npz new file mode 100644 index 000000000..f5a567250 Binary files /dev/null and b/tests/applications/qwen3_moe/reference/goldens/toy.npz differ diff --git a/tests/applications/qwen3_moe/reference/test.py b/tests/applications/qwen3_moe/reference/test.py new file mode 100644 index 000000000..484baa32d --- /dev/null +++ b/tests/applications/qwen3_moe/reference/test.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-MoE's graph against Hugging Face's Qwen3MoE in float32 over the +same random weights at a toy shape, its logits recorded in ``goldens/``, +device-free: the shared decoder's tests (``Decoder``). +""" + +import dataclasses + +import pytest + +from iron.applications.qwen3_moe.model import Qwen3Moe, Qwen3MoeConfig +from iron.tests.support.llama_reference import Decoder +from iron.tests.support.qwen3_moe import TOY, names + + +class TestQwen3Moe(Decoder): + model, config, model_type, name = Qwen3Moe, TOY, "qwen3_moe", "toy" + recorded = dict(num_experts=64) + # One whose bf16 router picks the experts float32 does at every step. + seed = 0 + + @classmethod + def names(cls): + assert isinstance(cls.config, Qwen3MoeConfig), "Qwen3-MoE's shape" + return names(cls.config) + + +@pytest.mark.parametrize( + "change", + [ + dict(norm_topk_prob=False), + dict(decoder_sparse_step=2), + dict(mlp_only_layers=(1,)), + ], +) +def test_unnormalized_weights_or_a_dense_layer_is_refused(change): + with pytest.raises(ValueError, match="Qwen3MoeConfig"): + dataclasses.replace(TOY, **change) diff --git a/tests/applications/qwen3_moe/test.py b/tests/applications/qwen3_moe/test.py new file mode 100644 index 000000000..039acf016 --- /dev/null +++ b/tests/applications/qwen3_moe/test.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-MoE on the NPU at a toy shape, on drawn weights: a prompt chunk and +decode steps, each call's chosen experts staged by the host between the +router and the experts (``iron.on_host``), held to the graph's own +reference on the host. Needs no checkpoint; the decoders suite runs the +real 30B nightly. +""" + +import dataclasses + +import numpy as np +import pytest + +from iron.applications.qwen3_moe.model import Qwen3Moe +from iron.lm.generation import divergence +from iron.tests.support.llama_model import checkpoint +from iron.tests.support.qwen3_moe import TOY, names + +SEED = 7 +# A chunk's router logits are flm.ProjectionGEMM's (N a multiple of 64), and +# a bounded row of its ids a 16-byte unit: 64 experts, 4 a token, not TOY's 8 and 2. +MOE = dataclasses.replace(TOY, num_experts=64, num_experts_per_tok=4) + +pytestmark = pytest.mark.supported_devices("npu2") + + +def test_paged_experts_follow_the_graph_reference(): + """A 37-token prompt chunk, then three decode steps: each call's next-token + distribution within KL 0.1 of the graph's reference, as long_context's deep + checks hold a decoder to it, its top token the same. + """ + model = Qwen3Moe(MOE, checkpoint(MOE, names(MOE)).tree("random", SEED)) + model.random_weights = SEED + model.load_versions() + reference = model.cpu_reference() + tokens = np.random.default_rng(SEED).integers(0, MOE.vocab_size, 40).tolist() + for n in (37, 38, 39, 40): + kl, top1 = divergence(reference.logits(tokens[:n]), model.logits(tokens[:n])) + assert kl <= 0.1 and top1, f"{n} tokens: KL {kl}, top-1 {top1}" diff --git a/tests/applications/qwen3_reranker/reference/goldens/toy.npz b/tests/applications/qwen3_reranker/reference/goldens/toy.npz new file mode 100644 index 000000000..43b693f2f Binary files /dev/null and b/tests/applications/qwen3_reranker/reference/goldens/toy.npz differ diff --git a/tests/applications/qwen3_reranker/reference/test.py b/tests/applications/qwen3_reranker/reference/test.py new file mode 100644 index 000000000..6bfa2cd0a --- /dev/null +++ b/tests/applications/qwen3_reranker/reference/test.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Reranker's score against Hugging Face's Qwen3 in float32 over +Qwen3's toy, its logits recorded in ``goldens/``, device-free: the prompt +it asks, and "yes" over "no" after it. +""" + +import math +import zlib + +import numpy as np + +from iron.applications.qwen3_reranker.model import PROMPT, Qwen3Reranker +from iron.lm import OnHost +from iron.lm.golden import HuggingFace +from tests.applications.qwen3.reference import test as qwen3 + + +class Words: + """A word a token, by its checksum: a tokenizer within the toy's vocab.""" + + def encode(self, text: str) -> list[int]: + return [zlib.crc32(w.encode()) % qwen3.TOY.vocab_size for w in text.split()] + + +def test_the_score_is_yes_over_no_after_the_prompt(npu2, golden): + weights = qwen3.TestQwen3.weights() + reference = golden( + "toy", + lambda: HuggingFace.build("qwen3", qwen3.TOY, weights), + qwen3.TOY, + weights, + ) + model = type("Host", (OnHost, Qwen3Reranker), {"profile": None})(qwen3.TOY, weights) + model.tokenizer = type("Tokenizer", (), {"encoding": Words()})() + query, document = "what lights the north sky", "charged particles from the sun" + got = model.score(query, document) + text = PROMPT.format( + instruction="Given a web search query, retrieve relevant passages " + "that answer the query", + query=query, + document=document, + ) + logits = reference.logits(Words().encode(text)) + [yes], [no] = Words().encode("yes"), Words().encode("no") + want = 1 / (1 + math.exp(logits[no] - logits[yes])) + assert np.isclose(got, want, atol=1e-2), (got, want) diff --git a/tests/applications/qwen3_reranker/test.py b/tests/applications/qwen3_reranker/test.py new file mode 100644 index 000000000..d58edca0c --- /dev/null +++ b/tests/applications/qwen3_reranker/test.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen3-Reranker 0.6B on the NPU, from the real checkpoint: ``score()`` +over its model card's queries and documents, and over a document of two +prompt chunks, each against the score transformers' float32 logits give +(``golden.npz`` beside the checkpoint), from one compile. +""" + +import math + +import aie.utils as aie_utils +import pytest + +from iron.applications.qwen3_reranker import Qwen3Reranker +from iron.applications.qwen3_reranker.model import INSTRUCTION, PROMPT +from iron.lm.generation import divergence +from iron.lm.model import prompt_text +from iron.lm.testing import load_from_flags, weights_dir + +WEIGHTS = weights_dir("qwen3-reranker-0.6b") / "model.safetensors" + +pytestmark = pytest.mark.supported_devices("npu2") + +# The model card's queries and documents. +CAPITAL = "What is the capital of China?" +GRAVITY = "Explain gravity" +BEIJING = "The capital of China is Beijing." +FORCE = ( + "Gravity is a force that attracts two bodies towards each other. It gives " + "weight to physical objects and is responsible for the movement of planets " + "around the sun." +) +# 12000 characters of prompt.txt: the prompt is two chunks. +LONG = prompt_text(12000) + +PAIRS = { + "capital_beijing": (CAPITAL, BEIJING), + "capital_force": (CAPITAL, FORCE), + "gravity_beijing": (GRAVITY, BEIJING), + "gravity_force": (GRAVITY, FORCE), + "gravity_long": (GRAVITY, LONG), +} + +# The score's error, and its input's: |logit("yes") - logit("no")| apart +# from transformers'. The next-token KL beside them is held to Qwen3's. +MAX_SCORE_ERROR = 0.02 +MAX_MARGIN_ERROR = 0.5 +MAX_KL = 0.2 + + +@pytest.fixture(scope="module") +def model(request): + """The model, compiled and loaded once for every pair.""" + yield load_from_flags(request, Qwen3Reranker, WEIGHTS)() + aie_utils.cleanup_npu_runtime() + + +@pytest.fixture(scope="module") +def reference(model): + return model.cpu_reference() + + +def margin(logits, encode) -> float: + """How far "yes" leads "no" as the answer's first token.""" + [yes], [no] = encode("yes"), encode("no") + return float(logits[yes] - logits[no]) + + +@pytest.mark.parametrize("pair", list(PAIRS)) +def test_qwen3_reranker_0_6b_scores_as_transformers( + model, reference, pair, record_property +): + query, document = PAIRS[pair] + encode = model.tokenizer.encoding.encode + tokens = encode( + PROMPT.format(instruction=INSTRUCTION, query=query, document=document) + ) + want = reference.logits(tokens) + expected = 1 / (1 + math.exp(-margin(want, encode))) + got = model.score(query, document) + logits = model.logits(tokens) + error = abs(margin(logits, encode) - margin(want, encode)) + kl, _ = divergence(want, logits) + record_property("Score", got) + record_property("ExpectedScore", expected) + record_property("MarginError", error) + record_property("KL", kl) + assert abs(got - expected) <= MAX_SCORE_ERROR, f"score {got} against {expected}" + assert error <= MAX_MARGIN_ERROR, f"yes over no {error} from transformers'" + assert kl <= MAX_KL, f"KL {kl} > {MAX_KL}"