From 16abd573ead4ecee2066229ac8f61bbbb98f7768 Mon Sep 17 00:00:00 2001 From: Vegard Berget Date: Wed, 2 Sep 2026 10:47:54 +0200 Subject: [PATCH 1/3] feat: add kernels/ -- in-tree AIE kernel sources for granite (MIT) This repository has no kernel source: src/xclbins/ holds 38 directories of prebuilt blobs with no record of what produced them and no way to rebuild them when the toolchain moves. An .xclbin is only valid for the toolchain that built it, so a committed binary rots silently, and the rot is invisible until a user hits it -- which is what issue #669's "will never work above 1.0.1" is. kernels/ proposes source plus a recorded toolchain fingerprint as the representation that survives a version bump. CONVENTION.md states the rules and is meant to be evaluated on its own; granite is the first worked example, not the point. kernels/CONVENTION.md the proposal kernels/README.md prerequisites, build, and what is not done yet kernels/LICENSE MIT, matching LICENSE_RUNTIME.txt kernels/common/ host-side helpers shared across families kernels/granite/aie/ 5 headers + 7 entry points -- the device C++ kernels/granite/iron/ 6 IRON designs -- placement and data movement kernels/granite/geometry.json model dimensions, so the build needs no weights No compiled artefact is committed here and none ever should be. Nothing under src/ is added, moved or renamed by this commit. 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 layout is the derivation already documented in src/include/models/granite/q4nx_host.hpp. --- kernels/CONVENTION.md | 88 ++++ kernels/LICENSE | 21 + kernels/README.md | 81 ++++ kernels/common/q4nx.py | 197 +++++++++ kernels/granite/README.md | 94 +++++ kernels/granite/aie/granite_attention.h | 238 +++++++++++ kernels/granite/aie/granite_attn_block_h.cc | 38 ++ kernels/granite/aie/granite_attn_finish_h.cc | 28 ++ kernels/granite/aie/granite_elementwise.h | 127 ++++++ kernels/granite/aie/granite_gemv.h | 302 ++++++++++++++ kernels/granite/aie/granite_qgemv_g0.cc | 18 + kernels/granite/aie/granite_qgemv_g1.cc | 18 + kernels/granite/aie/granite_qkv_rope.h | 90 ++++ kernels/granite/aie/granite_rmsnorm.h | 100 +++++ kernels/granite/aie/granite_rmsnorm_ip.cc | 17 + kernels/granite/aie/granite_rope_ip.cc | 19 + kernels/granite/aie/granite_swiglu_ip.cc | 14 + kernels/granite/geometry.json | 12 + kernels/granite/iron/granite_attn_o.py | 407 +++++++++++++++++++ kernels/granite/iron/granite_gemv.py | 399 ++++++++++++++++++ kernels/granite/iron/granite_gemv32.py | 297 ++++++++++++++ kernels/granite/iron/granite_norm_gemv.py | 330 +++++++++++++++ kernels/granite/iron/granite_qkv_wide.py | 360 ++++++++++++++++ kernels/granite/iron/granite_swiglu_down.py | 300 ++++++++++++++ kernels/requirements.txt | 12 + 25 files changed, 3607 insertions(+) create mode 100644 kernels/CONVENTION.md create mode 100644 kernels/LICENSE create mode 100644 kernels/README.md create mode 100644 kernels/common/q4nx.py create mode 100644 kernels/granite/README.md create mode 100644 kernels/granite/aie/granite_attention.h create mode 100644 kernels/granite/aie/granite_attn_block_h.cc create mode 100644 kernels/granite/aie/granite_attn_finish_h.cc create mode 100644 kernels/granite/aie/granite_elementwise.h create mode 100644 kernels/granite/aie/granite_gemv.h create mode 100644 kernels/granite/aie/granite_qgemv_g0.cc create mode 100644 kernels/granite/aie/granite_qgemv_g1.cc create mode 100644 kernels/granite/aie/granite_qkv_rope.h create mode 100644 kernels/granite/aie/granite_rmsnorm.h create mode 100644 kernels/granite/aie/granite_rmsnorm_ip.cc create mode 100644 kernels/granite/aie/granite_rope_ip.cc create mode 100644 kernels/granite/aie/granite_swiglu_ip.cc create mode 100644 kernels/granite/geometry.json create mode 100644 kernels/granite/iron/granite_attn_o.py create mode 100644 kernels/granite/iron/granite_gemv.py create mode 100644 kernels/granite/iron/granite_gemv32.py create mode 100644 kernels/granite/iron/granite_norm_gemv.py create mode 100644 kernels/granite/iron/granite_qkv_wide.py create mode 100644 kernels/granite/iron/granite_swiglu_down.py create mode 100644 kernels/requirements.txt 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/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 From 6ce7a7992bab46a22c7f2623ba10bcfe10d4d3d9 Mon Sep 17 00:00:00 2001 From: Vegard Berget Date: Wed, 2 Sep 2026 10:49:08 +0200 Subject: [PATCH 2/3] feat: build_kernels.py -- build granite designs from IRON source One entry point, usable with no CMake at all: python kernels/build_kernels.py --check-toolchain python kernels/build_kernels.py --family granite --out build/kernels python kernels/build_kernels.py --family granite --validate It drives the build_artifact(geometry, **kwargs) entry point each design exposes, so no model weights are needed: iron.jit keys its cache on argument shapes and dtypes rather than contents, and an artefact built from zeros is bit-identical to one built from real weights. The build therefore runs from a clean checkout with nothing but the toolchain, which is what makes the CMake option safe to enable in a container. It does not try to activate the toolchain. Activation does not survive a subprocess, and iron_env.cmd is Windows-only and machine-specific. It imports aie.iron in the current interpreter and, if that fails, reports exactly what was missing and where -- which is what makes the CMake error message useful. It locates the JIT output by set difference over the cache directory rather than by matching strings in the emitted MLIR. A set difference cannot be wrong and does not depend on IRON internals that change between versions. Every build writes manifest.json with the package versions, the aiecc flags, the geometry hash and a sha256 per artefact, so a consumer can refuse a mismatched pair rather than dispatch it. --- kernels/build_kernels.py | 252 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 252 insertions(+) create mode 100644 kernels/build_kernels.py 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()) From ad33b78eabc1a9310ae990493fd4f375310176a4 Mon Sep 17 00:00:00 2001 From: Vegard Berget Date: Wed, 2 Sep 2026 10:49:08 +0200 Subject: [PATCH 3/3] feat: optional FLM_BUILD_KERNELS integration, default OFF Appended at the end of src/CMakeLists.txt, after FLM_SHARE_DESTINATION has settled in the non-WIN32 branches. With the option OFF -- the default -- nothing inside the block executes: no target, no dependency, no install rule, no compile flag. The build is identical to a tree in which kernels/ does not exist, and `flm` never depends on the kernel target in either state, so a kernel build failure cannot block the binary. No file under src/ is added, moved or renamed, so no glob result changes and the flm link line is unchanged. No workflow under .github/ is touched. With it ON, --check-toolchain runs at CONFIGURE time, so a missing toolchain is one sentence naming both fixes rather than an aiecc traceback halfway through a build: FLM_BUILD_KERNELS=ON but the IRON toolchain is not usable: the `aie` package is not importable from (No module named 'aie'). Activate the mlir-aie/IRON environment before configuring, or pass -DFLM_KERNELS_PYTHON=. Set -DFLM_BUILD_KERNELS=OFF (the default) to build flm without it. The toolchain is not modelled in CMake on purpose: 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, so it cannot rot. Built artefacts are laid out as /.xclbin and install through the existing xclbins destination. --- src/CMakeLists.txt | 67 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) 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()