Add Qwen3, Qwen3-Embedding, Qwen3-Reranker and Qwen3-MoE on iron.lm

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.
 

iron/applications/granite/README.md

@@ -0,0 +1,41 @@
+<!--
+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
+
+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.

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

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

iron/applications/granite/profiles/npu1.json

@@ -0,0 +1,2 @@
+{"entries": [
+]}

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}
+]}

iron/applications/granite_hybrid/README.md

@@ -0,0 +1,66 @@
+<!--
+SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# 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.

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