Skip to content

Add a plugin mechanism for operator overrides - #743

Open
andrej wants to merge 19 commits into
ROCm:mainfrom
andrej:andrej/replace-gemm
Open

andrej wants to merge 19 commits into
ROCm:mainfrom
andrej:andrej/replace-gemm

Conversation

@andrej

@andrej andrej commented Sep 18, 2026

Copy link
Copy Markdown

Adds a plugin mechanism that allows users to override individual operators. This allows open-source contributors to experiment with dropping in individual custom kernels while harnessing the existing infrastructure.

The plugin mechanism consists of:

  • Common infrastructure, public in this repository: flm_plugin.hpp (plugin loading, flm::plugin_context, the FLM_PLUGIN macro), npu_utils/op_override.hpp (op_override, op_result, op_call) and npu_utils/op_registry.hpp (op_registry and its key matching).
  • Operator names, per-model and public: each model's header names the operators it dispatches, e.g. gemma4e_ops::op::dequant_qkv and the gemma4e_ops::key() that composes a layer index with one. A plugin includes
    that header and hooks the keys it wants.
  • The binding of those names to the engine's dispatch sites, which lives in the per-model DLL/SO (separate PR in internal repo).
  • The plugin itself, example below, compiled as a standalone DLL/SO and loaded at runtime via an environment variable.

Declaring operators is opt-in per model: causal_lm::ops() returns nullptr by default, and Gemma4 E2B/E4B is the only engine that declares any so far.

If FLM_PLUGIN is unset, no operator is overridden and behavior is unchanged from "stock".

An override does not have to serve every call it is bound to. Returning flm::op_result::decline() falls back to the default implementation for that call, so a hook can take only the shapes it handles and leave the rest alone.

Example:

#include <memory>

#include "flm_plugin.hpp"                   // op_override, plugin_context, FLM_PLUGIN
#include "models/gemma4e/gemma4e_npu.hpp"   // gemma4e_ops::key, gemma4e_ops::op

class my_dequant : public flm::op_override {
public:
    flm::op_result create_run(const flm::op_call& call) override {
        // ...
    }
};

void register_my_plugin(const flm::plugin_context& ctx) {
    int layer = 0;
    ctx.ops->override_op(gemma4e_ops::key(layer, gemma4e_ops::op::dequant_qkv),
                         std::make_shared<my_dequant>());
}

FLM_PLUGIN(register_my_plugin)

Then start with:

FLM_PLUGIN=/path/to/my_plugin.so flm serve gemma4-it:e2b

This PR also ships an example plugin, src/plugins/iron_gemm, which demonstrates how to override the dequant and matrix multiplication operators of Gemma4 E2B/E4B prefill. src/plugins/README.md documents it.

andrej and others added 19 commits September 18, 2026 12:08
npu_app::load_insts runs a prebuilt transaction binary in place of a sequence
the host generates. The instructions then travel as a buffer argument, so this
path uses a plain xrt::kernel and the app generates no sequence of its own.
That is what lets an operator compiled outside this repo run here.

The artifacts are one xclbin and twelve instruction streams covering every
prefill GEMM shape of Gemma4-E2B. M, K and N are runtime parameters of the
device body, so the shapes share the xclbin and each costs about 38 KB.

The rebuilt engine libraries dispatch the 14 apps mm.xclbin used to serve to
that xclbin instead, reading weights a converter packed to bfp16ebs8 ahead of
time. Prefill of a 247-token prompt takes 702 ms rather than 959 ms.

The packed weights are not here: they are a per-model sidecar of 2.09 GB, built
from model.q4nx.
An engine declares the operations it dispatches; a plugin named in FLM_PLUGIN
binds an override to the ones it wants and runs them however it likes. The
mechanism is header only, so a plugin builds against the public tree and needs
neither flm nor any engine library rebuilt.

An override returns a run for the caller to wait on, an empty result when it
already did the work, or declines, in which case the engine runs its own
implementation for that one call. Declining is per dispatch, so an override
serves the shapes it has and leaves the rest alone.

