Conversation
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>
Author
|
The CI failure appears to be pre-existing: same error occurs on main. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
flm_plugin.hpp(plugin loading,flm::plugin_context, theFLM_PLUGINmacro),npu_utils/op_override.hpp(op_override,op_result,op_call) andnpu_utils/op_registry.hpp(op_registryand its key matching).gemma4e_ops::op::dequant_qkvand thegemma4e_ops::key()that composes a layer index with one. A plugin includesthat header and hooks the keys it wants.
Declaring operators is opt-in per model:
causal_lm::ops()returnsnullptrby default, and Gemma4 E2B/E4B is the only engine that declares any so far.If
FLM_PLUGINis 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:
Then start with:
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.