Skip to content

feat: granite 4.2 - kernels - #701

Open
vegah wants to merge 3 commits into
ROCm:mainfrom
vegah:feat/kernels
Open

feat: granite 4.2 - kernels#701
vegah wants to merge 3 commits into
ROCm:mainfrom
vegah:feat/kernels

Conversation

@vegah

@vegah vegah commented Sep 2, 2026

Copy link
Copy Markdown

Motivation

One question, and the answer can be no: do you want AIE kernel source in
this repository? If not, say so and I will keep it out of tree — the work is not
wasted either way, and I would rather ask once than keep guessing.

src/xclbins/ holds 38 directories of prebuilt binaries 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 — #669's "will never
work above 1.0.1"
is that, observed from a user's side.

Source plus a recorded toolchain fingerprint is the representation that
survives a version bump. That is the proposal. kernels/CONVENTION.md states it
so it can be judged on its own; the granite kernels are the worked example, not
the point.

This is the second of two PRs. The first adds the granite model and a CPU host
engine. Neither depends on the other — this one adds no code that flm
links, and that engine needs nothing from here. Take either, both, or neither.

Technical Details

One new top-level directory and one guarded block appended to
src/CMakeLists.txt. That block is the only existing file this PR modifies.

kernels/CONVENTION.md      the proposal
kernels/build_kernels.py   the build, usable with no CMake at all
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

Checkable claims, if you would rather spot-check than read:

  • with FLM_BUILD_KERNELS=OFF — the default — the build is identical to a
    tree in which kernels/ does not exist.
    No target, no dependency, no
    install rule, no compile flag. 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 .github/workflows/*.yml is touched — CI cost and the release
    pipeline are unaffected
  • no compiled artefact is committed, and the convention makes that a rule
    rather than a habit

kernels/ is at the repository root rather than under src/ because
CMAKE_SOURCE_DIR is <repo>/src: root is outside the glob now and after any
future broadening. Relying on the .cc extension to stay outside it would be a
trap — the next contributor names one .cpp and it is silently compiled into
flm.

aie/ and iron/ are split because they are two review surfaces. aie/ is
what executes on the array and is the surface that matters for provenance;
iron/ only describes placement and data movement. Read aie/ first.

The 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
the CMake block free of toolchain knowledge, and keeps that script usable
standalone.

Licence: everything under kernels/ is MIT, matching
LICENSE_RUNTIME.txt, with SPDX in every file. Written from the public
MLIR-AIE/IRON examples and published AIE2P documentation — no shipped .xclbin
was disassembled and no closed component was reverse-engineered. If you would
prefer a different licence or location, say which and I will change it.

Not in this PR, on purpose: flm cannot dispatch these yet.
npu_utils_xrt.hpp builds its ELF from a control sequence assembled on the
host by npu_sequence in npu_instr_utils.hpp; IRON emits that same control
code as a prebuilt insts.bin at build time. Bridging them is roughly one
function. That is a runtime change, not a build change, and the two are worth
deciding separately.

Test Plan

The build has three states and all three matter, so all three were exercised on
top of main (6002e0f), Windows, Ryzen AI 9 HX 370:

  1. configure with the option absent (the default) — must leave no trace
  2. configure with the option on and a Python that lacks the toolchain — must
    fail at configure time with a usable message, not an aiecc traceback
    halfway through a build
  3. configure with the option on and the IRON environment — must build every
    artefact from a clean checkout with no model weights present, which is
    what makes the option safe to enable in a container

Then the artefacts themselves: build_kernels.py --validate runs each design
against a host reference built from the same bytes. That needs an NPU and the
model weights, so it is a local hardware check and CI is not asked to run it.

Test Result

cmake -DFLM_BUILD_KERNELS=OFF          (default)
  -- Configuring done / Generating done
  cache: FLM_BUILD_KERNELS:BOOL=OFF
  `cmake --build . --target help` lists no flm_kernels target at all

cmake -DFLM_BUILD_KERNELS=ON -DFLM_KERNELS_PYTHON=<python without the toolchain>
  CMake Error at CMakeLists.txt:929 (message):
    FLM_BUILD_KERNELS=ON but the IRON toolchain is not usable:
      the `aie` package is not importable from <python> (No module named 'aie').
      Activate the mlir-aie/IRON environment before configuring, or pass
      -DFLM_KERNELS_PYTHON=<path to the IRON python>.
    Set -DFLM_BUILD_KERNELS=OFF (the default) to build flm without it.

cmake -DFLM_BUILD_KERNELS=ON -DFLM_KERNELS_PYTHON=<IRON python>
  -- FLM kernel source build: ON -- mlir-aie 1.4.2.dev16+g7e00b57
  -- Configuring done / Generating done

python kernels/build_kernels.py --family granite --out build/kernels
  building norm_qkv_rope ... -> norm_qkv_rope.xclbin, norm_qkv_rope.insts.bin
  building attn_o        ... -> attn_o.xclbin, attn_o.insts.bin
  building norm_qkv      ... -> norm_qkv.xclbin, norm_qkv.insts.bin
  building norm_gate_up  ... -> norm_gate_up.xclbin, norm_gate_up.insts.bin
  building swiglu_down   ... -> swiglu_down.xclbin, swiglu_down.insts.bin
  manifest -> build/kernels/Granite-4.2-3B-NPU2/manifest.json
  OK

flm also builds unchanged with the option off, and the artefacts land laid
out as <Model>/<name>.xclbin so they install through the existing xclbins
rule.

manifest.json records the package versions, the aiecc flags, a geometry
hash and a sha256 per artefact, so a consumer can refuse a mismatched pair
rather than dispatch it into the NPU.

Per-design agreement against a host reference, for the record: cosine
1.00000000 for the GEMV groups under a one-hot activation, and > 0.9996 for the
fused blocks. For scale, the four fused dispatches measure 1744.7 µs for a
granite layer — 13.6 tok/s of device time for the stack, against 8.7 tok/s for
the CPU engine in the companion PR — but none of that is reachable from flm
until the dispatch path above exists.

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 ROCm#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.
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.
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 <python> (No module named 'aie').
    Activate the mlir-aie/IRON environment before configuring, or pass
    -DFLM_KERNELS_PYTHON=<path to the IRON 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 <Model>/<name>.xclbin and install through the
existing xclbins destination.
@Atomic-Germ

Copy link
Copy Markdown

AMD does not, but I do!!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants