diff --git a/.github/actions/lin-build/action.yml b/.github/actions/lin-build/action.yml index 62f4514c4..fa3bdb581 100644 --- a/.github/actions/lin-build/action.yml +++ b/.github/actions/lin-build/action.yml @@ -115,6 +115,17 @@ runs: echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" fi + - name: Build and run backend tests + if: inputs.variant == 'XRT-deb' + shell: bash + working-directory: ${{ github.workspace }} + run: | + for suite in iron_backend model_backend model_list_platform; do + cmake -S "src/test/$suite" -B "build-$suite" -DCMAKE_BUILD_TYPE=Release + cmake --build "build-$suite" --parallel 2 + ctest --test-dir "build-$suite" --output-on-failure + done + # Everything cached lives under .ci-cache in the workspace: a container # job's actions/cache resolves relative paths there, and an absolute # ${{ github.workspace }} would name the host's path, not the container's. diff --git a/docs/docs/models/llama.md b/docs/docs/models/llama.md index a85519d56..a3dc624ff 100644 --- a/docs/docs/models/llama.md +++ b/docs/docs/models/llama.md @@ -23,6 +23,20 @@ parent: Models flm run llama3.2:1b ``` +With an IRON-enabled build, `llama3.2-iron:1b` uses the same Instruct model in +BF16, with the tokenizer and chat template from FastFlowLM's package above. +Accept Meta's model license and provide a Hugging Face token (`HF_TOKEN` or +`hf auth login`), then run: + +```shell +flm pull llama3.2-iron:1b +flm run llama3.2-iron:1b --ctx-len 2048 +``` + +This downloads into `models/llama3.2-1b-instruct`. If you previously pulled the +base checkpoint, pull again after updating FastFlowLM; the new directory keeps +the two checkpoints separate. + --- ## 🧩 Model Card: [Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 789221bbe..b164bfc4c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -53,6 +53,11 @@ option(FLM_PORTABLE_BUILD "Build portable distribution with bundled runtime libr # It also fixes the NPU generation the build targets -- see utils/npu_platform.hpp. option(FLM_ENABLE_CORE "Reach kernels through ryzenai-corelib instead of FastFlowLM's own flow" OFF) +option(FLM_ENABLE_IRON "Run complete IRON graphs through the IRON server" OFF) + +if(FLM_ENABLE_IRON) + add_compile_definitions(FLM_ENABLE_IRON=1) +endif() if(FLM_ENABLE_CORE) # The corelib prefix lives outside this repository, so its headers have to @@ -221,10 +226,16 @@ include("${CMAKE_SOURCE_DIR}/common/chat/chat_sources.cmake") # ——————————————————————————————————————————————— file(GLOB SOURCES "src/*.cpp" "runner/*.cpp" "common/*.cpp" "common/*/*.cpp" "server/*.cpp" "pull/*.cpp" ) file(GLOB HEADERS "include/*.hpp" "runner/*.hpp" "common/*.hpp" "common/*/*.hpp" "server/*.hpp" "pull/*.hpp") +include(common/models/models_sources.cmake) +if(FLM_ENABLE_IRON) + list(APPEND SOURCES ${FLM_MODELS_IRON_SOURCES}) +else() + list(FILTER SOURCES EXCLUDE REGEX ".*/common/iron/.*\\.cpp$") +endif() list(FILTER SOURCES EXCLUDE REGEX ".*/common/core/.*\\.cpp$") -# Model sources live two levels deeper than the globs above reach. Only the core -# half exists as loose sources; it is built separately, below. +# Model sources live two levels deeper than the globs above reach. IRON links +# into flm. Core builds as a separate library below. # Exclude files that depend on missing libraries for Linux if(NOT WIN32) @@ -705,7 +716,7 @@ if(WIN32) target_link_libraries(flm PRIVATE ntdll wsock32 ws2_32 Bcrypt iphlpapi userenv psapi - crypt32 secur32 advapi32 normaliz wldap32 + crypt32 secur32 advapi32 normaliz wldap32 shell32 ) else() target_link_libraries(flm PRIVATE Threads::Threads) diff --git a/src/common/AutoModel/builtin_backends.cpp b/src/common/AutoModel/builtin_backends.cpp index 3cf755365..e9799987b 100644 --- a/src/common/AutoModel/builtin_backends.cpp +++ b/src/common/AutoModel/builtin_backends.cpp @@ -8,6 +8,10 @@ #include "AutoModel/flm_backend.hpp" #include "AutoModel/model_backend.hpp" +#if defined(FLM_ENABLE_IRON) +#include "models/llama3/iron/llama3_iron_backend.hpp" +#endif + #if defined(FLM_ENABLE_CORE) #include "models/gemma4/core/dev04/gemma4_core_backend.hpp" #include "models/llama/core/dev04/llama_core_backend.hpp" @@ -63,6 +67,12 @@ void register_builtin_backends(BackendRegistry& registry) { RegisterFlm(registry, "nanbeige"); RegisterFlm(registry, "phi4"); +#if defined(FLM_ENABLE_IRON) + registry.register_backend("llama3", flm::backend::kIronBackendId, + flm::llama3::iron_factory(), + flm::llama3::iron_traits()); +#endif + #if defined(FLM_ENABLE_CORE) registry.register_backend("gemma4e", flm::backend::kCoreBackendId, flm::gemma4::core_factory(), diff --git a/src/common/iron/causal_lm.cpp b/src/common/iron/causal_lm.cpp new file mode 100644 index 000000000..b8b802f07 --- /dev/null +++ b/src/common/iron/causal_lm.cpp @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "iron/causal_lm.hpp" + +#include +#include +#include +#include + +namespace flm::iron { +namespace { + +std::uint32_t ValidateContextLength(std::uint32_t length) { + if (length == 0) { + throw std::out_of_range("IRON context length must be positive"); + } + return length; +} + +} // namespace + +CausalLM::CausalLM(const std::string& model_path, std::string model, + std::uint32_t context_length) + : worker_(model_path, std::move(model), context_length), + max_length_(context_length) {} + +buffer CausalLM::forward(int id) { + tokens_.push_back(id); + return run(); +} + +buffer CausalLM::prefill(std::vector& ids, void*) { + tokens_.insert(tokens_.end(), ids.begin(), ids.end()); + return run(); +} + +void CausalLM::set_context_length(int length) { + if (length < 0 || static_cast(length) > tokens_.size()) { + throw std::out_of_range("Invalid IRON context length"); + } + tokens_.resize(static_cast(length)); +} + +void CausalLM::load_weights(Q4NX&) { + throw std::logic_error("The IRON server loads safetensors"); +} + +void CausalLM::update_max_length(std::uint32_t length) { + max_length_ = length; +} + +void CausalLM::clear_context() { + tokens_.clear(); + worker_.reset(); +} + +buffer CausalLM::get_k_cache(int, int) { return {}; } + +buffer CausalLM::get_v_cache(int, int) { return {}; } + +int CausalLM::get_current_context_length() { + return static_cast(tokens_.size()); +} + +int CausalLM::checkpoint() { return static_cast(tokens_.size()); } + +int CausalLM::restore() { return 0; } + +bool CausalLM::poisoned() const noexcept { return worker_.poisoned(); } + +buffer CausalLM::run() { + if (tokens_.empty() || tokens_.size() > max_length_) { + throw std::out_of_range("IRON token history exceeds the context length"); + } + auto words = worker_.logits(tokens_); + buffer output(words.size()); + std::memcpy(output.data(), words.data(), words.size() * sizeof(words[0])); + return output; +} + +CausalBackend::CausalBackend(const flm::backend::BackendContext& context, + std::string model) + : engine_(context.model_path, std::move(model), + ValidateContextLength(context.context_length)) {} + +::causal_lm& CausalBackend::engine() { return engine_; } + +std::string CausalBackend::id() const { + return flm::backend::kIronBackendId; +} + +std::string CausalBackend::detail() const { return "IRON server"; } + +bool CausalBackend::supports_preemption() const { return false; } + +bool CausalBackend::poisoned() const noexcept { return engine_.poisoned(); } + +flm::backend::BackendFactory causal_factory(std::string model) { + return [model = std::move(model)]( + const flm::backend::BackendContext& context) { + return std::make_unique(context, model); + }; +} + +} // namespace flm::iron \ No newline at end of file diff --git a/src/common/iron/embedding.cpp b/src/common/iron/embedding.cpp new file mode 100644 index 000000000..7c24df1f8 --- /dev/null +++ b/src/common/iron/embedding.cpp @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "iron/embedding.hpp" + +#include + +namespace flm::iron { + +void Embedding::load_model(std::string model_path, json, + bool enable_preemption) { + if (enable_preemption) { + throw std::invalid_argument( + "The IRON embedding backend does not support preemption"); + } + worker_ = std::make_unique(model_path, worker_model_); + is_model_loaded = true; + this->model_path = std::move(model_path); +} + +std::vector Embedding::embed(std::string& text, + embedding_task_type_t task_type) { + if (!worker_) { + throw std::runtime_error("The IRON embedding model is not loaded"); + } + return worker_->embed(text, task_type == task_document ? "document" : "query", + dimensions_); +} + +} // namespace flm::iron \ No newline at end of file diff --git a/src/common/iron/worker.cpp b/src/common/iron/worker.cpp new file mode 100644 index 000000000..bf77b5aa4 --- /dev/null +++ b/src/common/iron/worker.cpp @@ -0,0 +1,207 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "iron/worker.hpp" + +#define BOOST_PROCESS_USE_STD_FS +#if __has_include() +#define BOOST_PROCESS_VERSION 1 +#include +#include +#include +#include +#else +#include +#endif + +#if !defined(_WIN32) +#include +#endif +#include +#include + +namespace flm::iron { +namespace { + +constexpr std::uint32_t kMaxHeader = 1 << 20; +constexpr std::uint32_t kProtocolVersion = 1; +constexpr std::uint32_t kCapabilities = + static_cast(Capability::CausalLM) | + static_cast(Capability::Embeddings); +// Missing handshake fields mean the original protocol; never extend this mask. +constexpr std::uint32_t kLegacyCapabilities = 0b11; + +std::string PythonExecutable() { + const char* value = std::getenv("FLM_IRON_PYTHON"); + if (value && *value) return value; +#if defined(_WIN32) + return "python"; +#else + return "python3"; +#endif +} + +std::string WorkerModule() { + const char* value = std::getenv("FLM_IRON_WORKER_MODULE"); + return value && *value ? value : "iron.lm_server"; +} + +std::filesystem::path ResolveExecutable(const std::string& executable) { + const std::filesystem::path path(executable); + return path.has_parent_path() ? path : boost::process::search_path(path); +} + +void IgnoreBrokenPipeSignal() { +#if !defined(_WIN32) + static const bool ignored = std::signal(SIGPIPE, SIG_IGN) != SIG_ERR; + if (!ignored) throw std::runtime_error("Cannot ignore SIGPIPE"); +#endif +} + +} // namespace + +struct Worker::Impl { + boost::process::opstream input; + boost::process::ipstream output; + boost::process::child process; + + Impl(const std::string& python, const std::string& module) + : process(ResolveExecutable(python), + boost::process::args({"-m", module}), + boost::process::std_in < input, + boost::process::std_out > output) {} + + void close_pipes() noexcept { + try { input.pipe().close(); } catch (...) {} + try { output.pipe().close(); } catch (...) {} + } + + void stop() noexcept { + close_pipes(); + std::error_code error; + if (process.running(error)) process.terminate(error); + process.wait(error); + } +}; + +Worker::Worker(const std::filesystem::path& model_path, std::string model, + std::uint32_t context_length) { + IgnoreBrokenPipeSignal(); + impl_ = std::make_unique(PythonExecutable(), WorkerModule()); + try { + nlohmann::json command = {{"command", "load_model"}, + {"protocol_version", kProtocolVersion}, + {"capabilities", kCapabilities}, + {"model", std::move(model)}, + {"model_path", model_path.string()}}; + if (context_length) command["context_length"] = context_length; + const auto response = request(command); + if (response.value("protocol_version", 1u) != kProtocolVersion) { + throw std::runtime_error("Unsupported IRON server protocol version"); + } + capabilities_ = response.value("capabilities", kLegacyCapabilities) & kCapabilities; + } catch (...) { + impl_->stop(); + throw; + } +} + +Worker::~Worker() { + if (impl_) { + try { + request({{"command", "close"}}); + } catch (...) { + impl_->stop(); + return; + } + impl_->close_pipes(); + std::error_code error; + impl_->process.wait(error); + } +} + +void Worker::write_all(const void* data, std::size_t size) { + impl_->input.write(static_cast(data), size); + if (!impl_->input) throw std::runtime_error("IRON server write failed"); +} + +void Worker::read_all(void* data, std::size_t size) { + impl_->output.read(static_cast(data), size); + if (!impl_->output) throw std::runtime_error("IRON server exited"); +} + +nlohmann::json Worker::request(const nlohmann::json& command, + const void* payload, + std::size_t payload_size) { + if (poisoned_) throw std::runtime_error("IRON server is unusable; reload the model"); + try { + const std::string header = command.dump(); + const std::uint32_t header_size = static_cast(header.size()); + write_all(&header_size, sizeof(header_size)); + write_all(header.data(), header.size()); + if (payload_size) write_all(payload, payload_size); + impl_->input.flush(); + if (!impl_->input) throw std::runtime_error("IRON server write failed"); + + std::uint32_t response_size = 0; + read_all(&response_size, sizeof(response_size)); + if (response_size == 0 || response_size > kMaxHeader) { + throw std::runtime_error("IRON server returned an invalid header"); + } + std::string response(response_size, '\0'); + read_all(response.data(), response.size()); + auto result = nlohmann::json::parse(response); + if (!result.value("ok", false)) { + throw std::runtime_error(result.value("error", "IRON server failed")); + } + return result; + } catch (...) { + poisoned_ = true; + throw; + } +} + +std::vector Worker::logits(const std::vector& tokens) { + if (!supports(Capability::CausalLM)) + throw std::runtime_error("IRON server does not support causal inference"); + auto response = request({{"command", "logits"}, {"count", tokens.size()}}, + tokens.data(), tokens.size() * sizeof(int)); + try { + const std::size_t count = response.at("count").get(); + std::vector values(count); + read_all(values.data(), values.size() * sizeof(values[0])); + return values; + } catch (...) { + poisoned_ = true; + throw; + } +} + +void Worker::reset() { + if (!supports(Capability::CausalLM)) + throw std::runtime_error("IRON server does not support causal inference"); + request({{"command", "reset"}}); +} + +std::vector Worker::embed(const std::string& text, + const std::string& task, + std::uint32_t dimensions) { + if (!supports(Capability::Embeddings)) + throw std::runtime_error("IRON server does not support embeddings"); + auto response = request({{"command", "embed"}, + {"task", task}, + {"dimensions", dimensions}, + {"size", text.size()}}, + text.data(), text.size()); + try { + const std::size_t count = response.at("count").get(); + std::vector values(count); + read_all(values.data(), values.size() * sizeof(values[0])); + return values; + } catch (...) { + poisoned_ = true; + throw; + } +} + +} // namespace flm::iron diff --git a/src/common/models/llama3/iron/llama3_iron_backend.cpp b/src/common/models/llama3/iron/llama3_iron_backend.cpp new file mode 100644 index 000000000..d07cad852 --- /dev/null +++ b/src/common/models/llama3/iron/llama3_iron_backend.cpp @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "models/llama3/iron/llama3_iron_backend.hpp" + +#include "iron/causal_lm.hpp" + +namespace flm::llama3 { + +flm::backend::BackendFactory iron_factory() { + return flm::iron::causal_factory("llama"); +} + +} // namespace flm::llama3 \ No newline at end of file diff --git a/src/common/models/models_sources.cmake b/src/common/models/models_sources.cmake index c709a6da3..16c2768ac 100644 --- a/src/common/models/models_sources.cmake +++ b/src/common/models/models_sources.cmake @@ -1,8 +1,10 @@ # Per-model sources, split by where the kernels come from and which silicon they # run on. Those two axes are independent, so both appear in the path: # /// - e.g. phi4/core/dev04/ -# The core sources are built into flm_core when FLM_ENABLE_CORE is on. +# The core sources are built into flm_core when FLM_ENABLE_CORE is on. IRON +# sources are built into flm when FLM_ENABLE_IRON is on. # FastFlowLM's own flow has no per-model sources here: those engines ship as # prebuilt libraries under lib//. # Adding a model means adding the folder, not editing this file. file(GLOB FLM_MODELS_CORE_SOURCES "${CMAKE_CURRENT_LIST_DIR}/*/core/*/*.cpp") +file(GLOB FLM_MODELS_IRON_SOURCES "${CMAKE_CURRENT_LIST_DIR}/*/iron/*.cpp") diff --git a/src/home_install.sh b/src/home_install.sh index 1e6198c57..c6953118c 100755 --- a/src/home_install.sh +++ b/src/home_install.sh @@ -9,6 +9,7 @@ # ./home_install.sh # build (if needed) + install to ~/flm_exe # ./home_install.sh --no-build # install an existing build/ tree only # ./home_install.sh --core # ryzenai-corelib backend, to scratch +# ./home_install.sh --iron # IRON backend, using the active IRON environment # FLM_PREFIX=/path ./home_install.sh # # --core installs to /scratch/$USER/flm_exe_core rather than ~/flm_exe: a core @@ -74,10 +75,12 @@ detect_xrt_root() { DO_BUILD=1 WANT_CORE=0 +WANT_IRON=0 for arg in "$@"; do case "$arg" in --no-build) DO_BUILD=0 ;; --core) WANT_CORE=1 ;; + --iron) WANT_IRON=1 ;; -h|--help) grep '^#' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//' exit 0 ;; @@ -85,6 +88,16 @@ for arg in "$@"; do esac done +IRON_PYTHON="" +if [[ "$WANT_IRON" -eq 1 ]]; then + IRON_PYTHON="$(command -v python3 || true)" + if [[ -z "$IRON_PYTHON" ]] || ! "$IRON_PYTHON" -c 'import iron.lm_server' >/dev/null 2>&1; then + echo "[home_install] ERROR: --iron requires an environment with the IRON server (iron.lm_server) installed." >&2 + echo " Activate the IRON environment before running this installer." >&2 + exit 1 + fi +fi + if [[ "$WANT_CORE" -eq 1 && -z "$PRESET_FROM_ENV" ]]; then PRESET="linux-core-on" fi @@ -139,6 +152,12 @@ needs_configure() { local wanted=OFF [[ "$WANT_CORE" -eq 1 ]] && wanted=ON [[ "$cached" == "$wanted" ]] || return 0 + + cached=OFF + grep -q '^FLM_ENABLE_IRON:BOOL=ON' "$BUILD_DIR/CMakeCache.txt" && cached=ON + wanted=OFF + [[ "$WANT_IRON" -eq 1 ]] && wanted=ON + [[ "$cached" == "$wanted" ]] || return 0 fi return 1 } @@ -150,6 +169,7 @@ if [[ "$DO_BUILD" -eq 1 ]]; then if [[ -n "${RYZENAI_CORELIB_ROOT:-}" ]]; then configure_args+=(-DRYZENAI_CORELIB_ROOT="$RYZENAI_CORELIB_ROOT") fi + [[ "$WANT_IRON" -eq 1 ]] && configure_args+=(-DFLM_ENABLE_IRON=ON) # --fresh discards whatever a previous failed or differently-configured # run left behind. Without it CMake reloads those entries and the # reconfigure inherits the state it is meant to replace. @@ -184,6 +204,14 @@ if [[ ! -d "$BUILD_DIR" ]]; then exit 1 fi +if [[ "$WANT_IRON" -eq 1 ]] && + ! grep -q '^FLM_ENABLE_IRON:BOOL=ON' "$BUILD_DIR/CMakeCache.txt"; then + echo "[home_install] ERROR: --iron, but $BUILD_DIR was not configured with" >&2 + echo " FLM_ENABLE_IRON=ON. Drop --no-build, or point BUILD_DIR at" >&2 + echo " an IRON-enabled build tree." >&2 + exit 1 +fi + # ---- install into the home prefix ----------------------------------------- # This runs the same install rules as a system install, just relocated. # It populates: @@ -345,6 +373,10 @@ fi ENV_SCRIPT="$FLM_PREFIX/flm_env.sh" echo "[home_install] writing env script: $ENV_SCRIPT" +IRON_ENV_BLOCK="" +if [[ "$WANT_IRON" -eq 1 ]]; then + IRON_ENV_BLOCK="export FLM_IRON_PYTHON=\"$IRON_PYTHON\"" +fi cat > "$ENV_SCRIPT" <> get_auto_embedding_model(const std::string& model_tag, flm_rt::device* npu_device_inst) { +inline std::pair> get_auto_embedding_model(const std::string& model_tag, flm_rt::device* npu_device_inst, const std::string& backend = "") { #ifndef FASTFLOWLM_LINUX_LIMITED_MODELS static std::unordered_set gemma_embed_tags = { @@ -28,7 +35,18 @@ inline std::pair> get_auto_embe std::unique_ptr auto_embedding_engine = nullptr; std::string new_model_tag = complete_simple_embedding_tag(model_tag); - if (gemma_embed_tags.count(model_tag)) // tag + if (new_model_tag == "embeddinggemma-2:740m") { +#if defined(FLM_ENABLE_IRON) + if (backend.empty() || backend == flm::backend::kIronBackendId) + auto_embedding_engine = std::make_unique( + npu_device_inst, "embeddinggemma-2:740m", "embedding_gemma2", + 768); + else +#endif + throw std::runtime_error( + "EmbeddingGemma 2 requires the IRON backend"); + } + else if (gemma_embed_tags.count(model_tag)) // tag auto_embedding_engine = std::make_unique(npu_device_inst); else { new_model_tag = "embed-gemma:300m"; // No arguments, use default tag diff --git a/src/include/AutoModel/model_backend.hpp b/src/include/AutoModel/model_backend.hpp index 848e5071e..3c71a7d90 100644 --- a/src/include/AutoModel/model_backend.hpp +++ b/src/include/AutoModel/model_backend.hpp @@ -47,6 +47,8 @@ inline constexpr const char* kFlmBackendId = "flm"; /// build still runs every flm model with the flm kernels. inline constexpr const char* kCoreBackendId = "core"; +inline constexpr const char* kIronBackendId = "iron"; + /// \brief the entry key naming the kernel flow an entry's artifacts need /// \note Written onto the entry by model_list, which derives it from the tag, /// rather than read from model_list.json: the tag is what decides, and diff --git a/src/include/iron/causal_lm.hpp b/src/include/iron/causal_lm.hpp new file mode 100644 index 000000000..7973abf09 --- /dev/null +++ b/src/include/iron/causal_lm.hpp @@ -0,0 +1,62 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "AutoModel/model_backend.hpp" +#include "iron/worker.hpp" + +#include +#include +#include + +namespace flm::iron { + +class CausalLM final : public ::causal_lm { +public: + CausalLM(const std::string& model_path, std::string model, + std::uint32_t context_length); + + buffer forward(int id) override; + buffer prefill(std::vector& ids, void*) override; + void set_context_length(int length) override; + void load_weights(Q4NX&) override; + void update_max_length(std::uint32_t length) override; + void clear_context() override; + buffer get_k_cache(int, int) override; + buffer get_v_cache(int, int) override; + int get_current_context_length() override; + int checkpoint() override; + int restore() override; + bool poisoned() const noexcept; + +private: + buffer run(); + + Worker worker_; + std::vector tokens_; + std::uint32_t max_length_; +}; + +class CausalBackend final : public flm::backend::ModelBackend { +public: + CausalBackend(const flm::backend::BackendContext& context, + std::string model); + + ::causal_lm& engine() override; + std::string id() const override; + std::string detail() const override; + bool supports_preemption() const override; + bool poisoned() const noexcept override; + +private: + CausalLM engine_; +}; + +flm::backend::BackendFactory causal_factory(std::string model); + +inline flm::backend::BackendTraits causal_traits() { + return flm::backend::BackendTraits{false, false, 0}; +} + +} // namespace flm::iron \ No newline at end of file diff --git a/src/include/iron/embedding.hpp b/src/include/iron/embedding.hpp new file mode 100644 index 000000000..3431ac0fe --- /dev/null +++ b/src/include/iron/embedding.hpp @@ -0,0 +1,28 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "AutoEmbeddingModel/auto_embedding_model.hpp" +#include "iron/worker.hpp" + +namespace flm::iron { + +class Embedding final : public AutoEmbeddingModel { +public: + Embedding(flm_rt::device* device, std::string model_tag, + std::string worker_model, std::uint32_t dimensions) + : AutoEmbeddingModel(device, std::move(model_tag)), + worker_model_(std::move(worker_model)), dimensions_(dimensions) {} + + void load_model(std::string model_path, json, bool enable_preemption) override; + std::vector embed(std::string& text, + embedding_task_type_t task_type) override; + +private: + std::string worker_model_; + std::uint32_t dimensions_; + std::unique_ptr worker_; +}; + +} // namespace flm::iron \ No newline at end of file diff --git a/src/include/iron/worker.hpp b/src/include/iron/worker.hpp new file mode 100644 index 000000000..c861023d7 --- /dev/null +++ b/src/include/iron/worker.hpp @@ -0,0 +1,52 @@ +#pragma once +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "nlohmann/json.hpp" + +#include +#include +#include +#include +#include + +namespace flm::iron { + +// Wire protocol bits; keep in sync with IRON's iron/lm_server.py. +enum class Capability : std::uint32_t { + CausalLM = 1 << 0, // logits and reset + Embeddings = 1 << 1, +}; + +class Worker { +public: + Worker(const std::filesystem::path& model_path, std::string model, + std::uint32_t context_length = 0); + ~Worker(); + + Worker(const Worker&) = delete; + Worker& operator=(const Worker&) = delete; + + std::vector logits(const std::vector& tokens); + std::vector embed(const std::string& text, const std::string& task, + std::uint32_t dimensions = 768); + void reset(); + bool supports(Capability capability) const noexcept { + return (capabilities_ & static_cast(capability)) != 0; + } + bool poisoned() const noexcept { return poisoned_; } + +private: + nlohmann::json request(const nlohmann::json& command, + const void* payload = nullptr, + std::size_t payload_size = 0); + void write_all(const void* data, std::size_t size); + void read_all(void* data, std::size_t size); + + struct Impl; + std::unique_ptr impl_; + bool poisoned_ = false; + std::uint32_t capabilities_ = 0; +}; + +} // namespace flm::iron \ No newline at end of file diff --git a/src/include/model_list.hpp b/src/include/model_list.hpp index 225ddbfd1..cbc222d2d 100644 --- a/src/include/model_list.hpp +++ b/src/include/model_list.hpp @@ -225,6 +225,7 @@ class model_list { if (model_type == "whisper-v3") continue; else if (model_type == "embed-gemma") continue; else if (model_type == "d3") continue; // decision models answer /v1/systemone only + else if (model_type == "embeddinggemma-2") continue; for (const auto& [size, model_info] : model_subset.items()) { nlohmann::json model_entry = { {"name", model_type + ":" + size}, @@ -256,6 +257,7 @@ class model_list { if (model_type == "whisper-v3") continue; else if (model_type == "embed-gemma") continue; else if (model_type == "d3") continue; // decision models answer /v1/systemone only + else if (model_type == "embeddinggemma-2") continue; for (const auto& [size, model_info] : model_subset.items()) { // id uses the same "type:size" convention; created uses current epoch seconds nlohmann::json model_entry = { @@ -474,17 +476,24 @@ class model_list { /// \brief the kernel flow that runs a tag, read off the tag itself /// \param model_type the family part of the tag, e.g. "phi4-mini-it-core" - /// \return "core" for a family ending in "-core", "flm" otherwise + /// \return the backend id encoded by the family name /// \note The name is the mechanism, not a label on top of one. A /// corelib model is packaged differently from its FastFlowLM /// namesake -- different weights, different quantization -- so it /// is a different model to pull and a different tag to ask for, /// and there is nothing left for a catalog key to disambiguate. static std::string backend_for_family(const std::string& model_type) { - static const std::string suffix = "-core"; - if (model_type.size() > suffix.size() && - model_type.compare(model_type.size() - suffix.size(), - suffix.size(), suffix) == 0) { + if (model_type == "embeddinggemma-2") return "iron"; + static const std::string iron_suffix = "-iron"; + if (model_type.size() > iron_suffix.size() && + model_type.compare(model_type.size() - iron_suffix.size(), + iron_suffix.size(), iron_suffix) == 0) { + return "iron"; + } + static const std::string core_suffix = "-core"; + if (model_type.size() > core_suffix.size() && + model_type.compare(model_type.size() - core_suffix.size(), + core_suffix.size(), core_suffix) == 0) { return "core"; } return "flm"; diff --git a/src/include/models/llama3/iron/llama3_iron_backend.hpp b/src/include/models/llama3/iron/llama3_iron_backend.hpp new file mode 100644 index 000000000..a18fb8195 --- /dev/null +++ b/src/include/models/llama3/iron/llama3_iron_backend.hpp @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "AutoModel/model_backend.hpp" +#include "iron/causal_lm.hpp" + +namespace flm::llama3 { + +flm::backend::BackendFactory iron_factory(); + +inline flm::backend::BackendTraits iron_traits() { + return flm::iron::causal_traits(); +} + +} // namespace flm::llama3 \ No newline at end of file diff --git a/src/model_info.json b/src/model_info.json index a698aa7aa..04642fcc2 100644 --- a/src/model_info.json +++ b/src/model_info.json @@ -1918,6 +1918,32 @@ "path": "tokenizer_config.json" } ], + "embeddinggemma-2:740m": [ + { + "type": "file", + "path": "config.json", + "size": 4455, + "oid": "dc524b2cbef566d66d8533c67dc68a6a8f5d340a" + }, + { + "type": "file", + "path": "model.safetensors", + "size": 1488915288, + "sha256": "197a32965d4b1105faf060417baa899e193fb73cd401f42ec9295234d5553d79" + }, + { + "type": "file", + "path": "tokenizer.json", + "size": 32170510, + "sha256": "4d777ef5bdc1aa36227abdfb77c3e49e7b9c892d16e1b6bda41c393504828be4" + }, + { + "type": "file", + "path": "tokenizer_config.json", + "size": 1599, + "oid": "a9afd12fb3ca2c418fd405320ace85ad46748d03" + } + ], "whisper-v3:turbo": [ { "type": "file", @@ -2464,6 +2490,32 @@ "path": "vision_weight.q4nx" } ], + "llama3.2-iron:1b": [ + { + "type": "file", + "size": 1019, + "sha256": "a53d38c545333679d523863f60a45f04166e52d15793864da3464712810f8bd9", + "path": "config.json" + }, + { + "type": "file", + "size": 2471645608, + "sha256": "1ff795ff6a07e6a68085d206fb84417da2f083f68391c2843cd2b8ac6df8538f", + "path": "model.safetensors" + }, + { + "type": "file", + "size": 9085657, + "sha256": "79e3e522635f3171300913bb421464a87de6222182a0570b9b2ccba2a964b2b4", + "path": "tokenizer.json" + }, + { + "type": "file", + "size": 4331, + "sha256": "01b0d6dd3c391ffc2facc01592ac07f25e8c4212d3450b8103042e8a5fcbafa8", + "path": "tokenizer_config.json" + } + ], "llama3.2:1b": [ { "type": "file", diff --git a/src/model_list.json b/src/model_list.json index f645763f8..12b5472d9 100644 --- a/src/model_list.json +++ b/src/model_list.json @@ -965,6 +965,52 @@ "footprint": 0.62 } }, + "embeddinggemma-2": { + "740m": { + "supported_platforms": ["aie2p"], + "name": "embeddinggemma-2", + "url": "https://huggingface.co/google/embeddinggemma-2", + "file_url": "https://huggingface.co/api/models/google/embeddinggemma-2/tree/main", + "size": 1488915288, + "flm_min_version": "1.0.5", + "files": [ + "config.json", + "model.safetensors", + "tokenizer.json", + "tokenizer_config.json" + ], + "file_sources": { + "config.json": { + "url": "https://huggingface.co/google/embeddinggemma-2", + "revision": "914f7f89142e33e77833254d9c9b90c3cef7303b" + }, + "model.safetensors": { + "url": "https://huggingface.co/google/embeddinggemma-2", + "revision": "914f7f89142e33e77833254d9c9b90c3cef7303b" + }, + "tokenizer.json": { + "url": "https://huggingface.co/google/embeddinggemma-2", + "revision": "914f7f89142e33e77833254d9c9b90c3cef7303b" + }, + "tokenizer_config.json": { + "url": "https://huggingface.co/google/embeddinggemma-2", + "revision": "914f7f89142e33e77833254d9c9b90c3cef7303b" + } + }, + "default_context_length": 2048, + "details": { + "format": "safetensors", + "family": "embeddinggemma-2", + "think": false, + "parameter_size": "740M", + "quantization_level": "BF16" + }, + "label": [ + "embeddings" + ], + "footprint": 1.5 + } + }, "d3": { "edge": { "supported_platforms": ["aie2p"], @@ -1436,6 +1482,46 @@ "footprint": 2.7 } }, + "llama3.2-iron": { + "1b": { + "supported_platforms": ["aie2p"], + "name": "llama3.2-1b-instruct", + "url": "https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/resolve/9213176726f574b556790deb65791e0c5aa438b6", + "file_url": "https://huggingface.co/api/models/meta-llama/Llama-3.2-1B-Instruct/tree/9213176726f574b556790deb65791e0c5aa438b6", + "use_hf_token": true, + "size": 2480736615, + "flm_min_version": "1.0.7", + "default_context_length": 131072, + "max_prefill_len": 4096, + "files": [ + "config.json", + "model.safetensors", + "tokenizer.json", + "tokenizer_config.json" + ], + "file_sources": { + "config.json": { + "url": "https://huggingface.co/FastFlowLM/Llama-3.2-1B-NPU2", + "revision": "7bb88fcd5653534fe065b911123d10294b62a8d1" + }, + "tokenizer.json": { + "url": "https://huggingface.co/FastFlowLM/Llama-3.2-1B-NPU2", + "revision": "7bb88fcd5653534fe065b911123d10294b62a8d1" + }, + "tokenizer_config.json": { + "url": "https://huggingface.co/FastFlowLM/Llama-3.2-1B-NPU2", + "revision": "7bb88fcd5653534fe065b911123d10294b62a8d1" + } + }, + "details": { + "family": "llama3", + "think": false, + "parameter_size": "1B", + "quantization_level": "BF16" + }, + "footprint": 2.5 + } + }, "llama3.1": { "8b": { "supported_platforms": ["aie2p"], diff --git a/src/server/rest_handler.cpp b/src/server/rest_handler.cpp index 3109b7ade..70c8e03e4 100644 --- a/src/server/rest_handler.cpp +++ b/src/server/rest_handler.cpp @@ -525,13 +525,19 @@ RestHandler::RestHandler(model_list& models, ModelDownloader& downloader, progra this->prefix_min_gain = args.prefix_min_gain; // Initialize chat bot with default model + const bool embeddinggemma_2 = + default_model_tag == "embeddinggemma-2" || + default_model_tag.starts_with("embeddinggemma-2:"); #ifndef FASTFLOWLM_LINUX_LIMITED_MODELS if (this->asr) { std::string whisper_tag = "whisper-v3:turbo"; ensure_asr_model_loaded(whisper_tag); } - if (this->embed) { - std::string embed_tag = "embed-gemma:300m"; + if (this->embed || embeddinggemma_2) { + this->embed = true; + std::string embed_tag = embeddinggemma_2 + ? default_model_tag + : "embed-gemma:300m"; ensure_embed_model_loaded(embed_tag); } #else @@ -547,7 +553,7 @@ RestHandler::RestHandler(model_list& models, ModelDownloader& downloader, progra ensure_decision_model_loaded(args.decision); } - if (default_model_tag != "model-faker") { + if (default_model_tag != "model-faker" && !embeddinggemma_2) { if (!supported_models.is_model_supported(default_model_tag)) { header_print("Warning", "Default model tag '" << default_model_tag << "' is not supported. Falling back to 'llama3.2:1b'."); this->default_model_tag = "llama3.2:1b"; @@ -779,7 +785,8 @@ void RestHandler::ensure_embed_model_loaded(const std::string& model_tag) { this->embed = false; return; } - auto [embedding_model_tag, auto_embedding_engine] = get_auto_embedding_model(ensure_tag, this->npu_device_inst); + auto [embedding_model_tag, auto_embedding_engine] = + get_auto_embedding_model(ensure_tag, this->npu_device_inst, this->backend); this->auto_embedding_engine = std::move(auto_embedding_engine); auto [new_embedding_model_tag, embedding_model_info] = this->supported_models.get_model_info(embedding_model_tag); std::string embedding_model_path = this->supported_models.get_model_path(new_embedding_model_tag); diff --git a/src/test/iron_backend/CMakeLists.txt b/src/test/iron_backend/CMakeLists.txt new file mode 100644 index 000000000..6fb693f52 --- /dev/null +++ b/src/test/iron_backend/CMakeLists.txt @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT + +cmake_minimum_required(VERSION 3.22) +project(iron_backend_tests LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(FLM_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../..") +find_path(BOOST_PROCESS_INCLUDE_DIR NAMES boost/process.hpp + HINTS "$ENV{BOOST_ROOT}" "C:/dev/boost_1_88_0" REQUIRED) + +add_executable(test_iron_worker + test_worker.cpp + "${FLM_SOURCE_DIR}/common/iron/worker.cpp") +target_include_directories(test_iron_worker PRIVATE + "${FLM_SOURCE_DIR}/include" + "${BOOST_PROCESS_INCLUDE_DIR}") +if(WIN32) + target_link_libraries(test_iron_worker PRIVATE shell32) +endif() + +enable_testing() +add_test(NAME test_iron_worker COMMAND test_iron_worker) +set_tests_properties(test_iron_worker PROPERTIES + TIMEOUT 30 + ENVIRONMENT + "FLM_IRON_PYTHON=$") diff --git a/src/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp new file mode 100644 index 000000000..743802c7f --- /dev/null +++ b/src/test/iron_backend/test_worker.cpp @@ -0,0 +1,173 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "iron/worker.hpp" + +#include "../phi4_core/test_support.hpp" +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +#include +#include +#endif + +namespace { + +void SetEnvironment(const char* name, const char* value) { +#if defined(_WIN32) + _putenv_s(name, value); +#else + setenv(name, value, 1); +#endif +} + +// The test executable doubles as the server, using the same launch arguments. +int ServeFake() { +#if defined(_WIN32) + _setmode(_fileno(stdin), _O_BINARY); + _setmode(_fileno(stdout), _O_BINARY); +#endif + auto read = [](void* data, std::size_t size) { + TEST_REQUIRE(std::cin.read(static_cast(data), size)); + }; + auto send = [](const nlohmann::json& header, + const void* payload = nullptr, std::size_t size = 0) { + const auto data = header.dump(); + const std::uint32_t length = data.size(); + std::cout.write(reinterpret_cast(&length), sizeof(length)); + std::cout.write(data.data(), data.size()); + if (size) std::cout.write(static_cast(payload), size); + std::cout.flush(); + }; + std::string model; + std::uint32_t length; + while (std::cin.read(reinterpret_cast(&length), sizeof(length))) { + std::string header(length, '\0'); + read(header.data(), header.size()); + const auto command = nlohmann::json::parse(header); + const auto name = command.at("command").get(); + if (name == "load_model") { + model = command.at("model").get(); + if (model == "llama") TEST_REQUIRE(command.at("context_length") == 32); + else TEST_REQUIRE(!command.contains("context_length")); + nlohmann::json reply = {{"ok", true}}; + if (model != "legacy") { + TEST_REQUIRE(command.at("protocol_version") == 1); + TEST_REQUIRE(command.at("capabilities") == 3); + reply["protocol_version"] = model == "wrong_version" ? 2 : 1; + reply["capabilities"] = model == "embedding_only" ? 2 : + model == "causal_only" ? 1 : + model == "future" ? 7 : 3; + } + send(reply); + if (model == "dies") return 0; + } else if (name == "logits") { + std::vector tokens(command.at("count").get()); + read(tokens.data(), tokens.size() * sizeof(tokens[0])); + if (model == "bad_count") { + send({{"ok", true}}); + } else if (model == "bad_header") { + const std::uint32_t zero = 0; + std::cout.write(reinterpret_cast(&zero), sizeof(zero)); + std::cout.flush(); + } else if (model == "error") { + send({{"ok", false}, {"error", "model failed"}}); + } else if (model == "truncated") { + send({{"ok", true}, {"count", 2}}, "x", 1); + return 0; + } else { + const std::uint16_t values[] = { + static_cast(tokens.back()), + static_cast(tokens.size())}; + send({{"ok", true}, {"count", 2}}, values, sizeof(values)); + } + } else if (name == "embed") { + std::string text(command.at("size").get(), '\0'); + read(text.data(), text.size()); + const float values[] = {static_cast(text.size()), + command.at("dimensions").get()}; + send({{"ok", true}, {"count", 2}}, values, sizeof(values)); + } else { + TEST_REQUIRE(name == "reset" || name == "close"); + send({{"ok", true}}); + if (name == "close") return 0; + } + } + return 0; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc > 1) { + return argc == 3 && std::string(argv[1]) == "-m" && + std::string(argv[2]) == "fake_server" + ? ServeFake() : 1; + } + SetEnvironment("FLM_IRON_WORKER_MODULE", "fake_server"); + + flm::iron::Worker llama("/unused", "llama", 32); + const auto logits = llama.logits(std::vector{4, 10, 26}); + TEST_REQUIRE(logits == std::vector({26, 3})); + llama.reset(); + TEST_REQUIRE(!llama.poisoned()); + + flm::iron::Worker embedding("/unused", "embedding_gemma2"); + const auto result = embedding.embed("hello", "query", 128); + TEST_REQUIRE(result == std::vector({5.0f, 128.0f})); + + using flm::iron::Capability; + for (const char* model : {"legacy", "future"}) { + flm::iron::Worker compatible("/unused", model); + TEST_REQUIRE(compatible.supports(Capability::CausalLM)); + TEST_REQUIRE(compatible.supports(Capability::Embeddings)); + TEST_REQUIRE(!compatible.supports(static_cast(4))); + TEST_REQUIRE(compatible.logits({26}) == std::vector({26, 1})); + } + flm::iron::Worker embedding_only("/unused", "embedding_only"); + TEST_REQUIRE(!embedding_only.supports(Capability::CausalLM)); + RequireContains(RequireThrows([&] { embedding_only.logits({1}); }), "support"); + RequireContains(RequireThrows([&] { embedding_only.reset(); }), "support"); + TEST_REQUIRE(!embedding_only.poisoned()); + TEST_REQUIRE(embedding_only.embed("hello", "query", 128) == result); + + flm::iron::Worker causal_only("/unused", "causal_only"); + TEST_REQUIRE(!causal_only.supports(Capability::Embeddings)); + RequireContains(RequireThrows([&] { causal_only.embed("hello", "query"); }), "support"); + TEST_REQUIRE(!causal_only.poisoned()); + causal_only.reset(); + RequireContains(RequireThrows([] { + flm::iron::Worker incompatible("/unused", "wrong_version"); + }), "protocol version"); + + for (const char* model : {"bad_count", "bad_header", "error", "truncated"}) { + flm::iron::Worker broken("/unused", model); + RequireThrows([&] { broken.logits({1}); }); + TEST_REQUIRE(broken.poisoned()); + RequireContains(RequireThrows([&] { broken.reset(); }), "reload"); + } + + bool pipe_failed = false; + try { + flm::iron::Worker exits("/unused", "dies"); + exits.reset(); + } catch (const std::runtime_error&) { + pipe_failed = true; + } + TEST_REQUIRE(pipe_failed); + + bool failed = false; + SetEnvironment("FLM_IRON_WORKER_MODULE", "missing_worker"); + try { + flm::iron::Worker dead("/unused", "llama", 32); + } catch (const std::runtime_error&) { + failed = true; + } + TEST_REQUIRE(failed); + return 0; +} diff --git a/src/test/model_backend/CMakeLists.txt b/src/test/model_backend/CMakeLists.txt index 45689b977..7c32c38cb 100644 --- a/src/test/model_backend/CMakeLists.txt +++ b/src/test/model_backend/CMakeLists.txt @@ -20,7 +20,8 @@ else() if(XRT_FOUND) set(XRT_INCLUDE_DIR "${XRT_INCLUDE_DIRS}") else() - set(XRT_INCLUDE_DIR /opt/xilinx/xrt/include CACHE PATH "Where XRT headers live") + find_path(XRT_INCLUDE_DIR NAMES xrt/xrt_bo.h + HINTS /opt/xilinx/xrt/include REQUIRED) endif() endif() diff --git a/src/test/model_backend/test_model_backend.cpp b/src/test/model_backend/test_model_backend.cpp index df55b04f7..498a53849 100644 --- a/src/test/model_backend/test_model_backend.cpp +++ b/src/test/model_backend/test_model_backend.cpp @@ -21,6 +21,7 @@ using flm::backend::BackendRegistry; using flm::backend::BackendTraits; using flm::backend::kFlmBackendId; using flm::backend::kCoreBackendId; +using flm::backend::kIronBackendId; using flm::backend::ModelBackend; using flm::backend::resolve_backend_id; @@ -119,6 +120,16 @@ void test_traits_are_kept_per_backend() { TEST_REQUIRE(core.max_context_length == 4096); } +void test_iron_backend_registration() { + auto registry = MakeRegistry(); + registry.register_backend("llama3", kIronBackendId, StubFactory("iron"), + BackendTraits{false, false, 0}); + const auto traits = registry.traits("llama3", kIronBackendId); + TEST_REQUIRE(!traits.needs_npu_xclbin); + TEST_REQUIRE(!traits.supports_preemption); + TEST_REQUIRE(traits.max_context_length == 0); +} + void test_backend_defaults() { StubBackend backend(kFlmBackendId); TEST_REQUIRE(backend.detail().empty()); @@ -283,6 +294,7 @@ void test_resolution_rejects_with_a_readable_message() { int main() { RunTest(test_register_and_create, "register and create"); RunTest(test_traits_are_kept_per_backend, "traits are kept per backend"); + RunTest(test_iron_backend_registration, "iron backend registration"); RunTest(test_backend_defaults, "backend policy defaults"); RunTest(test_duplicate_registration_is_rejected, "duplicate registration is rejected"); RunTest(test_replace_backend_is_the_test_seam, "replace_backend is the test seam"); diff --git a/src/test/model_list_platform/test_model_list_platform.cpp b/src/test/model_list_platform/test_model_list_platform.cpp index ffd328518..da028aa62 100644 --- a/src/test/model_list_platform/test_model_list_platform.cpp +++ b/src/test/model_list_platform/test_model_list_platform.cpp @@ -19,6 +19,8 @@ namespace { constexpr const char* kCatalogPath = FLM_TEST_MODEL_LIST_PATH; constexpr const char* kPhiTag = "phi4-mini-it:4b"; constexpr const char* kCoreTag = "phi4-mini-it-core:4b"; +constexpr const char* kEmbeddingGemma2Tag = "embeddinggemma-2:740m"; +constexpr const char* kIronLlamaTag = "llama3.2-iron:1b"; /// \brief every generation the build knows about, for the catalog sweep const std::set kAllPlatforms = { @@ -116,6 +118,41 @@ void test_a_corelib_build_offers_the_corelib_phi4() { TEST_REQUIRE(corelib_on_aie2p.get_model_info(kPhiTag).second.at("backend") == "flm"); } +void test_embeddinggemma_2_requires_iron() { + auto stock = open_stock(); + TEST_REQUIRE(!stock.is_model_supported(kEmbeddingGemma2Tag)); + + auto iron = open_catalog("aie2p", {"flm", "iron"}); + TEST_REQUIRE(iron.is_model_supported(kEmbeddingGemma2Tag)); + TEST_REQUIRE(iron.rectify_model_tag("embeddinggemma-2") == kEmbeddingGemma2Tag); + + const auto [tag, info] = iron.get_model_info(kEmbeddingGemma2Tag); + TEST_REQUIRE(tag == kEmbeddingGemma2Tag); + TEST_REQUIRE(info.at("backend") == "iron"); + TEST_REQUIRE(info.at("details").at("family") == "embeddinggemma-2"); +} + +void test_bf16_llama_requires_iron() { + auto stock = open_stock(); + TEST_REQUIRE(!stock.is_model_supported(kIronLlamaTag)); + + auto iron = open_catalog("aie2p", {"flm", "iron"}); + TEST_REQUIRE(iron.is_model_supported(kIronLlamaTag)); + TEST_REQUIRE(iron.rectify_model_tag("llama3.2-iron") == kIronLlamaTag); + + const auto [tag, info] = iron.get_model_info(kIronLlamaTag); + TEST_REQUIRE(tag == kIronLlamaTag); + TEST_REQUIRE(info.at("backend") == "iron"); + TEST_REQUIRE(info.at("name") == "llama3.2-1b-instruct"); + TEST_REQUIRE(info.at("url").get().starts_with( + "https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/resolve/")); + TEST_REQUIRE(info.at("details").at("family") == "llama3"); + TEST_REQUIRE(info.at("details").at("quantization_level") == "BF16"); + TEST_REQUIRE(info.at("files").at(1) == "model.safetensors"); + TEST_REQUIRE(info.at("use_hf_token") == true); + TEST_REQUIRE(!info.at("file_sources").contains("model.safetensors")); +} + void test_a_catalog_can_filter_down_to_nothing() { // The other corner of the same rule, and the one that reads like a bug if // it is not stated: a build with only the FastFlowLM kernels, on silicon @@ -141,15 +178,15 @@ void test_a_catalog_can_filter_down_to_nothing() { } void test_the_tag_name_decides_the_backend() { - // The rule, on a catalog small enough to see: a family ending in "-core" is - // corelib's, anything else is the FastFlowLM kernels'. Nothing else in the - // entry is consulted, which is why the two never contend for a tag. + // The rule, on a catalog small enough to see: a family ending in "-core" or + // "-iron" names that backend; anything else uses FastFlowLM's kernels. const auto path = write_catalog( "flm_backend_model_list.json", {{"model_path", "models"}, {"models", {{"demo", {{"1b", {{"name", "Demo"}}}}}, - {"demo-core", {{"1b", {{"name", "Demo-core"}}}}}}}}); + {"demo-core", {{"1b", {{"name", "Demo-core"}}}}}, + {"demo-iron", {{"1b", {{"name", "Demo-iron"}}}}}}}}); std::string list_path = path.string(); std::string exe_dir = "."; @@ -159,6 +196,7 @@ void test_the_tag_name_decides_the_backend() { model_list stock(list_path, exe_dir, "aie2p", {"flm"}); TEST_REQUIRE(stock.is_model_supported("demo:1b")); TEST_REQUIRE(!stock.is_model_supported("demo-core:1b")); + TEST_REQUIRE(!stock.is_model_supported("demo-iron:1b")); TEST_REQUIRE(stock.get_model_info("demo:1b").second.at("backend") == "flm"); // Both lookups used to dereference the pruned entry. A pruned family keeps // its own name rather than resolving to a size, and an unknown tag comes @@ -175,6 +213,10 @@ void test_the_tag_name_decides_the_backend() { TEST_REQUIRE(corelib.get_model_info("demo-core:1b").second.at("backend") == "core"); TEST_REQUIRE(corelib.get_model_info("demo:1b").second.at("backend") == "flm"); + model_list iron(list_path, exe_dir, "aie2p", {"flm", "iron"}); + TEST_REQUIRE(iron.is_model_supported("demo-iron:1b")); + TEST_REQUIRE(iron.get_model_info("demo-iron:1b").second.at("backend") == "iron"); + std::filesystem::remove(path); } @@ -324,6 +366,10 @@ int main() { "stock build offers phi4 with its own artifacts"); RunTest(test_a_corelib_build_offers_the_corelib_phi4, "a corelib build offers the corelib phi4"); + RunTest(test_embeddinggemma_2_requires_iron, + "EmbeddingGemma 2 requires iron"); + RunTest(test_bf16_llama_requires_iron, + "BF16 Llama requires iron"); RunTest(test_the_tag_name_decides_the_backend, "the tag name decides the backend"); RunTest(test_a_catalog_can_filter_down_to_nothing,