Add a plugin mechanism for operator overrides

.gitignore

@@ -5,3 +5,8 @@ __pycache__
 _site
 .jekyll-metadata
 *.csv
+
+# iron_gemm's own artifacts (src/plugins/iron_gemm/tools/build_artifacts.py):
+# built by the user from IRON, not shipped.
+src/xclbins/*/FLM_GEMM_*
+src/xclbins/*/FLM_DequantBFP_*

src/CMakeLists.txt

@@ -898,3 +898,13 @@ 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}")
+
+# ———————————————————————————————————————————————
+# Operator plugins
+# ———————————————————————————————————————————————
+# Reference overrides, built on the public API only. Off by default because they
+# need artifacts and packed weights that are produced out of tree.
+option(FLM_BUILD_PLUGINS "Build the reference operator plugins" OFF)
+if(FLM_BUILD_PLUGINS)
+    add_subdirectory(plugins/iron_gemm)
+endif()

src/common/AutoModel/automodel.cpp

@@ -6,6 +6,7 @@
 /// \note This is a source file for the auto_model class
 
 #include "AutoModel/automodel.hpp"
+#include "flm_plugin.hpp"
 
 
 AutoModel::AutoModel(flm_rt::device* npu_device_inst, std::string current_model) {
@@ -138,6 +139,11 @@ void AutoModel::_shared_load_model(std::string model_path, json model_info, int
         exit(1);
     }
     this->npu = std::make_unique<npu_xclbin_manager>(npu_device::device_npu2, this->npu_device_inst, enable_preemption);
+    // Plugins bind overrides on npu->hooks; every model engine below resolves
+    // them in its own constructor, so plugins must load before any engine exists.
+    const std::string xclbin_path = utils::path_join(this->lm_config->exec_path, "xclbins", this->lm_config->model_name);
+    flm::plugin_context plugin_ctx{ this->npu.get(), this->model_path.c_str(), xclbin_path.c_str() };
+    flm::load_plugins_from_env(plugin_ctx);
     this->enable_preemption = enable_preemption;
     // Single-turn models (e.g. dedicated translation models) don't support arbitrary
     // context length overrides, so always fall back to the model's own default.

src/include/flm_plugin.hpp

@@ -0,0 +1,113 @@
+/// \file flm_plugin.hpp
+/// \brief Loading user supplied operator overrides from a shared library.
+/// \note  A plugin is a shared library exporting flm_plugin_abi_version() and
+///        flm_plugin_register(). Point FLM_PLUGIN at it and flm loads it right
+///        after the NPU device is ready and before any model engine is
+///        constructed, which is the window in which a plugin may register
+///        overrides and allocate its own weights.
+///
+/// A minimal plugin:
+/// \code
+///   #include "flm_plugin.hpp"
+///
+///   static void register_overrides(const flm::plugin_context& ctx) {
+///       ctx.npu->hooks.override_op<my_op_func_t>("mlp.up_proj",
+///           [ctx](bytes& out, bytes& in, bytes& weights, int layer, int m) {
+///               return flm::hook_result<ert_cmd_state>(...);
+///           });
+///   }
+///   FLM_PLUGIN(register_overrides)
+/// \endcode
+///
+/// \note A plugin and flm exchange C++ objects, so both must be built with the
+///       same compiler and standard library. The ABI version guards against
+///       loading a plugin built for a different override interface; it cannot
+///       detect a toolchain mismatch.
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "npu_utils/npu_utils.hpp"
+#include "utils/debug_utils.hpp"
+
+#ifdef _WIN32
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
+namespace flm {
+
+inline constexpr int plugin_abi_version = 1;
+
+/// \brief What a plugin is given when it registers.
+struct plugin_context {
+    npu_xclbin_manager* npu;   ///< the one device manager; also holds the hook registry (npu->hooks)
+    const char* model_path;    ///< directory holding the model's weights
+    const char* xclbin_path;   ///< directory holding this model's xclbins, where a plugin's belong too
+};
+
+using plugin_abi_version_fn = int (*)();
+using plugin_register_fn = void (*)(const plugin_context*);
+
+#ifdef _WIN32
+inline constexpr char plugin_path_separator = ';';
+#else
+inline constexpr char plugin_path_separator = ':';
+#endif
+
+// Load one plugin and let it register its overrides.
+/// \throws std::runtime_error if the library cannot be loaded, lacks the entry
+///         points, or was built against a different override interface.
+inline void load_plugin(const plugin_context& ctx, const std::string& path) {
+#ifdef _WIN32
+    HMODULE handle = LoadLibraryA(path.c_str());
+    if (handle == nullptr) throw std::runtime_error("failed to load plugin: " + path);
+    auto abi = reinterpret_cast<plugin_abi_version_fn>(GetProcAddress(handle, "flm_plugin_abi_version"));
+    auto reg = reinterpret_cast<plugin_register_fn>(GetProcAddress(handle, "flm_plugin_register"));
+#else
+    // RTLD_GLOBAL so a plugin's typeinfo and the host's agree, which dynamic_cast
+    // across the boundary depends on. The handle is deliberately never closed:
+    // the overrides it registered outlive registration.
+    void* handle = dlopen(path.c_str(), RTLD_NOW | RTLD_GLOBAL);
+    if (handle == nullptr) throw std::runtime_error("failed to load plugin " + path + ": " + dlerror());
+    auto abi = reinterpret_cast<plugin_abi_version_fn>(dlsym(handle, "flm_plugin_abi_version"));
+    auto reg = reinterpret_cast<plugin_register_fn>(dlsym(handle, "flm_plugin_register"));
+#endif
+    if (abi == nullptr || reg == nullptr) {
+        throw std::runtime_error("plugin " + path + " does not export flm_plugin_abi_version/flm_plugin_register");
+    }
+    const int found = abi();
+    if (found != plugin_abi_version) {
+        throw std::runtime_error("plugin " + path + " targets override ABI " + std::to_string(found)
+                                 + ", this build provides " + std::to_string(plugin_abi_version));
+    }
+    reg(&ctx);
+}
+
+// Load every plugin named in FLM_PLUGIN, separated by ':' (';' on Windows).
+inline void load_plugins_from_env(const plugin_context& ctx) {
+    const char* env = std::getenv("FLM_PLUGIN");
+    if (env == nullptr || *env == '\0') return;
+    const std::string spec(env);
+    size_t begin = 0;
+    while (begin <= spec.size()) {
+        const size_t end = std::min(spec.find(plugin_path_separator, begin), spec.size());
+        const std::string path = spec.substr(begin, end - begin);
+        if (!path.empty()) {
+            load_plugin(ctx, path);
+            header_print_g("info", "Loaded operator plugin: " << path);
+        }
+        begin = end + 1;
+    }
+}
+
+}  // namespace flm
+
+// Defines a plugin's entry points around a registration function.
+#define FLM_PLUGIN(register_function)                                            \
+    extern "C" int flm_plugin_abi_version() { return flm::plugin_abi_version; }  \
+    extern "C" void flm_plugin_register(const flm::plugin_context* ctx) {        \
+        (register_function)(*ctx);                                               \
+    }

src/include/models/gemma4e/gemma4e_npu.hpp

@@ -19,6 +19,192 @@
 #include <immintrin.h>  // For AVX intrinsics
 #endif
 
+// The operations this engine declares, and how their keys are spelled. Both
+// belong to the model, not to the plugin API: another engine may have no
+// layers, or a nest of them, and names its own operations.
+namespace gemma4e_ops {
+
+// One shape per call-site convention. A hook always returns
+// flm::hook_result<R>: a real result, skip() (the operation already
+// happened, nothing further to do), or defer() (run the engine's own
+// implementation). R is ert_cmd_state for a blocking call site, xrt::run
+// for one that builds a run to start/wait later.
+using proj_sig_t = flm::hook_result<ert_cmd_state>(bytes& out, bytes& in, bytes& weights, int64_t layer, int64_t padded);
+using proj_async_sig_t = flm::hook_result<xrt::run>(bytes& out, bytes& in, bytes& weights, int64_t layer, int64_t padded);
+using attn_core_sig_t = flm::hook_result<ert_cmd_state>(bytes& out, bytes& q, bytes& kv_cache, int64_t layer, int64_t padded, int64_t l_begin_chunked);
+using dequant_sig_t = flm::hook_result<ert_cmd_state>(bytes& dequantized, bytes& quantized, int64_t layer, int64_t padded);
+using decode_layer_sig_t = flm::hook_result<ert_cmd_state>(bytes& hidden_state_inout, bytes& proj_weights,
+    bytes& rms_weights, bytes& rope_rms_weights, bytes& kv_cache, int64_t layer, int64_t context_len);
+using decode_layer_async_sig_t = flm::hook_result<xrt::run>(bytes& hidden_state_inout, bytes& proj_weights,
+    bytes& rms_weights, bytes& rope_rms_weights, bytes& kv_cache, int64_t layer, int64_t context_len);
+using lm_head_sig_t = flm::hook_result<xrt::run>(bytes& logits, bytes& lm_head_weights, bytes& hidden_state);
+using audio_conv1d_sig_t = flm::hook_result<ert_cmd_state>(bytes& out, bytes& in, bytes& weights, int64_t layer);
+
+// The vision and audio encoders' own generated sequences take (in, weights,
+// out) -- the reverse of the decode path's (out, in, weights) -- so these get
+// their own shapes rather than reusing proj_sig_t/proj_async_sig_t.
+using layer_proj_sig_t = flm::hook_result<ert_cmd_state>(bytes& in, bytes& weights, bytes& out, int64_t layer, int64_t padded);
+using layer_proj_async_sig_t = flm::hook_result<xrt::run>(bytes& in, bytes& weights, bytes& out, int64_t layer, int64_t padded);
+// A one-shot embedding/projection step: no layer loop, so no layer or padded argument.
+using embed_sig_t = flm::hook_result<ert_cmd_state>(bytes& in, bytes& weights, bytes& out);
+// Vision attention takes q/k/v as three separate buffers, not one kv_cache.
+using vision_attn_core_sig_t = flm::hook_result<xrt::run>(bytes& out, bytes& q, bytes& k, bytes& v, int64_t layer, int64_t padded);
+// The background weight-preload run's sequence takes no buffer arguments at all.
+using preload_sig_t = flm::hook_result<xrt::run>();
+
+// Every operation, one name and one typedef each: sliding-window and global
+// attention run different kernels, and a skip layer's mlp is a different
+// (double-wide) sequence, so each gets its own rather than sharing one
+// distinguished by argument. decode.layer additionally has a plain and an
+// ".async" name, since the engine dispatches it both ways depending on
+// whether NPU preemption is enabled.
+namespace op {
+using q_swa_proj_func_t = proj_sig_t;
+using q_global_proj_func_t = proj_sig_t;
+using k_swa_proj_func_t = proj_async_sig_t;
+using k_global_proj_func_t = proj_async_sig_t;
+using v_swa_proj_func_t = proj_async_sig_t;
+using v_global_proj_func_t = proj_async_sig_t;
+using o_swa_proj_func_t = proj_sig_t;
+using o_global_proj_func_t = proj_sig_t;
+using swa_attn_core_func_t = attn_core_sig_t;
+using global_attn_core_func_t = attn_core_sig_t;
+inline constexpr std::string_view q_swa_proj = "self_attn.q_proj.swa";
+inline constexpr std::string_view q_global_proj = "self_attn.q_proj.global";
+inline constexpr std::string_view k_swa_proj = "self_attn.k_proj.swa";
+inline constexpr std::string_view k_global_proj = "self_attn.k_proj.global";
+inline constexpr std::string_view v_swa_proj = "self_attn.v_proj.swa";
+inline constexpr std::string_view v_global_proj = "self_attn.v_proj.global";
+inline constexpr std::string_view o_swa_proj = "self_attn.o_proj.swa";
+inline constexpr std::string_view o_global_proj = "self_attn.o_proj.global";
+inline constexpr std::string_view swa_attn_core = "self_attn.core.swa";
+inline constexpr std::string_view global_attn_core = "self_attn.core.global";
+
+using gate_proj_func_t = proj_sig_t;
+using gate_skip_proj_func_t = proj_sig_t;
+using up_proj_func_t = proj_sig_t;
+using up_skip_proj_func_t = proj_sig_t;
+using down_proj_func_t = proj_sig_t;
+using down_skip_proj_func_t = proj_sig_t;
+inline constexpr std::string_view gate_proj = "mlp.gate_proj";
+inline constexpr std::string_view gate_skip_proj = "mlp.gate_proj.skip";
+inline constexpr std::string_view up_proj = "mlp.up_proj";
+inline constexpr std::string_view up_skip_proj = "mlp.up_proj.skip";
+inline constexpr std::string_view down_proj = "mlp.down_proj";
+inline constexpr std::string_view down_skip_proj = "mlp.down_proj.skip";
+
+// One dequant call per gemma4e_layer_type_t, matching the engine's own
+// apps[4][matrix] table. Index with int(type) (e_gemma4e_swa_layer=0,
+// e_gemma4e_global_layer=1, e_gemma4e_swa_layer_skip=2, e_gemma4e_global_layer_skip=3).
+using dequant_func_t = dequant_sig_t;
+inline constexpr std::string_view dequant_qkv[4]  = { "dequant.qkv.swa",  "dequant.qkv.global",
+                                                      "dequant.qkv.swa_skip",  "dequant.qkv.global_skip" };
+inline constexpr std::string_view dequant_o[4]    = { "dequant.o.swa",    "dequant.o.global",
+                                                      "dequant.o.swa_skip",    "dequant.o.global_skip" };
+inline constexpr std::string_view dequant_gate[4] = { "dequant.gate.swa", "dequant.gate.global",
+                                                      "dequant.gate.swa_skip", "dequant.gate.global_skip" };
+inline constexpr std::string_view dequant_up[4]   = { "dequant.up.swa",   "dequant.up.global",
+                                                      "dequant.up.swa_skip",   "dequant.up.global_skip" };
+inline constexpr std::string_view dequant_down[4] = { "dequant.down.swa", "dequant.down.global",
+                                                      "dequant.down.swa_skip", "dequant.down.global_skip" };
+
+using swa_layer_func_t = decode_layer_sig_t;
+using global_layer_func_t = decode_layer_sig_t;
+using swa_skip_layer_func_t = decode_layer_sig_t;
+using global_skip_layer_func_t = decode_layer_sig_t;
+using swa_layer_async_func_t = decode_layer_async_sig_t;
+using global_layer_async_func_t = decode_layer_async_sig_t;
+using swa_skip_layer_async_func_t = decode_layer_async_sig_t;
+using global_skip_layer_async_func_t = decode_layer_async_sig_t;
+inline constexpr std::string_view swa_layer = "decode.layer.swa";
+inline constexpr std::string_view global_layer = "decode.layer.global";
+inline constexpr std::string_view swa_skip_layer = "decode.layer.swa_skip";
+inline constexpr std::string_view global_skip_layer = "decode.layer.global_skip";
+inline constexpr std::string_view swa_layer_async = "decode.layer.swa.async";
+inline constexpr std::string_view global_layer_async = "decode.layer.global.async";
+inline constexpr std::string_view swa_skip_layer_async = "decode.layer.swa_skip.async";
+inline constexpr std::string_view global_skip_layer_async = "decode.layer.global_skip.async";
+
+using lm_head_func_t = lm_head_sig_t;
+inline constexpr std::string_view lm_head = "lm_head";
+
+using audio_conv1d_func_t = audio_conv1d_sig_t;
+inline constexpr std::string_view audio_conv1d = "audio.conv1d";
+
+// The vision encoder: one patch-embedding step, then one attention block and
+// one mlp per hidden layer, then one projection into the language model's
+// embedding space.
+using vision_patch_embed_func_t = embed_sig_t;
+using vision_pos_embed_dim0_func_t = embed_sig_t;
+using vision_pos_embed_dim1_func_t = embed_sig_t;
+using vision_q_proj_func_t = layer_proj_async_sig_t;
+using vision_k_proj_func_t = layer_proj_async_sig_t;
+using vision_v_proj_func_t = layer_proj_async_sig_t;
+using vision_attn_core_func_t = vision_attn_core_sig_t;
+using vision_o_proj_func_t = layer_proj_sig_t;
+using vision_gate_proj_func_t = layer_proj_async_sig_t;
+using vision_up_proj_func_t = layer_proj_async_sig_t;
+using vision_down_proj_func_t = layer_proj_sig_t;
+using vision_to_language_proj_func_t = embed_sig_t;
+inline constexpr std::string_view vision_patch_embed = "vision.patch_embed";
+inline constexpr std::string_view vision_pos_embed_dim0 = "vision.pos_embed.dim0";
+inline constexpr std::string_view vision_pos_embed_dim1 = "vision.pos_embed.dim1";
+inline constexpr std::string_view vision_q_proj = "vision.q_proj";
+inline constexpr std::string_view vision_k_proj = "vision.k_proj";
+inline constexpr std::string_view vision_v_proj = "vision.v_proj";
+inline constexpr std::string_view vision_attn_core = "vision.attn_core";
+inline constexpr std::string_view vision_o_proj = "vision.o_proj";
+inline constexpr std::string_view vision_gate_proj = "vision.mlp.gate_proj";
+inline constexpr std::string_view vision_up_proj = "vision.mlp.up_proj";
+inline constexpr std::string_view vision_down_proj = "vision.mlp.down_proj";
+inline constexpr std::string_view vision_to_language_proj = "vision.to_language_proj";
+
+// The audio encoder: one sub-sample projection step, then one attention
+// block, one conv1d block (its own start/end pointwise projections; the
+// depthwise conv1d itself is audio_conv1d above) and one FFN pair per hidden
+// layer, then a pre-encode projection and one into the language model's
+// embedding space.
+using audio_sub_sample_proj_func_t = embed_sig_t;
+using audio_q_proj_func_t = layer_proj_async_sig_t;
+using audio_k_proj_func_t = layer_proj_async_sig_t;
+using audio_v_proj_func_t = layer_proj_async_sig_t;
+using audio_o_proj_func_t = layer_proj_sig_t;
+using audio_ffn_up_proj_func_t = layer_proj_async_sig_t;
+using audio_ffn_down_proj_func_t = layer_proj_sig_t;
+using audio_conv1d_start_proj_func_t = layer_proj_sig_t;
+using audio_conv1d_end_proj_func_t = layer_proj_sig_t;
+using audio_pre_encode_proj_func_t = embed_sig_t;
+using audio_to_language_proj_func_t = embed_sig_t;
+inline constexpr std::string_view audio_sub_sample_proj = "audio.sub_sample_proj";
+inline constexpr std::string_view audio_q_proj = "audio.q_proj";
+inline constexpr std::string_view audio_k_proj = "audio.k_proj";
+inline constexpr std::string_view audio_v_proj = "audio.v_proj";
+inline constexpr std::string_view audio_o_proj = "audio.o_proj";
+inline constexpr std::string_view audio_ffn_up_proj = "audio.ffn.up_proj";
+inline constexpr std::string_view audio_ffn_down_proj = "audio.ffn.down_proj";
+inline constexpr std::string_view audio_conv1d_start_proj = "audio.conv1d.start_proj";
+inline constexpr std::string_view audio_conv1d_end_proj = "audio.conv1d.end_proj";
+inline constexpr std::string_view audio_pre_encode_proj = "audio.pre_encode_proj";
+inline constexpr std::string_view audio_to_language_proj = "audio.to_language_proj";
+
+// The per-layer-input (PLI) path: a model-wide down projection run once per
+// batch, and a gate/up pair run at the end of every layer.
+using pli_down_proj_func_t = embed_sig_t;
+using pli_gate_proj_func_t = layer_proj_sig_t;
+using pli_up_proj_func_t = layer_proj_sig_t;
+inline constexpr std::string_view pli_down_proj = "pli.down_proj";
+inline constexpr std::string_view pli_gate_proj = "pli.gate_proj";
+inline constexpr std::string_view pli_up_proj = "pli.up_proj";
+
+// The background run that preloads the next layer's weights while the
+// current token executes. Its sequence takes no buffer arguments.
+using layer_pre_load_func_t = preload_sig_t;
+inline constexpr std::string_view layer_pre_load = "engine.layer_pre_load";
+}  // namespace op
+
+}  // namespace gemma4e_ops
+
+
 // some helper functions for convenience
 constexpr int GEMMA4E_IS_GLOBAL_MASK = 0x00000001;
 constexpr int GEMMA4E_IS_SKIP_MASK = 0x00000002;