diff --git a/kernels/CONVENTION.md b/kernels/CONVENTION.md new file mode 100644 index 00000000..bfa0f568 --- /dev/null +++ b/kernels/CONVENTION.md @@ -0,0 +1,88 @@ +# A convention for NPU kernel source in this repository + +This document proposes how contributed AIE kernel *source* lives in +FastFlowLM. It is deliberately separable from the kernels that come with it: +you can accept or reject the convention on its own terms, and `kernels/granite` +is just the first worked example. + +## What this is, and what it is not + +`kernels/` holds **source**. `src/xclbins/` continues to hold the shipped +binaries, and nothing here changes, moves, re-derives or replaces any of them. + +**No compiled artefact is ever committed under `kernels/`.** Artefacts are build +outputs. They land in the build directory and are installed from there, exactly +as the checked-in ones are. + +## Why source and not a binary + +An `.xclbin` is only valid for the toolchain that built it. A committed binary +carries no record of what produced it and cannot be rebuilt when the toolchain +moves; it rots silently, and the rot is invisible until a user hits it. Source +plus a recorded toolchain fingerprint is the only representation that survives +a version bump. + +That is the whole argument for this directory. The granite kernels are the +occasion, not the point. + +## Layout + +One directory per model family, named to match `src/xclbins//`: + +``` +kernels/ +├── CONVENTION.md this file +├── README.md how to build +├── LICENSE MIT +├── requirements.txt pinned toolchain packages +├── build_kernels.py the only entry point +├── common/ shared host-side helpers +└── / + ├── README.md geometry, measured numbers, what was rejected + ├── geometry.json model dimensions -- so the build needs no weights + ├── iron/ IRON Python: placement and data movement + └── aie/ C++ that runs on the AIE cores +``` + +**`iron/` and `aie/` are split because they are two different review surfaces.** +`aie/` is device C++: it is what actually executes, and it is the surface that +matters for provenance and for correctness. `iron/` is host Python that only +describes where things are placed and how data moves. **Read `aie/` first.** + +## Rules + +**Never a build dependency.** No family may become a dependency of `flm`. With +`FLM_BUILD_KERNELS=OFF` — the default — the build must be byte-identical to one +in which `kernels/` does not exist. + +**The build must not need model weights**, network access, or any path outside +the repository other than the toolchain itself. Shapes come from +`geometry.json`. This is what makes the option safe to enable in a container, +and it is why `geometry.json` exists at all. + +**No generated file is committed.** Generators live in `iron/`; their output +goes to the build directory. + +**The toolchain is declared, not assumed.** Every build writes a +`manifest.json` recording the exact package versions, compiler flags and a +sha256 per artefact. A consumer that finds a fingerprint mismatch must refuse +rather than dispatch a mismatched pair to the NPU. + +**Licence.** Everything under `kernels/` is MIT, matching +`LICENSE_RUNTIME.txt`, with `SPDX-License-Identifier: MIT` in every file. + +## Provenance + +The kernels here were written from the public MLIR-AIE/IRON examples and +published AIE2P documentation. **No shipped `.xclbin` was disassembled and no +closed component was reverse-engineered.** The q4nx container layout was +derived by inverting a published packer and cross-checked against published +model files. (The companion granite PR documents the same derivation in +`src/include/models/granite/q4nx_host.hpp`, if it has landed; this PR does not +depend on it.) + +## What a new family must ship + +A `README.md` with measured numbers against a host reference, a `geometry.json`, +and a `validate` path that can be run on hardware. A kernel with no reference +comparison is not a contribution; it is a claim. diff --git a/kernels/LICENSE b/kernels/LICENSE new file mode 100644 index 00000000..6c452600 --- /dev/null +++ b/kernels/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 the FastFlowLM kernel contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/kernels/README.md b/kernels/README.md new file mode 100644 index 00000000..85379fc3 --- /dev/null +++ b/kernels/README.md @@ -0,0 +1,81 @@ +# NPU kernel source + +Source for AIE kernels, built on request. See [CONVENTION.md](CONVENTION.md) +for what this directory is and the rules it follows. + +**Nothing here is built by default.** `FLM_BUILD_KERNELS` is `OFF`, and with it +off the `flm` build is identical to one in which this directory does not exist. + +## State + +| | | +|---|---| +| kernel source in tree | yes | +| builds from source | yes, `-DFLM_BUILD_KERNELS=ON` | +| validated against a host reference on NPU2 | yes — see `granite/README.md` | +| **dispatched by `flm`** | **not yet** — see *The gap* below | + +## Prerequisites + +The [mlir-aie](https://github.com/Xilinx/mlir-aie) / IRON toolchain, which +supplies the `aie` Python package, the `aiecc` driver and the Peano (LLVM-AIE) +backend. Versions are pinned in [requirements.txt](requirements.txt) and +recorded again in every build's `manifest.json`. + +Activate that environment first — the build does not try to activate it for +you, because activation does not survive a subprocess. + +## Build + +```shell +# with the IRON environment active +python kernels/build_kernels.py --check-toolchain +python kernels/build_kernels.py --family granite --out build/kernels +``` + +or through CMake: + +```shell +cmake -S src -B src/build -DFLM_BUILD_KERNELS=ON +cmake --build src/build --config Release +``` + +If the toolchain is not usable, configuring fails with one sentence naming both +fixes, rather than an `aiecc` traceback halfway through a build. Pass +`-DFLM_KERNELS_PYTHON=` to point at a specific interpreter. + +Output mirrors `src/xclbins//` so it installs through the rule that +already exists: + +``` +build/kernels/Granite-4.2-3B-NPU2/ + norm_qkv_rope.xclbin norm_qkv_rope.insts.bin norm_qkv_rope.json + attn_o.xclbin ... + manifest.json toolchain versions + sha256 per artefact +``` + +## Validate + +```shell +python kernels/build_kernels.py --family granite --validate +``` + +Runs each design against a host reference built from the same bytes. **This is +a hardware test**: it needs an NPU and the model weights, and CI is not asked to +run it. + +## The gap + +These artefacts are complete and numerically validated, but `flm` cannot +dispatch them yet, and the reason is specific. + +`src/include/npu_utils/npu_utils_xrt.hpp` builds its ELF from a control +sequence assembled **on the host** by `npu_sequence` +(`src/include/npu_utils/npu_instr_utils.hpp`). IRON emits that same control code +as a prebuilt `insts.bin` at build time. Bridging the two is roughly one +function: load `insts.bin` and hand it to `aiebu_assembler_get_elf` in place of +the assembled sequence. + +That is deliberately not here. **This is a build and convention change, not a +runtime change** — the two are worth deciding separately, and the runtime side +follows only if the convention is wanted. diff --git a/kernels/build_kernels.py b/kernels/build_kernels.py new file mode 100644 index 00000000..87ee4c19 --- /dev/null +++ b/kernels/build_kernels.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""Build FastFlowLM NPU kernels from source. See CONVENTION.md. + + python kernels/build_kernels.py --check-toolchain + python kernels/build_kernels.py --list + python kernels/build_kernels.py --family granite --out build/kernels + python kernels/build_kernels.py --family granite --validate + +Runs standalone, with no CMake. `src/CMakeLists.txt` shells out to it when +FLM_BUILD_KERNELS=ON, and to `--check-toolchain` at configure time so a missing +toolchain is one sentence rather than an aiecc traceback halfway through a +build. + +SPDX-License-Identifier: MIT +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shutil +import sys +from dataclasses import dataclass, field +from pathlib import Path + +HERE = Path(__file__).parent +JIT_CACHE = Path.home() / ".npu" / "cache" + + +@dataclass +class Design: + """One artefact: an IRON design plus the arguments that select its shape.""" + out: str + module: str + fn: str + kwargs: dict = field(default_factory=dict) + + +FAMILIES: dict[str, list[Design]] = { + "granite": [ + Design("norm_qkv_rope", "granite_qkv_wide", "granite_qkv_wide"), + Design("attn_o", "granite_attn_o", "granite_attn_o", {"seq": 64}), + Design("norm_qkv", "granite_norm_gemv", "granite_norm_gemv", {"tensor": "qkv"}), + Design("norm_gate_up", "granite_norm_gemv", "granite_norm_gemv", + {"tensor": "gate_up"}), + Design("swiglu_down", "granite_swiglu_down", "granite_swiglu_down"), + ], +} + +# The directory name must match src/xclbins// so the built artefacts +# install through the rule that already exists for the shipped ones. +MODEL_DIR = {"granite": "Granite-4.2-3B-NPU2"} + + +# ---------------------------------------------------------------- toolchain + +def probe_toolchain() -> tuple[bool, str]: + """Report whether this interpreter can build, and say precisely why not. + + Deliberately does not try to activate anything. `iron_env.cmd` is + Windows-only and machine-specific, and activating it in a subprocess would + not affect this process anyway. Activation is the caller's job; saying + exactly what is missing is ours -- that is what makes CMake's error message + useful. + """ + try: + import aie.iron # noqa: F401 + except ImportError as e: + return False, ( + f"the `aie` package is not importable from {sys.executable} ({e}). " + "Activate the mlir-aie/IRON environment before configuring, or pass " + "-DFLM_KERNELS_PYTHON=.") + try: + import importlib.metadata as md + ver = md.version("mlir_aie") + except Exception: + ver = "unknown" + try: + from aie.utils import config + hdr = Path(config.cxx_header_path()) + if not (hdr / "aie_kernels").is_dir(): + return False, (f"`aie` {ver} is importable but its kernel headers are " + f"missing at {hdr / 'aie_kernels'}.") + except Exception as e: + return False, f"`aie` {ver} is importable but unusable: {e!r}" + return True, f"mlir-aie {ver}, python {sys.version.split()[0]}" + + +def toolchain_manifest() -> dict: + import importlib.metadata as md + + def ver(p): + try: + return md.version(p) + except Exception: + return None + + return { + "mlir_aie": ver("mlir_aie"), + "ml_dtypes": ver("ml_dtypes"), + "numpy": ver("numpy"), + "python": sys.version.split()[0], + "aiecc_flags": ["--alloc-scheme=basic-sequential"], + } + + +# ------------------------------------------------------------------- build + +def _snapshot() -> set[Path]: + return set(JIT_CACHE.iterdir()) if JIT_CACHE.is_dir() else set() + + +def _new_cache_dir(before: set[Path]) -> Path | None: + """Find what the JIT just produced, by difference. + + Not by matching strings in the emitted MLIR: that is what + LLMNpuTest's export_design.py does, and its own docstring calls it fragile. + A set difference cannot be wrong, and it does not depend on IRON internals + that change between versions. If the design was already cached no directory + appears, and the caller falls back to the newest matching one. + """ + after = _snapshot() + fresh = [p for p in after - before if (p / "final.xclbin").is_file()] + if fresh: + return max(fresh, key=lambda p: (p / "final.xclbin").stat().st_mtime) + cached = [p for p in after if (p / "final.xclbin").is_file()] + return max(cached, key=lambda p: (p / "final.xclbin").stat().st_mtime) if cached else None + + +def build_family(family: str, out_root: Path, validate: bool) -> int: + designs = FAMILIES[family] + fam_dir = HERE / family + geometry = json.loads((fam_dir / "geometry.json").read_text(encoding="utf-8")) + + sys.path.insert(0, str(fam_dir / "iron")) + sys.path.insert(0, str(HERE / "common")) + # Generated entry points belong in the build tree, never beside tracked + # source. granite_gemv.py reads this. + gen = out_root / "_generated" + os.environ["GRANITE_GEN_DIR"] = str(gen) + + import aie.iron as iron + from aie.iron.device import from_name + iron.set_current_device(from_name("npu2", n_cols=None)) + + dest = out_root / MODEL_DIR[family] + dest.mkdir(parents=True, exist_ok=True) + entries, failures = [], 0 + + for d in designs: + mod = __import__(d.module) + builder = getattr(mod, "build_artifact", None) + if builder is None: + print(f" SKIP {d.out}: {d.module} has no build_artifact(geometry, **kw); " + f"see CONVENTION.md -- the build must not need model weights") + failures += 1 + continue + print(f" building {d.out} ...", flush=True) + before = _snapshot() + try: + builder(geometry, **d.kwargs) + except Exception as e: # noqa: BLE001 - report and continue + print(f" FAIL {d.out}: {e.__class__.__name__}: {e}") + failures += 1 + continue + src = _new_cache_dir(before) + if src is None: + print(f" FAIL {d.out}: no final.xclbin appeared in {JIT_CACHE}") + failures += 1 + continue + rec = {"design": d.out, "module": d.module, "kwargs": d.kwargs, "files": {}} + for name, suffix in (("final.xclbin", ".xclbin"), ("insts.bin", ".insts.bin")): + s = src / name + if not s.is_file(): + continue + t = dest / f"{d.out}{suffix}" + shutil.copy2(s, t) + rec["files"][t.name] = { + "bytes": t.stat().st_size, + "sha256": hashlib.sha256(t.read_bytes()).hexdigest(), + } + (dest / f"{d.out}.json").write_text(json.dumps(rec, indent=2) + "\n", + encoding="utf-8") + entries.append(rec) + print(f" -> {', '.join(rec['files'])}") + + if validate: + checker = getattr(mod, "validate_artifact", None) + if checker is None: + print(f" (no validate_artifact in {d.module})") + else: + ok = checker(geometry, **d.kwargs) + print(f" validate: {'PASS' if ok else 'FAIL'}") + failures += 0 if ok else 1 + + geom_hash = hashlib.sha256( + json.dumps(geometry, sort_keys=True).encode()).hexdigest()[:16] + (dest / "manifest.json").write_text(json.dumps({ + "family": family, + "model_dir": MODEL_DIR[family], + "geometry_sha256_16": geom_hash, + "toolchain": toolchain_manifest(), + "artifacts": entries, + }, indent=2) + "\n", encoding="utf-8") + print(f" manifest -> {dest / 'manifest.json'}") + return failures + + +# -------------------------------------------------------------------- main + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--check-toolchain", action="store_true", + help="report whether this interpreter can build, and exit") + ap.add_argument("--list", action="store_true", help="print the design map") + ap.add_argument("--family", choices=sorted(FAMILIES)) + ap.add_argument("--out", type=Path, default=Path("build/kernels")) + ap.add_argument("--validate", action="store_true", + help="also check each design against a host reference " + "(needs an NPU and the model weights)") + a = ap.parse_args(argv if argv is not None else sys.argv[1:]) + + if a.check_toolchain: + ok, msg = probe_toolchain() + print(msg) + return 0 if ok else 1 + + if a.list: + for fam, ds in FAMILIES.items(): + print(f"{fam} -> {MODEL_DIR[fam]}/") + for d in ds: + extra = f" {d.kwargs}" if d.kwargs else "" + print(f" {d.out:16} {d.module}.{d.fn}{extra}") + return 0 + + if not a.family: + ap.error("one of --check-toolchain, --list or --family is required") + + ok, msg = probe_toolchain() + if not ok: + print(f"toolchain unusable: {msg}", file=sys.stderr) + return 1 + print(f"toolchain: {msg}") + failures = build_family(a.family, a.out.resolve(), a.validate) + print("OK" if failures == 0 else f"{failures} failure(s)") + return 0 if failures == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kernels/common/q4nx.py b/kernels/common/q4nx.py new file mode 100644 index 00000000..56a83b5b --- /dev/null +++ b/kernels/common/q4nx.py @@ -0,0 +1,197 @@ +"""FastFlowLM `q4nx` container: reader and host-side reference dequantiser. + +`.q4nx` is a plain safetensors file (u64 header length, JSON header, data). The +weight tensors inside are pre-tiled for the NPU, shaped `[N/32][K/256][bytes]` +-- one tile is 32 output rows x 256 K, split into two row-blocks of 16. + + q4 5120 B [512 B bf16 d][512 B bf16 m][4096 B packed nibbles] + q8 8704 B [512 B bf16 scale][8192 B int8] + +Within a tile, with rb = row block (0..1), r = row in block (0..15), +k = 0..255, kb = k // 32: + + metadata index = kb * 32 + rb * 16 + r (256 entries per plane) + weight index i = rb * 4096 + k * 16 + r (8192 weights) + q4 byte, nibble = i >> 1, low nibble when i is even + + q4: w = code * d + m (GGUF Q4_1 semantics -- scale and minimum) + q8: w = code * scale (symmetric) + +This was not guessed. The q8 form was solved against ground truth: layer 0's +`ssm_alpha_proj` is stored twice in the same file, once quantised and once as +bf16, and under this mapping all 8192 codes of all four tiles reproduce the bf16 +weights to within one quantisation step. The q4 form then follows from three +independent checks -- `d` is everywhere positive, `m` everywhere negative, +mean(m/d) = -7.48 (the Q4_1 signature for a symmetric weight distribution), and +every group's codes span 0..15. + +The nibble parity -- whether the low nibble is the even or the odd weight index +-- could not be settled from the file alone, since it only swaps adjacent output +rows and every group stays a valid Q4_1 group either way. It is settled now, by +diffing against the upstream bf16 checkpoint the model was quantised from: +LOW_NIBBLE_IS_EVEN = True scores cosine 0.9975, False scores 0.0298. See +reference/check_weights.py. +""" + +from __future__ import annotations + +import json +import struct +from pathlib import Path +from typing import BinaryIO + +import numpy as np +from ml_dtypes import bfloat16 + +GROUP = 32 # weights per quantisation group +TILE_ROWS = 32 # output rows per tile +TILE_K = 256 # K per tile +ROW_BLOCK = 16 # rows per row-block; a tile holds two +N_META = TILE_ROWS * (TILE_K // GROUP) # 256 scale entries per plane +N_WEIGHTS = TILE_ROWS * TILE_K # 8192 + +Q4_TILE_BYTES = 2 * 2 * N_META + N_WEIGHTS // 2 # 5120 +Q8_TILE_BYTES = 2 * N_META + N_WEIGHTS # 8704 + +LOW_NIBBLE_IS_EVEN = True # measured, not assumed: see the module docstring + + +def read_header(f: BinaryIO) -> dict: + """Parse the safetensors header. Reads only the header, not the data.""" + (n,) = struct.unpack(" tuple[np.ndarray, np.ndarray, np.ndarray]: + """One 5120-byte q4 tile -> (nibbles uint8[4096], d bf16[256], m bf16[256]). + + Already planar in the file, so this is three slices. Convenient, because a + core has only 2 input DMA streams (trap 3b): d and m ride together in the + first 1024 bytes and the nibbles take the other stream. + """ + if len(raw) != Q4_TILE_BYTES: + raise ValueError(f"{len(raw)} bytes, expected a {Q4_TILE_BYTES} B q4 tile") + b = np.frombuffer(raw, dtype=np.uint8) + d = b[0:512].copy().view(bfloat16) + m = b[512:1024].copy().view(bfloat16) + return b[1024:].copy(), d, m + + +def unpack_nibbles(nib: np.ndarray) -> np.ndarray: + """uint8[4096] -> uint8[8192] raw 0..15 codes, in weight-index order.""" + out = np.empty(N_WEIGHTS, dtype=np.uint8) + lo, hi = nib & 0x0F, nib >> 4 + first, second = (lo, hi) if LOW_NIBBLE_IS_EVEN else (hi, lo) + out[0::2], out[1::2] = first, second + return out + + +def meta_index() -> np.ndarray: + """Metadata index for each of the 8192 weight positions: kb*32 + rb*16 + r.""" + i = np.arange(N_WEIGHTS) + rb, rest = divmod(i, 4096) + k, r = divmod(rest, ROW_BLOCK) + return (k // GROUP) * TILE_ROWS + rb * ROW_BLOCK + r + + +def dequant_q4(nib: np.ndarray, d: np.ndarray, m: np.ndarray) -> np.ndarray: + """The reference: w = code * d + m, accumulated in float32, returned bf16.""" + g = meta_index() + x = unpack_nibbles(nib).astype(np.float32) + return (x * d.astype(np.float32)[g] + m.astype(np.float32)[g]).astype(bfloat16) + + +def dequant_q8(raw: bytes) -> np.ndarray: + """One 8704-byte q8 tile -> bf16[8192] in weight-index order.""" + if len(raw) != Q8_TILE_BYTES: + raise ValueError(f"{len(raw)} bytes, expected a {Q8_TILE_BYTES} B q8 tile") + b = np.frombuffer(raw, dtype=np.uint8) + s = b[0:512].copy().view(bfloat16).astype(np.float32) + c = b[512:].view(np.int8).astype(np.float32) + return (c * s[meta_index()]).astype(bfloat16) + + +def to_matrix(flat: np.ndarray) -> np.ndarray: + """Weight-index order -> [32 rows][256 K], the natural view of a tile.""" + return flat.reshape(2, TILE_K, ROW_BLOCK).transpose(0, 2, 1).reshape(TILE_ROWS, TILE_K) + + +def pack_meta(d: np.ndarray, m: np.ndarray) -> np.ndarray: + """d and m as one bf16 buffer -- one input stream, not two (trap 3b).""" + return np.concatenate([np.asarray(d), np.asarray(m)]).astype(bfloat16) + + +# ---------------------------------------------------------------- whole file + + +def _q4_tiles(b: np.ndarray) -> np.ndarray: + """[T, 5120] uint8 -> [T, 8192] float32, weight-index order.""" + d = b[:, 0:512].copy().view(bfloat16).astype(np.float32) + m = b[:, 512:1024].copy().view(bfloat16).astype(np.float32) + nib = b[:, 1024:] + codes = np.empty((b.shape[0], N_WEIGHTS), dtype=np.uint8) + lo, hi = nib & 0x0F, nib >> 4 + first, second = (lo, hi) if LOW_NIBBLE_IS_EVEN else (hi, lo) + codes[:, 0::2], codes[:, 1::2] = first, second + g = meta_index() + return codes.astype(np.float32) * d[:, g] + m[:, g] + + +def _q8_tiles(b: np.ndarray) -> np.ndarray: + """[T, 8704] uint8 -> [T, 8192] float32, weight-index order.""" + s = b[:, 0:512].copy().view(bfloat16).astype(np.float32) + return b[:, 512:].view(np.int8).astype(np.float32) * s[:, meta_index()] + + +def _untile(flat: np.ndarray, n_t: int, k_t: int) -> np.ndarray: + """[n_t*k_t, 8192] -> [n_t*32, k_t*256], undoing the row-block interleave.""" + m = flat.reshape(n_t, k_t, 2, TILE_K, ROW_BLOCK).transpose(0, 2, 4, 1, 3) + return m.reshape(n_t * TILE_ROWS, k_t * TILE_K) + + +class Q4NX: + """A `.q4nx` model file. Tensors come back dequantised and un-tiled.""" + + _DT = {"BF16": bfloat16, "F32": np.float32, "F16": np.float16, "I8": np.int8} + + def __init__(self, path): + self.path = Path(path) + with self.path.open("rb") as f: + self.header = read_header(f) + self._data_start = f.tell() + self.header.pop("__metadata__", None) + + def __contains__(self, name: str) -> bool: + return name in self.header + + def names(self) -> list[str]: + return sorted(self.header) + + def raw(self, name: str) -> bytes: + first, last = self.header[name]["data_offsets"] + with self.path.open("rb") as f: + f.seek(self._data_start + first) + return f.read(last - first) + + def tensor(self, name: str, shape: tuple[int, int] | None = None) -> np.ndarray: + """Dequantised and un-tiled. `shape` trims the tile padding. + + A tile is 32 rows x 256 K, so a tensor whose N or K is not a multiple of + those is stored padded -- N=16 becomes 32 rows, the upper half a copy of + the lower. Pass the logical shape and the padding is dropped. + """ + e = self.header[name] + if e["dtype"] != "I8": # stored as-is: embeddings, norms, conv1d, biases + v = np.frombuffer(self.raw(name), dtype=self._DT[e["dtype"]]) + return v.reshape(e["shape"]) + + n_t, k_t, tile_bytes = e["shape"] + b = np.frombuffer(self.raw(name), dtype=np.uint8).reshape(n_t * k_t, tile_bytes) + if tile_bytes == Q4_TILE_BYTES: + flat = _q4_tiles(b) + elif tile_bytes == Q8_TILE_BYTES: + flat = _q8_tiles(b) + else: + raise ValueError(f"{name}: {tile_bytes} B/tile is neither q4 nor q8") + w = _untile(flat, n_t, k_t) + return w if shape is None else w[: shape[0], : shape[1]] diff --git a/kernels/granite/README.md b/kernels/granite/README.md new file mode 100644 index 00000000..11e64b5e --- /dev/null +++ b/kernels/granite/README.md @@ -0,0 +1,94 @@ +# granite kernels + +AIE kernels for IBM Granite 4.2 3B. A companion pull request adds a host +engine that runs the same model on the CPU; these kernels run its arithmetic on +the NPU instead. **Neither PR depends on the other** — this one adds no code +that `flm` links, and the engine needs nothing from here. + +## Geometry + +Hidden 2560, 40 layers, 40 query heads over 8 kv heads, head_dim 64, +intermediate 8192, vocab 100352, RoPE theta 1e7 with the half-split +(`rotate_half`) convention. All of it in [geometry.json](geometry.json), which +is what the build reads — **no model weights are needed to produce an +artefact.** + +Granite needs head_dim 64 at hidden 2560. Every shipped design at hidden >= 2560 +is head_dim 128, and head_dim cannot be padded, so nothing existing could be +reused for it. + +## The layer, in four dispatches + +| dispatch | ops | cores | device time | +|---|---|---|---| +| `norm_qkv_rope` | RMSNorm, q, k, v, RoPE | 28 | 286.8 µs | +| `attn_o` | attention (all 40 heads), o_proj | 8 + 16 | 330.2 µs | +| `norm_gate_up` | RMSNorm, gate, up | 32 | 669.2 µs | +| `swiglu_down` | SwiGLU, down | 20 | 458.5 µs | +| | | **layer** | **1744.7 µs** | + +×40 layers plus lm_head (3.66 ms, 43.9 GB/s) is **73.5 ms/token = 13.6 tok/s** +of device time, against 8.5 tok/s for the same work in nine dispatches, and +8.7 tok/s end to end for the CPU host engine. + +Measured on a Ryzen AI 9 HX 370, medians of repeated runs at `--iters 200`. +Every design checks against a host reference built from the same bytes; cosines +are in each design's own output (1.00000000 for the GEMV groups under a one-hot +activation, > 0.9996 for the fused blocks). + +## Why fuse the small ops and not the big ones + +The per-dispatch floor is about 200 µs regardless of size: RMSNorm on 2560 +values costs 244 µs, SwiGLU on 8192 costs 205, and RoPE cost **nothing** when it +moved inside a dispatch that was already running. Four GEMV groups move 49.2 MB +and cost 1.65 ms — nearly all real work. Five small ops move almost nothing and +cost 1.20 ms, of which ~1.0 ms is floor. + +So the fusions here all absorb a *small* op into a dispatch that was already +paying for bandwidth. + +## What was tried and rejected + +Worth recording, because it is the evidence for the shape above. + +**Fusing the whole MLP into one dispatch is a net loss.** Measured 2.32 ms at 8 +cores and 1.73 ms at 16, against 1.31 ms for the same three ops unfused across +20–32 cores. Even a perfect one-dispatch MLP at gate_up's 39.7 GB/s would be +1.19 ms. Fusion saves 0.40 ms of dispatch floor and costs more than that in +width, because fusion and width compete for the same scarce resources. + +**Two independent costs make wide fusion expensive.** Shim DMA channels: the +device has 16 each way, and private per-core streams at 16 cores want 17 in and +32 out. Routing through the memtile (one stream per column, `split()`/`join()`) +costs 5 and 8. And `PER_CALL` must divide the K-tile count of *every* matrix in +a dispatch, so fusing gate/up (10 tiles) with down (32) pins it at gcd = 2 and +halves the DMA element. + +**The same trade goes the other way for SwiGLU + down**, which is why it is in +the table above. Taking gate and up as the activation forces `per_call` from 4 +to 2, but down_proj has 32 K-tiles so that is still a long stream: 29.8 GB/s +becomes 28.8, a 3% cost against a 205 µs dispatch that moves no weights at all. +**The same decision, opposite outcomes, decided by K.** + +## Four hardware limits, in the order they bind + +None of these is derivable from bandwidth and L1 alone; each appeared only when +the placer ran. They are recorded in the sources at the point where they bind. + +| resource | budget | what it forces | +|---|---|---| +| shim MM2S / S2MM | 16 / 16 device-wide | streams per column, not per core | +| memtile DMA | ~6 in / 6 out per column | at most two split/join structures of four | +| compute tile DMA | 2 in / 2 out | at most two input streams per core | +| L1 | 62208 B | sets `per_call`, and so bandwidth | + +The compute-tile limit is why several kernels write **in place**: with weights +and one activation already using both input channels, a fused op has no third +stream, so the norm weight rides in the activation buffer and the result is +written back over its input. + +One more, which cost the most to find: **`aie::store_v` of 32 floats is a +128-byte operation**, so a per-head state array must be 128-byte aligned. A +72-float stride is 32-byte aligned and not enough — the second head's store +rounds down onto the first head's softmax denominator, and every head comes out +with the right direction and the wrong scale. diff --git a/kernels/granite/aie/granite_attention.h b/kernels/granite/aie/granite_attention.h new file mode 100644 index 00000000..66b22a6c --- /dev/null +++ b/kernels/granite/aie/granite_attention.h @@ -0,0 +1,238 @@ +#pragma once +//===- granite_attention.h ----------------------------------*- C++ -*-===// +// +// OpenFFLM -- granite GQA decode attention on the AIE core, one q head at a +// time, with an online (flash-style) softmax. +// SPDX-License-Identifier: MIT +// +// s[t] = scale * dot(q, K[t]) t over the whole KV cache +// p = softmax(s) +// out = sum_t p[t] * V[t] +// +// WHY ONLINE SOFTMAX AND NOT TWO PASSES +// ------------------------------------- +// The KV cache is the one tensor that grows: 8 kv heads x seq x 64 x 2 x 2 B is +// 2 MB at seq 1024, per layer, per token. It cannot live in a 64 KB L1, so it +// has to stream -- and a two-pass softmax would have to stream it twice, once +// for the max and once for the weights. The online form keeps a running max `m` +// and normaliser `l` and rescales the accumulator when the max moves, so the +// cache is read exactly once. +// +// m_new = max(m, max(s_block)) +// corr = exp(m - m_new) +// l = l*corr + sum(exp(s_block - m_new)) +// acc = acc*corr + sum_t exp(s[t] - m_new) * V[t] +// +// and `out = acc / l` once, after the last block. +// +// THE SCALE IS NOT OPTIONAL AND IS NOT GRANITE'S MULTIPLIER +// -------------------------------------------------------- +// q4nx-build folds `attention_multiplier` into q_proj as +// `q_proj *= attention_multiplier * sqrt(head_dim)`, precisely so that an engine +// applying the standard `head_dim ** -0.5` gets granite's intended result. So +// this kernel MUST still apply `head_dim ** -0.5` -- the fold assumes it. For +// head_dim 64 that is exactly 0.125, a power of two, so it is exact. +// +// TRAP CARRIED OVER FROM aie_kernels/aie2p/softmax.cc +// --------------------------------------------------- +// Its own comment warns: "The multiplication by log2e is very sensitive, +// casting it to bf16 before exponentiation leads to wrong output." `bf16_exp.cc` +// does exactly that (`broadcast(log2e)` rounds 1.44269504 to +// 1.4453125). Here the log2e scaling stays in fp32 and only the exp result is +// bf16. + +#include "aie_kernel_utils.h" +#include +#include + +#ifndef GRANITE_ATTN_HEAD_DIM +#define GRANITE_ATTN_HEAD_DIM 64 +#endif +#ifndef GRANITE_ATTN_BLOCK +#define GRANITE_ATTN_BLOCK 32 +#endif + +static constexpr unsigned kHD = GRANITE_ATTN_HEAD_DIM; +static constexpr unsigned kBlk = GRANITE_ATTN_BLOCK; // KV positions per call +static constexpr unsigned kHDV = 32; // vector width over head_dim +static constexpr unsigned kHDVecs = kHD / kHDV; // 2 for head_dim 64 + +// head_dim ** -0.5. Exact for any power-of-two head_dim. +#ifndef GRANITE_ATTN_SCALE +#define GRANITE_ATTN_SCALE 0.125f +#endif + +// log2(e) is folded straight into the score scale, so every score is already in +// log2 units and the softmax is pure exp2. This is what aie_kernels' softmax.cc +// does ("the max value scaled by log2e"), and it has two payoffs beyond speed: +// there is no per-element `* log2e` needing an fp32 vector multiply (which +// AIE2P does not have), and log2e never passes through bf16 -- the rounding +// that softmax.cc's own comment warns about and that bf16_exp.cc walks into. +static constexpr float kLog2e = 1.4426950408889634f; +static constexpr float kScaleL2 = GRANITE_ATTN_SCALE * kLog2e; + +// state layout, all float: acc[kHD] | m | l +// acc running unnormalised output +// m running max of the scores seen so far +// l running sum of exp(s - m) +// +// `kv` holds this block's K then this block's V, each [kBlk][kHD] bf16. +__attribute__((noinline)) inline void +granite_attn_block_impl(const bfloat16 *__restrict q, + const bfloat16 *__restrict kv, + float *__restrict state, unsigned n_t, unsigned first) { + event0(); + aie::set_rounding(aie::rounding_mode::conv_even); + + const bfloat16 *__restrict K = kv; + const bfloat16 *__restrict V = kv + (unsigned)(kBlk * kHD); + + aie::vector q0 = aie::load_v(q); + aie::vector q1 = aie::load_v(q + kHDV); + + // 1. scores for this block. Each is a bf16 dot product accumulated in fp32. + // Padded, because the exponentiation below runs over the whole block: a + // stale stack value at t >= n_t would exponentiate to garbage and, if large, + // would poison m_new for every later block. + float s[kBlk]; + for (unsigned t = n_t; t < kBlk; ++t) s[t] = -3.0e38f; + float m_blk = -3.0e38f; + for (unsigned t = 0; t < n_t; ++t) { + const bfloat16 *__restrict kt = K + t * kHD; + aie::accum a = aie::mul(q0, aie::load_v(kt)); + a = aie::mac(a, q1, aie::load_v(kt + kHDV)); + float v = aie::reduce_add(a.template to_vector()) * kScaleL2; + s[t] = v; + if (v > m_blk) m_blk = v; + } + + // 2. fold this block's max into the running one, and rescale what we have. + const float m_old = first ? -3.0e38f : state[kHD]; + const float l_old = first ? 0.0f : state[kHD + 1]; + const float m_new = m_blk > m_old ? m_blk : m_old; + + // Scores are already in log2 units, so the correction is a plain exp2 of a + // difference. aie::exp2 is a vector op with no scalar form, so this evaluates + // one lane of a broadcast -- once per block, against n_t dot products. + float corr = 0.0f; + if (!first) { + aie::vector cv = aie::broadcast(m_old - m_new); + corr = (float)aie::exp2(cv).get(0); + } + + // 3. p[t] = exp(s[t] - m_new), and the accumulator rescaled by corr. + float l_new = l_old * corr; + aie::accum acc[kHDVecs]; + for (unsigned v = 0; v < kHDVecs; ++v) { + if (first) { + acc[v] = aie::zeros(); + } else { + // acc *= corr. No fp32 vector multiplier on AIE2P, so this goes through + // a bf16 hi/lo split of the accumulator, as everywhere else here. + aie::vector a = aie::load_v(state + v * kHDV); + aie::accum t; + t.from_vector(a); + aie::vector hi = t.template to_vector(); + aie::vector lo = + aie::sub(t, hi).template to_vector(); + const bfloat16 c_hi = (bfloat16)corr; + const bfloat16 c_lo = (bfloat16)(corr - (float)c_hi); + aie::accum r = aie::zeros(); + r = aie::mac(r, hi, c_hi); + r = aie::mac(r, lo, c_hi); + r = aie::mac(r, hi, c_lo); + acc[v] = r; + } + } + + // Exponentiate the whole block at once. kBlk is a multiple of the vector + // width, so this is one exp2 per kHDV scores rather than one per score. + bfloat16 pbuf[kBlk]; + for (unsigned t0 = 0; t0 < kBlk; t0 += kHDV) { + aie::vector sv = aie::load_v(s + t0); + aie::vector d = + aie::sub(sv, aie::broadcast(m_new)); + aie::store_v(pbuf + t0, aie::exp2(d)); + } + + for (unsigned t = 0; t < n_t; ++t) { + const float p = (float)pbuf[t]; + l_new += p; + const bfloat16 p_b = pbuf[t]; + const bfloat16 *__restrict vt = V + t * kHD; + for (unsigned v = 0; v < kHDVecs; ++v) { + aie::vector vv = aie::load_v(vt + v * kHDV); + acc[v] = aie::mac(acc[v], vv, p_b); + } + } + + for (unsigned v = 0; v < kHDVecs; ++v) + aie::store_v(state + v * kHDV, acc[v].template to_vector()); + state[kHD] = m_new; + state[kHD + 1] = l_new; +#ifdef GRANITE_ATTN_DEBUG_SCORES + // Overwrite the accumulator with the raw scores so they can be compared + // element by element against numpy. Which scores are wrong, and by how much, + // is a fact; which stage is at fault has so far only been a guess. + for (unsigned t = 0; t < kBlk; ++t) state[t] = s[t]; + // and what the kernel actually received for q and K[0] -- which is what + // distinguishes 'the maths is wrong' from 'the data never arrived'. + // Via vector loads: scalar bf16 reads at a computed offset make the Peano + // backend fail with "immediate operand value -120 is not a multiple of 64". + // q only: state is kHD+2 floats, and writing K at state+kBlk+kHDV ran off + // the end of the buffer. K was already confirmed to arrive correctly. + { + aie::accum tq; + tq.from_vector(aie::load_v(q)); + aie::store_v(state + kBlk, tq.template to_vector()); + } +#endif + event1(); +} + +// Divide the accumulator by the normaliser, once, after the last block. +__attribute__((noinline)) inline void +granite_attn_finish_impl(const float *__restrict state, + bfloat16 *__restrict out) { + event0(); + const float inv = 1.0f / state[kHD + 1]; + const bfloat16 i_hi = (bfloat16)inv; + const bfloat16 i_lo = (bfloat16)(inv - (float)i_hi); + for (unsigned v = 0; v < kHDVecs; ++v) { + aie::accum t; + t.from_vector(aie::load_v(state + v * kHDV)); + aie::vector hi = t.template to_vector(); + aie::vector lo = aie::sub(t, hi).template to_vector(); + aie::accum r = aie::zeros(); + r = aie::mac(r, hi, i_hi); + r = aie::mac(r, lo, i_hi); + r = aie::mac(r, hi, i_lo); + aie::store_v(out + v * kHDV, r.template to_vector()); + } + event1(); +} + +// The two entry points are emitted into SEPARATE translation units, selected by +// these macros. IRON compiles the kernel source once per ExternalFunction, so +// pointing two ExternalFunctions at one .cc yields two objects that each define +// BOTH symbols and the link fails on duplicates -- the same trap granite_gemv.h +// records, walked into again from the other direction. The `impl` functions +// above are `inline`, so the shared body still merges into one COMDAT. +#ifdef GRANITE_ATTN_EMIT_BLOCK +extern "C" { +void granite_attn_block(const bfloat16 *__restrict q, + const bfloat16 *__restrict kv, float *__restrict state, + unsigned n_t, unsigned first) { + granite_attn_block_impl(q, kv, state, n_t, first); +} +} +#endif + +#ifdef GRANITE_ATTN_EMIT_FINISH +extern "C" { +void granite_attn_finish(const float *__restrict state, + bfloat16 *__restrict out) { + granite_attn_finish_impl(state, out); +} +} +#endif diff --git a/kernels/granite/aie/granite_attn_block_h.cc b/kernels/granite/aie/granite_attn_block_h.cc new file mode 100644 index 00000000..e05ce74c --- /dev/null +++ b/kernels/granite/aie/granite_attn_block_h.cc @@ -0,0 +1,38 @@ +// OpenFFLM -- attention block for one of a core's several query heads. +// SPDX-License-Identifier: MIT +// +// granite_attn_block takes pointers to ONE head's q and state. A core that owns +// a whole kv head owns the five query heads that share it (granite is GQA, 40 +// over 8), and IRON passes whole buffers, not offsets -- so the head index has +// to be an argument and the offsetting has to happen here. +// +// Separate translation unit for the reason granite_attention.h records: IRON +// compiles a kernel source once per ExternalFunction, so two entry points in one +// .cc link as duplicate symbols. +#include "granite_attention.h" + +// The per-head state is kHD + 2 = 66 floats, but the stride is 128, not 66 and +// not 72. The kernel stores the accumulator with aie::store_v of 32 floats -- +// vectors that are 128 BYTES wide -- so a head's state has to start on a +// 128-byte boundary. +// +// 72 floats = 288 bytes is 32-byte aligned, which is what a first attempt at +// this checked, and it is not enough: head 1's store rounded down to the +// nearest 128-byte boundary, float index 64, which is exactly where head 0 +// keeps its softmax max and denominator. Head 1 was overwriting head 0's +// divisor. +// +// The symptom named the cause once it was read properly: every head came out +// with cosine +-1.0 against the reference -- right direction, wrong scale -- +// and head 0 was wrong too, which no offset error of head 1's own could do. +// q_per = 1 passed because there was no head 1 to do the overwriting. +static constexpr unsigned kStStride = 128; // floats: 512 B, 128-B aligned + +extern "C" { +void granite_attn_block_h(const bfloat16 *__restrict q_all, + const bfloat16 *__restrict kv, float *__restrict st_all, + unsigned n_t, unsigned first, unsigned h) { + granite_attn_block_impl(q_all + h * kHD, kv, + st_all + h * kStStride, n_t, first); +} +} diff --git a/kernels/granite/aie/granite_attn_finish_h.cc b/kernels/granite/aie/granite_attn_finish_h.cc new file mode 100644 index 00000000..5d486589 --- /dev/null +++ b/kernels/granite/aie/granite_attn_finish_h.cc @@ -0,0 +1,28 @@ +// OpenFFLM -- attention epilogue for one of a core's several query heads. +// SPDX-License-Identifier: MIT +#include "granite_attention.h" + +// The per-head state is kHD + 2 = 66 floats, but the stride is 128, not 66 and +// not 72. The kernel stores the accumulator with aie::store_v of 32 floats -- +// vectors that are 128 BYTES wide -- so a head's state has to start on a +// 128-byte boundary. +// +// 72 floats = 288 bytes is 32-byte aligned, which is what a first attempt at +// this checked, and it is not enough: head 1's store rounded down to the +// nearest 128-byte boundary, float index 64, which is exactly where head 0 +// keeps its softmax max and denominator. Head 1 was overwriting head 0's +// divisor. +// +// The symptom named the cause once it was read properly: every head came out +// with cosine +-1.0 against the reference -- right direction, wrong scale -- +// and head 0 was wrong too, which no offset error of head 1's own could do. +// q_per = 1 passed because there was no head 1 to do the overwriting. +static constexpr unsigned kStStride = 128; // floats: 512 B, 128-B aligned + +extern "C" { +void granite_attn_finish_h(const float *__restrict st_all, + bfloat16 *__restrict out_all, unsigned h) { + granite_attn_finish_impl(st_all + h * kStStride, + out_all + h * kHD); +} +} diff --git a/kernels/granite/aie/granite_elementwise.h b/kernels/granite/aie/granite_elementwise.h new file mode 100644 index 00000000..24fb57e0 --- /dev/null +++ b/kernels/granite/aie/granite_elementwise.h @@ -0,0 +1,127 @@ +#pragma once +//===- granite_elementwise.h --------------------------------*- C++ -*-===// +// +// OpenFFLM -- granite's RoPE and SwiGLU on the AIE core. +// SPDX-License-Identifier: MIT +// +// WHY NOT THE aie_kernels/aie2p REFERENCES +// ---------------------------------------- +// Both are demos that fail open, in the same way `rms_norm.cc` does (see +// granite_rmsnorm.h): +// +// * **`rope.cc` uses the interleaved-pair convention.** It pairs element 0 with +// 1, 2 with 3 (`filter_even`/`filter_odd`, GPT-NeoX style). Llama and Granite +// use **half-split** `rotate_half`, pairing i with i + head_dim/2. The two +// are different rotations that produce identical magnitudes, so the output +// looks entirely reasonable and is wrong. +// +// * **`swiglu.cc`'s entry point hardcodes `input_size = 1024`.** The templated +// body takes a size, but the `extern "C"` wrapper passes a literal. Granite's +// intermediate is 8192, so it would compute one eighth of the vector and +// leave the remaining seven eighths as whatever was already in the buffer -- +// partially correct output, no error, no shape mismatch. +// +// The silu maths is worth keeping though: `sigmoid(x) = (tanh(x/2) + 1) / 2` is +// an identity, not an approximation (tanh(x/2) = 2*sigmoid(x) - 1). Only +// `aie::tanh`'s own implementation is approximate. + +#include "aie_kernel_utils.h" +#include +#include + +// granite-4.2-3B is head_dim 64, so each half is one 32-lane vector. Both halves +// of a head are then a single load, a single multiply and a single store. +#ifndef GRANITE_HEAD_HALF +#define GRANITE_HEAD_HALF 32 +#endif +static constexpr unsigned kHalf = GRANITE_HEAD_HALF; +static constexpr unsigned kHeadDim = 2 * kHalf; + +// RoPE, HALF-SPLIT (Llama/Granite `rotate_half`): +// +// y[i] = x[i] * cos[i] - x[i + half] * sin[i] +// y[i + half] = x[i + half] * cos[i] + x[i] * sin[i] +// +// `cs` holds cos[0..half) then sin[0..half) for this position -- the caller +// owns the position, so this kernel is stateless and the same code serves +// prefill and decode. +__attribute__((noinline)) inline void +granite_rope_impl(const bfloat16 *__restrict x, const bfloat16 *__restrict cs, + bfloat16 *__restrict y, unsigned n_heads) { + event0(); + aie::set_rounding(aie::rounding_mode::conv_even); + + const aie::vector c = aie::load_v(cs); + const aie::vector s = aie::load_v(cs + kHalf); + + for (unsigned h = 0; h < n_heads; ++h) { + const bfloat16 *__restrict xh = x + h * kHeadDim; + bfloat16 *__restrict yh = y + h * kHeadDim; + + aie::vector lo = aie::load_v(xh); + aie::vector hi = aie::load_v(xh + kHalf); + + // Each product is bf16 x bf16 into an fp32 accumulator, so the rotation + // carries full precision and only the final store rounds. + aie::accum lo_c = aie::mul(lo, c); + aie::accum hi_s = aie::mul(hi, s); + aie::accum hi_c = aie::mul(hi, c); + aie::accum lo_s = aie::mul(lo, s); + + aie::store_v(yh, aie::sub(lo_c, hi_s).template to_vector()); + aie::store_v(yh + kHalf, aie::add(hi_c, lo_s).template to_vector()); + } + event1(); +} + +// SwiGLU: y = silu(gate) * up, with `n` a RUNTIME argument. +// +// None of the three pointers is __restrict: granite_swiglu_ip.cc calls this +// with gate, up and y all inside ONE buffer -- up at gate + n, and the result +// written back over gate. That is what lets a fused SwiGLU+down_proj keep its +// activation at 2 x 8192 bf16 with no destination buffer, which is in turn what +// keeps it inside L1. Safe because the op is elementwise at the same index and +// each vector's loads precede its store. +// +// silu(x) = x * sigmoid(x), sigmoid(x) = (tanh(x/2) + 1) / 2 +__attribute__((noinline)) inline void +granite_swiglu_impl(const bfloat16 *gate, + const bfloat16 *up, bfloat16 *y, + unsigned n) { + event0(); + aie::set_rounding(aie::rounding_mode::conv_even); + + constexpr unsigned V = 32; + const aie::vector half = aie::broadcast((bfloat16)0.5f); + const aie::vector one = aie::broadcast((bfloat16)1.0f); + + for (unsigned i = 0; i < n; i += V) { + aie::vector g = aie::load_v(gate + i); + aie::vector u = aie::load_v(up + i); + + // Keep x/2 as the fp32 accumulator and feed tanh from there -- rounding it + // to bf16 first would throw away half the mantissa before the nonlinearity. + aie::accum gh = aie::mul(g, half); + aie::vector t = + aie::tanh(gh.template to_vector()); + aie::vector sig = + aie::mul(aie::add(t, one), half).template to_vector(); + + aie::vector silu = + aie::mul(g, sig).template to_vector(); + aie::store_v(y + i, aie::mul(silu, u).template to_vector()); + } + event1(); +} + +extern "C" { +void granite_rope(const bfloat16 *__restrict x, const bfloat16 *__restrict cs, + bfloat16 *__restrict y, unsigned n_heads) { + granite_rope_impl(x, cs, y, n_heads); +} +void granite_swiglu(const bfloat16 *__restrict gate, + const bfloat16 *__restrict up, bfloat16 *__restrict y, + unsigned n) { + granite_swiglu_impl(gate, up, y, n); +} +} diff --git a/kernels/granite/aie/granite_gemv.h b/kernels/granite/aie/granite_gemv.h new file mode 100644 index 00000000..663d2b43 --- /dev/null +++ b/kernels/granite/aie/granite_gemv.h @@ -0,0 +1,302 @@ +#pragma once +//===- granite_gemv.h --------------------------------------*- C++ -*-===// +// +// OpenFFLM -- W4A16 GEMV for granite-4.2-3B's lm_head, on the AIE core. +// SPDX-License-Identifier: MIT +// +// y[100352] = W[100352, 2560] @ x[2560], W in q4nx's **q4** form. +// +// WHY THIS EXISTS AND NOT lm_head.h +// --------------------------------- +// `lm_head.h` is W8A16: one scale per block, `w = code * scale`. Granite is +// stored Q4_1, which carries a **minimum** as well as a scale: +// +// w = code * d + m code 0..15, d and m bf16 per 32-wide K block +// +// and a half-width code. Both differences are in this file; the design around +// it is the same shape. +// +// WHY THE MINIMUM IS ALMOST FREE +// ------------------------------ +// Naively `m` is another per-element term. It is not, because it is constant +// across the 32 K of a block, so it factors out of the inner sum: +// +// y[r] = sum_k (code[k][r]*d[kb][r] + m[kb][r]) * x[k] +// = sum_kb { d[kb][r] * (sum_{k in kb} code[k][r]*x[k]) +// + m[kb][r] * (sum_{k in kb} x[k]) } +// +// The second term needs one **scalar** sum of x per K block -- 8 per tile, +// against 8192 MACs. The minimum costs essentially nothing. +// +// THE LAYOUT, AND WHY IT IS STILL CHEAP +// ------------------------------------- +// One q4 tile is 32 output rows x 256 K in 5120 bytes, two row-blocks of 16: +// +// d [256] bf16 at tile[0 : 512] index = kb*32 + rb*16 + r +// m [256] bf16 at tile[512 : 1024] same index +// nib[4096] B at tile[1024 : 5120] weight i = rb*4096 + k*16 + r, +// byte i>>1, low nibble when i is even +// +// The row index is the FASTEST axis, so 16 consecutive nibbles are 16 different +// rows at one k -- they need 16 different scales, and those are constant for a +// whole k-block. Load them once per (rb, kb) and the inner loop is a mask, a +// to_float with a shifted binary point, a zip and a mac. No gather. +// +// TRAPS OBSERVED (NpuEmbeddings CLAUDE.md) +// ---------------------------------------- +// - AIE default rounding is `floor`, a systematic downward bias baked into +// every weight. Set conv_even. +// - AIE2P has NO fp32 vector multiplier: `aie::mul(vector, +// vector)` compiles and returns **zero**, silently. So the fp32 +// partial is split into two bf16 halves and both are scaled; 8 + 8 mantissa +// bits land exactly in the fp32 accumulator. +// - `aie::downshift` on uint8 lowers to a deprecated intrinsic the build +// promotes to an error; mask and use to_float's shift argument instead, +// which also does the divide by 16 for free. + +#include "aie_kernel_utils.h" +#include +#include + +static constexpr unsigned kRowBlocks = 2; +static constexpr unsigned kRowsPerBlock = 16; +static constexpr unsigned kRows = 32; // output rows per tile +static constexpr unsigned kKBlocks = 8; // 32-wide K blocks per tile +static constexpr unsigned kKInBlock = 32; +static constexpr unsigned kTileK = 256; // K per tile +static constexpr unsigned kMetaEntries = 256; // d entries, and m entries +static constexpr unsigned kDBytes = 512; // bf16[256] +static constexpr unsigned kMetaBytes = 1024; // d then m +static constexpr unsigned kTileBytes = 5120; // whole q4 tile +// K tiles per entry point. Set per shape by the generator in granite_gemv.py: +// the largest divisor of K/256 that L1 can hold, so K = 2560 -> 5 (2 entry +// points) and down_proj's K = 8192 -> 4 (8 entry points). +#ifndef GRANITE_TILES_PER_CALL +#define GRANITE_TILES_PER_CALL 5 +#endif +static constexpr unsigned kTilesPerCall = GRANITE_TILES_PER_CALL; + +// Tokens processed per weight pass. THIS IS THE ONE KNOB THAT BEATS THE MEMORY +// BOUND: decode reads all 2.13 GB of weights per token, so B tokens sharing one +// pass divide the per-token traffic by B. It only applies to *independent* +// tokens -- concurrent requests, prefill, or speculative decoding -- since in a +// single autoregressive stream token t+1 needs token t. +// +// The cost is B times the arithmetic per byte, and at B = 1 this design is +// already at the DMA bound, so the compute becomes the wall almost immediately. +// See the task notes: useful, but bounded by this kernel's 16-lane MACs. +#ifndef GRANITE_BATCH +#define GRANITE_BATCH 1 +#endif +static constexpr unsigned kBatch = GRANITE_BATCH; + +// x is [kBatch][K], so token b starts at x + b*GRANITE_K. Only needed to stride +// between tokens; at kBatch == 1 it is never used. +#ifndef GRANITE_K +#define GRANITE_K 0 +#endif + +// `kt` selects which 256-wide slice of x this tile covers; `first` starts the +// accumulator rather than adding to it, so the K tiles chain without a separate +// zeroing pass over the output. +// +// Both are RUNTIME arguments and this is deliberately **noinline**. As a +// template with compile-time KT, every entry point instantiated its own copy of +// the body -- measured at 4736 bytes of .text each, against a core program +// memory of 16 KB, so even a single entry point plus the runtime overflowed. +// Emitted once, each entry point below is a 208-byte wrapper (measured). +// KT was only ever pointer arithmetic; it never needed to be compile-time. +__attribute__((noinline)) inline void gemv_q4_tile(const uint8_t *__restrict tile, + const bfloat16 *__restrict x, + unsigned kt, bool first, + float *__restrict y) { + event0(); + aie::set_rounding(aie::rounding_mode::conv_even); + + const bfloat16 *__restrict dp = (const bfloat16 *)tile; + const bfloat16 *__restrict mp = (const bfloat16 *)(tile + kDBytes); + const uint8_t *__restrict nib = tile + kMetaBytes; + const bfloat16 *__restrict xt = x + kt * kTileK; + + // One accumulator per row-block rather than one 32-lane accumulator for the + // tile: the two halves are updated independently, and keeping them apart + // avoids an insert/extract on every (kb, rb) step. Rows 0..15 are row-block 0 + // and 16..31 row-block 1, which is the order the host expects. + aie::accum acc[kRowBlocks][kBatch]; +#pragma clang loop unroll(full) + for (unsigned b = 0; b < kBatch; ++b) { + float *__restrict yb = y + b * kRows; + if (first) { + acc[0][b] = aie::zeros(); + acc[1][b] = aie::zeros(); + } else { + acc[0][b].from_vector(aie::load_v(yb)); + acc[1][b].from_vector(aie::load_v(yb + kRowsPerBlock)); + } + } + + // Program memory, not speed, is the binding constraint here: the entry + // points all link into ONE core program, so an unrolled body is multiplied by + // however many there are. Peano silently drops chess_* pragmas; clang's are + // honoured. +#pragma clang loop unroll(disable) + for (unsigned kb = 0; kb < kKBlocks; ++kb) { + // sum of x over this K block -- the whole cost of the Q4_1 minimum. + // + // Vectorised, and it matters far more than it looks: as 32 scalar adds this + // was ~256 scalar ops per tile, and scalar work does not overlap the vector + // pipeline. It is also pure redundancy -- xs depends only on x and kb, never + // on the weights, yet it is recomputed for every one of the 3136 tile-rows. + // Summing through an fp32 accumulator rather than in bf16 keeps it exact; + // a bf16 running sum over 32 terms would lose ~5 bits. + float xs[kBatch]; +#pragma clang loop unroll(full) + for (unsigned b = 0; b < kBatch; ++b) { + aie::accum xa; + xa.from_vector(aie::load_v(xt + b * GRANITE_K + kb * kKInBlock)); + xs[b] = aie::reduce_add(xa.template to_vector()); + } + +#pragma clang loop unroll(disable) + for (unsigned rb = 0; rb < kRowBlocks; ++rb) { + const unsigned g = kb * kRows + rb * kRowsPerBlock; + aie::vector d16 = aie::load_v(dp + g); + aie::vector m16 = aie::load_v(mp + g); + + const uint8_t *__restrict src = nib + rb * 2048 + kb * 256; + + // 16 lanes, one per row of this row-block, accumulated over the + // block's 32 k. + aie::accum part[kBatch]; +#pragma clang loop unroll(full) + for (unsigned b = 0; b < kBatch; ++b) + part[b] = aie::zeros(); + +#pragma clang loop unroll(disable) + for (unsigned kk = 0; kk < kKInBlock; kk += 8) { + aie::vector p = aie::load_v<64>(src + kk * 8); + + // High nibble is masked rather than shifted, and to_float's shift + // argument (the binary point) does the divide by 16 for free. + aie::vector flo = + aie::to_float(aie::bit_and((uint8_t)0x0F, p), 0); + aie::vector fhi = + aie::to_float(aie::bit_and((uint8_t)0xF0, p), 4); + + // Low nibble is the even weight index, so zipping at chunk size 1 + // restores weight order: [lo0, hi0, lo1, hi1, ...]. + auto [c0, c1] = aie::interleave_zip(flo, fhi, 1); + + // c0 covers k = kk..kk+3, c1 covers kk+4..kk+7, 16 rows each, so + // lane group j of c0 is k = kbase + j. + // + // `aie::mac` takes a SCALAR third operand, so x needs no broadcast + // vector -- building 8 broadcast vectors per iteration is pure code. + // `extract`'s lane index must be a COMPILE-TIME constant; with a loop + // variable it lowers to a dynamic shuffle chain. Written out, each + // extract is register selection and costs nothing. + // The weight vectors are decoded ONCE and reused by every token in the + // batch -- that reuse is the whole point of batching. Only the scalar + // activation changes per token. + const unsigned kbase = kb * kKInBlock + kk; +#pragma clang loop unroll(full) + for (unsigned b = 0; b < kBatch; ++b) { + const bfloat16 *__restrict xb = xt + b * GRANITE_K; + part[b] = aie::mac(part[b], c0.template extract(0), xb[kbase + 0]); + part[b] = aie::mac(part[b], c0.template extract(1), xb[kbase + 1]); + part[b] = aie::mac(part[b], c0.template extract(2), xb[kbase + 2]); + part[b] = aie::mac(part[b], c0.template extract(3), xb[kbase + 3]); + part[b] = aie::mac(part[b], c1.template extract(0), xb[kbase + 4]); + part[b] = aie::mac(part[b], c1.template extract(1), xb[kbase + 5]); + part[b] = aie::mac(part[b], c1.template extract(2), xb[kbase + 6]); + part[b] = aie::mac(part[b], c1.template extract(3), xb[kbase + 7]); + } + } + +#pragma clang loop unroll(full) + for (unsigned b = 0; b < kBatch; ++b) { + aie::accum sum16 = part[b]; + + // acc[row] += d*sum + m*xs. Both products must be bf16 x bf16: AIE2P has + // no fp32 vector multiplier, and aie_api compiles one and returns **zero** + // in silence. Splitting each fp32 value into two bf16 halves keeps 8 + 8 + // mantissa bits, which land exactly in the fp32 accumulator. + aie::vector hi = + sum16.template to_vector(); + aie::vector lo = + aie::sub(sum16, hi).template to_vector(); + acc[rb][b] = aie::mac(acc[rb][b], hi, d16); + acc[rb][b] = aie::mac(acc[rb][b], lo, d16); + + // Same split for the block's sum of x, which carries the minimum. + const bfloat16 xs_hi = (bfloat16)xs[b]; + const bfloat16 xs_lo = (bfloat16)(xs[b] - (float)xs_hi); + acc[rb][b] = aie::mac(acc[rb][b], m16, xs_hi); + acc[rb][b] = aie::mac(acc[rb][b], m16, xs_lo); + } + } + } + +#pragma clang loop unroll(full) + for (unsigned b = 0; b < kBatch; ++b) { + float *__restrict yb = y + b * kRows; + aie::store_v(yb, acc[0][b].template to_vector()); + aie::store_v(yb + kRowsPerBlock, acc[1][b].template to_vector()); + } + event1(); +} + +// Five K tiles per call. +// +// Program memory is the binding constraint, and it was found by measuring, not +// by reasoning: three separate theories (entry-point count, loop unrolling, +// runtime extract indices) were all wrong. Compiling the object directly and +// reading `llvm-objdump -h` settled it in seconds -- the body was 4736 B because +// it was a template on the K-tile index, so EVERY entry point instantiated its +// own copy and even one overflowed the core's 16 KB. +// +// KT was only ever pointer arithmetic (`x + KT*256`); it never needed to be +// compile-time. Runtime + noinline emits the body once per translation unit. +// +// It is `inline`, NOT `static`, and that distinction is what makes wide K +// affordable. `static` gives each entry point's translation unit a private copy; +// `inline` gives the body vague linkage, so the copies land in a COMDAT and the +// linker keeps exactly one (`llvm-objdump -t` shows the symbol as `w`). The cost +// of an entry point drops from a whole body to its 208-byte wrapper -- which is +// what lets down_proj's K = 8192 have its 8 entry points at all. +// +// The other side of the trade is L1: an element is now 5 x 5120 = 25600 B, and +// double-buffered that is 51200 B against the 63 KB budget (64 KB less ~1 KB of +// stack). Adding x (5120 B) and y still fits, and keeping depth 2 matters -- +// this kernel is bandwidth-bound, so the DMA must overlap the compute. +static inline void gemv_q4_group(const uint8_t *__restrict tiles, + const bfloat16 *__restrict x, unsigned group, + float *__restrict y) { +#pragma clang loop unroll(disable) + for (unsigned i = 0; i < kTilesPerCall; ++i) { + const unsigned kt = group * kTilesPerCall + i; + gemv_q4_tile(tiles + i * kTileBytes, x, kt, kt == 0, y); + } +} + +extern "C" { + +// One entry point per group of five K tiles. They live in separate translation +// units because IRON compiles the kernel source once per ExternalFunction: +// several functions in one .cc become several objects that each define every +// symbol, and the link fails on duplicates. +// +// Only the group index is needed -- whether to start or continue the accumulator +// follows from it (`kt == 0`), so it is not a second argument that could +// disagree with the first. +// The tiles-per-call variant is part of the SYMBOL, not just the file name. +// down_proj wants 4 tiles per call and everything else wants 5; if both variants +// produced `granite_gemv_k0`, a build cache keyed on anything other than the +// kernel source bytes could hand one shape the other's object -- which is not an +// error, just a different and wrong matmul. Distinct names make that impossible. +// The extra indirection is so GRANITE_TILES_PER_CALL expands before ## pastes. +#define GRANITE_GEMV_ENTRY__(P, B, N) void granite_gemv_p##P##b##B##_k##N(const uint8_t *__restrict t, const bfloat16 *__restrict x, float *__restrict y) { gemv_q4_group(t, x, N, y); } +#define GRANITE_GEMV_ENTRY_(P, B, N) GRANITE_GEMV_ENTRY__(P, B, N) +#define GRANITE_GEMV_ENTRY(N) GRANITE_GEMV_ENTRY_(GRANITE_TILES_PER_CALL, GRANITE_BATCH, N) + +} // extern "C" diff --git a/kernels/granite/aie/granite_qgemv_g0.cc b/kernels/granite/aie/granite_qgemv_g0.cc new file mode 100644 index 00000000..f4a46262 --- /dev/null +++ b/kernels/granite/aie/granite_qgemv_g0.cc @@ -0,0 +1,18 @@ +// OpenFFLM -- q_proj GEMV group 0, writing at a row offset so a core can +// accumulate its whole slice before the RoPE epilogue runs on it. +// One entry point per translation unit: IRON compiles the source once per +// ExternalFunction, so two here would each define both symbols. +// SPDX-License-Identifier: MIT +#define GRANITE_TILES_PER_CALL 5 +#define GRANITE_BATCH 1 +#define GRANITE_K 2560 +#include "granite_gemv.h" + +extern "C" { +void granite_qgemv_g0(const uint8_t *__restrict t, const bfloat16 *__restrict x, + float *__restrict y, unsigned row) { + // `row` selects the 32-float window inside the core's slice; without it every + // call would write y[0..31] and only the last tile-row would survive. + gemv_q4_group(t, x, 0, y + row * kRows); +} +} diff --git a/kernels/granite/aie/granite_qgemv_g1.cc b/kernels/granite/aie/granite_qgemv_g1.cc new file mode 100644 index 00000000..c4fbd4a1 --- /dev/null +++ b/kernels/granite/aie/granite_qgemv_g1.cc @@ -0,0 +1,18 @@ +// OpenFFLM -- q_proj GEMV group 1, writing at a row offset so a core can +// accumulate its whole slice before the RoPE epilogue runs on it. +// One entry point per translation unit: IRON compiles the source once per +// ExternalFunction, so two here would each define both symbols. +// SPDX-License-Identifier: MIT +#define GRANITE_TILES_PER_CALL 5 +#define GRANITE_BATCH 1 +#define GRANITE_K 2560 +#include "granite_gemv.h" + +extern "C" { +void granite_qgemv_g1(const uint8_t *__restrict t, const bfloat16 *__restrict x, + float *__restrict y, unsigned row) { + // `row` selects the 32-float window inside the core's slice; without it every + // call would write y[0..31] and only the last tile-row would survive. + gemv_q4_group(t, x, 1, y + row * kRows); +} +} diff --git a/kernels/granite/aie/granite_qkv_rope.h b/kernels/granite/aie/granite_qkv_rope.h new file mode 100644 index 00000000..92fba7b5 --- /dev/null +++ b/kernels/granite/aie/granite_qkv_rope.h @@ -0,0 +1,90 @@ +#pragma once +//===- granite_qkv_rope.h -----------------------------------*- C++ -*-===// +// +// OpenFFLM -- epilogue for a fused q/k/v projection: rotate q and k, pass v. +// SPDX-License-Identifier: MIT +// +// q, k and v all consume the same input x, so all three GEMVs share one +// dispatch with no gather. RoPE then applies to q and k but NOT to v, and a +// single core owns whole heads of each, so the whole epilogue is core-local. +// +// The accumulator holds, in order: q_heads*64 | k_heads*64 | v_len +// and the output has the same shape. Only the layout differs from +// granite_qrope.h, which is why this is a separate kernel rather than three +// calls with pointer offsets -- IRON kernels take whole buffers. +// +// cos/sin ride at the end of the x buffer (GRANITE_QROPE_XOFF): a compute tile +// has 2 input DMA channels, the weights take one and x the other, so a third +// stream for cos/sin does not exist. + +#include "aie_kernel_utils.h" +#include +#include +#include "granite_elementwise.h" + +#ifndef GRANITE_QROPE_XOFF +#define GRANITE_QROPE_XOFF 2560 +#endif + +// Rotate `n_heads` consecutive heads starting at `y` into `out`. +static inline void rope_heads(const float *__restrict y, + const aie::vector &c, + const aie::vector &s, + bfloat16 *__restrict out, unsigned n_heads) { + for (unsigned h = 0; h < n_heads; ++h) { + const float *__restrict yh = y + h * kHeadDim; + bfloat16 *__restrict oh = out + h * kHeadDim; + aie::accum alo, ahi; + alo.from_vector(aie::load_v(yh)); + ahi.from_vector(aie::load_v(yh + kHalf)); + aie::vector lo = alo.template to_vector(); + aie::vector hi = ahi.template to_vector(); + aie::store_v(oh, aie::sub(aie::mul(lo, c), aie::mul(hi, s)) + .template to_vector()); + aie::store_v(oh + kHalf, aie::add(aie::mul(hi, c), aie::mul(lo, s)) + .template to_vector()); + } +} + +// `acc` and `out` are NOT __restrict: granite_rope_ip.cc calls this with the +// output aliasing the accumulator, so the result narrows into the buffer it was +// read from. That is not a micro-optimisation -- it is 512 B of L1, and with a +// 62208 B budget it is the difference between norm+qkv+RoPE fitting in one +// dispatch at per_call 5 and not fitting at all. +// +// Safe because rope_heads loads BOTH halves of a head before storing either, +// and consecutive heads do not overlap: head h reads 64 floats at acc + 64h and +// writes 64 bfloat16 at the same base, i.e. the first half of what it just +// read. Verified by the design's own check against a host reference, not by +// inspection alone. +__attribute__((noinline)) inline void +granite_qkv_rope_impl(const float *acc, const bfloat16 *__restrict xcs, + bfloat16 *out, + unsigned q_heads, unsigned k_heads, unsigned v_len) { + event0(); + aie::set_rounding(aie::rounding_mode::conv_even); + const bfloat16 *__restrict cs = xcs + GRANITE_QROPE_XOFF; + const aie::vector c = aie::load_v(cs); + const aie::vector s = aie::load_v(cs + kHalf); + + rope_heads(acc, c, s, out, q_heads); + const unsigned qn = q_heads * kHeadDim; + rope_heads(acc + qn, c, s, out + qn, k_heads); + + // v is not rotated -- only narrowed. + const unsigned kn = qn + k_heads * kHeadDim; + for (unsigned i = 0; i < v_len; i += kHalf) { + aie::accum a; + a.from_vector(aie::load_v(acc + kn + i)); + aie::store_v(out + kn + i, a.template to_vector()); + } + event1(); +} + +extern "C" { +void granite_qkv_rope(const float *__restrict acc, + const bfloat16 *__restrict xcs, bfloat16 *__restrict out, + unsigned q_heads, unsigned k_heads, unsigned v_len) { + granite_qkv_rope_impl(acc, xcs, out, q_heads, k_heads, v_len); +} +} diff --git a/kernels/granite/aie/granite_rmsnorm.h b/kernels/granite/aie/granite_rmsnorm.h new file mode 100644 index 00000000..43648d19 --- /dev/null +++ b/kernels/granite/aie/granite_rmsnorm.h @@ -0,0 +1,100 @@ +#pragma once +//===- granite_rmsnorm.h ------------------------------------*- C++ -*-===// +// +// OpenFFLM -- Llama/Granite RMSNorm on the AIE core. +// SPDX-License-Identifier: MIT +// +// y[c] = x[c] * rsqrt(mean(x^2) + eps) * w[c] +// +// WHY NOT aie_kernels/aie2p/rms_norm.cc +// ------------------------------------- +// That kernel hardcodes `const float gamma = 1.0f` and never applies the +// per-channel weight tensor, and its `epsilon` is a `constexpr`. It normalises +// correctly and then throws away the learned scale -- output of the right +// magnitude and the wrong value, which no shape check catches. `cols` IS a +// runtime argument there, so only the weight was ever the problem. +// +// PRECISION, AND THE TRAP IT AVOIDS +// --------------------------------- +// AIE2P has no fp32 vector multiplier: `aie::mul(vector, vector)` +// compiles and returns **zero**, silently. So `x * w * inv` cannot be done in +// fp32 vectors. Instead: +// +// * x*w is a bf16 x bf16 product, which lands EXACTLY in an fp32 accumulator; +// * that fp32 partial is split into two bf16 halves (8 + 8 mantissa bits); +// * the scalar `inv` is split the same way; +// * three of the four cross terms are accumulated (hi*hi, lo*hi, hi*lo), +// giving ~2^-17 relative -- far below the bf16 output's own 2^-9, so the +// stored result is correctly rounded. +// +// Summing x^2 through an fp32 accumulator rather than a bf16 running sum is the +// other half of it: 2560 bf16 additions would lose ~6 bits. + +#include "aie_kernel_utils.h" +#include +#include + +#ifndef GRANITE_NORM_EPS +#define GRANITE_NORM_EPS 1e-5f +#endif + +static constexpr unsigned kNormVec = 32; + +// cols is a RUNTIME argument, so one build serves hidden 2560 and any other +// width the model happens to use. +// `x` and `y` are deliberately NOT __restrict: the fused prologue in +// granite_norm_gemv.py calls this with y == x, normalising in place inside the +// activation fifo's own element. That is what keeps the fused design's fixed L1 +// at 10240 B instead of 15360, which is in turn what lets the weight double +// buffer keep per_call = 5 -- and per_call is worth more than the aliasing +// optimisation this gives up. The two passes read x fully before writing y, so +// in-place is correct; it was only ever the promise that was wrong. +// `w` is not __restrict either, for a reason that only shows up on hardware: a +// compute tile has 2 input DMA channels, and the fused design already spends +// one on weights. So x and the norm weight cannot arrive as separate streams -- +// they share one fifo element, w sitting at x + cols. Disjoint ranges of one +// object do not alias in practice, but declaring otherwise would be another +// promise this code cannot keep. +__attribute__((noinline)) inline void +granite_rms_norm_impl(const bfloat16 *x, const bfloat16 *w, + bfloat16 *y, unsigned cols) { + event0(); + aie::set_rounding(aie::rounding_mode::conv_even); + + // sum of x^2, accumulated in fp32. + aie::accum sq = aie::zeros(); + for (unsigned i = 0; i < cols; i += kNormVec) { + aie::vector v = aie::load_v(x + i); + sq = aie::mac(sq, v, v); + } + const float sum_sq = aie::reduce_add(sq.template to_vector()); + + const float inv = aie::invsqrt(sum_sq / (float)cols + GRANITE_NORM_EPS); + const bfloat16 inv_hi = (bfloat16)inv; + const bfloat16 inv_lo = (bfloat16)(inv - (float)inv_hi); + + for (unsigned i = 0; i < cols; i += kNormVec) { + aie::vector xv = aie::load_v(x + i); + aie::vector wv = aie::load_v(w + i); + + // x*w exactly, in fp32. + aie::accum xw = aie::mul(xv, wv); + aie::vector hi = xw.template to_vector(); + aie::vector lo = + aie::sub(xw, hi).template to_vector(); + + aie::accum out = aie::zeros(); + out = aie::mac(out, hi, inv_hi); + out = aie::mac(out, lo, inv_hi); + out = aie::mac(out, hi, inv_lo); + aie::store_v(y + i, out.template to_vector()); + } + event1(); +} + +extern "C" { +void granite_rms_norm(const bfloat16 *__restrict x, const bfloat16 *__restrict w, + bfloat16 *__restrict y, unsigned cols) { + granite_rms_norm_impl(x, w, y, cols); +} +} diff --git a/kernels/granite/aie/granite_rmsnorm_ip.cc b/kernels/granite/aie/granite_rmsnorm_ip.cc new file mode 100644 index 00000000..defa19c3 --- /dev/null +++ b/kernels/granite/aie/granite_rmsnorm_ip.cc @@ -0,0 +1,17 @@ +// OpenFFLM -- RMSNorm in place over a combined [x | weight] buffer. +// SPDX-License-Identifier: MIT +// +// ONE pointer, because a compute tile has only 2 input DMA channels and the +// fused norm+GEMV design spends one of them on weights. The activation and the +// norm weight therefore share a single fifo element: x in the first `cols`, +// the weight in the second. The result is written back over x, which costs no +// extra L1 -- and that is what lets the weight double buffer keep per_call 5. +// +// See granite_rmsnorm.h for why neither pointer is __restrict. +#include "granite_rmsnorm.h" + +extern "C" { +void granite_rms_norm_ip(bfloat16 *xw, unsigned cols) { + granite_rms_norm_impl(xw, xw + cols, xw, cols); +} +} diff --git a/kernels/granite/aie/granite_rope_ip.cc b/kernels/granite/aie/granite_rope_ip.cc new file mode 100644 index 00000000..be81ca77 --- /dev/null +++ b/kernels/granite/aie/granite_rope_ip.cc @@ -0,0 +1,19 @@ +// OpenFFLM -- RoPE epilogue writing back into the accumulator it reads. +// SPDX-License-Identifier: MIT +// +// The fused norm+qkv+RoPE design has 128 B of L1 to spare, so the output cannot +// have a buffer of its own: the float32 accumulator IS the output fifo element, +// and the rotated result narrows into its first half. See the aliasing note in +// granite_qkv_rope.h. +// +// cos|sin sit after x and the norm weight, hence XOFF = 2 * K. +#define GRANITE_QROPE_XOFF 5120 +#include "granite_qkv_rope.h" + +extern "C" { +void granite_rope_ip(float *acc, const bfloat16 *__restrict xcs, + unsigned q_heads, unsigned k_heads, unsigned v_len) { + granite_qkv_rope_impl(acc, xcs, reinterpret_cast(acc), + q_heads, k_heads, v_len); +} +} diff --git a/kernels/granite/aie/granite_swiglu_ip.cc b/kernels/granite/aie/granite_swiglu_ip.cc new file mode 100644 index 00000000..c253855e --- /dev/null +++ b/kernels/granite/aie/granite_swiglu_ip.cc @@ -0,0 +1,14 @@ +// OpenFFLM -- SwiGLU in place over a combined [gate | up] buffer. +// SPDX-License-Identifier: MIT +// +// ONE pointer, because the fused SwiGLU+down_proj design has two input DMA +// channels per core and both are spoken for -- weights and activation. gate and +// up therefore share the activation element, and the result is written back +// over gate so no destination buffer is needed. See granite_elementwise.h. +#include "granite_elementwise.h" + +extern "C" { +void granite_swiglu_ip(bfloat16 *gu, unsigned n) { + granite_swiglu_impl(gu, gu + n, gu, n); +} +} diff --git a/kernels/granite/geometry.json b/kernels/granite/geometry.json new file mode 100644 index 00000000..9bdeb6da --- /dev/null +++ b/kernels/granite/geometry.json @@ -0,0 +1,12 @@ +{ + "hidden_size": 2560, + "intermediate_size": 8192, + "num_hidden_layers": 40, + "num_attention_heads": 40, + "num_key_value_heads": 8, + "head_dim": 64, + "rms_norm_eps": 1e-05, + "vocab_size": 100352, + "rope_theta": 10000000, + "attention_multiplier": 0.125 +} diff --git a/kernels/granite/iron/granite_attn_o.py b/kernels/granite/iron/granite_attn_o.py new file mode 100644 index 00000000..6e2fa13b --- /dev/null +++ b/kernels/granite/iron/granite_attn_o.py @@ -0,0 +1,407 @@ +r"""attention (40 heads) + o_proj in ONE dispatch, on disjoint cores. + +WHY +--- +0148 measured the per-dispatch floor at ~200 us: RMSNorm on 2560 values costs +244 us, SwiGLU on 8192 costs 205, and RoPE cost NOTHING when it moved inside an +existing dispatch. A granite layer is 6 dispatches and 2.337 ms, so about 1.2 ms +of it is the cost of asking. Removing a dispatch is worth ~200 us wherever it +can be done without losing cores. + +WHY DISJOINT CORES, AND WHY THE WHOLE BLOCK DOES NOT FIT +-------------------------------------------------------- +Fusing all of norm+qkv+RoPE, attention and o_proj would want 7 columns for qkv, +2 for attention and 5 for o_proj -- 14 of the 8 the array has. Cores would then +have to serve several phases, which means one ObjectFifo carrying different +payloads in different phases (the granite_mlp_full.py pattern) against a much +tighter L1. + +This pair fits disjointly: + + columns 0-4 o_proj 20 cores, 4 tile-rows each + columns 5-6 attention 8 cores, one kv head each + column 7 unused + +so no core does both, no fifo is reused, and the shim budget is 10 MM2S and +7 S2MM against 16 of each. + + call c:\dev\mlir-aie\iron_env.cmd + python designs\granite_gemv\granite_attn_o.py +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +import aie.iron as iron +from aie.iron import (Buffer, CompileTime, In, ObjectFifo, Out, Program, + Runtime, TaskGroup, Worker) +from aie.iron.controlflow import range_ +from aie.iron.device import from_name +from aie.iron.kernel import ExternalFunction +from aie.helpers.taplib import TensorTiler2D +from aie.utils.benchmark import run_iters + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parent.parent / "common")) +import q4nx # noqa: E402 +from granite_gemv import (AIE, MODEL, ROWS_PER_TILE, TILE_BYTES, TILE_K, # noqa: E402 + _include_dirs, ensure_entry_points, + projection_shape, reference, tiles_per_call) +from granite_gemv32 import (ROWS_PER_COL, permute_weights, # noqa: E402 + unpermute_y) + +SRC_BLOCK_H = str(AIE / "granite_attn_block_h.cc") +SRC_FINISH_H = str(AIE / "granite_attn_finish_h.cc") +HD = 64 # granite head_dim +BLK = 32 # KV positions per call, matches GRANITE_ATTN_BLOCK +ST_STRIDE = 128 # floats; 512 B, 128-byte aligned. granite_attn_block_h.cc +O_COLS = 4 # columns 0..3 -> o_proj, 16 cores, 5 tile-rows each +A_COLS = 4 # columns 4..7 -> attention +A_ROWS = 2 # attention cores per column; 4 x 2 = 8 kv heads + +# Why 4 x 2 and not 2 x 4: a memtile has about 6 DMA channels each way, and an +# attention column needs a split for q (1 in, 4 out), a split for kv (1 in, +# 4 out) and a join for the result (4 in, 1 out) -- 6 in and 9 out at four cores +# per column, which does not place: +# +# error: no MemTile has sufficient DMA capacity for 1 input/4 output channels +# +# At two cores per column it is 4 in and 5 out. o_proj gives up a column for it +# (20 cores -> 16), which costs little because o_proj is floor-dominated: 4.1 MB +# at 15.8 GB/s is nearly all fixed cost already. + + +@iron.jit(aiecc_flags=["--alloc-scheme=basic-sequential"]) +def granite_attn_o(a_q: In, a_kv: In, scratch_o: Out, scratch_i: In, + a_w: In, c_y: Out, *, seq: CompileTime[int], + tile_rows: CompileTime[int], k: CompileTime[int], + per_call: CompileTime[int], + q_per: CompileTime[int] = 5, + o_cols: CompileTime[int] = O_COLS, + a_cols: CompileTime[int] = A_COLS, + a_rows: CompileTime[int] = A_ROWS): + n_blk = (seq + BLK - 1) // BLK + n_kv = a_cols * a_rows + o_cores = o_cols * ROWS_PER_COL + per_core = tile_rows // o_cores + k_tiles = k // TILE_K + n_entry = k_tiles // per_call + call_bytes = per_call * TILE_BYTES + chunks = per_core * n_entry + + srcs = ensure_entry_points(n_entry, per_call, False, 1, k) + + # ---- attention side ---- + # No Q_PAD: 320 bf16 is 640 B, clear of the 128-byte transfer that arrives + # as zeros, and an unpadded output makes the scratch the plain head-ordered + # 2560-wide vector o_proj wants, with no holes to skip. + qc_ty = np.ndarray[(q_per * HD,), np.dtype[bfloat16]] + kv_ty = np.ndarray[(2 * BLK * HD,), np.dtype[bfloat16]] + st_ty = np.ndarray[(q_per * ST_STRIDE,), np.dtype[np.float32]] + ao_ty = np.ndarray[(q_per * HD,), np.dtype[bfloat16]] + + # ---- o_proj side ---- + w_l1_ty = np.ndarray[(call_bytes,), np.dtype[np.uint8]] + w_l2_ty = np.ndarray[(ROWS_PER_COL * call_bytes,), np.dtype[np.uint8]] + y_l1_ty = np.ndarray[(ROWS_PER_TILE,), np.dtype[np.float32]] + y_l2_ty = np.ndarray[(ROWS_PER_COL * ROWS_PER_TILE,), np.dtype[np.float32]] + x_ty = np.ndarray[(k,), np.dtype[bfloat16]] + + all_q_ty = np.ndarray[(n_kv * q_per * HD,), np.dtype[bfloat16]] + all_kv_ty = np.ndarray[(n_kv * n_blk * 2 * BLK * HD,), np.dtype[bfloat16]] + attn_ty = np.ndarray[(k,), np.dtype[bfloat16]] + w_ty = np.ndarray[(tile_rows * k_tiles * TILE_BYTES,), np.dtype[np.uint8]] + y_ty = np.ndarray[(tile_rows * ROWS_PER_TILE,), np.dtype[np.float32]] + + blk = ExternalFunction("granite_attn_block_h", source_file=SRC_BLOCK_H, + arg_types=[qc_ty, kv_ty, st_ty, np.int32, np.int32, + np.int32], + include_dirs=_include_dirs()) + fin = ExternalFunction("granite_attn_finish_h", source_file=SRC_FINISH_H, + arg_types=[st_ty, ao_ty, np.int32], + include_dirs=_include_dirs()) + gemv = [ExternalFunction(f"granite_gemv_p{per_call}b1_k{i}", + source_file=str(srcs[i]), + arg_types=[w_l1_ty, x_ty, y_l1_ty], + include_dirs=_include_dirs()) + for i in range(n_entry)] + + # Per COLUMN, not per core. Eight cores own eight distinct kv heads, so + # there is nothing to broadcast -- but giving each its own shim stream costs + # 8 q + 8 kv + 5 weights + 1 activation = 22 MM2S against the 16 the device + # has, and the placer says so: + # + # error: no ShimNOCTile has sufficient DMA capacity ... + # + # Through the memtile it is 2 + 2 + 5 + 1 = 10 MM2S and 2 + 5 = 7 S2MM. + q_l2_ty = np.ndarray[(a_rows * q_per * HD,), np.dtype[bfloat16]] + kv_l2_ty = np.ndarray[(a_rows * 2 * BLK * HD,), np.dtype[bfloat16]] + ao_l2_ty = np.ndarray[(a_rows * q_per * HD,), np.dtype[bfloat16]] + of_x = ObjectFifo(x_ty, name="aox", depth=1) + q_l3l2, kv_l3l2, ao_l2l3 = [], [], [] + q_cores, kv_cores, ao_cores = [], [], [] + for c in range(a_cols): + qf = ObjectFifo(q_l2_ty, name=f"aoqL2_{c}", depth=2) + q_l3l2.append(qf) + q_cores.append(qf.cons().split( + [r * q_per * HD for r in range(a_rows)], + obj_types=[qc_ty] * a_rows, + names=[f"aoq_{c}_{r}" for r in range(a_rows)])) + kf = ObjectFifo(kv_l2_ty, name=f"aokvL2_{c}", depth=2) + kv_l3l2.append(kf) + kv_cores.append(kf.cons().split( + [r * 2 * BLK * HD for r in range(a_rows)], + obj_types=[kv_ty] * a_rows, + names=[f"aokv_{c}_{r}" for r in range(a_rows)])) + af = ObjectFifo(ao_l2_ty, name=f"aoaL2_{c}", depth=2) + ao_l2l3.append(af) + ao_cores.append(af.prod().join( + [r * q_per * HD for r in range(a_rows)], + depths=[1] * a_rows, + obj_types=[ao_ty] * a_rows, + names=[f"aoa_{c}_{r}" for r in range(a_rows)])) + + w_l3l2, y_l2l3, w_cores, y_cores = [], [], [], [] + for c in range(o_cols): + wf = ObjectFifo(w_l2_ty, name=f"aowL2_{c}", depth=2) + w_l3l2.append(wf) + w_cores.append(wf.cons().split( + [r * call_bytes for r in range(ROWS_PER_COL)], + obj_types=[w_l1_ty] * ROWS_PER_COL, + names=[f"aow_{c}_{r}" for r in range(ROWS_PER_COL)])) + yf = ObjectFifo(y_l2_ty, name=f"aoyL2_{c}", depth=2) + y_l2l3.append(yf) + y_cores.append(yf.prod().join( + [r * ROWS_PER_TILE for r in range(ROWS_PER_COL)], + obj_types=[y_l1_ty] * ROWS_PER_COL, + names=[f"aoy_{c}_{r}" for r in range(ROWS_PER_COL)])) + + def attn_body(qi, kvi, oo, state, kb, fn): + qe = qi.acquire(1) + ke = kvi.acquire(1) + # Python range, not range_: the head loop is unrolled at build time. + for h in range(q_per): + kb(qe, ke, state, BLK, 1, h) + kvi.release(1) + if n_blk > 1: + for _ in range_(n_blk - 1): + ke = kvi.acquire(1) + for h in range(q_per): + kb(qe, ke, state, BLK, 0, h) + kvi.release(1) + oe = oo.acquire(1) + for h in range(q_per): + fn(state, oe, h) + oo.release(1) + qi.release(1) + + def gemv_body(win, xin, yout, *ks): + xe = xin.acquire(1) + for _ in range_(per_core): + ye = yout.acquire(1) + for fn in ks: + we = win.acquire(1) + fn(we, xe, ye) + win.release(1) + yout.release(1) + xin.release(1) + + workers = [] + for c in range(o_cols): + for r in range(ROWS_PER_COL): + workers.append(Worker( + gemv_body, + fn_args=[w_cores[c][r].cons(), of_x.cons(), + y_cores[c][r].prod(), *gemv], + stack_size=0xD00)) + for c in range(a_cols): + for r in range(a_rows): + state = Buffer(np.ndarray[(q_per * ST_STRIDE,), np.dtype[np.float32]], + name=f"aost{c}_{r}") + workers.append(Worker( + attn_body, + fn_args=[q_cores[c][r].cons(), kv_cores[c][r].cons(), + ao_cores[c][r].prod(), state, blk, fin], + stack_size=0xD00)) + + col_w = chunks * ROWS_PER_COL * call_bytes + col_y = per_core * ROWS_PER_COL * ROWS_PER_TILE + w_taps = TensorTiler2D.simple_tiler((1, o_cols * col_w), (1, col_w)) + y_taps = TensorTiler2D.simple_tiler((1, o_cols * col_y), (1, col_y)) + col_q = a_rows * q_per * HD + col_kv = a_rows * n_blk * 2 * BLK * HD + q_taps = TensorTiler2D.simple_tiler((1, a_cols * col_q), (1, col_q)) + kv_taps = TensorTiler2D.simple_tiler((1, a_cols * col_kv), (1, col_kv)) + a_taps = TensorTiler2D.simple_tiler((1, a_cols * col_q), (1, col_q)) + + def sequence(t_q, t_kv, t_so, t_si, t_w, t_y, qp, kvp, aoc, wp, xp, yc): + # Phase 1: attention. The o_proj cores have nothing to do and are simply + # not mentioned -- a worker with no traffic in a phase costs nothing. + tg1 = TaskGroup() + for c in range(a_cols): + qp[c].fill(t_q, tap=q_taps[c], group=tg1) + kvp[c].fill(t_kv, tap=kv_taps[c], group=tg1) + aoc[c].drain(t_so, tap=a_taps[c], wait=True, group=tg1) + tg1.finish() + # Phase 2: o_proj over the gathered attention output. A SEPARATE task + # group, for the reason granite_mlp_full.py records: one group spanning + # both phases cannot complete, because finish() would be waiting on a + # fill the core cannot consume until finish() returns. Deadlock, not an + # error. + tg2 = TaskGroup() + xp.fill(t_si, group=tg2) + for c in range(o_cols): + wp[c].fill(t_w, tap=w_taps[c], group=tg2) + yc[c].drain(t_y, tap=y_taps[c], wait=True, group=tg2) + tg2.finish() + + rt = Runtime(sequence, [all_q_ty, all_kv_ty, attn_ty, attn_ty, w_ty, y_ty, + [f.prod() for f in q_l3l2], + [f.prod() for f in kv_l3l2], + [f.cons() for f in ao_l2l3], + [f.prod() for f in w_l3l2], of_x.prod(), + [f.cons() for f in y_l2l3]]) + return Program(iron.get_current_device(), rt, workers=workers).resolve_program() + + +def main() -> int: + import argparse + ap = argparse.ArgumentParser() + ap.add_argument("--seq", type=int, default=64) + ap.add_argument("--iters", type=int, default=200) + a = ap.parse_args(sys.argv[1:]) + seq, n_blk = a.seq, (a.seq + BLK - 1) // BLK + + cfg = json.loads((MODEL / "config.json").read_text(encoding="utf-8")) + n_kv_model = cfg["num_key_value_heads"] + q_per = cfg["num_attention_heads"] // n_kv_model + scale = cfg["attention_multiplier"] + name = "model.layers.0.self_attn.o_proj.weight" + n, k = projection_shape(name, cfg) + tile_rows, k_tiles = n // ROWS_PER_TILE, k // TILE_K + n_kv = A_COLS * A_ROWS + assert n_kv == n_kv_model, ( + f"{A_COLS} x {A_ROWS} gives {n_kv} cores, model has {n_kv_model} kv heads") + o_cores = O_COLS * ROWS_PER_COL + assert tile_rows % o_cores == 0 + per_core = tile_rows // o_cores + per_call = tiles_per_call(k_tiles, 1, k) + n_entry = k_tiles // per_call + + f = q4nx.Q4NX(MODEL / "model.q4nx") + off, _ = f.header[name]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + off) + raw = fh.read(tile_rows * k_tiles * TILE_BYTES) + w = permute_weights(raw, O_COLS, per_core, n_entry, per_call * TILE_BYTES) + + rng = np.random.default_rng(0) + q = rng.standard_normal((n_kv, q_per, HD)).astype(np.float32).astype(bfloat16) + K = rng.standard_normal((n_kv, n_blk * BLK, HD)).astype(np.float32).astype(bfloat16) + V = rng.standard_normal((n_kv, n_blk * BLK, HD)).astype(np.float32).astype(bfloat16) + # A column's memtile object holds its FOUR cores' block i back to back, and + # split() hands each core its own 4096 bf16. So the layout is + # column-major, then block, then core -- not core-major as it is when every + # core has its own shim stream. + kv = np.concatenate([ + np.concatenate([ + np.concatenate([ + np.concatenate([K[col * A_ROWS + r, i * BLK:(i + 1) * BLK].reshape(-1), + V[col * A_ROWS + r, i * BLK:(i + 1) * BLK].reshape(-1)]) + for r in range(A_ROWS)]) + for i in range(n_blk)]) + for col in range(A_COLS)]) + + iron.set_current_device(from_name("npu2", n_cols=None)) + scratch = iron.zeros(k, dtype=bfloat16, device="npu") + c_y = iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu") + b = run_iters(granite_attn_o, + iron.tensor(q.reshape(-1), dtype=bfloat16, device="npu"), + iron.tensor(kv, dtype=bfloat16, device="npu"), + scratch, scratch, + iron.tensor(w, dtype=np.uint8, device="npu"), c_y, + seq=seq, tile_rows=tile_rows, k=k, per_call=per_call, + q_per=q_per, o_cols=O_COLS, a_cols=A_COLS, a_rows=A_ROWS, + warmup=1, iters=a.iters) + got = unpermute_y(c_y.numpy().copy(), O_COLS, per_core).astype(np.float64) + + # Reference: attention for all 40 heads, THEN o_proj on the concatenation. + # Both together, because checking the halves separately is what let a wrong + # composition pass earlier in this workstream. + attn = np.empty(n_kv * q_per * HD, np.float64) + for c in range(n_kv): + Kc, Vc = K[c, :seq].astype(np.float64), V[c, :seq].astype(np.float64) + for h in range(q_per): + s = (Kc @ q[c, h].astype(np.float64)) * scale + s -= s.max() + e = np.exp(s) + attn[(c * q_per + h) * HD:(c * q_per + h + 1) * HD] = (e / e.sum()) @ Vc + ref = reference(raw, attn.astype(np.float32).astype(bfloat16), + tile_rows, k_tiles).astype(np.float64) + + sc = scratch.numpy().astype(np.float64) + a_rel = np.abs(sc - attn).max() / (np.abs(attn).max() + 1e-30) + rel = np.abs(got - ref).max() / (np.abs(ref).max() + 1e-30) + g1, r1 = got.ravel(), ref.ravel() + cos = float(g1 @ r1 / (np.linalg.norm(g1) * np.linalg.norm(r1) + 1e-30)) + ok = cos > 0.999 and rel < 8e-2 and a_rel < 8e-2 + mb = len(raw) / 1e6 + us = b.npu.avg_us + print(f"attention (40 heads) + o_proj, ONE dispatch seq {seq}") + print(f" {O_COLS} cols o_proj ({o_cores} cores) + {A_COLS}x{A_ROWS} " + f"attention ({n_kv} cores), per_call {per_call}, {mb:.1f} MB") + print(f" cosine {cos:.8f} max rel err {rel:.3e}") + # The intermediate separately: if the scratch matches, phase 1 is right and + # any error is phase 2's. Guessing which half is at fault cost three builds + # on the standalone attention design. + print(f" attention scratch vs reference: {a_rel:.3e} " + f"{'(phase 1 OK)' if a_rel < 8e-2 else '(PHASE 1 WRONG)'}") + print(f" {us:.1f} us device {b.e2e.avg_us:.1f} us wall") + print(f" separate: attention 278 + o_proj 260 = 538 us") + print(f" [fused == attention then o_proj on the host] " + f"{'PASS' if ok else 'FAIL'}") + return 0 if ok else 1 + + +# -------------------------------------------------------------------------- +# build_artifact: produce the xclbin WITHOUT model weights. +# +# iron.jit keys its cache on argument shapes and dtypes, not contents, so an +# artefact built from zeros is bit-identical to one built from real weights. +# That is what lets the in-tree build run from a clean checkout with nothing +# but the toolchain -- see kernels/CONVENTION.md. main() below is unchanged and +# remains the developer path: it needs the model and checks the result. +# -------------------------------------------------------------------------- + + +def build_artifact(geometry: dict, seq: int = 64) -> None: + n_blk = (seq + BLK - 1) // BLK + n_kv_model = geometry["num_key_value_heads"] + q_per = geometry["num_attention_heads"] // n_kv_model + n, k = projection_shape("o_proj", geometry) + tile_rows, k_tiles = n // ROWS_PER_TILE, k // TILE_K + n_kv = A_COLS * A_ROWS + per_call = tiles_per_call(k_tiles, 1, k) + w_bytes = tile_rows * k_tiles * TILE_BYTES + + iron.set_current_device(from_name("npu2", n_cols=None)) + scratch = iron.zeros(k, dtype=bfloat16, device="npu") + granite_attn_o( + iron.zeros(n_kv * q_per * HD, dtype=bfloat16, device="npu"), + iron.zeros(n_kv * n_blk * 2 * BLK * HD, dtype=bfloat16, device="npu"), + scratch, scratch, + iron.zeros(w_bytes, dtype=np.uint8, device="npu"), + iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu"), + seq=seq, tile_rows=tile_rows, k=k, per_call=per_call, + q_per=q_per, o_cols=O_COLS, a_cols=A_COLS, a_rows=A_ROWS) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kernels/granite/iron/granite_gemv.py b/kernels/granite/iron/granite_gemv.py new file mode 100644 index 00000000..6c86c4b7 --- /dev/null +++ b/kernels/granite/iron/granite_gemv.py @@ -0,0 +1,399 @@ +r"""Any granite-4.2-3B projection on the NPU: y[N] = W[N, K] @ x[K], W in q4nx q4. + +Every matmul in granite is the same kernel at a different (N, K) -- the seven +projections plus lm_head differ only in shape, so one design covers the lot: + + q_proj 2560 x 2560 gate_proj 8192 x 2560 lm_head 100352 x 2560 + k_proj 512 x 2560 up_proj 8192 x 2560 + v_proj 512 x 2560 down_proj 2560 x 8192 <- the only K != 2560 + o_proj 2560 x 2560 + +**(N, K) cannot be read off the file.** q4nx stores a tiled shape +`[N/32 * K/256, 5120]`, and `gate_proj` and `down_proj` are BOTH `[2560, 5120]` +-- 256 tile-rows x 10 K-tiles against 80 x 32. The two factor differently and a +GEMV against the wrong factoring is silently a different (wrong) matmul, not an +error. So (N, K) comes from `config.json` and the product is checked against the +stored shape. + + call c:\dev\mlir-aie\iron_env.cmd + python designs\granite_gemv\granite_gemv.py --tensor lm_head.weight + python designs\granite_gemv\granite_gemv.py --all :: every shape +""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +import aie.iron as iron +from aie.iron import (CompileTime, In, ObjectFifo, Out, Program, Runtime, + TaskGroup, Worker) +from aie.iron.controlflow import range_ +from aie.iron.device import from_name +from aie.iron.kernel import ExternalFunction +from aie.helpers.taplib import TensorTiler2D +from aie.utils import config +from aie.utils.benchmark import run_iters + +HERE = Path(__file__).parent +AIE = HERE.parent / "aie" # the device C++ this design compiles +# Generated entry points go to a build directory, never beside tracked source. +# GRANITE_GEN_DIR is set by build_kernels.py; the fallback keeps a bare +# `python granite_gemv.py` working for development. +GEN = Path(os.environ.get("GRANITE_GEN_DIR", "")) or (HERE.parent / "_generated") +sys.path.insert(0, str(HERE.parent.parent / "common")) +import q4nx # noqa: E402 + +TILE_BYTES = q4nx.Q4_TILE_BYTES # 5120 +ROWS_PER_TILE = q4nx.TILE_ROWS # 32 +TILE_K = q4nx.TILE_K # 256 + +# L1 is the cap on how many K tiles one call may hold: an element is +# TILES_PER_CALL * 5120 B and it is double-buffered, against a 63 KB budget +# (64 KB less ~1 KB of stack). 5 * 5120 * 2 = 51200 B fits; 6 would not. +MAX_TILES_PER_CALL = 5 + +# one result object per token: 32 floats, double-buffered +TILE_ROWS_BYTES = 32 * 4 * 2 + +MODEL = Path.home() / ".cache" / "openfflm" / "Granite-4.2-3B-NPU2" + + +# L1 less the worker stack (0xD00). Set from evidence, not from the 63 KB rule +# of thumb: batch 2 with 5 tiles per call needs 61952 B of buffers and builds and +# runs, which a 63 KB - stack budget (61184) would have wrongly rejected -- and +# rejecting it is not a safe error, it silently halves the DMA element and cost +# ~20% throughput before the cause was spotted. +L1_BUDGET = 64 * 1024 - 0xD00 # 62208, verified against 61952 + + +def tiles_per_call(k_tiles: int, batch: int = 1, k: int = 0) -> int: + """Largest divisor of k_tiles whose double-buffered weights fit in L1. + + Entry points are nearly free since the kernel body has vague linkage (see + granite_gemv.h), so this maximises the DMA element: K = 2560 -> 5 tiles x 2 + entry points, K = 8192 -> 4 x 8. + + **Batching competes for the same L1.** The activation buffer is batch*K*2 + bytes and grows with the batch while the weights do not, so a batch that is + free in DMA terms still shrinks the weight element. At K = 2560 a batch of 4 + takes 20 KB of activations, which pushes 5 tiles (51200 B double-buffered) + over the budget -- measured as `allocated buffers exceeded available memory`, + at MLIR level, not as anything the kernel could report. + """ + fixed = batch * k * 2 + batch * TILE_ROWS_BYTES # activations + results + for d in range(min(MAX_TILES_PER_CALL, k_tiles), 0, -1): + if k_tiles % d == 0 and d * TILE_BYTES * 2 + fixed <= L1_BUDGET: + return d + return 1 + + +def projection_shape(name: str, cfg: dict) -> tuple[int, int]: + """(N, K) for a tensor name, from the config -- NOT from the stored shape.""" + h, i = cfg["hidden_size"], cfg["intermediate_size"] + hd = cfg.get("head_dim") or h // cfg["num_attention_heads"] + q = cfg["num_attention_heads"] * hd + kv = cfg["num_key_value_heads"] * hd + if name == "lm_head.weight": + return cfg["vocab_size"], h + leaf = name.rsplit(".", 2)[-2] if name.endswith(".weight") else name + return { + "q_proj": (q, h), "k_proj": (kv, h), "v_proj": (kv, h), "o_proj": (h, q), + "gate_proj": (i, h), "up_proj": (i, h), "down_proj": (h, i), + }[leaf] + + +NULL_SRC = """// OpenFFLM -- DMA probe, group {i} (batch {batch}). GENERATED. +// No arithmetic: measures what the weight stream alone sustains, which +// separates 'the kernel is slow' from 'the memory path is slow'. +// SPDX-License-Identifier: MIT +#include "granite_gemv.h" + +extern "C" {{ +void granite_gemv_p{per_call}b{batch}_k{i}(const uint8_t *__restrict t, + float *__restrict y) {{ + // One byte, so the load is not elided; the DMA has already moved the + // whole object by the time this runs. No activation argument: without an + // x stream the design spends all 16 shim MM2S channels on weights, which + // is what lets it reach 16 cores where the GEMV caps at 8. + y[0] += (float)t[0]; +}} +}} +""" + + +def ensure_entry_points(n_entry: int, per_call: int, + null: bool = False, batch: int = 1, + k: int = 0) -> list[Path]: + """Write one .cc per entry point. + + They must be separate translation units: IRON compiles the kernel source + once per ExternalFunction, so several entry points in one .cc become several + objects that each define every symbol and the link fails on duplicates. + The shared body costs nothing extra per TU -- it is `inline`, so the copies + merge into one COMDAT at link time. + """ + out = [] + for i in range(n_entry): + stem = "granite_null" if null else "granite_gemv" + GEN.mkdir(parents=True, exist_ok=True) + p = GEN / f"{stem}_p{per_call}b{batch}_k{i}.cc" + lo, hi = i * per_call, (i + 1) * per_call - 1 + if null: + src = NULL_SRC.format(i=i, per_call=per_call, batch=batch) + if not p.is_file() or p.read_text(encoding='utf-8') != src: + p.write_text(src, encoding='utf-8', newline=chr(10)) + out.append(p) + continue + src = ( + f"// OpenFFLM -- K tile group {i} (tiles {lo}..{hi}) of a granite " + f"q4nx GEMV.\n" + f"// GENERATED by granite_gemv.py -- edit the generator, not this.\n" + f"// See granite_gemv.h.\n" + f"// SPDX-License-Identifier: MIT\n" + f"#define GRANITE_TILES_PER_CALL {per_call}\n" + f"#define GRANITE_BATCH {batch}\n" + f"#define GRANITE_K {k}\n" + f'#include "granite_gemv.h"\n\n' + f'extern "C" {{\n' + f"GRANITE_GEMV_ENTRY({i})\n" + f"}}\n" + ) + # Only rewrite on change: IRON caches compiled kernels by source hash, + # and rewriting identical bytes would still invalidate mtime-based caches. + if not p.is_file() or p.read_text(encoding="utf-8") != src: + p.write_text(src, encoding="utf-8", newline="\n") + out.append(p) + return out + + +def _include_dirs() -> list[str]: + from aie.iron.kernels._common import _detect_arch, _include_dirs as base + + inc = base() + root = Path(config.cxx_header_path()) / "aie_kernels" + inc.append(str(root)) + inc.append(str(root / _detect_arch())) + # This tree's own device headers. Needed because the generated entry points + # live in a build directory, so their quoted #include cannot resolve + # relative to themselves. + inc.append(str(AIE)) + return inc + + +@iron.jit(aiecc_flags=["--alloc-scheme=basic-sequential"]) +def granite_gemv(w: In, x: In, y: Out, *, tile_rows: CompileTime[int], + k: CompileTime[int], n_cores: CompileTime[int] = 1, + null: CompileTime[bool] = False, + per_call: CompileTime[int] = 5): + # Explicit, not derived: see the note in granite_gemv32.py -- iron.jit's + # cache key never sees a value computed inside the generator. + k_tiles = k // TILE_K + n_entry = k_tiles // per_call + call_bytes = per_call * TILE_BYTES + row_bytes = k_tiles * TILE_BYTES + per_core = tile_rows // n_cores + + srcs = ensure_entry_points(n_entry, per_call, null, 1, k) + tile_ty = np.ndarray[(call_bytes,), np.dtype[np.uint8]] + x_ty = np.ndarray[(k,), np.dtype[bfloat16]] + acc_ty = np.ndarray[(ROWS_PER_TILE,), np.dtype[np.float32]] + w_ty = np.ndarray[(tile_rows * row_bytes,), np.dtype[np.uint8]] + y_ty = np.ndarray[(tile_rows * ROWS_PER_TILE,), np.dtype[np.float32]] + + kernels = [ + ExternalFunction( + f"granite_gemv_p{per_call}b1_k{i}", + source_file=str(srcs[i]), + arg_types=([tile_ty, acc_ty] if null else [tile_ty, x_ty, acc_ty]), + include_dirs=_include_dirs(), + ) + for i in range(n_entry) + ] + + of_w = [ObjectFifo(tile_ty, name=f"w{c}", depth=2) for c in range(n_cores)] + of_y = [ObjectFifo(acc_ty, name=f"y{c}", depth=2) for c in range(n_cores)] + # One activation, broadcast. Private copies would want one shim MM2S channel + # each and there are only 16 device-wide -- the weights need those. + of_x = ObjectFifo(x_ty, name="x", depth=1) + + def core_body(win, xin, yout, *ks): + # x is acquired once and held: the same activation feeds every tile-row. + xe = None if null else xin.acquire(1) + for _ in range_(per_core): + ye = yout.acquire(1) + for fn in ks: + we = win.acquire(1) + if null: + fn(we, ye) + else: + fn(we, xe, ye) + win.release(1) + yout.release(1) + if not null: + xin.release(1) + + workers = [ + Worker(core_body, + fn_args=[of_w[c].cons(), of_x.cons(), of_y[c].prod(), *kernels], + stack_size=0xD00) + for c in range(n_cores) + ] + + w_taps = TensorTiler2D.simple_tiler( + (1, tile_rows * row_bytes), (1, per_core * row_bytes)) + y_taps = TensorTiler2D.simple_tiler( + (1, tile_rows * ROWS_PER_TILE), (1, per_core * ROWS_PER_TILE)) + + def sequence(a_w, a_x, c_y, w_prods, x_prod, y_conss): + tg = TaskGroup() + x_prod.fill(a_x, group=tg) + for c in range(n_cores): + w_prods[c].fill(a_w, tap=w_taps[c], group=tg) + y_conss[c].drain(c_y, tap=y_taps[c], wait=True, group=tg) + tg.finish() + + rt = Runtime(sequence, + [w_ty, x_ty, y_ty, + [f.prod() for f in of_w], of_x.prod(), [f.cons() for f in of_y]]) + return Program(iron.get_current_device(), rt, workers=workers).resolve_program() + + +def reference(raw: bytes, x: np.ndarray, tile_rows: int, k_tiles: int, + chunk: int = 256) -> np.ndarray: + """The same GEMV on the host, from the same bytes, in float32. + + Chunked over tile-rows: lm_head dequantised to float32 in one piece would be + 100352 x 2560 x 4 B = 1.0 GB and pointless to hold at once. + """ + xf = x.astype(np.float32) + out = np.empty(tile_rows * ROWS_PER_TILE, np.float32) + b_all = np.frombuffer(raw, dtype=np.uint8).reshape(tile_rows * k_tiles, TILE_BYTES) + for lo in range(0, tile_rows, chunk): + n = min(chunk, tile_rows - lo) + b = b_all[lo * k_tiles:(lo + n) * k_tiles] + w = q4nx._untile(q4nx._q4_tiles(b), n, k_tiles) + out[lo * ROWS_PER_TILE:(lo + n) * ROWS_PER_TILE] = w.astype(np.float32) @ xf + return out + + +def run_one(f, cfg, name: str, cores: int, limit_rows: int | None, + xmode: str, iters: int, null: bool = False) -> bool: + n, k = projection_shape(name, cfg) + k_tiles = k // TILE_K + per_call = tiles_per_call(k_tiles) + row_bytes = k_tiles * TILE_BYTES + n_rows_all = n // ROWS_PER_TILE + + # Cross-check the config-derived factoring against what the file stores. + # gate_proj and down_proj share a stored shape and differ only here. + stored = f.header[name]["shape"] + if stored[0] != n_rows_all * k_tiles or stored[1] != TILE_BYTES: + print(f" [SKIP] {name}: config says N={n} K={k} -> " + f"{n_rows_all * k_tiles} x {TILE_BYTES}, file has {stored}") + return False + + tile_rows = n_rows_all if limit_rows is None else min(limit_rows, n_rows_all) + while cores > 1 and tile_rows % cores: + cores -= 1 + + first, _ = f.header[name]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + first) + raw = fh.read(tile_rows * row_bytes) + + if xmode == "ones": + x = np.ones(k, np.float32).astype(bfloat16) + elif xmode.startswith("onehot"): + x = np.zeros(k, np.float32) + x[int(xmode.split(":")[1]) if ":" in xmode else 0] = 1.0 + x = x.astype(bfloat16) + else: + x = np.random.default_rng(0).standard_normal(k).astype(np.float32).astype(bfloat16) + + a_w = iron.tensor(np.frombuffer(raw, dtype=np.uint8).copy(), dtype=np.uint8, device="npu") + a_x = iron.tensor(x, dtype=bfloat16, device="npu") + c_y = iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu") + bench = run_iters(granite_gemv, a_w, a_x, c_y, tile_rows=tile_rows, k=k, + n_cores=cores, null=null, per_call=per_call, + warmup=1, iters=iters) + mb = tile_rows * row_bytes / 1e6 + if null: + # No arithmetic ran, so there is nothing to check: bandwidth only. + us = bench.npu.avg_us + print(f" DMA {name:44} {mb:6.1f}MB cores={cores:2} " + f"{us / 1000:7.2f}ms {mb / us * 1e3:5.1f}GB/s" + f" <- weight stream, no compute") + return True + + got = c_y.numpy().copy() + ref = reference(raw, x, tile_rows, k_tiles) + + # float64: both of these are reductions over up to 100352 terms, and in + # float32 the dot and the two norms accumulate in different orders, so the + # ratio drifts from 1 even when the vectors are bit-identical. Measured: + # x = ones reproduced the reference exactly (max rel err 0.0) and the float32 + # cosine still read 0.99999988 -- the metric's own rounding, reported as if + # it were the kernel's error. + g64, r64 = got.astype(np.float64), ref.astype(np.float64) + rel = np.abs(g64 - r64).max() / (np.abs(r64).max() + 1e-30) + cos = float(g64 @ r64 / (np.linalg.norm(g64) * np.linalg.norm(r64) + 1e-30)) + ok = cos > 0.9999999 and rel < 1e-4 + + us = bench.npu.avg_us if bench.npu is not None else float("nan") + print(f" {'PASS' if ok else 'FAIL'} {name:44} N={n:6} K={k:5} " + f"kt={k_tiles:2}/{per_call} rows={tile_rows * ROWS_PER_TILE:6} " + f"cores={cores:2} {mb:6.1f}MB {us / 1000:7.2f}ms " + f"{mb / us * 1e3:5.1f}GB/s cos={cos:.8f} rel={rel:.2e}") + return ok + + +def main(argv: list[str]) -> int: + import argparse + + ap = argparse.ArgumentParser() + ap.add_argument("--tensor", default="lm_head.weight") + ap.add_argument("--all", action="store_true", + help="every distinct projection shape, layer 0 + lm_head") + ap.add_argument("--cores", type=int, default=8) + ap.add_argument("--tile-rows", type=int, default=None, + help="cap tile-rows (default: the whole tensor)") + ap.add_argument("--x", default="random", help="random | ones | onehot:") + ap.add_argument("--iters", type=int, default=5) + ap.add_argument("--null", action="store_true", + help="DMA probe: stream the weights, do no arithmetic") + a = ap.parse_args(argv[1:]) + + if not (MODEL / "model.q4nx").is_file(): + raise SystemExit(f"model not found: {MODEL / 'model.q4nx'}") + cfg = json.loads((MODEL / "config.json").read_text(encoding="utf-8")) + f = q4nx.Q4NX(MODEL / "model.q4nx") + + # Trap: without this IRON silently falls back to aie2/NPU1 -- no error, wrong + # mac_dims, halved shim DMA burst. + iron.set_current_device(from_name("npu2", n_cols=None)) + + if a.all: + names = [f"model.layers.0.self_attn.{p}.weight" + for p in ("q_proj", "k_proj", "v_proj", "o_proj")] + names += [f"model.layers.0.mlp.{p}.weight" + for p in ("gate_proj", "up_proj", "down_proj")] + names.append("lm_head.weight") + else: + names = [a.tensor] + + print(f"granite q4nx GEMV on NPU ({a.x} activation)") + results = [run_one(f, cfg, n, a.cores, a.tile_rows, a.x, a.iters, a.null) + for n in names] + n_ok = sum(results) + print(f"\n{n_ok}/{len(results)} shapes match the host GEMV") + return 0 if n_ok == len(results) else 1 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/kernels/granite/iron/granite_gemv32.py b/kernels/granite/iron/granite_gemv32.py new file mode 100644 index 00000000..957d590d --- /dev/null +++ b/kernels/granite/iron/granite_gemv32.py @@ -0,0 +1,297 @@ +r"""The same granite q4nx GEMV, but on all 32 cores via the memtile leg. + +WHY THIS EXISTS +--------------- +`granite_gemv.py` gives every core its own shim stream, which costs n+1 MM2S and +n S2MM channels. There are 16 of each device-wide, so it tops out at **8 cores** +and 16 dies at placement with `no ShimNOCTile has sufficient DMA capacity`. +npu2 is 8 columns x 6 rows (row 0 shim, row 1 memtile, rows 2-5 compute) = **32 +compute cores**, so three quarters of the array is unreachable that way. + +Measured, and this is the reason to bother: + + NPU DMA alone, no arithmetic 3.47 ms 46.3 GB/s + NPU GEMV, 8 cores 8.08 ms 19.9 GB/s + per-core compute ~2.5 GB/s, perfectly linear + +The weight path sustains 46.3 GB/s and the kernel uses 19.9 of them, so this is +compute-bound and ~18 cores saturates it. The bandwidth is available long before +the cores are. + +THE SHAPE OF THE FIX +-------------------- +One shim stream per **column** into the memtile, split four ways to that column's +compute cores, and joined back on the way out: + + shim MM2S x1 -> memtile -> split -> 4 cores (per column) + shim S2MM x1 <- memtile <- join <- 4 cores + +That is 8 + 1 MM2S (weights + broadcast x) and 8 S2MM for 32 cores, comfortably +inside 16/16. + +WHY THE WEIGHTS ARE PERMUTED ON THE HOST +---------------------------------------- +`split()` hands child i the i-th slice of each parent object, so the DDR stream +must arrive as [core0 chunk k][core1 chunk k][core2 chunk k][core3 chunk k]. +With each core owning a contiguous block of tile-rows, that is a strided 3-D tap +whose innermost run is 25600 B -- close enough to the BD size limits to be a +liability. Permuting host-side makes each column's stream plain contiguous +instead, and it is a **one-time cost at model load**, not per token: the weights +are uploaded once and streamed for every token thereafter. + + call c:\dev\mlir-aie\iron_env.cmd + python designs\granite_gemv\granite_gemv32.py --cols 8 :: 32 cores + python designs\granite_gemv\granite_gemv32.py --cols 8 --null :: DMA only +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +import aie.iron as iron +from aie.iron import (CompileTime, In, ObjectFifo, Out, Program, Runtime, + TaskGroup, Worker) +from aie.iron.controlflow import range_ +from aie.iron.device import from_name +from aie.iron.kernel import ExternalFunction +from aie.helpers.taplib import TensorTiler2D + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parent.parent / "common")) + +import q4nx # noqa: E402 +from granite_gemv import (MODEL, ROWS_PER_TILE, TILE_BYTES, TILE_K, # noqa: E402 + _include_dirs, ensure_entry_points, projection_shape, + reference, tiles_per_call) + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + +def permute_weights(raw: bytes, n_cols: int, per_core: int, n_entry: int, + call_bytes: int) -> np.ndarray: + """Reorder so each column's stream is contiguous in the order split() wants. + + In: core-major, each core's tile-rows contiguous. + Out: per column, chunk-major then core -- [k][r] -- which is exactly the + layout `split()` consumes, one parent object per k. + """ + a = np.frombuffer(raw, dtype=np.uint8) + # (col, row_in_col, chunk, bytes) -> (col, chunk, row_in_col, bytes) + a = a.reshape(n_cols, ROWS_PER_COL, per_core * n_entry, call_bytes) + return np.ascontiguousarray(a.transpose(0, 2, 1, 3)).reshape(-1) + + +def unpermute_y(y: np.ndarray, n_cols: int, per_core: int, + batch: int = 1) -> np.ndarray: + """Inverse of the above for the joined output. + + On the wire it is [col][t][row_in_col][token][32]; the caller wants one + contiguous result vector per token, so the token axis comes out front. + Returns (batch, tile_rows * 32). + """ + a = y.reshape(n_cols, per_core, ROWS_PER_COL, batch, ROWS_PER_TILE) + a = a.transpose(3, 0, 2, 1, 4) # [token][col][row][t][32] + return np.ascontiguousarray(a).reshape(batch, -1) + + +@iron.jit(aiecc_flags=["--alloc-scheme=basic-sequential"]) +def granite_gemv32(w: In, x: In, y: Out, *, tile_rows: CompileTime[int], + k: CompileTime[int], n_cols: CompileTime[int] = 8, + null: CompileTime[bool] = False, + batch: CompileTime[int] = 1, + per_call: CompileTime[int] = 5): + # per_call MUST be an explicit argument, not derived here. iron.jit's cache + # key hashes the call's arguments and nothing else (trap 7d), so a derived + # per_call is invisible to it: two runs that differ only in per_call collide + # on one cache entry and the second silently gets the first one's xclbin, + # while the host permutes for its own value. That is not a crash -- it + # returned cosine 0.208 against the reference, a plausible-looking wrong + # answer, and it cost a confusing debug cycle to find. + k_tiles = k // TILE_K + n_entry = k_tiles // per_call + call_bytes = per_call * TILE_BYTES + n_cores = n_cols * ROWS_PER_COL + per_core = tile_rows // n_cores + chunks = per_core * n_entry # parent objects per column + + srcs = ensure_entry_points(n_entry, per_call, null, batch, k) + + w_l1_ty = np.ndarray[(call_bytes,), np.dtype[np.uint8]] + w_l2_ty = np.ndarray[(ROWS_PER_COL * call_bytes,), np.dtype[np.uint8]] + # One weight pass serves `batch` tokens, so every activation and result + # buffer carries the token axis; the weights do not change size at all -- + # that asymmetry is the entire point of batching. + y_l1_ty = np.ndarray[(batch * ROWS_PER_TILE,), np.dtype[np.float32]] + y_l2_ty = np.ndarray[(ROWS_PER_COL * batch * ROWS_PER_TILE,), + np.dtype[np.float32]] + x_ty = np.ndarray[(batch * k,), np.dtype[bfloat16]] + w_ty = np.ndarray[(tile_rows * k_tiles * TILE_BYTES,), np.dtype[np.uint8]] + y_ty = np.ndarray[(tile_rows * batch * ROWS_PER_TILE,), np.dtype[np.float32]] + + kernels = [ + ExternalFunction( + f"granite_gemv_p{per_call}b{batch}_k{i}", + source_file=str(srcs[i]), + arg_types=[w_l1_ty, x_ty, y_l1_ty], + include_dirs=_include_dirs(), + ) + for i in range(n_entry) + ] + + # One activation for the whole array. 32 private copies are impossible -- + # there are only 16 shim MM2S channels and the weights need 8 of them. + of_x = ObjectFifo(x_ty, name="x", depth=1) + + w_l3l2, y_l2l3, w_cores, y_cores = [], [], [], [] + for c in range(n_cols): + wf = ObjectFifo(w_l2_ty, name=f"wL2_{c}", depth=2) + w_l3l2.append(wf) + w_cores.append(wf.cons().split( + [r * call_bytes for r in range(ROWS_PER_COL)], + obj_types=[w_l1_ty] * ROWS_PER_COL, + names=[f"w_{c}_{r}" for r in range(ROWS_PER_COL)], + )) + yf = ObjectFifo(y_l2_ty, name=f"yL2_{c}", depth=2) + y_l2l3.append(yf) + y_cores.append(yf.prod().join( + [r * batch * ROWS_PER_TILE for r in range(ROWS_PER_COL)], + obj_types=[y_l1_ty] * ROWS_PER_COL, + names=[f"y_{c}_{r}" for r in range(ROWS_PER_COL)], + )) + + def core_body(win, xin, yout, *ks): + xe = xin.acquire(1) + for _ in range_(per_core): + ye = yout.acquire(1) + for fn in ks: + we = win.acquire(1) + fn(we, xe, ye) + win.release(1) + yout.release(1) + xin.release(1) + + workers = [ + Worker(core_body, + fn_args=[w_cores[c][r].cons(), of_x.cons(), + y_cores[c][r].prod(), *kernels], + stack_size=0xD00) + for c in range(n_cols) for r in range(ROWS_PER_COL) + ] + + # After the host-side permutation each column's weights are one contiguous + # run, so this is a plain split rather than a strided gather. + col_w = chunks * ROWS_PER_COL * call_bytes + col_y = per_core * ROWS_PER_COL * batch * ROWS_PER_TILE + w_taps = TensorTiler2D.simple_tiler((1, n_cols * col_w), (1, col_w)) + y_taps = TensorTiler2D.simple_tiler((1, n_cols * col_y), (1, col_y)) + + def sequence(a_w, a_x, c_y, w_prods, x_prod, y_conss): + tg = TaskGroup() + x_prod.fill(a_x, group=tg) + for c in range(n_cols): + w_prods[c].fill(a_w, tap=w_taps[c], group=tg) + y_conss[c].drain(c_y, tap=y_taps[c], wait=True, group=tg) + tg.finish() + + rt = Runtime(sequence, + [w_ty, x_ty, y_ty, + [f.prod() for f in w_l3l2], of_x.prod(), + [f.cons() for f in y_l2l3]]) + return Program(iron.get_current_device(), rt, workers=workers).resolve_program() + + +def main(argv: list[str]) -> int: + import argparse + from aie.utils.benchmark import run_iters + + ap = argparse.ArgumentParser() + ap.add_argument("--tensor", default="lm_head.weight") + ap.add_argument("--cols", type=int, default=8, help="columns; cores = 4x") + ap.add_argument("--tile-rows", type=int, default=None) + ap.add_argument("--x", default="random") + ap.add_argument("--iters", type=int, default=5) + ap.add_argument("--batch", type=int, default=1, + help="tokens per weight pass (independent tokens only)") + ap.add_argument("--null", action="store_true", + help="DMA probe: stream the weights, do no arithmetic") + a = ap.parse_args(argv[1:]) + + cfg = json.loads((MODEL / "config.json").read_text(encoding="utf-8")) + f = q4nx.Q4NX(MODEL / "model.q4nx") + n, k = projection_shape(a.tensor, cfg) + k_tiles = k // TILE_K + per_call = tiles_per_call(k_tiles, a.batch, k) # passed in explicitly + n_entry = k_tiles // per_call + call_bytes = per_call * TILE_BYTES + row_bytes = k_tiles * TILE_BYTES + + n_cores = a.cols * ROWS_PER_COL + tile_rows = n // ROWS_PER_TILE if a.tile_rows is None else a.tile_rows + if tile_rows % n_cores: + tile_rows -= tile_rows % n_cores + per_core = tile_rows // n_cores + + first, _ = f.header[a.tensor]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + first) + raw = fh.read(tile_rows * row_bytes) + + if a.x == "ones": + x = np.ones(k, np.float32).astype(bfloat16) + elif a.x.startswith("onehot"): + x = np.zeros(k, np.float32) + x[int(a.x.split(":")[1]) if ":" in a.x else 0] = 1.0 + x = x.astype(bfloat16) + else: + x = np.random.default_rng(0).standard_normal( + a.batch * k).astype(np.float32).astype(bfloat16) + if a.batch > 1 and x.size == k: + x = np.tile(x, a.batch) + + iron.set_current_device(from_name("npu2", n_cols=None)) + w_perm = permute_weights(raw, a.cols, per_core, n_entry, call_bytes) + + a_w = iron.tensor(w_perm, dtype=np.uint8, device="npu") + a_x = iron.tensor(x, dtype=bfloat16, device="npu") + c_y = iron.zeros(tile_rows * a.batch * ROWS_PER_TILE, dtype=np.float32, + device="npu") + bench = run_iters(granite_gemv32, a_w, a_x, c_y, tile_rows=tile_rows, k=k, + n_cols=a.cols, null=a.null, batch=a.batch, + per_call=per_call, warmup=1, iters=a.iters) + + mb = tile_rows * row_bytes / 1e6 + us = bench.npu.avg_us + head = (f"{a.tensor} N={n} K={k} {tile_rows * ROWS_PER_TILE} rows " + f"{a.cols} cols x {ROWS_PER_COL} = {n_cores} cores {mb:.1f} MB" + + (f" batch {a.batch}" if a.batch > 1 else "")) + if a.null: + print(f"DMA {head}\n {us / 1000:.2f} ms {mb / us * 1e3:.1f} GB/s" + f" <- weight stream, no compute") + return 0 + + got = unpermute_y(c_y.numpy().copy(), a.cols, per_core, a.batch) + # Check EVERY token, not just the first: a batched kernel that ignored the + # token index would still reproduce token 0 perfectly. + rel, cos, ok = 0.0, 1.0, True + for b in range(a.batch): + ref = reference(raw, x[b * k:(b + 1) * k], tile_rows, k_tiles) + g64, r64 = got[b].astype(np.float64), ref.astype(np.float64) + rel = max(rel, np.abs(g64 - r64).max() / (np.abs(r64).max() + 1e-30)) + cos = min(cos, float(g64 @ r64 / + (np.linalg.norm(g64) * np.linalg.norm(r64) + 1e-30))) + ok = cos > 0.9999999 and rel < 1e-4 + print(f"{'PASS' if ok else 'FAIL'} {head}\n" + f" {us / 1000:.2f} ms {mb / us * 1e3:.1f} GB/s " + f"({2.0 * a.batch * tile_rows * ROWS_PER_TILE * k / us / 1e3:.1f} GFLOP/s) " + f"cos={cos:.8f} rel={rel:.2e}") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/kernels/granite/iron/granite_norm_gemv.py b/kernels/granite/iron/granite_norm_gemv.py new file mode 100644 index 00000000..d7443daf --- /dev/null +++ b/kernels/granite/iron/granite_norm_gemv.py @@ -0,0 +1,330 @@ +r"""RMSNorm folded into a GEMV as a per-core prologue. ONE dispatch, full width. + +WHY THIS AND NOT MORE FUSION OF THE BIG OPS +------------------------------------------- +0148 measured where a granite layer's time actually goes: the four GEMV groups +cost 1.65 ms and move 49.2 MB, while five small ops (2x RMSNorm, RoPE, SwiGLU, +attention) cost 1.20 ms and move almost nothing. ~1.0 ms of that is pure +per-dispatch floor -- about 200 us each, paid regardless of size. + +Fusing the BIG ops does not help: even a perfect one-dispatch MLP at gate_up's +39.7 GB/s is 1.19 ms against 1.31 ms for the same three ops unfused across +20-32 cores. Measured, in granite_mlp_wide.py. The big ops are already work, +not waiting. + +So the lever is the small ones, and RMSNorm is the easiest because it is +elementwise and has no head structure to keep core-local. + +THE TRICK: REDUNDANT, NOT DISTRIBUTED +------------------------------------- +Every core's GEMV consumes the WHOLE activation, so every core needs the whole +normalised vector. Computing the norm once and broadcasting it would need a +cross-core barrier inside the dispatch. Computing it **redundantly on every +core** needs nothing: it is 2560 multiply-accumulates against 8192 MACs per +weight tile, and it removes the dispatch outright. + +AND IN PLACE, WHICH IS THE PART THAT DECIDES IT +----------------------------------------------- +The activation fifo carries `x` then the norm weight, 2 x 2560 bf16. The +prologue normalises into the first half. A separate destination buffer would +cost another 5120 B of L1, pushing the fixed cost from 10240 to 15360 -- and +with a 62208 B budget that drops the weight double buffer from per_call 5 to +per_call 2 (the only divisors of 10 tiles), which costs more bandwidth than the +saved dispatch is worth. See the __restrict note in granite_rmsnorm.h. + + call c:\dev\mlir-aie\iron_env.cmd + python designs\granite_gemv\granite_norm_gemv.py + python designs\granite_gemv\granite_norm_gemv.py --tensor gate_up +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +import aie.iron as iron +from aie.iron import (CompileTime, In, ObjectFifo, Out, Program, Runtime, + TaskGroup, Worker) +from aie.iron.controlflow import range_ +from aie.iron.device import from_name +from aie.iron.kernel import ExternalFunction +from aie.helpers.taplib import TensorTiler2D + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parent.parent / "common")) + +import q4nx # noqa: E402 +from granite_gemv import (AIE, MODEL, ROWS_PER_TILE, TILE_BYTES, TILE_K, # noqa: E402 + _include_dirs, ensure_entry_points, projection_shape, + reference, tiles_per_call) + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + +def permute_weights(raw: bytes, n_cols: int, per_core: int, n_entry: int, + call_bytes: int) -> np.ndarray: + """Reorder so each column's stream is contiguous in the order split() wants. + + In: core-major, each core's tile-rows contiguous. + Out: per column, chunk-major then core -- [k][r] -- which is exactly the + layout `split()` consumes, one parent object per k. + """ + a = np.frombuffer(raw, dtype=np.uint8) + # (col, row_in_col, chunk, bytes) -> (col, chunk, row_in_col, bytes) + a = a.reshape(n_cols, ROWS_PER_COL, per_core * n_entry, call_bytes) + return np.ascontiguousarray(a.transpose(0, 2, 1, 3)).reshape(-1) + + +def unpermute_y(y: np.ndarray, n_cols: int, per_core: int, + batch: int = 1) -> np.ndarray: + """Inverse of the above for the joined output. + + On the wire it is [col][t][row_in_col][token][32]; the caller wants one + contiguous result vector per token, so the token axis comes out front. + Returns (batch, tile_rows * 32). + """ + a = y.reshape(n_cols, per_core, ROWS_PER_COL, batch, ROWS_PER_TILE) + a = a.transpose(3, 0, 2, 1, 4) # [token][col][row][t][32] + return np.ascontiguousarray(a).reshape(batch, -1) + + + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + +@iron.jit(aiecc_flags=["--alloc-scheme=basic-sequential"]) +def granite_norm_gemv(w: In, xw: In, y: Out, *, + tile_rows: CompileTime[int], k: CompileTime[int], + n_cols: CompileTime[int] = 8, + per_call: CompileTime[int] = 5): + # per_call is an explicit argument for the reason granite_gemv32 records: + # iron.jit's cache key hashes the call's arguments and nothing else, so a + # value derived in here is invisible to it and two configs collide. + k_tiles = k // TILE_K + n_entry = k_tiles // per_call + call_bytes = per_call * TILE_BYTES + n_cores = n_cols * ROWS_PER_COL + per_core = tile_rows // n_cores + chunks = per_core * n_entry + + srcs = ensure_entry_points(n_entry, per_call, False, 1, k) + + w_l1_ty = np.ndarray[(call_bytes,), np.dtype[np.uint8]] + w_l2_ty = np.ndarray[(ROWS_PER_COL * call_bytes,), np.dtype[np.uint8]] + y_l1_ty = np.ndarray[(ROWS_PER_TILE,), np.dtype[np.float32]] + y_l2_ty = np.ndarray[(ROWS_PER_COL * ROWS_PER_TILE,), np.dtype[np.float32]] + # x and the norm weight share ONE fifo: a compute tile has 2 input DMA + # channels and the weights already take one, so a second activation stream + # is not available at any L1 cost. Two 5120 B buffers and one 10240 B buffer + # cost the same anyway. + xw_ty = np.ndarray[(2 * k,), np.dtype[bfloat16]] + w_ty = np.ndarray[(tile_rows * k_tiles * TILE_BYTES,), np.dtype[np.uint8]] + y_ty = np.ndarray[(tile_rows * ROWS_PER_TILE,), np.dtype[np.float32]] + + kernels = [ExternalFunction(f"granite_gemv_p{per_call}b1_k{i}", + source_file=str(srcs[i]), + arg_types=[w_l1_ty, xw_ty, y_l1_ty], + include_dirs=_include_dirs()) + for i in range(n_entry)] + rmsnorm = ExternalFunction("granite_rms_norm_ip", + source_file=str(AIE / "granite_rmsnorm_ip.cc"), + arg_types=[xw_ty, np.int32], + include_dirs=_include_dirs()) + + of_x = ObjectFifo(xw_ty, name="nx", depth=1) + + w_l3l2, y_l2l3, w_cores, y_cores = [], [], [], [] + for c in range(n_cols): + wf = ObjectFifo(w_l2_ty, name=f"nwL2_{c}", depth=2) + w_l3l2.append(wf) + w_cores.append(wf.cons().split( + [r * call_bytes for r in range(ROWS_PER_COL)], + obj_types=[w_l1_ty] * ROWS_PER_COL, + names=[f"nw_{c}_{r}" for r in range(ROWS_PER_COL)])) + yf = ObjectFifo(y_l2_ty, name=f"nyL2_{c}", depth=2) + y_l2l3.append(yf) + y_cores.append(yf.prod().join( + [r * ROWS_PER_TILE for r in range(ROWS_PER_COL)], + obj_types=[y_l1_ty] * ROWS_PER_COL, + names=[f"ny_{c}_{r}" for r in range(ROWS_PER_COL)])) + + def core_body(win, xin, yout, norm, *ks): + xe = xin.acquire(1) + # Redundantly on every core, in place. Every core's GEMV reads the whole + # activation anyway, so a broadcast would need a barrier inside the + # dispatch; this needs nothing. + norm(xe, k) + for _ in range_(per_core): + ye = yout.acquire(1) + for fn in ks: + we = win.acquire(1) + fn(we, xe, ye) + win.release(1) + yout.release(1) + xin.release(1) + + workers = [ + Worker(core_body, + fn_args=[w_cores[c][r].cons(), of_x.cons(), + y_cores[c][r].prod(), rmsnorm, *kernels], + stack_size=0xD00) + for c in range(n_cols) for r in range(ROWS_PER_COL) + ] + + col_w = chunks * ROWS_PER_COL * call_bytes + col_y = per_core * ROWS_PER_COL * ROWS_PER_TILE + w_taps = TensorTiler2D.simple_tiler((1, n_cols * col_w), (1, col_w)) + y_taps = TensorTiler2D.simple_tiler((1, n_cols * col_y), (1, col_y)) + + def sequence(a_w, a_x, c_y, w_prods, x_prod, y_conss): + tg = TaskGroup() + x_prod.fill(a_x, group=tg) + for c in range(n_cols): + w_prods[c].fill(a_w, tap=w_taps[c], group=tg) + y_conss[c].drain(c_y, tap=y_taps[c], wait=True, group=tg) + tg.finish() + + rt = Runtime(sequence, + [w_ty, xw_ty, y_ty, + [f.prod() for f in w_l3l2], of_x.prod(), + [f.cons() for f in y_l2l3]]) + return Program(iron.get_current_device(), rt, workers=workers).resolve_program() + + +def main(argv: list[str] | None = None) -> int: + import argparse + from aie.utils.benchmark import run_iters + ap = argparse.ArgumentParser() + ap.add_argument("--tensor", default="qkv", + choices=("qkv", "gate_up", "o")) + ap.add_argument("--cols", type=int, default=0, + help="0 = widest that divides the tile-rows") + ap.add_argument("--iters", type=int, default=200) + a = ap.parse_args((argv or sys.argv)[1:]) + + cfg = json.loads((MODEL / "config.json").read_text(encoding="utf-8")) + groups = {"qkv": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + "gate_up": ["mlp.gate_proj", "mlp.up_proj"], + "o": ["self_attn.o_proj"]}[a.tensor] + # o_proj takes the attention output, not a normalised hidden state -- it is + # here only to show the norm prologue is shape-agnostic, and its reference + # normalises too so the comparison stays honest. + norm_name = ("model.layers.0.post_attention_layernorm.weight" + if a.tensor == "gate_up" + else "model.layers.0.input_layernorm.weight") + + f = q4nx.Q4NX(MODEL / "model.q4nx") + raws, tile_rows, k = [], 0, None + for g in groups: + nm = f"model.layers.0.{g}.weight" + n, kk = projection_shape(nm, cfg) + k = kk + rows = n // ROWS_PER_TILE + off, _ = f.header[nm]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + off) + raws.append(fh.read(rows * (kk // TILE_K) * TILE_BYTES)) + tile_rows += rows + raw = b"".join(raws) + k_tiles = k // TILE_K + + off, _ = f.header[norm_name]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + off) + nwv = np.frombuffer(fh.read(k * 2), dtype=bfloat16) + + # The array has 8 columns of 4, but a group only fills the ones its + # tile-rows divide over: qkv is 112 rows = 16 x 7, so 7 columns and not 8. + # This is the same cap that pins o_proj and down_proj (80 rows) to 20 cores. + cols = a.cols or max(c for c in range(1, 9) if tile_rows % (c * ROWS_PER_COL) == 0) + n_cores = cols * ROWS_PER_COL + per_call = tiles_per_call(k_tiles, 1, k) + # The norm weight adds a second k-wide fifo, so the fixed L1 is 2*k*2 rather + # than k*2. Check it here rather than discover it as a build failure. + fixed = 2 * k * 2 # the combined [x | norm weight] element + assert per_call * TILE_BYTES * 2 + fixed <= 64 * 1024 - 0xD00, ( + f"per_call {per_call} plus the norm weight overflows L1") + n_entry = k_tiles // per_call + per_core = tile_rows // n_cores + w = permute_weights(raw, cols, per_core, n_entry, per_call * TILE_BYTES) + + rng = np.random.default_rng(0) + x = rng.standard_normal(k).astype(np.float32).astype(bfloat16) + + iron.set_current_device(from_name("npu2", n_cols=None)) + c_y = iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu") + b = run_iters(granite_norm_gemv, + iron.tensor(w, dtype=np.uint8, device="npu"), + iron.tensor(np.concatenate([x, nwv]), dtype=bfloat16, + device="npu"), + c_y, tile_rows=tile_rows, k=k, n_cols=cols, + per_call=per_call, warmup=1, iters=a.iters) + got = unpermute_y(c_y.numpy().copy(), cols, per_core).astype(np.float64) + + # Reference: normalise, THEN the GEMV -- both together. Checking the two + # halves separately is what let a wrong composition pass in 0147. + xf = x.astype(np.float32) + inv = 1.0 / np.sqrt((xf * xf).mean() + cfg["rms_norm_eps"]) + h = (xf * inv * nwv.astype(np.float32)).astype(bfloat16) + ref = reference(raw, h, tile_rows, k_tiles).astype(np.float64) + + rel = np.abs(got - ref).max() / (np.abs(ref).max() + 1e-30) + g1, r1 = got.ravel(), ref.ravel() + cos = float(g1 @ r1 / (np.linalg.norm(g1) * np.linalg.norm(r1) + 1e-30)) + mb = len(raw) / 1e6 + us = b.npu.avg_us + ok = cos > 0.9999 and rel < 8e-3 + print(f"RMSNorm + {a.tensor} GEMV, ONE dispatch " + f"{cols} cols x {ROWS_PER_COL} = {n_cores} cores, per_call {per_call}") + print(f" {tile_rows * ROWS_PER_TILE} rows K={k} {mb:.1f} MB") + print(f" cosine {cos:.8f} max rel err {rel:.3e}") + print(f" {us:.1f} us device {b.e2e.avg_us:.1f} us wall " + f"{mb / us * 1e3:.1f} GB/s") + print(f" separate would be RMSNorm 244 us + GEMV -> compare against 0148") + print(f" [fused == normalise then GEMV on the host] " + f"{'PASS' if ok else 'FAIL'}") + return 0 if ok else 1 + + +# -------------------------------------------------------------------------- +# build_artifact: produce the xclbin WITHOUT model weights. +# +# iron.jit keys its cache on argument shapes and dtypes, not contents, so an +# artefact built from zeros is bit-identical to one built from real weights. +# That is what lets the in-tree build run from a clean checkout with nothing +# but the toolchain -- see kernels/CONVENTION.md. main() below is unchanged and +# remains the developer path: it needs the model and checks the result. +# -------------------------------------------------------------------------- + + +def build_artifact(geometry: dict, tensor: str = "qkv") -> None: + groups = {"qkv": ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"], + "gate_up": ["mlp.gate_proj", "mlp.up_proj"], + "o": ["self_attn.o_proj"]}[tensor] + tile_rows, k = 0, None + for g in groups: + n, kk = projection_shape(g.rsplit(".", 1)[-1], geometry) + k = kk + tile_rows += n // ROWS_PER_TILE + k_tiles = k // TILE_K + cols = max(c for c in range(1, 9) if tile_rows % (c * ROWS_PER_COL) == 0) + per_call = tiles_per_call(k_tiles, 1, k) + n_cores = cols * ROWS_PER_COL + per_core = tile_rows // n_cores + w_bytes = tile_rows * k_tiles * TILE_BYTES + + iron.set_current_device(from_name("npu2", n_cols=None)) + granite_norm_gemv( + iron.zeros(w_bytes, dtype=np.uint8, device="npu"), + iron.zeros(2 * k, dtype=bfloat16, device="npu"), + iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu"), + tile_rows=tile_rows, k=k, n_cols=cols, per_call=per_call) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kernels/granite/iron/granite_qkv_wide.py b/kernels/granite/iron/granite_qkv_wide.py new file mode 100644 index 00000000..c8706e85 --- /dev/null +++ b/kernels/granite/iron/granite_qkv_wide.py @@ -0,0 +1,360 @@ +r"""RMSNorm + q + k + v + RoPE. FOUR ops, ONE dispatch, 28 cores. + +WHY THIS BREAKS THE 8-CORE CAP granite_qkv.py HIT +-------------------------------------------------- +granite_qkv.py fuses the same four ops and is stuck at 8 cores. The reason +looked structural: granite has 8 kv heads, RoPE is head-local, so a core must +own whole heads -- and there are only 8 to go round. + +But that follows from ITS layout, where every core gets a slice of q AND k AND +v. Concatenating the three along N instead gives each core a slice of exactly +one of them, and then the arithmetic is different: k is 16 tile-rows over 4 +dedicated cores = 4 tile-rows = 2 whole heads each. The cap was in the design, +not in the model. + + 112 tile-rows = 80 (q) + 16 (k) + 16 (v), over 7 columns x 4 rows: + columns 0-4 -> q, 2 heads per core, rotated + column 5 -> k, 2 heads per core, rotated + column 6 -> v, 128 values per core, NOT rotated, only narrowed + +Three roles, so three core bodies, chosen per worker by column. The RoPE kernel +already takes (q_heads, k_heads, v_len) and does the right thing for each. + +L1, WHICH IS WHAT DECIDES THE SHAPE OF THIS +------------------------------------------- + weights, per_call 5, double buffered 51200 + activation [x | norm weight | cos,sin] 10368 + accumulator 512 + ----- + 62080 against a 62208 budget + +128 bytes spare. That is why the float32 accumulator IS the output fifo element +and the rotated result narrows into its first half (granite_rope_ip.cc): a +separate 256 B output element would put it 128 B over, and dropping per_call to +2 -- the only other divisor of 10 K-tiles -- would cost more bandwidth than all +three fused ops save. + + call c:\dev\mlir-aie\iron_env.cmd + python designs\granite_gemv\granite_qkv_wide.py +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +import aie.iron as iron +from aie.iron import (CompileTime, In, ObjectFifo, Out, Program, Runtime, + TaskGroup, Worker) +from aie.iron.controlflow import range_ +from aie.iron.device import from_name +from aie.iron.kernel import ExternalFunction +from aie.helpers.taplib import TensorTiler2D + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parent.parent / "common")) + +import q4nx # noqa: E402 +from granite_gemv import (AIE, MODEL, ROWS_PER_TILE, TILE_BYTES, TILE_K, # noqa: E402 + _include_dirs, ensure_entry_points, projection_shape, + reference, tiles_per_call) + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + +def permute_weights(raw: bytes, n_cols: int, per_core: int, n_entry: int, + call_bytes: int) -> np.ndarray: + """Reorder so each column's stream is contiguous in the order split() wants. + + In: core-major, each core's tile-rows contiguous. + Out: per column, chunk-major then core -- [k][r] -- which is exactly the + layout `split()` consumes, one parent object per k. + """ + a = np.frombuffer(raw, dtype=np.uint8) + # (col, row_in_col, chunk, bytes) -> (col, chunk, row_in_col, bytes) + a = a.reshape(n_cols, ROWS_PER_COL, per_core * n_entry, call_bytes) + return np.ascontiguousarray(a.transpose(0, 2, 1, 3)).reshape(-1) + + +def unpermute_y(y: np.ndarray, n_cols: int, per_core: int, + batch: int = 1) -> np.ndarray: + """Inverse of the above for the joined output. + + On the wire it is [col][t][row_in_col][token][32]; the caller wants one + contiguous result vector per token, so the token axis comes out front. + Returns (batch, tile_rows * 32). + """ + a = y.reshape(n_cols, per_core, ROWS_PER_COL, batch, ROWS_PER_TILE) + a = a.transpose(3, 0, 2, 1, 4) # [token][col][row][t][32] + return np.ascontiguousarray(a).reshape(batch, -1) + + + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + + +ROWS_PER_COL = 4 +HD = 64 # head_dim +PER_CALL = 5 + + +@iron.jit(aiecc_flags=["--alloc-scheme=basic-sequential"]) +def granite_qkv_wide(w: In, xw: In, out: Out, *, + q_rows: CompileTime[int], kv_rows: CompileTime[int], + k: CompileTime[int], n_cols: CompileTime[int] = 7): + k_tiles = k // TILE_K + n_entry = k_tiles // PER_CALL + call_bytes = PER_CALL * TILE_BYTES + n_cores = n_cols * ROWS_PER_COL + tile_rows = q_rows + 2 * kv_rows + per_core = tile_rows // n_cores + slice_f = per_core * ROWS_PER_TILE # 128 floats = 2 heads + heads_per_core = slice_f // HD + chunks = per_core * n_entry + + # Which columns own q, which own k, which own v. Whole columns, because the + # memtile split hands a column's stream to its own four cores. + q_cols = q_rows // (per_core * ROWS_PER_COL) + k_cols = kv_rows // (per_core * ROWS_PER_COL) + + w_l1_ty = np.ndarray[(call_bytes,), np.dtype[np.uint8]] + w_l2_ty = np.ndarray[(ROWS_PER_COL * call_bytes,), np.dtype[np.uint8]] + # The output element is float32 and doubles as the accumulator; the RoPE + # epilogue narrows into its first half. See granite_rope_ip.cc. + y_l1_ty = np.ndarray[(slice_f,), np.dtype[np.float32]] + y_l2_ty = np.ndarray[(ROWS_PER_COL * slice_f,), np.dtype[np.float32]] + act_ty = np.ndarray[(2 * k + 2 * (HD // 2),), np.dtype[bfloat16]] + w_ty = np.ndarray[(tile_rows * k_tiles * TILE_BYTES,), np.dtype[np.uint8]] + out_ty = np.ndarray[(tile_rows * ROWS_PER_TILE,), np.dtype[np.float32]] + + gemv = [ExternalFunction(f"granite_qgemv_g{i}", + source_file=str(AIE / f"granite_qgemv_g{i}.cc"), + arg_types=[w_l1_ty, act_ty, y_l1_ty, np.int32], + include_dirs=_include_dirs()) + for i in range(n_entry)] + rmsnorm = ExternalFunction("granite_rms_norm_ip", + source_file=str(AIE / "granite_rmsnorm_ip.cc"), + arg_types=[act_ty, np.int32], + include_dirs=_include_dirs()) + rope = ExternalFunction("granite_rope_ip", + source_file=str(AIE / "granite_rope_ip.cc"), + arg_types=[y_l1_ty, act_ty, np.int32, np.int32, + np.int32], + include_dirs=_include_dirs()) + + of_x = ObjectFifo(act_ty, name="qx", depth=1) + + w_l3l2, y_l2l3, w_cores, y_cores = [], [], [], [] + for c in range(n_cols): + wf = ObjectFifo(w_l2_ty, name=f"qwL2_{c}", depth=2) + w_l3l2.append(wf) + w_cores.append(wf.cons().split( + [r * call_bytes for r in range(ROWS_PER_COL)], + obj_types=[w_l1_ty] * ROWS_PER_COL, + names=[f"qw_{c}_{r}" for r in range(ROWS_PER_COL)])) + yf = ObjectFifo(y_l2_ty, name=f"qyL2_{c}", depth=2) + y_l2l3.append(yf) + # depths=1: a core produces exactly ONE output element per dispatch, so + # the default double buffer buys no overlap and costs 512 B of L1 -- + # which is more than the 128 B this design has spare. Measured as a + # build failure, not guessed: + # qy_6_3_buff_0 and _1, 512 bytes each -> Basic sequential allocation failed + y_cores.append(yf.prod().join( + [r * slice_f for r in range(ROWS_PER_COL)], + depths=[1] * ROWS_PER_COL, + obj_types=[y_l1_ty] * ROWS_PER_COL, + names=[f"qy_{c}_{r}" for r in range(ROWS_PER_COL)])) + + def make_body(qh, kh, vl): + """One body per role. The counts are baked per worker, not branched on.""" + def body(win, xin, yout, norm, rp, *ks): + xe = xin.acquire(1) + norm(xe, k) # redundant, in place, on every core + ye = yout.acquire(1) + for r in range_(per_core): + for fn in ks: + we = win.acquire(1) + fn(we, xe, ye, r) + win.release(1) + rp(ye, xe, qh, kh, vl) # narrows into ye's first half + yout.release(1) + xin.release(1) + return body + + workers = [] + for c in range(n_cols): + if c < q_cols: + body = make_body(heads_per_core, 0, 0) + elif c < q_cols + k_cols: + body = make_body(0, heads_per_core, 0) + else: + body = make_body(0, 0, slice_f) # v is narrowed, never rotated + for r in range(ROWS_PER_COL): + workers.append(Worker( + body, + fn_args=[w_cores[c][r].cons(), of_x.cons(), + y_cores[c][r].prod(), rmsnorm, rope, *gemv], + stack_size=0xD00)) + + col_w = chunks * ROWS_PER_COL * call_bytes + col_y = ROWS_PER_COL * slice_f + w_taps = TensorTiler2D.simple_tiler((1, n_cols * col_w), (1, col_w)) + y_taps = TensorTiler2D.simple_tiler((1, n_cols * col_y), (1, col_y)) + + def sequence(a_w, a_x, c_y, w_prods, x_prod, y_conss): + tg = TaskGroup() + x_prod.fill(a_x, group=tg) + for c in range(n_cols): + w_prods[c].fill(a_w, tap=w_taps[c], group=tg) + y_conss[c].drain(c_y, tap=y_taps[c], wait=True, group=tg) + tg.finish() + + rt = Runtime(sequence, [w_ty, act_ty, out_ty, + [f.prod() for f in w_l3l2], of_x.prod(), + [f.cons() for f in y_l2l3]]) + return Program(iron.get_current_device(), rt, workers=workers).resolve_program() + + +def main(argv: list[str] | None = None) -> int: + import argparse + from aie.utils.benchmark import run_iters + ap = argparse.ArgumentParser() + ap.add_argument("--iters", type=int, default=200) + ap.add_argument("--pos", type=int, default=17) + a = ap.parse_args((argv or sys.argv)[1:]) + + cfg = json.loads((MODEL / "config.json").read_text(encoding="utf-8")) + names = [f"model.layers.0.self_attn.{p}.weight" + for p in ("q_proj", "k_proj", "v_proj")] + nq, k = projection_shape(names[0], cfg) + nkv, _ = projection_shape(names[1], cfg) + k_tiles = k // TILE_K + q_rows, kv_rows = nq // ROWS_PER_TILE, nkv // ROWS_PER_TILE + tile_rows = q_rows + 2 * kv_rows + + f = q4nx.Q4NX(MODEL / "model.q4nx") + raws = [] + for nm, rows in zip(names, (q_rows, kv_rows, kv_rows)): + off, _ = f.header[nm]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + off) + raws.append(fh.read(rows * k_tiles * TILE_BYTES)) + raw = b"".join(raws) + + off, _ = f.header["model.layers.0.input_layernorm.weight"]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + off) + nwv = np.frombuffer(fh.read(k * 2), dtype=bfloat16) + + n_cols = 7 + n_cores = n_cols * ROWS_PER_COL + assert tile_rows % n_cores == 0 + per_core = tile_rows // n_cores + n_entry = k_tiles // PER_CALL + slice_f = per_core * ROWS_PER_TILE + w = permute_weights(raw, n_cols, per_core, n_entry, PER_CALL * TILE_BYTES) + + half = HD // 2 + theta = float(cfg.get("rope_theta") or cfg["rope_parameters"]["rope_theta"]) + inv_f = 1.0 / (theta ** (np.arange(0, half, dtype=np.float64) * 2.0 / HD)) + ang = a.pos * inv_f + rng = np.random.default_rng(0) + act = np.zeros(2 * k + HD, np.float32) + act[:k] = rng.standard_normal(k) + act[k:2 * k] = nwv.astype(np.float32) + act[2 * k:2 * k + half] = np.cos(ang) + act[2 * k + half:2 * k + HD] = np.sin(ang) + act_bf = act.astype(bfloat16) + + iron.set_current_device(from_name("npu2", n_cols=None)) + c_y = iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu") + b = run_iters(granite_qkv_wide, + iron.tensor(w, dtype=np.uint8, device="npu"), + iron.tensor(act_bf, dtype=bfloat16, device="npu"), c_y, + q_rows=q_rows, kv_rows=kv_rows, k=k, n_cols=n_cols, + warmup=1, iters=a.iters) + # Each core's 512 B element holds its 128 bf16 results in the first 256 B -- + # the accumulator it narrowed into. Read the bytes, not the floats. + raw_y = c_y.numpy().tobytes() + got = np.concatenate([ + np.frombuffer(raw_y[i * slice_f * 4: i * slice_f * 4 + slice_f * 2], + dtype=bfloat16) for i in range(n_cores)]).astype(np.float64) + + xf = act[:k] + invn = 1.0 / np.sqrt((xf * xf).mean() + cfg["rms_norm_eps"]) + h = (xf * invn * nwv.astype(np.float32)).astype(bfloat16) + co, si = np.cos(ang), np.sin(ang) + + def rope(y): + hh = y.reshape(-1, HD) + r = np.empty_like(hh) + r[:, :half] = hh[:, :half] * co - hh[:, half:] * si + r[:, half:] = hh[:, half:] * co + hh[:, :half] * si + return r.reshape(-1) + + q = reference(raws[0], h, q_rows, k_tiles).astype(np.float64) + kk = reference(raws[1], h, kv_rows, k_tiles).astype(np.float64) + v = reference(raws[2], h, kv_rows, k_tiles).astype(np.float64) + ref = np.concatenate([rope(q), rope(kk), v]) + + rel = np.abs(got - ref).max() / (np.abs(ref).max() + 1e-30) + cos_ = float(got @ ref / (np.linalg.norm(got) * np.linalg.norm(ref) + 1e-30)) + # v must NOT be rotated: checked separately, or a kernel that rotated + # everything would still pass on the q and k majority. + vn = 2 * kv_rows * ROWS_PER_TILE + v_rel = np.abs(got[-vn // 2:] - ref[-vn // 2:]).max() / ( + np.abs(ref[-vn // 2:]).max() + 1e-30) + mb = len(raw) / 1e6 + us = b.npu.avg_us + ok = cos_ > 0.9999 and rel < 8e-3 + print(f"RMSNorm + q + k + v + RoPE, ONE dispatch " + f"{n_cols} cols x {ROWS_PER_COL} = {n_cores} cores, per_call {PER_CALL}") + print(f" {q_rows * ROWS_PER_TILE} q + {kv_rows * ROWS_PER_TILE} k + " + f"{kv_rows * ROWS_PER_TILE} v K={k} {mb:.1f} MB") + print(f" cosine {cos_:.8f} max rel err {rel:.3e} " + f"v (unrotated) {v_rel:.3e}") + print(f" {us:.1f} us device {b.e2e.avg_us:.1f} us wall " + f"{mb / us * 1e3:.1f} GB/s") + print(f" separate: RMSNorm 244 + qkv 290 + RoPE 226 = 760 us " + f"(norm+qkv fused was 365 + 226 = 591)") + print(f" [fused == norm, then three GEMVs, then RoPE on q and k] " + f"{'PASS' if ok else 'FAIL'}") + return 0 if ok else 1 + + +# -------------------------------------------------------------------------- +# build_artifact: produce the xclbin WITHOUT model weights. +# +# iron.jit keys its cache on argument shapes and dtypes, not contents, so an +# artefact built from zeros is bit-identical to one built from real weights. +# That is what lets the in-tree build run from a clean checkout with nothing +# but the toolchain -- see kernels/CONVENTION.md. main() below is unchanged and +# remains the developer path: it needs the model and checks the result. +# -------------------------------------------------------------------------- + + +def build_artifact(geometry: dict) -> None: + nq, k = projection_shape("q_proj", geometry) + nkv, _ = projection_shape("k_proj", geometry) + q_rows, kv_rows = nq // ROWS_PER_TILE, nkv // ROWS_PER_TILE + tile_rows = q_rows + 2 * kv_rows + k_tiles = k // TILE_K + n_cols = 7 + w_bytes = tile_rows * k_tiles * TILE_BYTES + + iron.set_current_device(from_name("npu2", n_cols=None)) + granite_qkv_wide( + iron.zeros(w_bytes, dtype=np.uint8, device="npu"), + iron.zeros(2 * k + HD, dtype=bfloat16, device="npu"), + iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu"), + q_rows=q_rows, kv_rows=kv_rows, k=k, n_cols=n_cols) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kernels/granite/iron/granite_swiglu_down.py b/kernels/granite/iron/granite_swiglu_down.py new file mode 100644 index 00000000..1bc79e71 --- /dev/null +++ b/kernels/granite/iron/granite_swiglu_down.py @@ -0,0 +1,300 @@ +r"""SwiGLU folded into down_proj as a per-core prologue. ONE dispatch. + +WHY +--- +0148 measured the per-dispatch floor at ~200 us and SwiGLU at 205 us for 8192 +values -- almost pure floor, since it moves no weights at all. Folding it into +the dispatch that consumes its output removes that entirely. + +THE COST, WHICH IS REAL +----------------------- +down_proj has K = 8192, so its activation is 8192 bf16. Taking gate and up +instead doubles that to 32768 B, and the L1 budget is 62208: + + per_call 1 10240 + 32768 = 43008 fits + per_call 2 20480 + 32768 = 53248 fits + per_call 4 40960 + 32768 = 73728 over + +Unfused, down_proj runs at per_call 4. Fused it must drop to 2, which halves the +DMA element and costs bandwidth. The trade is one whole dispatch against that, +and it is worth measuring rather than assuming -- granite_mlp_wide.py is the +case where exactly this trade came out NEGATIVE. + +IN PLACE, AND WHY THAT IS FORCED +-------------------------------- +A core has two input DMA channels and both are taken: weights and activation. +gate and up therefore share one element, up at gate + 8192, and the SwiGLU +result is written back over gate so the GEMV can read it as x with no third +buffer. See granite_swiglu_ip.cc. + + call c:\dev\mlir-aie\iron_env.cmd + python designs\granite_gemv\granite_swiglu_down.py +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + +import aie.iron as iron +from aie.iron import (CompileTime, In, ObjectFifo, Out, Program, Runtime, + TaskGroup, Worker) +from aie.iron.controlflow import range_ +from aie.iron.device import from_name +from aie.iron.kernel import ExternalFunction +from aie.helpers.taplib import TensorTiler2D + +HERE = Path(__file__).parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parent.parent / "common")) + +import q4nx # noqa: E402 +from granite_gemv import (AIE, MODEL, ROWS_PER_TILE, TILE_BYTES, TILE_K, # noqa: E402 + _include_dirs, ensure_entry_points, projection_shape, + reference, tiles_per_call) + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + +def permute_weights(raw: bytes, n_cols: int, per_core: int, n_entry: int, + call_bytes: int) -> np.ndarray: + """Reorder so each column's stream is contiguous in the order split() wants. + + In: core-major, each core's tile-rows contiguous. + Out: per column, chunk-major then core -- [k][r] -- which is exactly the + layout `split()` consumes, one parent object per k. + """ + a = np.frombuffer(raw, dtype=np.uint8) + # (col, row_in_col, chunk, bytes) -> (col, chunk, row_in_col, bytes) + a = a.reshape(n_cols, ROWS_PER_COL, per_core * n_entry, call_bytes) + return np.ascontiguousarray(a.transpose(0, 2, 1, 3)).reshape(-1) + + +def unpermute_y(y: np.ndarray, n_cols: int, per_core: int, + batch: int = 1) -> np.ndarray: + """Inverse of the above for the joined output. + + On the wire it is [col][t][row_in_col][token][32]; the caller wants one + contiguous result vector per token, so the token axis comes out front. + Returns (batch, tile_rows * 32). + """ + a = y.reshape(n_cols, per_core, ROWS_PER_COL, batch, ROWS_PER_TILE) + a = a.transpose(3, 0, 2, 1, 4) # [token][col][row][t][32] + return np.ascontiguousarray(a).reshape(batch, -1) + + + +ROWS_PER_COL = 4 # compute rows per column on npu2 (array rows 2..5) + + +@iron.jit(aiecc_flags=["--alloc-scheme=basic-sequential"]) +def granite_swiglu_down(w: In, xw: In, y: Out, *, + tile_rows: CompileTime[int], k: CompileTime[int], + n_cols: CompileTime[int] = 8, + per_call: CompileTime[int] = 5): + # per_call is an explicit argument for the reason granite_gemv32 records: + # iron.jit's cache key hashes the call's arguments and nothing else, so a + # value derived in here is invisible to it and two configs collide. + k_tiles = k // TILE_K + n_entry = k_tiles // per_call + call_bytes = per_call * TILE_BYTES + n_cores = n_cols * ROWS_PER_COL + per_core = tile_rows // n_cores + chunks = per_core * n_entry + + srcs = ensure_entry_points(n_entry, per_call, False, 1, k) + + w_l1_ty = np.ndarray[(call_bytes,), np.dtype[np.uint8]] + w_l2_ty = np.ndarray[(ROWS_PER_COL * call_bytes,), np.dtype[np.uint8]] + y_l1_ty = np.ndarray[(ROWS_PER_TILE,), np.dtype[np.float32]] + y_l2_ty = np.ndarray[(ROWS_PER_COL * ROWS_PER_TILE,), np.dtype[np.float32]] + # x and the norm weight share ONE fifo: a compute tile has 2 input DMA + # channels and the weights already take one, so a second activation stream + # is not available at any L1 cost. Two 5120 B buffers and one 10240 B buffer + # cost the same anyway. + xw_ty = np.ndarray[(2 * k,), np.dtype[bfloat16]] + w_ty = np.ndarray[(tile_rows * k_tiles * TILE_BYTES,), np.dtype[np.uint8]] + y_ty = np.ndarray[(tile_rows * ROWS_PER_TILE,), np.dtype[np.float32]] + + kernels = [ExternalFunction(f"granite_gemv_p{per_call}b1_k{i}", + source_file=str(srcs[i]), + arg_types=[w_l1_ty, xw_ty, y_l1_ty], + include_dirs=_include_dirs()) + for i in range(n_entry)] + swiglu = ExternalFunction("granite_swiglu_ip", + source_file=str(AIE / "granite_swiglu_ip.cc"), + arg_types=[xw_ty, np.int32], + include_dirs=_include_dirs()) + + of_x = ObjectFifo(xw_ty, name="nx", depth=1) + + w_l3l2, y_l2l3, w_cores, y_cores = [], [], [], [] + for c in range(n_cols): + wf = ObjectFifo(w_l2_ty, name=f"nwL2_{c}", depth=2) + w_l3l2.append(wf) + w_cores.append(wf.cons().split( + [r * call_bytes for r in range(ROWS_PER_COL)], + obj_types=[w_l1_ty] * ROWS_PER_COL, + names=[f"nw_{c}_{r}" for r in range(ROWS_PER_COL)])) + yf = ObjectFifo(y_l2_ty, name=f"nyL2_{c}", depth=2) + y_l2l3.append(yf) + y_cores.append(yf.prod().join( + [r * ROWS_PER_TILE for r in range(ROWS_PER_COL)], + obj_types=[y_l1_ty] * ROWS_PER_COL, + names=[f"ny_{c}_{r}" for r in range(ROWS_PER_COL)])) + + def core_body(win, xin, yout, sw, *ks): + xe = xin.acquire(1) + # Redundantly on every core, in place. down_proj's K is the whole + # intermediate, so every core needs all 8192 values anyway -- computing + # them per core needs no barrier and no communication. + sw(xe, k) + for _ in range_(per_core): + ye = yout.acquire(1) + for fn in ks: + we = win.acquire(1) + fn(we, xe, ye) + win.release(1) + yout.release(1) + xin.release(1) + + workers = [ + Worker(core_body, + fn_args=[w_cores[c][r].cons(), of_x.cons(), + y_cores[c][r].prod(), swiglu, *kernels], + stack_size=0xD00) + for c in range(n_cols) for r in range(ROWS_PER_COL) + ] + + col_w = chunks * ROWS_PER_COL * call_bytes + col_y = per_core * ROWS_PER_COL * ROWS_PER_TILE + w_taps = TensorTiler2D.simple_tiler((1, n_cols * col_w), (1, col_w)) + y_taps = TensorTiler2D.simple_tiler((1, n_cols * col_y), (1, col_y)) + + def sequence(a_w, a_x, c_y, w_prods, x_prod, y_conss): + tg = TaskGroup() + x_prod.fill(a_x, group=tg) + for c in range(n_cols): + w_prods[c].fill(a_w, tap=w_taps[c], group=tg) + y_conss[c].drain(c_y, tap=y_taps[c], wait=True, group=tg) + tg.finish() + + rt = Runtime(sequence, + [w_ty, xw_ty, y_ty, + [f.prod() for f in w_l3l2], of_x.prod(), + [f.cons() for f in y_l2l3]]) + return Program(iron.get_current_device(), rt, workers=workers).resolve_program() + + +def main(argv: list[str] | None = None) -> int: + import argparse + from aie.utils.benchmark import run_iters + ap = argparse.ArgumentParser() + ap.add_argument("--cols", type=int, default=0) + ap.add_argument("--iters", type=int, default=200) + a = ap.parse_args((argv or sys.argv)[1:]) + + cfg = json.loads((MODEL / "config.json").read_text(encoding="utf-8")) + name = "model.layers.0.mlp.down_proj.weight" + n, k = projection_shape(name, cfg) # 2560 x 8192 + tile_rows, k_tiles = n // ROWS_PER_TILE, k // TILE_K + + f = q4nx.Q4NX(MODEL / "model.q4nx") + off, _ = f.header[name]["data_offsets"] + with f.path.open("rb") as fh: + fh.seek(f._data_start + off) + raw = fh.read(tile_rows * k_tiles * TILE_BYTES) + + cols = a.cols or max(c for c in range(1, 9) + if tile_rows % (c * ROWS_PER_COL) == 0) + n_cores = cols * ROWS_PER_COL + per_core = tile_rows // n_cores + + # per_call from the ACTUAL fixed cost, which is 2*k*2 here because the + # activation carries gate and up: tiles_per_call() assumes k*2 and would + # return 4, which does not fit. + L1 = 64 * 1024 - 0xD00 + fixed = 2 * k * 2 + per_call = max(d for d in range(1, k_tiles + 1) + if k_tiles % d == 0 and d * TILE_BYTES * 2 + fixed <= L1) + n_entry = k_tiles // per_call + w = permute_weights(raw, cols, per_core, n_entry, per_call * TILE_BYTES) + + rng = np.random.default_rng(0) + gate = rng.standard_normal(k).astype(np.float32).astype(bfloat16) + up = rng.standard_normal(k).astype(np.float32).astype(bfloat16) + gu = np.concatenate([gate, up]) + + iron.set_current_device(from_name("npu2", n_cols=None)) + c_y = iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu") + b = run_iters(granite_swiglu_down, + iron.tensor(w, dtype=np.uint8, device="npu"), + iron.tensor(gu, dtype=bfloat16, device="npu"), c_y, + tile_rows=tile_rows, k=k, n_cols=cols, per_call=per_call, + warmup=1, iters=a.iters) + got = unpermute_y(c_y.numpy().copy(), cols, per_core).astype(np.float64) + + # Reference: SwiGLU then the GEMV, together. The kernel narrows the SwiGLU + # result to bf16 before the matmul because that is what the buffer holds, so + # the reference must too -- comparing against an fp64 intermediate would + # charge the kernel for a rounding the storage format dictates. + g64 = gate.astype(np.float64) + h = ((g64 / (1.0 + np.exp(-g64))) * up.astype(np.float64)) + h_bf = h.astype(np.float32).astype(bfloat16) + ref = reference(raw, h_bf, tile_rows, k_tiles).astype(np.float64) + + rel = np.abs(got - ref).max() / (np.abs(ref).max() + 1e-30) + g1, r1 = got.ravel(), ref.ravel() + cos = float(g1 @ r1 / (np.linalg.norm(g1) * np.linalg.norm(r1) + 1e-30)) + mb = len(raw) / 1e6 + us = b.npu.avg_us + # aie::tanh sets the floor, as granite_mlp.py records: ~8e-3 on the SwiGLU + # alone, and down_proj adds little on top. + ok = cos > 0.999 and rel < 5e-2 + print(f"SwiGLU + down_proj, ONE dispatch {cols} cols x {ROWS_PER_COL} = " + f"{n_cores} cores, per_call {per_call}") + print(f" {n} rows K={k} {mb:.1f} MB") + print(f" cosine {cos:.8f} max rel err {rel:.3e}") + print(f" {us:.1f} us device {b.e2e.avg_us:.1f} us wall " + f"{mb / us * 1e3:.1f} GB/s") + print(f" separate: SwiGLU 205 + down 440 = 645 us " + f"(down alone runs per_call 4, this must use {per_call})") + print(f" [fused == SwiGLU then down on the host] {'PASS' if ok else 'FAIL'}") + return 0 if ok else 1 + + +# -------------------------------------------------------------------------- +# build_artifact: produce the xclbin WITHOUT model weights. +# +# iron.jit keys its cache on argument shapes and dtypes, not contents, so an +# artefact built from zeros is bit-identical to one built from real weights. +# That is what lets the in-tree build run from a clean checkout with nothing +# but the toolchain -- see kernels/CONVENTION.md. main() below is unchanged and +# remains the developer path: it needs the model and checks the result. +# -------------------------------------------------------------------------- + + +def build_artifact(geometry: dict) -> None: + n, k = projection_shape("down_proj", geometry) + tile_rows, k_tiles = n // ROWS_PER_TILE, k // TILE_K + cols = max(c for c in range(1, 9) if tile_rows % (c * ROWS_PER_COL) == 0) + L1 = 64 * 1024 - 0xD00 + fixed = 2 * k * 2 + per_call = max(d for d in range(1, k_tiles + 1) + if k_tiles % d == 0 and d * TILE_BYTES * 2 + fixed <= L1) + w_bytes = tile_rows * k_tiles * TILE_BYTES + + iron.set_current_device(from_name("npu2", n_cols=None)) + granite_swiglu_down( + iron.zeros(w_bytes, dtype=np.uint8, device="npu"), + iron.zeros(2 * k, dtype=bfloat16, device="npu"), + iron.zeros(tile_rows * ROWS_PER_TILE, dtype=np.float32, device="npu"), + tile_rows=tile_rows, k=k, n_cols=cols, per_call=per_call) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/kernels/requirements.txt b/kernels/requirements.txt new file mode 100644 index 00000000..d56e3d03 --- /dev/null +++ b/kernels/requirements.txt @@ -0,0 +1,12 @@ +# The toolchain these kernels were built and validated against. +# +# mlir-aie ships the `aie` Python package, the aiecc driver and the Peano +# (LLVM-AIE) backend. It is not on PyPI under this name; install it from the +# project's own wheels or a source build, then activate that environment before +# configuring with -DFLM_BUILD_KERNELS=ON. See README.md. +# +# An .xclbin is only valid for the toolchain that produced it, which is why +# these are pinned and why every build records them again in manifest.json. +mlir-aie==1.4.2.dev16+g7e00b57 +ml_dtypes==0.6.0 +numpy==2.5.1 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index cc85d041..addf8fe8 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -880,3 +880,70 @@ install(FILES model_info.json DESTINATION "${FLM_SHARE_DESTINATION}") # xclbins, which are loaded by shared libraries need to be in location # relative to the executable, so we install them relative to the binary. install(DIRECTORY xclbins DESTINATION "${FLM_SHARE_DESTINATION}") + +# --------------------------------------------------------------------------- +# Optional: build NPU kernels from source. See ../kernels/CONVENTION.md. +# +# OFF by default, and when OFF this block adds no target, no dependency, no +# install rule and no compile flag -- the build is identical to a tree in which +# kernels/ does not exist. The `flm` target never depends on it in either +# state, so a kernel build failure cannot block the binary. +# +# The AIE toolchain is deliberately not modelled in CMake. There is no +# find_package for it; aiecc is invoked by the IRON runtime from inside Python; +# and the dependency graph is not representable, because a design's entry +# points do not exist until the design has run. Shelling out to one script +# keeps this block free of toolchain knowledge, and keeps that script usable +# with no CMake at all. +# --------------------------------------------------------------------------- +option(FLM_BUILD_KERNELS "Build NPU kernels from source in ../kernels (needs the IRON toolchain)" OFF) + +if(FLM_BUILD_KERNELS) + set(FLM_KERNELS_DIR "${CMAKE_SOURCE_DIR}/../kernels" CACHE PATH + "Kernel source tree") + set(FLM_KERNELS_PYTHON "" CACHE FILEPATH + "Python from an activated IRON environment (empty = found on PATH)") + + if(FLM_KERNELS_PYTHON) + set(_flm_kpy "${FLM_KERNELS_PYTHON}") + else() + find_package(Python3 3.10 COMPONENTS Interpreter) + if(NOT Python3_Interpreter_FOUND) + message(FATAL_ERROR + "FLM_BUILD_KERNELS=ON but no Python 3.10+ was found. Activate the " + "IRON environment before configuring, or pass " + "-DFLM_KERNELS_PYTHON=. See " + "kernels/README.md. Leave FLM_BUILD_KERNELS=OFF to skip this.") + endif() + set(_flm_kpy "${Python3_EXECUTABLE}") + endif() + + # Fail at CONFIGURE time with the script's own one-line diagnostic, rather + # than at build time with an aiecc traceback. + execute_process( + COMMAND "${_flm_kpy}" "${FLM_KERNELS_DIR}/build_kernels.py" --check-toolchain + RESULT_VARIABLE _flm_k_rc + OUTPUT_VARIABLE _flm_k_out + ERROR_VARIABLE _flm_k_out + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT _flm_k_rc EQUAL 0) + message(FATAL_ERROR + "FLM_BUILD_KERNELS=ON but the IRON toolchain is not usable:\n" + " ${_flm_k_out}\n" + "Set -DFLM_BUILD_KERNELS=OFF (the default) to build flm without it.") + endif() + message(STATUS "FLM kernel source build: ON -- ${_flm_k_out}") + + add_custom_target(flm_kernels ALL + COMMAND "${_flm_kpy}" "${FLM_KERNELS_DIR}/build_kernels.py" + --family granite --out "${CMAKE_BINARY_DIR}/kernels" + WORKING_DIRECTORY "${FLM_KERNELS_DIR}" + COMMENT "Building NPU kernels from source (IRON/aiecc)" + VERBATIM USES_TERMINAL) + + # Built artefacts install exactly like the checked-in ones, into the same + # tree, because build_kernels.py lays them out as /.xclbin. + install(DIRECTORY "${CMAKE_BINARY_DIR}/kernels/" + DESTINATION "${FLM_SHARE_DESTINATION}/xclbins" + PATTERN "_generated" EXCLUDE) +endif()