plugins/flm_gemm runs Gemma4 E2B's prefill projections on the IRON FLMGEMM
operator and switches off the dequantization that fed them. Prefill of a
247-token prompt: 672 ms against 714 ms for the engine's own GEMM reading
dequantized weights.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Each model constructs its own engine, so there is no shared step between that
and reading the weights — which is the one window in which an override may be
registered. _load_engine_weights() becomes that step: it loads the plugins,
reads the weights and releases the file, and every model calls it immediately
after constructing its engine.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
A role constant that nothing declares is a key an override can bind to and
never fire on, which is the failure the declared registry exists to prevent.
gemma4e dispatches lm_head through a run it builds once, not through the
override path.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The plugin ran the GEMM from weights a converter had packed offline, which meant
a 1.94 GiB sidecar every model had to ship and keep in step. IRON's DequantBFP
now runs where the engine's dequant.xclbin ran, per layer per chunk, writing
bfp16ebs8 straight into the packed order the GEMM reads. The weights stay 4-bit
in DRAM and the staging comes to 75 MiB.

Its instruction streams are built with qw_layout=engine. The engine's loader
interleaves pairs of 32-row block-rows, and that is byte for byte the order the
operator's fill would otherwise gather out of the weights file, so in engine mode
the descriptor reads the layer's buffer straight through. q, k and v go in one
dispatch at their combined width and the GEMMs take their share through a
sub-buffer, which is what the shipped mm expresses as a weight_offset.

FLM_GEMM_MODE keeps the offline configurations, now selectable from one binary
rather than an instrumented engine build. Against 939.3 ms for the stock engine
on a 247-token prompt: 874.7 ms per chunk, 733.6 ms with bf16 weights resident
and the shipped kernel, 669.9 ms with bfp16 weights and this one. So dropping the
dequant from the timed path is worth 205.7 ms and the kernel a further 63.7 ms.

FLM_DEQUANT_VERIFY compares every buffer the dequant writes against the sidecar
for the same tensor: 205 checks on E2B, every projection of every layer.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The entry point opened on the mechanism in the abstract, which is the wrong end
for someone deciding whether any of this is for them. It now says what the
example plugin overrides and what that buys, shows the prefill medians for each
combination of its two options, and only then walks the mechanism -- entirely
through lines from that plugin rather than a sketch.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The medians were pasted into the plot script, so the chart could drift from the
benchmark without anyone noticing. It now reads bench3.sh's own
"<label>: median <ms>" lines off stdin and keys the bars by label.

Redrawn from one session rather than measurements taken days apart: 929 ms
stock, 883 ms per chunk, 726 ms with bf16 weights on the shipped kernel, 670 ms
with bfp16 on the IRON one. The numbers move within the 2% run-to-run spread;
the conclusions do not.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Same 205 tensors byte-identical to the sidecar, and prefill unmoved at 881.7 ms
against 882.9 before -- inside the run-to-run spread, which is what a change
described as simplifying the runtime sequence should look like.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… set

The GEMM lineage this plugin was built against had moved on: its artifact stem
now leads with the configuration tag rather than trailing it, and the tag itself
gained ck and mc fields. Discovery no longer spells that tag out. It takes it
from whichever FLM_GEMM xclbin is present -- a configuration's stem never
carries _M<digits>, a stream's always does -- and keeps the file each shape came
from instead of rebuilding the name to load it.

A layer the plugin would otherwise serve but whose shape has no instruction
stream now throws at load, naming the layer, the projection and the shape. It
used to shrink coverage silently, which leaves the model correct and merely
slower: the failure mode nothing reports. An M nothing was built for stays a
run-time fallback, since that is a chunk length rather than a broken build.

Rebuilt both operators from IRON 3b9e4bb1b and re-measured all four
combinations in one session. Against 922.0 ms stock: 880.3 per chunk, 732.3 with
bf16 weights on the shipped kernel, 673.7 with bfp16 on this one. Still 205 of
205 tensors byte-identical to the sidecar.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The per-plugin README repeated the root's intro, its sidecar instructions and
its performance numbers, which put the measurements in three places -- the
chart, that table, and the internal reproduction doc -- all of which had to be
edited on every re-measure.

What was only there splits in two. The artifact layout and the environment
variables are what a user needs, and move to the root. The rest -- the engine
block order, the gate/up interleave, the q/k/v fusion, how shapes and offsets
are derived -- is implementation, and is already commented at the code it
describes, where it is likelier to stay true.

