Add a plugin mechanism for operator overrides

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,29 @@
 /// \note This is a source file for the auto_model class
 
 #include "AutoModel/automodel.hpp"
+#include "flm_plugin.hpp"
+
+
+void AutoModel::_load_operator_plugins() {
+    flm::op_registry* ops = this->lm_engine ? this->lm_engine->ops() : nullptr;
+    if (ops == nullptr) return;
+    const std::string xclbin_path = utils::path_join(
+        this->lm_config->exec_path, "xclbins", this->lm_config->model_name);
+    flm::plugin_context ctx{
+        ops,
+        this->npu.get(),
+        this->model_path.c_str(),
+        xclbin_path.c_str(),
+    };
+    flm::load_plugins_from_env(ctx);
+}
+
+
+void AutoModel::_load_engine_weights() {
+    this->_load_operator_plugins();
+    this->lm_engine->load_weights(*this->q4nx);
+    this->q4nx.reset();
+}
 
 
 AutoModel::AutoModel(flm_rt::device* npu_device_inst, std::string current_model) {

src/common/AutoModel/modeling_gemma3.cpp

@@ -16,11 +16,7 @@ void Gemma3::load_model(std::string model_path, json model_info, int default_con
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == gemma
     this->lm_engine = std::make_unique<gemma_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_gemma3_text.cpp

@@ -18,11 +18,7 @@ void Gemma3_Text_Only::load_model(std::string model_path, json model_info, int d
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == gemma_text_only
     this->lm_engine = std::make_unique<gemma_text_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_gemma4_12b.cpp

@@ -409,10 +409,7 @@ void Gemma4_12B::load_model(std::string model_path, json model_info, int default
 
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     this->lm_engine = std::make_unique<gemma4_12b_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-    // free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     // The soft token budget is a preprocessing constant (processor_config.json
     // image_seq_length / image_processor.max_soft_tokens), read by the engine
     // along with the rest of the image front end parameters.

src/common/AutoModel/modeling_gemma4e.cpp

@@ -654,10 +654,7 @@ void Gemma4e::load_model(std::string model_path, json model_info, int default_co
 
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     this->create_engine();
-
-    this->lm_engine->load_weights(*this->q4nx);
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_gpt_oss.cpp

@@ -14,8 +14,7 @@ void GPT_OSS::load_model(std::string model_path, json model_info, int default_co
     this->_shared_load_model(model_path, model_info, default_context_length, enable_preemption);
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     this->lm_engine = std::make_unique<gpt_oss_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-    this->lm_engine->load_weights(*this->q4nx);
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->tokenizer = std::make_unique<Tokenizer>(model_path);
 
     this->setup_tokenizer(model_path);

src/common/AutoModel/modeling_hunyuan.cpp

@@ -24,11 +24,7 @@ void Hunyuan::load_model(std::string model_path, json model_info, int default_co
 
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     this->lm_engine = std::make_unique<hunyuan_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    // free the mmap'd weights immediately
-    this->q4nx.reset();
+    this->_load_engine_weights();
 
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);

src/common/AutoModel/modeling_lfm2.cpp

@@ -18,11 +18,7 @@ void LFM2::load_model(std::string model_path, json model_info, int default_conte
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == llama
     this->lm_engine = std::make_unique<lfm2_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();
@@ -273,11 +269,7 @@ void LFM2_5_TK::load_model(std::string model_path, json model_info, int default_
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == llama
     this->lm_engine = std::make_unique<lfm2_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_llama3.cpp

@@ -16,11 +16,7 @@ void Llama3::load_model(std::string model_path, json model_info, int default_con
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == llama
     this->lm_engine = std::make_unique<llama_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
 
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
@@ -98,11 +94,7 @@ void DeepSeek_r1_8b::load_model(std::string model_path, json model_info, int def
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == llama
     this->lm_engine = std::make_unique<llama_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_nanbeige.cpp

@@ -16,11 +16,7 @@ void Nanbeige::load_model(std::string model_path, json model_info, int default_c
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == nanbeige
     this->lm_engine = std::make_unique<nanbeige_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
 
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);

src/common/AutoModel/modeling_phi4.cpp

@@ -16,10 +16,7 @@ void Phi4::load_model(std::string model_path, json model_info, int default_conte
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == phi4
     this->lm_engine = std::make_unique<phi4_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
 
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);

src/common/AutoModel/modeling_qwen2.cpp

@@ -18,11 +18,7 @@ void Qwen2::load_model(std::string model_path, json model_info, int default_cont
 
     // lm_config->get<std::string>("model_type", "") == qwen2
     this->lm_engine = std::make_unique<qwen2_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_qwen2vl.cpp

@@ -20,10 +20,7 @@ void Qwen2VL::load_model(std::string model_path, json model_info, int default_co
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // lm_config->get<std::string>("model_type", "") == qwen2
     this->lm_engine = std::make_unique<qwen2vl_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_qwen3.cpp

@@ -18,11 +18,7 @@ void Qwen3::load_model(std::string model_path, json model_info, int default_cont
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // lm_config->get<std::string>("model_type", "") == qwen3
     this->lm_engine = std::make_unique<qwen3_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();
@@ -316,11 +312,7 @@ void Qwen3_IT::load_model(std::string model_path, json model_info, int default_c
 
     // lm_config->get<std::string>("model_type", "") == qwen3
     this->lm_engine = std::make_unique<qwen3_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();
@@ -466,11 +458,7 @@ void Qwen3_TK::load_model(std::string model_path, json model_info, int default_c
 
     // lm_config->get<std::string>("model_type", "") == qwen3
     this->lm_engine = std::make_unique<qwen3_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();
@@ -687,11 +675,7 @@ void DeepSeek_r1_0528_8b::load_model(std::string model_path, json model_info, in
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // model_type == llama
     this->lm_engine = std::make_unique<qwen3_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_qwen3_5vl.cpp

@@ -19,10 +19,7 @@ void Qwen3_5VL::load_model(std::string model_path, json model_info, int default_
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // lm_config->get<std::string>("model_type", "") == qwen3
     this->lm_engine = std::make_unique<qwen3_5vl_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_qwen3_6_moe.cpp

@@ -19,10 +19,7 @@ void Qwen3_6_MOE::load_model(std::string model_path, json model_info, int defaul
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // lm_config->get<std::string>("model_type", "") == qwen3
     this->lm_engine = std::make_unique<qwen3_6_moe_npu>(*this->lm_config, this->npu.get(), this->MAX_L);
-
-    this->lm_engine->load_weights(*this->q4nx);
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/common/AutoModel/modeling_qwen3vl.cpp

@@ -27,10 +27,7 @@ void Qwen3VL::load_model(std::string model_path, json model_info, int default_co
     this->q4nx = std::make_unique<Q4NX>(this->model_path);
     // lm_config->get<std::string>("model_type", "") == qwen3
     this->create_engine();
-
-    this->lm_engine->load_weights(*this->q4nx);
-    //free the q4nx
-    this->q4nx.reset();
+    this->_load_engine_weights();
     this->lm_engine->clear_context();
     this->setup_tokenizer(model_path);
     this->sampler.reset();

src/include/AutoModel/automodel.hpp

@@ -215,6 +215,17 @@ class AutoModel {
 
 
 	void _shared_load_model(std::string model_path, json model_info, int default_context_length = -1, bool enable_preemption = false);
+
+	/// \brief Give the plugins named in FLM_PLUGIN a chance to override operations
+	/// \note Call once the engine exists and before its weights are loaded: the
+	///       engine declares its operations in its constructor, and an override
+	///       must be in place before anything dispatches through it.
+	void _load_operator_plugins();
+
+	/// \brief Bring the engine up: load plugins, read the weights, release the file
+	/// \note Every model calls this immediately after constructing its engine.
+	///       Keeping it in one place is what makes a plugin apply to all of them.
+	void _load_engine_weights();
 	nlohmann::json _shared_setup_tokenizer(std::string model_path);
 
 	/// \brief Insert tokens into the model

src/include/causal_lm.hpp

@@ -11,12 +11,21 @@
 #include "utils/utils.hpp"
 #include "buffer.hpp"
 
+class npu_app;
+namespace flm {
+    template <typename App> class op_registry_t;
+}
+
 /// \brief causal_lm class
 class causal_lm {
 public:
     causal_lm(){}
     virtual ~causal_lm(){}
 
+    /// \brief The operations of this engine that a plugin may override
+    /// \return nullptr if the engine declares none
+    virtual flm::op_registry_t<npu_app>* ops() { return nullptr; }
+
     /// \brief forward the causal_lm
     /// \param ids the ids
     /// \return the output