The plugin's own header also still advertised a default for FLM_GEMM_CONFIG,
which stopped being a default when discovery started reading the tag off the
xclbin.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
dequantize.py and sidecar.py existed in both trees, and the internal copies were
the older ones. sidecar.py was a library with a single caller, so it folds into
dequantize.py; its self-check went with it, having compared against dumps from
an instrumented engine build that no longer exists. FLM_DEQUANT_VERIFY is the
check now.

build_artifacts.py replaces the pair of build scripts that lived internally, and
builds both operators' shapes in one pass. The reproduction section no longer
reaches for a benchmark script in another repository: it is one block from
building the operators to serving the model.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The README framed the plugin against a baseline it called "stock" and described
its GEMM as the faster kernel. Both are now named for what they are: the
measured engine is FastFlowLM v1.0.5, which is the version the numbers were
taken on, and the override is the IRON bfp16 GEMM. The chart carries the same
two names.

The options are described by what they do and what they cost -- disk and memory
for prefill time, bfp16 rather than bf16 weights -- and the four measurements
are left to speak for themselves.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The plugin is named for what it brings, not for the engine it plugs into, so
flm_gemm becomes iron_gemm and its variables follow: IRON_GEMM_MODE,
IRON_GEMM_OFF, IRON_GEMM_CONFIG, IRON_DEQUANT_VERIFY. FLM_PLUGIN stays, being
the engine's.

The FLM_ prefixes on the artifact filenames are unchanged: they come from each
IRON operator's own naming, not from here. The header comment saying so was also
still describing the pre-rebase artifact stem, which put the shape before the
configuration rather than after.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
no_op_override, is_overridden and all_overridden had no callers. The first was
how a plugin was meant to switch an operator off; iron_gemm does that from its
own hook instead, because the decision is per dispatch. The other two were how
an engine could ask what had been taken over, which the engine stopped needing
once the dequant staging was allocated unconditionally.

plugin_context::model_name went the same way: unused, and one filename() call
from model_path if a plugin wants it.

streams_ and stream_files_ were a set and a map with identical keys, filled
together and queried separately. The map is enough.

list_ops() stays, having no caller here but being how a plugin discovers what a
model declares.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The plugin API named the hook points: a role vocabulary in the shared header, a
key spelled "layers.<i>.<role>", and an op_call carrying a layer index. All
three were Gemma4e's decomposition promoted to universal, which would have made
a second engine either adopt someone else's names or silently collide with them
on a shared key.

The framework now takes a key it never parses, a name it only hands back, and a
dense index it only uses as a subscript into each app's override table. That
index exists because one npu_app serves many dispatch sites -- q_swa_proj is one
app across 28 layers -- so a dispatch has to say which site it is; it is not a
layer as far as anything generic is concerned.

gemma4e_ops::op and gemma4e_ops::key move to the model's own header, beside the
comment that already documented each operation's arguments. An engine with one
layer declares index 0; one with nested blocks spells whatever key it likes.

205 of 205 tensors still byte-identical to the sidecar.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Mostly artifacts, but not only. The plugin inferred which layers share a kv
cache from the MLP being double-wide there, which holds on E2B and not on E4B:
its MLP is 10240 on all 42 layers, so every layer looked non-skip and the 18
that share a cache would have read q, o, gate, up and down at offsets shifted by
the k and v the engine never wrote. num_kv_shared_layers says it outright, so
that is what decides it now.

build_artifacts.py derives its shape set from the model's own weights, the same
way the plugin does, rather than listing E2B's. It reproduces E2B's 12 GEMM and
10 dequant shapes exactly and gives E4B 9 and 8.

E4B binds all 294 projections. 24 of 24 tensors byte-identical to the sidecar
across layers 0, 5, 24 and 29 -- sliding-window and global, kv-shared and not --
with k and v correctly absent on the shared ones.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The prefill path upstream now runs through gemma4e_prefill_context's
dequant/attn/mlp blocks, so the dispatch sites the registry declares
moved with it. This is that engine, rebuilt.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The README explained iron_gemm without ever saying what the pieces of the
mechanism are or how to build the smallest thing that uses it. Says both,
with a standalone example and the two build flags that are not optional.

The binding snippet still spelled keys with op_key and flm::role, which
went away when each model took ownership of naming its own operations.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@andrej

andrej commented Sep 18, 2026

Copy link
Copy Markdown
Author

The CI failure appears to be pre-existing: same error occurs on main.

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.

1 participant