From bea82f1243981230ad0d15129ce72a67bf194d92 Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 7 Oct 2026 12:28:45 -0600 Subject: [PATCH 01/19] Add end-to-end IRON backend --- docs/iron-backend.md | 58 ++++++ python/flm_iron/__init__.py | 4 + python/flm_iron/worker.py | 91 ++++++++++ src/CMakeLists.txt | 17 +- .../modeling_iron_embedding.cpp | 24 +++ src/common/AutoModel/builtin_backends.cpp | 10 ++ src/common/iron/worker.cpp | 169 ++++++++++++++++++ .../llama3/iron/llama3_iron_backend.cpp | 103 +++++++++++ src/common/models/models_sources.cmake | 4 +- .../all_embedding_model.hpp | 12 +- .../modeling_iron_embedding.hpp | 20 +++ src/include/AutoModel/model_backend.hpp | 2 + src/include/iron/worker.hpp | 42 +++++ .../llama3/iron/llama3_iron_backend.hpp | 16 ++ src/server/rest_handler.cpp | 3 +- src/test/iron_backend/CMakeLists.txt | 19 ++ src/test/iron_backend/fake_worker.py | 50 ++++++ src/test/iron_backend/test_worker.cpp | 33 ++++ src/test/model_backend/test_model_backend.cpp | 12 ++ 19 files changed, 684 insertions(+), 5 deletions(-) create mode 100644 docs/iron-backend.md create mode 100644 python/flm_iron/__init__.py create mode 100644 python/flm_iron/worker.py create mode 100644 src/common/AutoEmbeddingModel/modeling_iron_embedding.cpp create mode 100644 src/common/iron/worker.cpp create mode 100644 src/common/models/llama3/iron/llama3_iron_backend.cpp create mode 100644 src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp create mode 100644 src/include/iron/worker.hpp create mode 100644 src/include/models/llama3/iron/llama3_iron_backend.hpp create mode 100644 src/test/iron_backend/CMakeLists.txt create mode 100644 src/test/iron_backend/fake_worker.py create mode 100644 src/test/iron_backend/test_worker.cpp diff --git a/docs/iron-backend.md b/docs/iron-backend.md new file mode 100644 index 000000000..182acfa41 --- /dev/null +++ b/docs/iron-backend.md @@ -0,0 +1,58 @@ + + +# IRON backend + +The `iron` backend runs a complete IRON graph in a persistent Python process. +FastFlowLM handles tokenization, chat templates, sampling, and HTTP requests. +The worker retains the IRON model, compiled images, weights, and device state. + +## Build + +Configure FastFlowLM with `FLM_ENABLE_IRON=ON`. Set these variables when you run +the source build: + +```bash +source /opt/xilinx/xrt/setup.sh +export FLM_IRON_PYTHON=/scratch/roesti/IRON/.venv/bin/python +export PYTHONPATH=/scratch/roesti/FastFlowLM/python +``` + +An installed build places `flm_iron` under `share/flm/python`. Add that directory +to `PYTHONPATH`. + +## Llama 3.2 1B + +The FastFlowLM package supplies `config.json`, tokenizer files, and chat +metadata. IRON loads a separate checkpoint directory that contains +`model.safetensors` and `tokenizer.model`: + +```bash +export FLM_IRON_MODEL_PATH=/path/to/llama3.2-1b +flm run llama3.2:1b --backend iron --ctx-len 32768 +``` + +The context length must be a multiple of 2048 and cannot exceed 32768. +Preemption is unavailable. + +## EmbeddingGemma + +IRON loads a directory that contains `model.safetensors` and `tokenizer.json`: + +```bash +export FLM_IRON_EMBEDDING_PATH=/path/to/embeddinggemma-2 +flm serve llama3.2:1b --embed 1 --backend iron +``` + +The embedding endpoint passes text and its query or document task to the IRON +encoder. The worker returns the normalized 768-element float32 embedding. + +## Protocol + +FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary +payloads over pipes. Llama requests contain the complete token history and +return bfloat16 logits. Embedding requests contain UTF-8 text and return +float32 values. Worker errors poison the backend, so FastFlowLM requires a model +reload after a process or protocol failure. \ No newline at end of file diff --git a/python/flm_iron/__init__.py b/python/flm_iron/__init__.py new file mode 100644 index 000000000..82f81270c --- /dev/null +++ b/python/flm_iron/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT + +"""FastFlowLM transport for complete IRON graphs.""" \ No newline at end of file diff --git a/python/flm_iron/worker.py b/python/flm_iron/worker.py new file mode 100644 index 000000000..99f975d56 --- /dev/null +++ b/python/flm_iron/worker.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT + +import json +import os +import struct +import sys +from dataclasses import replace +from pathlib import Path + +import numpy as np +from ml_dtypes import bfloat16 + + +def read_exact(size): + data = bytearray() + while len(data) < size: + chunk = sys.stdin.buffer.read(size - len(data)) + if not chunk: + raise EOFError + data.extend(chunk) + return bytes(data) + + +def receive(): + size = struct.unpack("=I", read_exact(4))[0] + return json.loads(read_exact(size)) + + +def send(header, payload=b""): + data = json.dumps(header, separators=(",", ":")).encode() + sys.stdout.buffer.write(struct.pack("=I", len(data))) + sys.stdout.buffer.write(data) + sys.stdout.buffer.write(payload) + sys.stdout.buffer.flush() + + +def load_llama(command): + from iron.lm.llama3.model import LLAMA_3_2_1B, Runner + + root = Path(command["model_path"]) + weights = root / "model.safetensors" + tokenizer = root / "tokenizer.model" + config = replace(LLAMA_3_2_1B, max_seq_len=command["context_length"]) + return Runner(weights, tokenizer, config).npu() + + +def load_embedding(command): + from iron.lm.embeddinggemma2.encoder import Encoder + + return Encoder(command["model_path"], max_tokens=command["context_length"]) + + +def main(): + model = None + while True: + command = receive() + try: + name = command["command"] + if name == "load_llama": + model = load_llama(command) + send({"ok": True}) + elif name == "load_embedding": + model = load_embedding(command) + send({"ok": True}) + elif name == "logits": + tokens = np.frombuffer(read_exact(command["count"] * 4), np.int32) + logits = np.asarray(model.logits(tokens), dtype=bfloat16) + words = logits.view(np.uint16).astype(" +#include + +void IronEmbedding::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"); + } + const char* override_path = std::getenv("FLM_IRON_EMBEDDING_PATH"); + if (override_path && *override_path) model_path = override_path; + worker_ = std::make_unique(model_path, 512, "embedding"); + is_model_loaded = true; + this->model_path = std::move(model_path); +} + +std::vector IronEmbedding::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"); +} \ No newline at end of file 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/worker.cpp b/src/common/iron/worker.cpp new file mode 100644 index 000000000..8e26f9888 --- /dev/null +++ b/src/common/iron/worker.cpp @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "iron/worker.hpp" + +#if !defined(_WIN32) +#include +#include +#include +#endif + +#include +#include +#include +#include +#include + +namespace flm::iron { +namespace { + +constexpr std::uint32_t kMaxHeader = 1 << 20; + +std::string PythonExecutable() { + const char* value = std::getenv("FLM_IRON_PYTHON"); + return value && *value ? value : "python3"; +} + +std::string WorkerModule() { + const char* value = std::getenv("FLM_IRON_WORKER_MODULE"); + return value && *value ? value : "flm_iron.worker"; +} + +} // namespace + +Worker::Worker(const std::filesystem::path& model_path, + std::uint32_t context_length, std::string model) { +#if defined(_WIN32) + throw std::runtime_error("The IRON backend requires POSIX process support"); +#else + int sockets[2]; + if (socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) != 0) { + throw std::runtime_error("Cannot create the IRON worker socket"); + } + process_ = fork(); + if (process_ < 0) throw std::runtime_error("Cannot start the IRON worker"); + if (process_ == 0) { + dup2(sockets[1], STDIN_FILENO); + dup2(sockets[1], STDOUT_FILENO); + close(sockets[0]); + close(sockets[1]); + const std::string python = PythonExecutable(); + const std::string module = WorkerModule(); + execlp(python.c_str(), python.c_str(), "-m", module.c_str(), + static_cast(nullptr)); + _exit(127); + } + close(sockets[1]); + input_ = output_ = sockets[0]; + try { + request({{"command", model == "embedding" ? "load_embedding" : "load_llama"}, + {"model_path", model_path.string()}, + {"context_length", context_length}}); + } catch (...) { + close(input_); + waitpid(process_, nullptr, 0); + input_ = output_ = process_ = -1; + throw; + } +#endif +} + +Worker::~Worker() { +#if !defined(_WIN32) + if (input_ >= 0) { + try { request({{"command", "close"}}); } catch (...) {} + close(input_); + } + if (process_ > 0) waitpid(process_, nullptr, 0); +#endif +} + +void Worker::write_all(const void* data, std::size_t size) { +#if !defined(_WIN32) + const auto* cursor = static_cast(data); + while (size) { + const ssize_t written = send(input_, cursor, size, MSG_NOSIGNAL); + if (written <= 0) throw std::runtime_error("IRON worker write failed"); + cursor += written; + size -= static_cast(written); + } +#endif +} + +void Worker::read_all(void* data, std::size_t size) { +#if !defined(_WIN32) + auto* cursor = static_cast(data); + while (size) { + const ssize_t received = read(output_, cursor, size); + if (received <= 0) throw std::runtime_error("IRON worker exited"); + cursor += received; + size -= static_cast(received); + } +#endif +} + +nlohmann::json Worker::request(const nlohmann::json& command, + const void* payload, + std::size_t payload_size) { + 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); + + 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 worker 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 worker failed")); + } + return result; + } catch (...) { + poisoned_ = true; + throw; + } +} + +std::vector Worker::logits(const std::vector& tokens) { + auto response = request({{"command", "logits"}, {"count", tokens.size()}}, + tokens.data(), tokens.size() * sizeof(int)); + const std::size_t count = response.at("count").get(); + std::vector values(count); + try { + read_all(values.data(), values.size() * sizeof(values[0])); + } catch (...) { + poisoned_ = true; + throw; + } + return values; +} + +void Worker::reset() { request({{"command", "reset"}}); } + +std::vector Worker::embed(const std::string& text, + const std::string& task, + std::uint32_t dimensions) { + auto response = request({{"command", "embed"}, + {"task", task}, + {"dimensions", dimensions}, + {"size", text.size()}}, + text.data(), text.size()); + const std::size_t count = response.at("count").get(); + std::vector values(count); + try { + read_all(values.data(), values.size() * sizeof(values[0])); + } catch (...) { + poisoned_ = true; + throw; + } + return values; +} + +} // namespace flm::iron \ No newline at end of file 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..54c9aaa53 --- /dev/null +++ b/src/common/models/llama3/iron/llama3_iron_backend.cpp @@ -0,0 +1,103 @@ +// 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/worker.hpp" + +#include +#include +#include +#include +#include + +namespace flm::llama3 { +namespace { + +class IronLlama final : public causal_lm { +public: + IronLlama(const std::string& model_path, std::uint32_t context_length) + : worker_(model_path, context_length), max_length_(context_length) {} + + buffer forward(int id) override { + tokens_.push_back(id); + return run(); + } + + buffer prefill(std::vector& ids, void*) override { + tokens_.insert(tokens_.end(), ids.begin(), ids.end()); + return run(); + } + + void set_context_length(int length) override { + if (length < 0 || static_cast(length) > tokens_.size()) { + throw std::out_of_range("Invalid IRON context length"); + } + tokens_.resize(static_cast(length)); + } + void load_weights(Q4NX&) override { + throw std::logic_error("IRON loads safetensors in its worker"); + } + void update_max_length(std::uint32_t length) override { max_length_ = length; } + void clear_context() override { + tokens_.clear(); + worker_.reset(); + } + buffer get_k_cache(int, int) override { return {}; } + buffer get_v_cache(int, int) override { return {}; } + int get_current_context_length() override { return static_cast(tokens_.size()); } + int checkpoint() override { return static_cast(tokens_.size()); } + int restore() override { return 0; } + bool poisoned() const noexcept { return worker_.poisoned(); } + +private: + buffer 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; + } + + flm::iron::Worker worker_; + std::vector tokens_; + std::uint32_t max_length_; +}; + +class IronBackend final : public flm::backend::ModelBackend { +public: + explicit IronBackend(const flm::backend::BackendContext& context) + : engine_(CheckpointPath(context.model_path), + ValidateContextLength(context.context_length)) {} + causal_lm& engine() override { return engine_; } + std::string id() const override { return flm::backend::kIronBackendId; } + std::string detail() const override { return "IRON graph worker"; } + bool supports_preemption() const override { return false; } + bool poisoned() const noexcept override { return engine_.poisoned(); } + +private: + static std::string CheckpointPath(const std::string& fallback) { + const char* path = std::getenv("FLM_IRON_MODEL_PATH"); + return path && *path ? path : fallback; + } + static std::uint32_t ValidateContextLength(std::uint32_t length) { + if (length == 0 || length > 32768 || length % 2048 != 0) { + throw std::out_of_range( + "IRON Llama context length must be a multiple of 2048 in 2048..32768"); + } + return length; + } + IronLlama engine_; +}; + +} // namespace + +flm::backend::BackendFactory iron_factory() { + return [](const flm::backend::BackendContext& context) { + return std::make_unique(context); + }; +} + +} // 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/include/AutoEmbeddingModel/all_embedding_model.hpp b/src/include/AutoEmbeddingModel/all_embedding_model.hpp index 4e0c5da9f..3755d4fc3 100644 --- a/src/include/AutoEmbeddingModel/all_embedding_model.hpp +++ b/src/include/AutoEmbeddingModel/all_embedding_model.hpp @@ -6,8 +6,13 @@ /// \note This is a header file for get_auto_embedding_model func #pragma once +#include "AutoModel/model_backend.hpp" + #ifndef FASTFLOWLM_LINUX_LIMITED_MODELS #include "modeling_gemma_embedding.hpp" +#if defined(FLM_ENABLE_IRON) +#include "modeling_iron_embedding.hpp" +#endif #endif inline std::string complete_simple_embedding_tag(std::string model_tag) { @@ -18,7 +23,7 @@ inline std::string complete_simple_embedding_tag(std::string model_tag) { } -inline std::pair> 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,6 +33,11 @@ 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 defined(FLM_ENABLE_IRON) + if (backend == flm::backend::kIronBackendId) + auto_embedding_engine = std::make_unique(npu_device_inst); + else + #endif if (gemma_embed_tags.count(model_tag)) // tag auto_embedding_engine = std::make_unique(npu_device_inst); else { diff --git a/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp b/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp new file mode 100644 index 000000000..f1f81c5dd --- /dev/null +++ b/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp @@ -0,0 +1,20 @@ +// 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" + +class IronEmbedding final : public AutoEmbeddingModel { +public: + explicit IronEmbedding(flm_rt::device* device) + : AutoEmbeddingModel(device, "embed-gemma:300m") {} + + 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::unique_ptr worker_; +}; \ No newline at end of file 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/worker.hpp b/src/include/iron/worker.hpp new file mode 100644 index 000000000..ebe49876d --- /dev/null +++ b/src/include/iron/worker.hpp @@ -0,0 +1,42 @@ +#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 + +namespace flm::iron { + +class Worker { +public: + Worker(const std::filesystem::path& model_path, std::uint32_t context_length, + std::string model = "llama"); + ~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 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); + + int input_ = -1; + int output_ = -1; + int process_ = -1; + bool poisoned_ = false; +}; + +} // namespace flm::iron \ No newline at end of file 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..7e1c13881 --- /dev/null +++ b/src/include/models/llama3/iron/llama3_iron_backend.hpp @@ -0,0 +1,16 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#pragma once + +#include "AutoModel/model_backend.hpp" + +namespace flm::llama3 { + +flm::backend::BackendFactory iron_factory(); + +inline flm::backend::BackendTraits iron_traits() { + return flm::backend::BackendTraits{false, false, 32768}; +} + +} // namespace flm::llama3 \ No newline at end of file diff --git a/src/server/rest_handler.cpp b/src/server/rest_handler.cpp index 3109b7ade..cc779cd41 100644 --- a/src/server/rest_handler.cpp +++ b/src/server/rest_handler.cpp @@ -779,7 +779,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..d7518f304 --- /dev/null +++ b/src/test/iron_backend/CMakeLists.txt @@ -0,0 +1,19 @@ +# 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}/../..") + +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") + +enable_testing() +add_test(NAME test_iron_worker COMMAND test_iron_worker) +set_tests_properties(test_iron_worker PROPERTIES + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_LIST_DIR}") \ No newline at end of file diff --git a/src/test/iron_backend/fake_worker.py b/src/test/iron_backend/fake_worker.py new file mode 100644 index 000000000..076a8d5d4 --- /dev/null +++ b/src/test/iron_backend/fake_worker.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT + +import json +import struct +import sys + + +def read_exact(size): + data = bytearray() + while len(data) < size: + chunk = sys.stdin.buffer.read(size - len(data)) + if not chunk: + raise EOFError + data.extend(chunk) + return bytes(data) + + +def receive(): + size = struct.unpack("=I", read_exact(4))[0] + return json.loads(read_exact(size)) + + +def send(header, payload=b""): + data = json.dumps(header).encode() + sys.stdout.buffer.write(struct.pack("=I", len(data)) + data + payload) + sys.stdout.buffer.flush() + + +def main(): + while True: + command = receive() + name = command["command"] + if name.startswith("load_") or name == "reset": + send({"ok": True}) + elif name == "logits": + tokens = struct.unpack(f"={command['count']}i", read_exact(command["count"] * 4)) + send({"ok": True, "count": 2}, struct.pack("=HH", tokens[-1], len(tokens))) + elif name == "embed": + text = read_exact(command["size"]).decode() + send({"ok": True, "count": 2}, struct.pack("=ff", len(text), command["dimensions"])) + elif name == "exit": + return + elif name == "close": + send({"ok": True}) + return + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/src/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp new file mode 100644 index 000000000..7f902268e --- /dev/null +++ b/src/test/iron_backend/test_worker.cpp @@ -0,0 +1,33 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: MIT + +#include "iron/worker.hpp" + +#include +#include +#include +#include + +int main() { + setenv("FLM_IRON_WORKER_MODULE", "fake_worker", 1); + + flm::iron::Worker llama("/unused", 32); + const auto logits = llama.logits(std::vector{4, 7, 9}); + assert(logits == std::vector({9, 3})); + llama.reset(); + assert(!llama.poisoned()); + + flm::iron::Worker embedding("/unused", 512, "embedding"); + const auto result = embedding.embed("hello", "query", 128); + assert(result == std::vector({5.0f, 128.0f})); + + bool failed = false; + setenv("FLM_IRON_WORKER_MODULE", "missing_worker", 1); + try { + flm::iron::Worker dead("/unused", 32); + } catch (const std::runtime_error&) { + failed = true; + } + assert(failed); + return 0; +} \ No newline at end of file diff --git a/src/test/model_backend/test_model_backend.cpp b/src/test/model_backend/test_model_backend.cpp index df55b04f7..888b1134a 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, 32768}); + const auto traits = registry.traits("llama3", kIronBackendId); + TEST_REQUIRE(!traits.needs_npu_xclbin); + TEST_REQUIRE(!traits.supports_preemption); + TEST_REQUIRE(traits.max_context_length == 32768); +} + 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"); From 1bd97d8716ce2625aba978b527247b1e6e3b707c Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 7 Oct 2026 14:05:47 -0600 Subject: [PATCH 02/19] Fix IRON worker device setup --- docs/iron-backend.md | 23 ++++++++++++++++------- python/flm_iron/worker.py | 5 ++++- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index 182acfa41..fe87b5639 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -17,7 +17,7 @@ the source build: ```bash source /opt/xilinx/xrt/setup.sh export FLM_IRON_PYTHON=/scratch/roesti/IRON/.venv/bin/python -export PYTHONPATH=/scratch/roesti/FastFlowLM/python +export PYTHONPATH=/scratch/roesti/FastFlowLM/python:${PYTHONPATH:-} ``` An installed build places `flm_iron` under `share/flm/python`. Add that directory @@ -31,11 +31,20 @@ metadata. IRON loads a separate checkpoint directory that contains ```bash export FLM_IRON_MODEL_PATH=/path/to/llama3.2-1b -flm run llama3.2:1b --backend iron --ctx-len 32768 +flm serve llama3.2:1b --backend iron --ctx-len 2048 + +curl -H 'Content-Type: application/json' \ + -d '{"model":"llama3.2:1b","messages":[{"role":"user","content":"Hello"}],"max_tokens":8}' \ + http://127.0.0.1:52625/v1/chat/completions ``` The context length must be a multiple of 2048 and cannot exceed 32768. -Preemption is unavailable. +XRT locks the model arena. Increase the process memlock limit before using a +context whose arena exceeds the current limit. Preemption is unavailable. + +The FastFlowLM tokenizer and chat metadata must name the same model variant as +the IRON checkpoint. In particular, use Llama 3.2 1B Instruct metadata with an +Instruct checkpoint, and base-model metadata with the base checkpoint. ## EmbeddingGemma @@ -52,7 +61,7 @@ encoder. The worker returns the normalized 768-element float32 embedding. ## Protocol FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary -payloads over pipes. Llama requests contain the complete token history and -return bfloat16 logits. Embedding requests contain UTF-8 text and return -float32 values. Worker errors poison the backend, so FastFlowLM requires a model -reload after a process or protocol failure. \ No newline at end of file +payloads over a Unix socket pair. Llama requests contain the complete token +history and return bfloat16 logits. Embedding requests contain UTF-8 text and +return float32 values. Worker errors poison the backend, so FastFlowLM requires +a model reload after a process or protocol failure. diff --git a/python/flm_iron/worker.py b/python/flm_iron/worker.py index 99f975d56..27b6227d5 100644 --- a/python/flm_iron/worker.py +++ b/python/flm_iron/worker.py @@ -9,6 +9,8 @@ from pathlib import Path import numpy as np +import aie.utils as aie_utils +from aie.iron.device import from_name from ml_dtypes import bfloat16 @@ -52,6 +54,7 @@ def load_embedding(command): def main(): + aie_utils.set_current_device(from_name("npu2", n_cols=8)) model = None while True: command = receive() @@ -88,4 +91,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() From 61104820b95e72d67c0c80f22bbd5e291cfde2ef Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 7 Oct 2026 14:34:25 -0600 Subject: [PATCH 03/19] Add EmbeddingGemma 2 IRON model --- docs/iron-backend.md | 13 ++++-- .../all_embedding_model.hpp | 18 +++++--- .../modeling_iron_embedding.hpp | 2 +- src/include/model_list.hpp | 5 +- src/model_info.json | 26 +++++++++++ src/model_list.json | 46 +++++++++++++++++++ src/server/rest_handler.cpp | 12 +++-- .../test_model_list_platform.cpp | 17 +++++++ 8 files changed, 124 insertions(+), 15 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index fe87b5639..be724679c 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -46,13 +46,18 @@ The FastFlowLM tokenizer and chat metadata must name the same model variant as the IRON checkpoint. In particular, use Llama 3.2 1B Instruct metadata with an Instruct checkpoint, and base-model metadata with the base checkpoint. -## EmbeddingGemma +## EmbeddingGemma 2 -IRON loads a directory that contains `model.safetensors` and `tokenizer.json`: +EmbeddingGemma 2 is available only in builds that enable the IRON backend. Pull +the model from its public repository and serve it as an embedding-only model: ```bash -export FLM_IRON_EMBEDDING_PATH=/path/to/embeddinggemma-2 -flm serve llama3.2:1b --embed 1 --backend iron +flm pull embeddinggemma-2 +flm serve embeddinggemma-2 --backend iron + +curl -H 'Content-Type: application/json' \ + -d '{"model":"embeddinggemma-2","input":"Hello"}' \ + http://127.0.0.1:52625/v1/embeddings ``` The embedding endpoint passes text and its query or document task to the IRON diff --git a/src/include/AutoEmbeddingModel/all_embedding_model.hpp b/src/include/AutoEmbeddingModel/all_embedding_model.hpp index 3755d4fc3..c44f997eb 100644 --- a/src/include/AutoEmbeddingModel/all_embedding_model.hpp +++ b/src/include/AutoEmbeddingModel/all_embedding_model.hpp @@ -18,6 +18,8 @@ inline std::string complete_simple_embedding_tag(std::string model_tag) { if (model_tag == "embed-gemma:300m") return "embed-gemma:300m"; + else if (model_tag == "embeddinggemma-2") + return "embeddinggemma-2:740m"; else return model_tag; } @@ -33,12 +35,16 @@ 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 defined(FLM_ENABLE_IRON) - if (backend == flm::backend::kIronBackendId) - auto_embedding_engine = std::make_unique(npu_device_inst); - else - #endif - 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); + 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/AutoEmbeddingModel/modeling_iron_embedding.hpp b/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp index f1f81c5dd..a492dd423 100644 --- a/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp +++ b/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp @@ -9,7 +9,7 @@ class IronEmbedding final : public AutoEmbeddingModel { public: explicit IronEmbedding(flm_rt::device* device) - : AutoEmbeddingModel(device, "embed-gemma:300m") {} + : AutoEmbeddingModel(device, "embeddinggemma-2:740m") {} void load_model(std::string model_path, json, bool enable_preemption) override; std::vector embed(std::string& text, diff --git a/src/include/model_list.hpp b/src/include/model_list.hpp index 225ddbfd1..fed0a5697 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,13 +476,14 @@ 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) { + if (model_type == "embeddinggemma-2") return "iron"; static const std::string suffix = "-core"; if (model_type.size() > suffix.size() && model_type.compare(model_type.size() - suffix.size(), diff --git a/src/model_info.json b/src/model_info.json index a698aa7aa..9cc19e04f 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", diff --git a/src/model_list.json b/src/model_list.json index f645763f8..13ceec19d 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"], diff --git a/src/server/rest_handler.cpp b/src/server/rest_handler.cpp index cc779cd41..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"; 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..8e834c8af 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,7 @@ 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"; /// \brief every generation the build knows about, for the catalog sweep const std::set kAllPlatforms = { @@ -116,6 +117,20 @@ 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_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 @@ -324,6 +339,8 @@ 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_the_tag_name_decides_the_backend, "the tag name decides the backend"); RunTest(test_a_catalog_can_filter_down_to_nothing, From 469884fd14bb888d048e17d4d368ba22f60d2e65 Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 28 Oct 2026 11:23:37 -0600 Subject: [PATCH 04/19] Port IRON backend to graph model APIs --- docs/iron-backend.md | 13 ++++++----- python/flm_iron/worker.py | 23 ++++++++----------- .../modeling_iron_embedding.cpp | 2 +- .../llama3/iron/llama3_iron_backend.cpp | 5 ++-- .../llama3/iron/llama3_iron_backend.hpp | 2 +- src/test/model_backend/test_model_backend.cpp | 4 ++-- 6 files changed, 22 insertions(+), 27 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index be724679c..1faea7f3e 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -25,9 +25,9 @@ to `PYTHONPATH`. ## Llama 3.2 1B -The FastFlowLM package supplies `config.json`, tokenizer files, and chat -metadata. IRON loads a separate checkpoint directory that contains -`model.safetensors` and `tokenizer.model`: +The FastFlowLM package supplies chat metadata and the tokenizer for its HTTP +API. IRON loads a separate checkpoint directory. The directory contains +`config.json`, `tokenizer.json`, and one or more `model*.safetensors` files: ```bash export FLM_IRON_MODEL_PATH=/path/to/llama3.2-1b @@ -38,9 +38,10 @@ curl -H 'Content-Type: application/json' \ http://127.0.0.1:52625/v1/chat/completions ``` -The context length must be a multiple of 2048 and cannot exceed 32768. -XRT locks the model arena. Increase the process memlock limit before using a -context whose arena exceeds the current limit. Preemption is unavailable. +The context length must be positive and cannot exceed the checkpoint's +`max_position_embeddings`. XRT locks the model arena. Increase the process +memlock limit before using a context whose arena exceeds the current limit. +Preemption is unavailable. The FastFlowLM tokenizer and chat metadata must name the same model variant as the IRON checkpoint. In particular, use Llama 3.2 1B Instruct metadata with an diff --git a/python/flm_iron/worker.py b/python/flm_iron/worker.py index 27b6227d5..6002454ee 100644 --- a/python/flm_iron/worker.py +++ b/python/flm_iron/worker.py @@ -2,14 +2,11 @@ # SPDX-License-Identifier: MIT import json -import os import struct import sys -from dataclasses import replace -from pathlib import Path -import numpy as np import aie.utils as aie_utils +import numpy as np from aie.iron.device import from_name from ml_dtypes import bfloat16 @@ -38,19 +35,17 @@ def send(header, payload=b""): def load_llama(command): - from iron.lm.llama3.model import LLAMA_3_2_1B, Runner + from iron.applications.llama import Llama - root = Path(command["model_path"]) - weights = root / "model.safetensors" - tokenizer = root / "tokenizer.model" - config = replace(LLAMA_3_2_1B, max_seq_len=command["context_length"]) - return Runner(weights, tokenizer, config).npu() + return Llama.load( + command["model_path"], max_seq_len=command["context_length"] + ) def load_embedding(command): - from iron.lm.embeddinggemma2.encoder import Encoder + from iron.applications.embedding_gemma2 import EmbeddingGemma - return Encoder(command["model_path"], max_tokens=command["context_length"]) + return EmbeddingGemma.load(command["model_path"]) def main(): @@ -72,12 +67,12 @@ def main(): words = logits.view(np.uint16).astype("(model_path, 512, "embedding"); + worker_ = std::make_unique(model_path, 2048, "embedding"); is_model_loaded = true; this->model_path = std::move(model_path); } diff --git a/src/common/models/llama3/iron/llama3_iron_backend.cpp b/src/common/models/llama3/iron/llama3_iron_backend.cpp index 54c9aaa53..1208648b9 100644 --- a/src/common/models/llama3/iron/llama3_iron_backend.cpp +++ b/src/common/models/llama3/iron/llama3_iron_backend.cpp @@ -83,9 +83,8 @@ class IronBackend final : public flm::backend::ModelBackend { return path && *path ? path : fallback; } static std::uint32_t ValidateContextLength(std::uint32_t length) { - if (length == 0 || length > 32768 || length % 2048 != 0) { - throw std::out_of_range( - "IRON Llama context length must be a multiple of 2048 in 2048..32768"); + if (length == 0) { + throw std::out_of_range("IRON Llama context length must be positive"); } return length; } diff --git a/src/include/models/llama3/iron/llama3_iron_backend.hpp b/src/include/models/llama3/iron/llama3_iron_backend.hpp index 7e1c13881..80b3b9d38 100644 --- a/src/include/models/llama3/iron/llama3_iron_backend.hpp +++ b/src/include/models/llama3/iron/llama3_iron_backend.hpp @@ -10,7 +10,7 @@ namespace flm::llama3 { flm::backend::BackendFactory iron_factory(); inline flm::backend::BackendTraits iron_traits() { - return flm::backend::BackendTraits{false, false, 32768}; + return flm::backend::BackendTraits{false, false, 0}; } } // namespace flm::llama3 \ No newline at end of file diff --git a/src/test/model_backend/test_model_backend.cpp b/src/test/model_backend/test_model_backend.cpp index 888b1134a..498a53849 100644 --- a/src/test/model_backend/test_model_backend.cpp +++ b/src/test/model_backend/test_model_backend.cpp @@ -123,11 +123,11 @@ void test_traits_are_kept_per_backend() { void test_iron_backend_registration() { auto registry = MakeRegistry(); registry.register_backend("llama3", kIronBackendId, StubFactory("iron"), - BackendTraits{false, false, 32768}); + 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 == 32768); + TEST_REQUIRE(traits.max_context_length == 0); } void test_backend_defaults() { From 72313491a7bc8672a368c993b8270d9a0a18f9f5 Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 28 Oct 2026 16:07:42 -0600 Subject: [PATCH 05/19] Use IRON graph model server --- docs/iron-backend.md | 18 ++-- python/flm_iron/__init__.py | 4 - python/flm_iron/worker.py | 89 ------------------- src/CMakeLists.txt | 6 +- .../modeling_iron_embedding.cpp | 2 +- src/common/iron/worker.cpp | 14 +-- .../llama3/iron/llama3_iron_backend.cpp | 3 +- src/include/iron/worker.hpp | 4 +- src/test/iron_backend/fake_worker.py | 9 +- src/test/iron_backend/test_worker.cpp | 6 +- 10 files changed, 35 insertions(+), 120 deletions(-) delete mode 100644 python/flm_iron/__init__.py delete mode 100644 python/flm_iron/worker.py diff --git a/docs/iron-backend.md b/docs/iron-backend.md index 1faea7f3e..69bb3e545 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -5,7 +5,8 @@ SPDX-License-Identifier: MIT # IRON backend -The `iron` backend runs a complete IRON graph in a persistent Python process. +The `iron` backend runs a complete IRON graph through IRON's persistent Python +server. FastFlowLM handles tokenization, chat templates, sampling, and HTTP requests. The worker retains the IRON model, compiled images, weights, and device state. @@ -17,11 +18,11 @@ the source build: ```bash source /opt/xilinx/xrt/setup.sh export FLM_IRON_PYTHON=/scratch/roesti/IRON/.venv/bin/python -export PYTHONPATH=/scratch/roesti/FastFlowLM/python:${PYTHONPATH:-} +export PYTHONPATH=/scratch/roesti/IRON:${PYTHONPATH:-} ``` -An installed build places `flm_iron` under `share/flm/python`. Add that directory -to `PYTHONPATH`. +`FLM_IRON_PYTHON` must name an environment with IRON installed. `PYTHONPATH` is +needed only when that environment uses an IRON source checkout. ## Llama 3.2 1B @@ -67,7 +68,8 @@ encoder. The worker returns the normalized 768-element float32 embedding. ## Protocol FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary -payloads over a Unix socket pair. Llama requests contain the complete token -history and return bfloat16 logits. Embedding requests contain UTF-8 text and -return float32 values. Worker errors poison the backend, so FastFlowLM requires -a model reload after a process or protocol failure. +payloads over a Unix socket pair. The load request names a model in +`iron.lm.server.MODELS`. Llama requests contain the complete token history and +return bfloat16 logits. EmbeddingGemma 2 requests contain UTF-8 text and return +float32 values. Worker errors poison the backend, so FastFlowLM requires a model +reload after a process or protocol failure. diff --git a/python/flm_iron/__init__.py b/python/flm_iron/__init__.py deleted file mode 100644 index 82f81270c..000000000 --- a/python/flm_iron/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: MIT - -"""FastFlowLM transport for complete IRON graphs.""" \ No newline at end of file diff --git a/python/flm_iron/worker.py b/python/flm_iron/worker.py deleted file mode 100644 index 6002454ee..000000000 --- a/python/flm_iron/worker.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: MIT - -import json -import struct -import sys - -import aie.utils as aie_utils -import numpy as np -from aie.iron.device import from_name -from ml_dtypes import bfloat16 - - -def read_exact(size): - data = bytearray() - while len(data) < size: - chunk = sys.stdin.buffer.read(size - len(data)) - if not chunk: - raise EOFError - data.extend(chunk) - return bytes(data) - - -def receive(): - size = struct.unpack("=I", read_exact(4))[0] - return json.loads(read_exact(size)) - - -def send(header, payload=b""): - data = json.dumps(header, separators=(",", ":")).encode() - sys.stdout.buffer.write(struct.pack("=I", len(data))) - sys.stdout.buffer.write(data) - sys.stdout.buffer.write(payload) - sys.stdout.buffer.flush() - - -def load_llama(command): - from iron.applications.llama import Llama - - return Llama.load( - command["model_path"], max_seq_len=command["context_length"] - ) - - -def load_embedding(command): - from iron.applications.embedding_gemma2 import EmbeddingGemma - - return EmbeddingGemma.load(command["model_path"]) - - -def main(): - aie_utils.set_current_device(from_name("npu2", n_cols=8)) - model = None - while True: - command = receive() - try: - name = command["command"] - if name == "load_llama": - model = load_llama(command) - send({"ok": True}) - elif name == "load_embedding": - model = load_embedding(command) - send({"ok": True}) - elif name == "logits": - tokens = np.frombuffer(read_exact(command["count"] * 4), np.int32) - logits = np.asarray(model.logits(tokens), dtype=bfloat16) - words = logits.view(np.uint16).astype("(model_path, 2048, "embedding"); + worker_ = std::make_unique(model_path, "embedding_gemma2"); is_model_loaded = true; this->model_path = std::move(model_path); } diff --git a/src/common/iron/worker.cpp b/src/common/iron/worker.cpp index 8e26f9888..2f28258bb 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -27,13 +27,13 @@ std::string PythonExecutable() { std::string WorkerModule() { const char* value = std::getenv("FLM_IRON_WORKER_MODULE"); - return value && *value ? value : "flm_iron.worker"; + return value && *value ? value : "iron.lm.server"; } } // namespace -Worker::Worker(const std::filesystem::path& model_path, - std::uint32_t context_length, std::string model) { +Worker::Worker(const std::filesystem::path& model_path, std::string model, + std::uint32_t context_length) { #if defined(_WIN32) throw std::runtime_error("The IRON backend requires POSIX process support"); #else @@ -57,9 +57,11 @@ Worker::Worker(const std::filesystem::path& model_path, close(sockets[1]); input_ = output_ = sockets[0]; try { - request({{"command", model == "embedding" ? "load_embedding" : "load_llama"}, - {"model_path", model_path.string()}, - {"context_length", context_length}}); + nlohmann::json command = {{"command", "load_model"}, + {"model", std::move(model)}, + {"model_path", model_path.string()}}; + if (context_length) command["context_length"] = context_length; + request(command); } catch (...) { close(input_); waitpid(process_, nullptr, 0); diff --git a/src/common/models/llama3/iron/llama3_iron_backend.cpp b/src/common/models/llama3/iron/llama3_iron_backend.cpp index 1208648b9..a2725ba43 100644 --- a/src/common/models/llama3/iron/llama3_iron_backend.cpp +++ b/src/common/models/llama3/iron/llama3_iron_backend.cpp @@ -17,7 +17,8 @@ namespace { class IronLlama final : public causal_lm { public: IronLlama(const std::string& model_path, std::uint32_t context_length) - : worker_(model_path, context_length), max_length_(context_length) {} + : worker_(model_path, "llama", context_length), + max_length_(context_length) {} buffer forward(int id) override { tokens_.push_back(id); diff --git a/src/include/iron/worker.hpp b/src/include/iron/worker.hpp index ebe49876d..c2b6d05d7 100644 --- a/src/include/iron/worker.hpp +++ b/src/include/iron/worker.hpp @@ -13,8 +13,8 @@ namespace flm::iron { class Worker { public: - Worker(const std::filesystem::path& model_path, std::uint32_t context_length, - std::string model = "llama"); + Worker(const std::filesystem::path& model_path, std::string model, + std::uint32_t context_length = 0); ~Worker(); Worker(const Worker&) = delete; diff --git a/src/test/iron_backend/fake_worker.py b/src/test/iron_backend/fake_worker.py index 076a8d5d4..a36d89e86 100644 --- a/src/test/iron_backend/fake_worker.py +++ b/src/test/iron_backend/fake_worker.py @@ -31,7 +31,14 @@ def main(): while True: command = receive() name = command["command"] - if name.startswith("load_") or name == "reset": + if name == "load_model": + assert command["model"] in {"embedding_gemma2", "llama"} + if command["model"] == "llama": + assert command["context_length"] == 32 + else: + assert "context_length" not in command + send({"ok": True}) + elif name == "reset": send({"ok": True}) elif name == "logits": tokens = struct.unpack(f"={command['count']}i", read_exact(command["count"] * 4)) diff --git a/src/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp index 7f902268e..d20616fea 100644 --- a/src/test/iron_backend/test_worker.cpp +++ b/src/test/iron_backend/test_worker.cpp @@ -11,20 +11,20 @@ int main() { setenv("FLM_IRON_WORKER_MODULE", "fake_worker", 1); - flm::iron::Worker llama("/unused", 32); + flm::iron::Worker llama("/unused", "llama", 32); const auto logits = llama.logits(std::vector{4, 7, 9}); assert(logits == std::vector({9, 3})); llama.reset(); assert(!llama.poisoned()); - flm::iron::Worker embedding("/unused", 512, "embedding"); + flm::iron::Worker embedding("/unused", "embedding_gemma2"); const auto result = embedding.embed("hello", "query", 128); assert(result == std::vector({5.0f, 128.0f})); bool failed = false; setenv("FLM_IRON_WORKER_MODULE", "missing_worker", 1); try { - flm::iron::Worker dead("/unused", 32); + flm::iron::Worker dead("/unused", "llama", 32); } catch (const std::runtime_error&) { failed = true; } From a428b9906a00d86cce331fe1dff6bfec38d188cf Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 28 Oct 2026 16:19:29 -0600 Subject: [PATCH 06/19] Run the top-level IRON server --- docs/iron-backend.md | 2 +- src/common/iron/worker.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index 69bb3e545..e4d4a608d 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -69,7 +69,7 @@ encoder. The worker returns the normalized 768-element float32 embedding. FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary payloads over a Unix socket pair. The load request names a model in -`iron.lm.server.MODELS`. Llama requests contain the complete token history and +`iron.server.MODELS`. Llama requests contain the complete token history and return bfloat16 logits. EmbeddingGemma 2 requests contain UTF-8 text and return float32 values. Worker errors poison the backend, so FastFlowLM requires a model reload after a process or protocol failure. diff --git a/src/common/iron/worker.cpp b/src/common/iron/worker.cpp index 2f28258bb..f11e4e208 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -27,7 +27,7 @@ std::string PythonExecutable() { std::string WorkerModule() { const char* value = std::getenv("FLM_IRON_WORKER_MODULE"); - return value && *value ? value : "iron.lm.server"; + return value && *value ? value : "iron.server"; } } // namespace From 91e7fb8457583c53237773d964f6293a651ff1ae Mon Sep 17 00:00:00 2001 From: andrej Date: Wed, 28 Oct 2026 16:41:14 -0600 Subject: [PATCH 07/19] Rename the IRON language model server --- docs/iron-backend.md | 2 +- src/common/iron/worker.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index e4d4a608d..16e811e6d 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -69,7 +69,7 @@ encoder. The worker returns the normalized 768-element float32 embedding. FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary payloads over a Unix socket pair. The load request names a model in -`iron.server.MODELS`. Llama requests contain the complete token history and +`iron.applications.MODELS`. Llama requests contain the complete token history and return bfloat16 logits. EmbeddingGemma 2 requests contain UTF-8 text and return float32 values. Worker errors poison the backend, so FastFlowLM requires a model reload after a process or protocol failure. diff --git a/src/common/iron/worker.cpp b/src/common/iron/worker.cpp index f11e4e208..56928ff0b 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -27,7 +27,7 @@ std::string PythonExecutable() { std::string WorkerModule() { const char* value = std::getenv("FLM_IRON_WORKER_MODULE"); - return value && *value ? value : "iron.server"; + return value && *value ? value : "iron.lm_server"; } } // namespace From 6ae2b2d49b7f9fbebead3fd84cc2e1cd9158b0dd Mon Sep 17 00:00:00 2001 From: andrej Date: Thu, 29 Oct 2026 10:11:28 -0600 Subject: [PATCH 08/19] Use Boost.Process for the IRON worker --- docs/iron-backend.md | 10 +-- src/CMakeLists.txt | 2 +- src/common/iron/worker.cpp | 123 ++++++++++++++------------ src/include/iron/worker.hpp | 6 +- src/test/iron_backend/CMakeLists.txt | 13 ++- src/test/iron_backend/fake_worker.py | 9 +- src/test/iron_backend/test_worker.cpp | 29 +++++- 7 files changed, 117 insertions(+), 75 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index 16e811e6d..52c43e68d 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -68,8 +68,8 @@ encoder. The worker returns the normalized 768-element float32 embedding. ## Protocol FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary -payloads over a Unix socket pair. The load request names a model in -`iron.applications.MODELS`. Llama requests contain the complete token history and -return bfloat16 logits. EmbeddingGemma 2 requests contain UTF-8 text and return -float32 values. Worker errors poison the backend, so FastFlowLM requires a model -reload after a process or protocol failure. +payloads through the worker's standard input and output pipes. The load request +names a model in `iron.applications.MODELS`. Llama requests contain the complete +token history and return bfloat16 logits. EmbeddingGemma 2 requests contain UTF-8 +text and return float32 values. Worker errors poison the backend, so FastFlowLM +requires a model reload after a process or protocol failure. diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7fe688abf..cb0a2538f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -714,7 +714,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/iron/worker.cpp b/src/common/iron/worker.cpp index 56928ff0b..42d3e3a9c 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -3,16 +3,13 @@ #include "iron/worker.hpp" +#define BOOST_PROCESS_USE_STD_FS +#include + #if !defined(_WIN32) -#include -#include -#include +#include #endif - -#include -#include #include -#include #include namespace flm::iron { @@ -22,7 +19,12 @@ constexpr std::uint32_t kMaxHeader = 1 << 20; std::string PythonExecutable() { const char* value = std::getenv("FLM_IRON_PYTHON"); - return value && *value ? value : "python3"; + if (value && *value) return value; +#if defined(_WIN32) + return "python"; +#else + return "python3"; +#endif } std::string WorkerModule() { @@ -30,32 +32,48 @@ std::string WorkerModule() { 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 -Worker::Worker(const std::filesystem::path& model_path, std::string model, - std::uint32_t context_length) { -#if defined(_WIN32) - throw std::runtime_error("The IRON backend requires POSIX process support"); -#else - int sockets[2]; - if (socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) != 0) { - throw std::runtime_error("Cannot create the IRON worker socket"); +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 (...) {} } - process_ = fork(); - if (process_ < 0) throw std::runtime_error("Cannot start the IRON worker"); - if (process_ == 0) { - dup2(sockets[1], STDIN_FILENO); - dup2(sockets[1], STDOUT_FILENO); - close(sockets[0]); - close(sockets[1]); - const std::string python = PythonExecutable(); - const std::string module = WorkerModule(); - execlp(python.c_str(), python.c_str(), "-m", module.c_str(), - static_cast(nullptr)); - _exit(127); + + void stop() noexcept { + close_pipes(); + std::error_code error; + if (process.running(error)) process.terminate(error); + process.wait(error); } - close(sockets[1]); - input_ = output_ = sockets[0]; +}; + +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"}, {"model", std::move(model)}, @@ -63,46 +81,33 @@ Worker::Worker(const std::filesystem::path& model_path, std::string model, if (context_length) command["context_length"] = context_length; request(command); } catch (...) { - close(input_); - waitpid(process_, nullptr, 0); - input_ = output_ = process_ = -1; + impl_->stop(); throw; } -#endif } Worker::~Worker() { -#if !defined(_WIN32) - if (input_ >= 0) { - try { request({{"command", "close"}}); } catch (...) {} - close(input_); + if (impl_) { + try { + request({{"command", "close"}}); + } catch (...) { + impl_->stop(); + return; + } + impl_->close_pipes(); + std::error_code error; + impl_->process.wait(error); } - if (process_ > 0) waitpid(process_, nullptr, 0); -#endif } void Worker::write_all(const void* data, std::size_t size) { -#if !defined(_WIN32) - const auto* cursor = static_cast(data); - while (size) { - const ssize_t written = send(input_, cursor, size, MSG_NOSIGNAL); - if (written <= 0) throw std::runtime_error("IRON worker write failed"); - cursor += written; - size -= static_cast(written); - } -#endif + impl_->input.write(static_cast(data), size); + if (!impl_->input) throw std::runtime_error("IRON worker write failed"); } void Worker::read_all(void* data, std::size_t size) { -#if !defined(_WIN32) - auto* cursor = static_cast(data); - while (size) { - const ssize_t received = read(output_, cursor, size); - if (received <= 0) throw std::runtime_error("IRON worker exited"); - cursor += received; - size -= static_cast(received); - } -#endif + impl_->output.read(static_cast(data), size); + if (!impl_->output) throw std::runtime_error("IRON worker exited"); } nlohmann::json Worker::request(const nlohmann::json& command, @@ -114,6 +119,8 @@ nlohmann::json Worker::request(const nlohmann::json& command, 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 worker write failed"); std::uint32_t response_size = 0; read_all(&response_size, sizeof(response_size)); diff --git a/src/include/iron/worker.hpp b/src/include/iron/worker.hpp index c2b6d05d7..ce078be25 100644 --- a/src/include/iron/worker.hpp +++ b/src/include/iron/worker.hpp @@ -6,6 +6,7 @@ #include #include +#include #include #include @@ -33,9 +34,8 @@ class Worker { void write_all(const void* data, std::size_t size); void read_all(void* data, std::size_t size); - int input_ = -1; - int output_ = -1; - int process_ = -1; + struct Impl; + std::unique_ptr impl_; bool poisoned_ = false; }; diff --git a/src/test/iron_backend/CMakeLists.txt b/src/test/iron_backend/CMakeLists.txt index d7518f304..14d176e27 100644 --- a/src/test/iron_backend/CMakeLists.txt +++ b/src/test/iron_backend/CMakeLists.txt @@ -7,13 +7,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) +find_package(Python3 REQUIRED COMPONENTS Interpreter) 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") +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 - ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_LIST_DIR}") \ No newline at end of file + ENVIRONMENT + "FLM_IRON_PYTHON=${Python3_EXECUTABLE};PYTHONPATH=${CMAKE_CURRENT_LIST_DIR}") \ No newline at end of file diff --git a/src/test/iron_backend/fake_worker.py b/src/test/iron_backend/fake_worker.py index a36d89e86..d0cf52232 100644 --- a/src/test/iron_backend/fake_worker.py +++ b/src/test/iron_backend/fake_worker.py @@ -29,15 +29,20 @@ def send(header, payload=b""): def main(): while True: - command = receive() + try: + command = receive() + except EOFError: + return name = command["command"] if name == "load_model": - assert command["model"] in {"embedding_gemma2", "llama"} + assert command["model"] in {"dies", "embedding_gemma2", "llama"} if command["model"] == "llama": assert command["context_length"] == 32 else: assert "context_length" not in command send({"ok": True}) + if command["model"] == "dies": + return elif name == "reset": send({"ok": True}) elif name == "logits": diff --git a/src/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp index d20616fea..e96f49dda 100644 --- a/src/test/iron_backend/test_worker.cpp +++ b/src/test/iron_backend/test_worker.cpp @@ -8,12 +8,24 @@ #include #include +namespace { + +void SetEnvironment(const char* name, const char* value) { +#if defined(_WIN32) + _putenv_s(name, value); +#else + setenv(name, value, 1); +#endif +} + +} // namespace + int main() { - setenv("FLM_IRON_WORKER_MODULE", "fake_worker", 1); + SetEnvironment("FLM_IRON_WORKER_MODULE", "fake_worker"); flm::iron::Worker llama("/unused", "llama", 32); - const auto logits = llama.logits(std::vector{4, 7, 9}); - assert(logits == std::vector({9, 3})); + const auto logits = llama.logits(std::vector{4, 10, 26}); + assert(logits == std::vector({26, 3})); llama.reset(); assert(!llama.poisoned()); @@ -21,8 +33,17 @@ int main() { const auto result = embedding.embed("hello", "query", 128); assert(result == std::vector({5.0f, 128.0f})); + bool pipe_failed = false; + try { + flm::iron::Worker exits("/unused", "dies"); + exits.reset(); + } catch (const std::runtime_error&) { + pipe_failed = true; + } + assert(pipe_failed); + bool failed = false; - setenv("FLM_IRON_WORKER_MODULE", "missing_worker", 1); + SetEnvironment("FLM_IRON_WORKER_MODULE", "missing_worker"); try { flm::iron::Worker dead("/unused", "llama", 32); } catch (const std::runtime_error&) { From 92525638313ee2e737946e560a9b87ed8a383e59 Mon Sep 17 00:00:00 2001 From: andrej Date: Thu, 29 Oct 2026 10:55:36 -0600 Subject: [PATCH 09/19] Register BF16 Llama for the IRON backend --- docs/iron-backend.md | 12 +++--- src/include/model_list.hpp | 14 +++++-- src/model_info.json | 26 ++++++++++++ src/model_list.json | 40 +++++++++++++++++++ .../test_model_list_platform.cpp | 35 ++++++++++++++-- 5 files changed, 113 insertions(+), 14 deletions(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index 52c43e68d..eb780e76f 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -26,16 +26,16 @@ needed only when that environment uses an IRON source checkout. ## Llama 3.2 1B -The FastFlowLM package supplies chat metadata and the tokenizer for its HTTP -API. IRON loads a separate checkpoint directory. The directory contains -`config.json`, `tokenizer.json`, and one or more `model*.safetensors` files: +Accept the license for `meta-llama/Llama-3.2-1B`, then provide a Hugging Face +token through `HF_TOKEN`, `HUGGING_FACE_HUB_TOKEN` or `HF_TOKEN_PATH`. Pull and +serve the BF16 model registered for IRON: ```bash -export FLM_IRON_MODEL_PATH=/path/to/llama3.2-1b -flm serve llama3.2:1b --backend iron --ctx-len 2048 +flm pull llama3.2-iron:1b +flm serve llama3.2-iron:1b --ctx-len 2048 curl -H 'Content-Type: application/json' \ - -d '{"model":"llama3.2:1b","messages":[{"role":"user","content":"Hello"}],"max_tokens":8}' \ + -d '{"model":"llama3.2-iron:1b","messages":[{"role":"user","content":"Hello"}],"max_tokens":8}' \ http://127.0.0.1:52625/v1/chat/completions ``` diff --git a/src/include/model_list.hpp b/src/include/model_list.hpp index fed0a5697..cbc222d2d 100644 --- a/src/include/model_list.hpp +++ b/src/include/model_list.hpp @@ -484,10 +484,16 @@ class model_list { /// and there is nothing left for a catalog key to disambiguate. static std::string backend_for_family(const std::string& model_type) { if (model_type == "embeddinggemma-2") return "iron"; - static const std::string suffix = "-core"; - if (model_type.size() > suffix.size() && - model_type.compare(model_type.size() - suffix.size(), - suffix.size(), suffix) == 0) { + 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/model_info.json b/src/model_info.json index 9cc19e04f..5b620b5b2 100644 --- a/src/model_info.json +++ b/src/model_info.json @@ -2490,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": "68a2e4be76fa709455a60272fba8e512c02d81c46e6c671cc9449e374fd6809a", + "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 13ceec19d..25e35e254 100644 --- a/src/model_list.json +++ b/src/model_list.json @@ -1482,6 +1482,46 @@ "footprint": 2.7 } }, + "llama3.2-iron": { + "1b": { + "supported_platforms": ["aie2p"], + "name": "llama3.2-1b", + "url": "https://huggingface.co/meta-llama/Llama-3.2-1B", + "file_url": "https://huggingface.co/api/models/meta-llama/Llama-3.2-1B/tree/main", + "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/test/model_list_platform/test_model_list_platform.cpp b/src/test/model_list_platform/test_model_list_platform.cpp index 8e834c8af..5cb6e612e 100644 --- a/src/test/model_list_platform/test_model_list_platform.cpp +++ b/src/test/model_list_platform/test_model_list_platform.cpp @@ -20,6 +20,7 @@ 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 = { @@ -131,6 +132,25 @@ void test_embeddinggemma_2_requires_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"); + 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 @@ -156,15 +176,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 = "."; @@ -174,6 +194,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 @@ -190,6 +211,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); } @@ -341,6 +366,8 @@ int main() { "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, From 8a9de352fd7dfe8797385526069f70445af4132f Mon Sep 17 00:00:00 2001 From: andrej Date: Thu, 29 Oct 2026 11:55:39 -0600 Subject: [PATCH 10/19] Add IRON home installer mode --- docs/iron-backend.md | 12 ++++++++++++ src/home_install.sh | 34 +++++++++++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/docs/iron-backend.md b/docs/iron-backend.md index eb780e76f..c1054a153 100644 --- a/docs/iron-backend.md +++ b/docs/iron-backend.md @@ -24,6 +24,18 @@ export PYTHONPATH=/scratch/roesti/IRON:${PYTHONPATH:-} `FLM_IRON_PYTHON` must name an environment with IRON installed. `PYTHONPATH` is needed only when that environment uses an IRON source checkout. +For a user install, activate the IRON environment and let the installer record +its Python executable: + +```bash +source /path/to/iron/.venv/bin/activate +./src/home_install.sh --iron +source ~/flm_exe/flm_env.sh +``` + +The generated environment script configures FastFlowLM, XRT, and the IRON +worker. An editable IRON installation also removes the need for `PYTHONPATH`. + ## Llama 3.2 1B Accept the license for `meta-llama/Llama-3.2-1B`, then provide a Hugging Face diff --git a/src/home_install.sh b/src/home_install.sh index 1e6198c57..790e48e9b 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 Python # 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 a Python environment that can import iron.lm_server." >&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" < Date: Thu, 29 Oct 2026 15:26:23 -0600 Subject: [PATCH 11/19] Check the Rust toolchain during configuration --- docs/docs/install_lin.md | 14 ++++++++++++-- docs/linux-getting-started.md | 8 ++++++-- src/CMakeLists.txt | 17 +++++++++++++++++ 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/docs/docs/install_lin.md b/docs/docs/install_lin.md index 62ae23706..63db68835 100644 --- a/docs/docs/install_lin.md +++ b/docs/docs/install_lin.md @@ -166,12 +166,22 @@ If `flm validate` passes but `flm run` fails with `No such device with index '0' ### Building from Source -1. Clone the repository: +1. Install the build dependencies and select a Rust toolchain: + ```sh + sudo apt install ninja libedit-dev libxrt-dev uuid-dev libdrm-dev rustup + rustup default stable + ``` + + FastFlowLM builds its tokenizer with Rust. If `cargo` exists but reports + that no default toolchain is configured, `rustup default stable` repairs the + installation. + +2. Clone the repository: ```sh git clone https://github.com/ROCm/FastFlowLM.git cd FastFlowLM ``` -2. Build: +3. Build: ```sh cd src cmake --preset linux-default diff --git a/docs/linux-getting-started.md b/docs/linux-getting-started.md index 20d7b2dcf..4693ec126 100644 --- a/docs/linux-getting-started.md +++ b/docs/linux-getting-started.md @@ -121,10 +121,14 @@ In short: if 1.1 firmware breaks probing on stock 6.19, do not keep forcing the 1. Ensure all required development packages are installed: ```sh - sudo apt install ninja - sudo apt install libedit-dev libxrt-dev uuid-dev libdrm-dev + sudo apt install ninja libedit-dev libxrt-dev uuid-dev libdrm-dev rustup + rustup default stable ``` + FastFlowLM builds its tokenizer with Rust. If `cargo` exists but reports + that no default toolchain is configured, `rustup default stable` repairs the + installation. + 2. Clone the repository and pull all submodules: ```sh git clone --recursive https://github.com/ROCm/FastFlowLM.git diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index cb0a2538f..b0ac5589c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -200,6 +200,23 @@ endif() # ——————————————————————————————————————————————— # Add tokenizers-cpp subproject # ——————————————————————————————————————————————— +find_program(CARGO_EXECUTABLE cargo REQUIRED) +execute_process( + COMMAND "${CARGO_EXECUTABLE}" --version + RESULT_VARIABLE CARGO_VERSION_RESULT + OUTPUT_VARIABLE CARGO_VERSION_OUTPUT + ERROR_VARIABLE CARGO_VERSION_ERROR + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_STRIP_TRAILING_WHITESPACE) +if(NOT CARGO_VERSION_RESULT EQUAL 0) + message(FATAL_ERROR + "Cargo was found at ${CARGO_EXECUTABLE}, but it cannot run:\n" + "${CARGO_VERSION_ERROR}\n" + "Install a Rust toolchain. If rustup manages Cargo, run: " + "rustup default stable") +endif() +message(STATUS "Cargo: ${CARGO_VERSION_OUTPUT} (${CARGO_EXECUTABLE})") + # Ensure C++17 is available for tokenizers-cpp if(CMAKE_CXX_STANDARD LESS 17) set(CMAKE_CXX_STANDARD 17) From bbeca746c555f012b749620f6efe16a6710b83eb Mon Sep 17 00:00:00 2001 From: andrej Date: Thu, 29 Oct 2026 16:29:35 -0600 Subject: [PATCH 12/19] Limit Rust setup docs to the source guide --- docs/docs/install_lin.md | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/docs/docs/install_lin.md b/docs/docs/install_lin.md index 63db68835..62ae23706 100644 --- a/docs/docs/install_lin.md +++ b/docs/docs/install_lin.md @@ -166,22 +166,12 @@ If `flm validate` passes but `flm run` fails with `No such device with index '0' ### Building from Source -1. Install the build dependencies and select a Rust toolchain: - ```sh - sudo apt install ninja libedit-dev libxrt-dev uuid-dev libdrm-dev rustup - rustup default stable - ``` - - FastFlowLM builds its tokenizer with Rust. If `cargo` exists but reports - that no default toolchain is configured, `rustup default stable` repairs the - installation. - -2. Clone the repository: +1. Clone the repository: ```sh git clone https://github.com/ROCm/FastFlowLM.git cd FastFlowLM ``` -3. Build: +2. Build: ```sh cd src cmake --preset linux-default From 33d18d94ac61bdb53a07dc992dd026af8a115b10 Mon Sep 17 00:00:00 2001 From: andrej Date: Thu, 29 Oct 2026 16:53:54 -0600 Subject: [PATCH 13/19] Share IRON model adapters --- docs/iron-backend.md | 87 -------------- docs/linux-getting-started.md | 8 +- src/CMakeLists.txt | 19 +--- .../modeling_iron_embedding.cpp | 24 ---- src/common/iron/causal_lm.cpp | 106 ++++++++++++++++++ src/common/iron/embedding.cpp | 30 +++++ .../llama3/iron/llama3_iron_backend.cpp | 93 +-------------- .../all_embedding_model.hpp | 6 +- .../modeling_iron_embedding.hpp | 20 ---- src/include/iron/causal_lm.hpp | 62 ++++++++++ src/include/iron/embedding.hpp | 28 +++++ .../llama3/iron/llama3_iron_backend.hpp | 3 +- 12 files changed, 238 insertions(+), 248 deletions(-) delete mode 100644 docs/iron-backend.md delete mode 100644 src/common/AutoEmbeddingModel/modeling_iron_embedding.cpp create mode 100644 src/common/iron/causal_lm.cpp create mode 100644 src/common/iron/embedding.cpp delete mode 100644 src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp create mode 100644 src/include/iron/causal_lm.hpp create mode 100644 src/include/iron/embedding.hpp diff --git a/docs/iron-backend.md b/docs/iron-backend.md deleted file mode 100644 index c1054a153..000000000 --- a/docs/iron-backend.md +++ /dev/null @@ -1,87 +0,0 @@ - - -# IRON backend - -The `iron` backend runs a complete IRON graph through IRON's persistent Python -server. -FastFlowLM handles tokenization, chat templates, sampling, and HTTP requests. -The worker retains the IRON model, compiled images, weights, and device state. - -## Build - -Configure FastFlowLM with `FLM_ENABLE_IRON=ON`. Set these variables when you run -the source build: - -```bash -source /opt/xilinx/xrt/setup.sh -export FLM_IRON_PYTHON=/scratch/roesti/IRON/.venv/bin/python -export PYTHONPATH=/scratch/roesti/IRON:${PYTHONPATH:-} -``` - -`FLM_IRON_PYTHON` must name an environment with IRON installed. `PYTHONPATH` is -needed only when that environment uses an IRON source checkout. - -For a user install, activate the IRON environment and let the installer record -its Python executable: - -```bash -source /path/to/iron/.venv/bin/activate -./src/home_install.sh --iron -source ~/flm_exe/flm_env.sh -``` - -The generated environment script configures FastFlowLM, XRT, and the IRON -worker. An editable IRON installation also removes the need for `PYTHONPATH`. - -## Llama 3.2 1B - -Accept the license for `meta-llama/Llama-3.2-1B`, then provide a Hugging Face -token through `HF_TOKEN`, `HUGGING_FACE_HUB_TOKEN` or `HF_TOKEN_PATH`. Pull and -serve the BF16 model registered for IRON: - -```bash -flm pull llama3.2-iron:1b -flm serve llama3.2-iron:1b --ctx-len 2048 - -curl -H 'Content-Type: application/json' \ - -d '{"model":"llama3.2-iron:1b","messages":[{"role":"user","content":"Hello"}],"max_tokens":8}' \ - http://127.0.0.1:52625/v1/chat/completions -``` - -The context length must be positive and cannot exceed the checkpoint's -`max_position_embeddings`. XRT locks the model arena. Increase the process -memlock limit before using a context whose arena exceeds the current limit. -Preemption is unavailable. - -The FastFlowLM tokenizer and chat metadata must name the same model variant as -the IRON checkpoint. In particular, use Llama 3.2 1B Instruct metadata with an -Instruct checkpoint, and base-model metadata with the base checkpoint. - -## EmbeddingGemma 2 - -EmbeddingGemma 2 is available only in builds that enable the IRON backend. Pull -the model from its public repository and serve it as an embedding-only model: - -```bash -flm pull embeddinggemma-2 -flm serve embeddinggemma-2 --backend iron - -curl -H 'Content-Type: application/json' \ - -d '{"model":"embeddinggemma-2","input":"Hello"}' \ - http://127.0.0.1:52625/v1/embeddings -``` - -The embedding endpoint passes text and its query or document task to the IRON -encoder. The worker returns the normalized 768-element float32 embedding. - -## Protocol - -FastFlowLM and the worker exchange length-prefixed JSON headers and raw binary -payloads through the worker's standard input and output pipes. The load request -names a model in `iron.applications.MODELS`. Llama requests contain the complete -token history and return bfloat16 logits. EmbeddingGemma 2 requests contain UTF-8 -text and return float32 values. Worker errors poison the backend, so FastFlowLM -requires a model reload after a process or protocol failure. diff --git a/docs/linux-getting-started.md b/docs/linux-getting-started.md index 4693ec126..20d7b2dcf 100644 --- a/docs/linux-getting-started.md +++ b/docs/linux-getting-started.md @@ -121,14 +121,10 @@ In short: if 1.1 firmware breaks probing on stock 6.19, do not keep forcing the 1. Ensure all required development packages are installed: ```sh - sudo apt install ninja libedit-dev libxrt-dev uuid-dev libdrm-dev rustup - rustup default stable + sudo apt install ninja + sudo apt install libedit-dev libxrt-dev uuid-dev libdrm-dev ``` - FastFlowLM builds its tokenizer with Rust. If `cargo` exists but reports - that no default toolchain is configured, `rustup default stable` repairs the - installation. - 2. Clone the repository and pull all submodules: ```sh git clone --recursive https://github.com/ROCm/FastFlowLM.git diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index b0ac5589c..c1c6cbfb9 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -200,23 +200,6 @@ endif() # ——————————————————————————————————————————————— # Add tokenizers-cpp subproject # ——————————————————————————————————————————————— -find_program(CARGO_EXECUTABLE cargo REQUIRED) -execute_process( - COMMAND "${CARGO_EXECUTABLE}" --version - RESULT_VARIABLE CARGO_VERSION_RESULT - OUTPUT_VARIABLE CARGO_VERSION_OUTPUT - ERROR_VARIABLE CARGO_VERSION_ERROR - OUTPUT_STRIP_TRAILING_WHITESPACE - ERROR_STRIP_TRAILING_WHITESPACE) -if(NOT CARGO_VERSION_RESULT EQUAL 0) - message(FATAL_ERROR - "Cargo was found at ${CARGO_EXECUTABLE}, but it cannot run:\n" - "${CARGO_VERSION_ERROR}\n" - "Install a Rust toolchain. If rustup manages Cargo, run: " - "rustup default stable") -endif() -message(STATUS "Cargo: ${CARGO_VERSION_OUTPUT} (${CARGO_EXECUTABLE})") - # Ensure C++17 is available for tokenizers-cpp if(CMAKE_CXX_STANDARD LESS 17) set(CMAKE_CXX_STANDARD 17) @@ -246,6 +229,8 @@ file(GLOB HEADERS "include/*.hpp" "runner/*.hpp" "common/*.hpp" "common/*/*.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$") diff --git a/src/common/AutoEmbeddingModel/modeling_iron_embedding.cpp b/src/common/AutoEmbeddingModel/modeling_iron_embedding.cpp deleted file mode 100644 index 42e0863f2..000000000 --- a/src/common/AutoEmbeddingModel/modeling_iron_embedding.cpp +++ /dev/null @@ -1,24 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -// SPDX-License-Identifier: MIT - -#include "AutoEmbeddingModel/modeling_iron_embedding.hpp" - -#include -#include - -void IronEmbedding::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"); - } - const char* override_path = std::getenv("FLM_IRON_EMBEDDING_PATH"); - if (override_path && *override_path) model_path = override_path; - worker_ = std::make_unique(model_path, "embedding_gemma2"); - is_model_loaded = true; - this->model_path = std::move(model_path); -} - -std::vector IronEmbedding::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"); -} \ No newline at end of file diff --git a/src/common/iron/causal_lm.cpp b/src/common/iron/causal_lm.cpp new file mode 100644 index 000000000..082aba567 --- /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("IRON loads safetensors in its worker"); +} + +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 graph worker"; } + +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/models/llama3/iron/llama3_iron_backend.cpp b/src/common/models/llama3/iron/llama3_iron_backend.cpp index a2725ba43..d07cad852 100644 --- a/src/common/models/llama3/iron/llama3_iron_backend.cpp +++ b/src/common/models/llama3/iron/llama3_iron_backend.cpp @@ -3,101 +3,12 @@ #include "models/llama3/iron/llama3_iron_backend.hpp" -#include "iron/worker.hpp" - -#include -#include -#include -#include -#include +#include "iron/causal_lm.hpp" namespace flm::llama3 { -namespace { - -class IronLlama final : public causal_lm { -public: - IronLlama(const std::string& model_path, std::uint32_t context_length) - : worker_(model_path, "llama", context_length), - max_length_(context_length) {} - - buffer forward(int id) override { - tokens_.push_back(id); - return run(); - } - - buffer prefill(std::vector& ids, void*) override { - tokens_.insert(tokens_.end(), ids.begin(), ids.end()); - return run(); - } - - void set_context_length(int length) override { - if (length < 0 || static_cast(length) > tokens_.size()) { - throw std::out_of_range("Invalid IRON context length"); - } - tokens_.resize(static_cast(length)); - } - void load_weights(Q4NX&) override { - throw std::logic_error("IRON loads safetensors in its worker"); - } - void update_max_length(std::uint32_t length) override { max_length_ = length; } - void clear_context() override { - tokens_.clear(); - worker_.reset(); - } - buffer get_k_cache(int, int) override { return {}; } - buffer get_v_cache(int, int) override { return {}; } - int get_current_context_length() override { return static_cast(tokens_.size()); } - int checkpoint() override { return static_cast(tokens_.size()); } - int restore() override { return 0; } - bool poisoned() const noexcept { return worker_.poisoned(); } - -private: - buffer 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; - } - - flm::iron::Worker worker_; - std::vector tokens_; - std::uint32_t max_length_; -}; - -class IronBackend final : public flm::backend::ModelBackend { -public: - explicit IronBackend(const flm::backend::BackendContext& context) - : engine_(CheckpointPath(context.model_path), - ValidateContextLength(context.context_length)) {} - causal_lm& engine() override { return engine_; } - std::string id() const override { return flm::backend::kIronBackendId; } - std::string detail() const override { return "IRON graph worker"; } - bool supports_preemption() const override { return false; } - bool poisoned() const noexcept override { return engine_.poisoned(); } - -private: - static std::string CheckpointPath(const std::string& fallback) { - const char* path = std::getenv("FLM_IRON_MODEL_PATH"); - return path && *path ? path : fallback; - } - static std::uint32_t ValidateContextLength(std::uint32_t length) { - if (length == 0) { - throw std::out_of_range("IRON Llama context length must be positive"); - } - return length; - } - IronLlama engine_; -}; - -} // namespace flm::backend::BackendFactory iron_factory() { - return [](const flm::backend::BackendContext& context) { - return std::make_unique(context); - }; + return flm::iron::causal_factory("llama"); } } // namespace flm::llama3 \ No newline at end of file diff --git a/src/include/AutoEmbeddingModel/all_embedding_model.hpp b/src/include/AutoEmbeddingModel/all_embedding_model.hpp index c44f997eb..4419d6e1e 100644 --- a/src/include/AutoEmbeddingModel/all_embedding_model.hpp +++ b/src/include/AutoEmbeddingModel/all_embedding_model.hpp @@ -11,7 +11,7 @@ #ifndef FASTFLOWLM_LINUX_LIMITED_MODELS #include "modeling_gemma_embedding.hpp" #if defined(FLM_ENABLE_IRON) -#include "modeling_iron_embedding.hpp" +#include "iron/embedding.hpp" #endif #endif @@ -38,7 +38,9 @@ inline std::pair> get_auto_embe 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); + auto_embedding_engine = std::make_unique( + npu_device_inst, "embeddinggemma-2:740m", "embedding_gemma2", + 768); else #endif throw std::runtime_error( diff --git a/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp b/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp deleted file mode 100644 index a492dd423..000000000 --- a/src/include/AutoEmbeddingModel/modeling_iron_embedding.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// 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" - -class IronEmbedding final : public AutoEmbeddingModel { -public: - explicit IronEmbedding(flm_rt::device* device) - : AutoEmbeddingModel(device, "embeddinggemma-2:740m") {} - - 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::unique_ptr worker_; -}; \ No newline at end of file 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/models/llama3/iron/llama3_iron_backend.hpp b/src/include/models/llama3/iron/llama3_iron_backend.hpp index 80b3b9d38..a18fb8195 100644 --- a/src/include/models/llama3/iron/llama3_iron_backend.hpp +++ b/src/include/models/llama3/iron/llama3_iron_backend.hpp @@ -4,13 +4,14 @@ #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::backend::BackendTraits{false, false, 0}; + return flm::iron::causal_traits(); } } // namespace flm::llama3 \ No newline at end of file From f10d0ba04d155d447796665406dd22713df31388 Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 30 Oct 2026 09:22:47 -0600 Subject: [PATCH 14/19] fix(iron): validate worker failures and run backend tests in CI --- .github/actions/lin-build/action.yml | 11 +++++++++++ src/common/iron/worker.cpp | 23 ++++++++++++++++------- src/test/iron_backend/CMakeLists.txt | 3 ++- src/test/iron_backend/fake_worker.py | 17 ++++++++++++++--- src/test/iron_backend/test_worker.cpp | 21 ++++++++++++++------- src/test/model_backend/CMakeLists.txt | 3 ++- 6 files changed, 59 insertions(+), 19 deletions(-) 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/src/common/iron/worker.cpp b/src/common/iron/worker.cpp index 42d3e3a9c..edbf4918c 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -4,7 +4,15 @@ #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 @@ -113,6 +121,7 @@ void Worker::read_all(void* data, std::size_t size) { nlohmann::json Worker::request(const nlohmann::json& command, const void* payload, std::size_t payload_size) { + if (poisoned_) throw std::runtime_error("IRON worker is unusable; reload the model"); try { const std::string header = command.dump(); const std::uint32_t header_size = static_cast(header.size()); @@ -143,15 +152,15 @@ nlohmann::json Worker::request(const nlohmann::json& command, std::vector Worker::logits(const std::vector& tokens) { auto response = request({{"command", "logits"}, {"count", tokens.size()}}, tokens.data(), tokens.size() * sizeof(int)); - const std::size_t count = response.at("count").get(); - std::vector values(count); 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; } - return values; } void Worker::reset() { request({{"command", "reset"}}); } @@ -164,15 +173,15 @@ std::vector Worker::embed(const std::string& text, {"dimensions", dimensions}, {"size", text.size()}}, text.data(), text.size()); - const std::size_t count = response.at("count").get(); - std::vector values(count); 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; } - return values; } -} // namespace flm::iron \ No newline at end of file +} // namespace flm::iron diff --git a/src/test/iron_backend/CMakeLists.txt b/src/test/iron_backend/CMakeLists.txt index 14d176e27..a881e7d70 100644 --- a/src/test/iron_backend/CMakeLists.txt +++ b/src/test/iron_backend/CMakeLists.txt @@ -24,5 +24,6 @@ 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=${Python3_EXECUTABLE};PYTHONPATH=${CMAKE_CURRENT_LIST_DIR}") \ No newline at end of file + "FLM_IRON_PYTHON=${Python3_EXECUTABLE};PYTHONPATH=${CMAKE_CURRENT_LIST_DIR}") diff --git a/src/test/iron_backend/fake_worker.py b/src/test/iron_backend/fake_worker.py index d0cf52232..2968f240f 100644 --- a/src/test/iron_backend/fake_worker.py +++ b/src/test/iron_backend/fake_worker.py @@ -35,7 +35,7 @@ def main(): return name = command["command"] if name == "load_model": - assert command["model"] in {"dies", "embedding_gemma2", "llama"} + model = command["model"] if command["model"] == "llama": assert command["context_length"] == 32 else: @@ -47,7 +47,18 @@ def main(): send({"ok": True}) elif name == "logits": tokens = struct.unpack(f"={command['count']}i", read_exact(command["count"] * 4)) - send({"ok": True, "count": 2}, struct.pack("=HH", tokens[-1], len(tokens))) + if model == "bad_count": + send({"ok": True}) + elif model == "bad_header": + sys.stdout.buffer.write(struct.pack("=I", 0)) + sys.stdout.buffer.flush() + elif model == "error": + send({"ok": False, "error": "model failed"}) + elif model == "truncated": + send({"ok": True, "count": 2}, b"x") + return + else: + send({"ok": True, "count": 2}, struct.pack("=HH", tokens[-1], len(tokens))) elif name == "embed": text = read_exact(command["size"]).decode() send({"ok": True, "count": 2}, struct.pack("=ff", len(text), command["dimensions"])) @@ -59,4 +70,4 @@ def main(): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp index e96f49dda..fa825efcb 100644 --- a/src/test/iron_backend/test_worker.cpp +++ b/src/test/iron_backend/test_worker.cpp @@ -3,7 +3,7 @@ #include "iron/worker.hpp" -#include +#include "../phi4_core/test_support.hpp" #include #include #include @@ -25,13 +25,20 @@ int main() { flm::iron::Worker llama("/unused", "llama", 32); const auto logits = llama.logits(std::vector{4, 10, 26}); - assert(logits == std::vector({26, 3})); + TEST_REQUIRE(logits == std::vector({26, 3})); llama.reset(); - assert(!llama.poisoned()); + TEST_REQUIRE(!llama.poisoned()); flm::iron::Worker embedding("/unused", "embedding_gemma2"); const auto result = embedding.embed("hello", "query", 128); - assert(result == std::vector({5.0f, 128.0f})); + TEST_REQUIRE(result == std::vector({5.0f, 128.0f})); + + 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 { @@ -40,7 +47,7 @@ int main() { } catch (const std::runtime_error&) { pipe_failed = true; } - assert(pipe_failed); + TEST_REQUIRE(pipe_failed); bool failed = false; SetEnvironment("FLM_IRON_WORKER_MODULE", "missing_worker"); @@ -49,6 +56,6 @@ int main() { } catch (const std::runtime_error&) { failed = true; } - assert(failed); + TEST_REQUIRE(failed); return 0; -} \ No newline at end of file +} 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() From 5aa1a98b6d29c88aecdcf6172c31e887b31ee59a Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 30 Oct 2026 09:24:43 -0600 Subject: [PATCH 15/19] test: match catalog expectations to the selected platform --- src/test/model_list_platform/test_model_list_platform.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 5cb6e612e..44da73d09 100644 --- a/src/test/model_list_platform/test_model_list_platform.cpp +++ b/src/test/model_list_platform/test_model_list_platform.cpp @@ -366,7 +366,7 @@ int main() { "a corelib build offers the corelib phi4"); RunTest(test_embeddinggemma_2_requires_iron, "EmbeddingGemma 2 requires iron"); - RunTest(test_bf16_llama_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"); From 538818bf58bfffa7e063d75db2b91300ff311ba6 Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 30 Oct 2026 11:41:30 -0600 Subject: [PATCH 16/19] docs: refer to the IRON server consistently --- src/CMakeLists.txt | 2 +- src/common/iron/causal_lm.cpp | 4 ++-- src/common/iron/worker.cpp | 12 ++++++------ src/home_install.sh | 4 ++-- 4 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c1c6cbfb9..b164bfc4c 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -53,7 +53,7 @@ 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 Python server" 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) diff --git a/src/common/iron/causal_lm.cpp b/src/common/iron/causal_lm.cpp index 082aba567..b8b802f07 100644 --- a/src/common/iron/causal_lm.cpp +++ b/src/common/iron/causal_lm.cpp @@ -43,7 +43,7 @@ void CausalLM::set_context_length(int length) { } void CausalLM::load_weights(Q4NX&) { - throw std::logic_error("IRON loads safetensors in its worker"); + throw std::logic_error("The IRON server loads safetensors"); } void CausalLM::update_max_length(std::uint32_t length) { @@ -90,7 +90,7 @@ std::string CausalBackend::id() const { return flm::backend::kIronBackendId; } -std::string CausalBackend::detail() const { return "IRON graph worker"; } +std::string CausalBackend::detail() const { return "IRON server"; } bool CausalBackend::supports_preemption() const { return false; } diff --git a/src/common/iron/worker.cpp b/src/common/iron/worker.cpp index edbf4918c..135faaa84 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -110,18 +110,18 @@ Worker::~Worker() { 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 worker write failed"); + 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 worker exited"); + 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 worker is unusable; reload the model"); + 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()); @@ -129,18 +129,18 @@ nlohmann::json Worker::request(const nlohmann::json& command, 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 worker write failed"); + 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 worker returned an invalid header"); + 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 worker failed")); + throw std::runtime_error(result.value("error", "IRON server failed")); } return result; } catch (...) { diff --git a/src/home_install.sh b/src/home_install.sh index 790e48e9b..c6953118c 100755 --- a/src/home_install.sh +++ b/src/home_install.sh @@ -9,7 +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 Python +# ./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 @@ -92,7 +92,7 @@ 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 a Python environment that can import iron.lm_server." >&2 + 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 From aa6fc1d85b7c9a01b5d1b9141901534a73b74bcc Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 30 Oct 2026 11:45:06 -0600 Subject: [PATCH 17/19] test: replace IRON Python stub with a C++ server --- src/test/iron_backend/CMakeLists.txt | 3 +- src/test/iron_backend/fake_worker.py | 73 ----------------------- src/test/iron_backend/test_worker.cpp | 83 ++++++++++++++++++++++++++- 3 files changed, 82 insertions(+), 77 deletions(-) delete mode 100644 src/test/iron_backend/fake_worker.py diff --git a/src/test/iron_backend/CMakeLists.txt b/src/test/iron_backend/CMakeLists.txt index a881e7d70..6fb693f52 100644 --- a/src/test/iron_backend/CMakeLists.txt +++ b/src/test/iron_backend/CMakeLists.txt @@ -9,7 +9,6 @@ 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) -find_package(Python3 REQUIRED COMPONENTS Interpreter) add_executable(test_iron_worker test_worker.cpp @@ -26,4 +25,4 @@ add_test(NAME test_iron_worker COMMAND test_iron_worker) set_tests_properties(test_iron_worker PROPERTIES TIMEOUT 30 ENVIRONMENT - "FLM_IRON_PYTHON=${Python3_EXECUTABLE};PYTHONPATH=${CMAKE_CURRENT_LIST_DIR}") + "FLM_IRON_PYTHON=$") diff --git a/src/test/iron_backend/fake_worker.py b/src/test/iron_backend/fake_worker.py deleted file mode 100644 index 2968f240f..000000000 --- a/src/test/iron_backend/fake_worker.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: MIT - -import json -import struct -import sys - - -def read_exact(size): - data = bytearray() - while len(data) < size: - chunk = sys.stdin.buffer.read(size - len(data)) - if not chunk: - raise EOFError - data.extend(chunk) - return bytes(data) - - -def receive(): - size = struct.unpack("=I", read_exact(4))[0] - return json.loads(read_exact(size)) - - -def send(header, payload=b""): - data = json.dumps(header).encode() - sys.stdout.buffer.write(struct.pack("=I", len(data)) + data + payload) - sys.stdout.buffer.flush() - - -def main(): - while True: - try: - command = receive() - except EOFError: - return - name = command["command"] - if name == "load_model": - model = command["model"] - if command["model"] == "llama": - assert command["context_length"] == 32 - else: - assert "context_length" not in command - send({"ok": True}) - if command["model"] == "dies": - return - elif name == "reset": - send({"ok": True}) - elif name == "logits": - tokens = struct.unpack(f"={command['count']}i", read_exact(command["count"] * 4)) - if model == "bad_count": - send({"ok": True}) - elif model == "bad_header": - sys.stdout.buffer.write(struct.pack("=I", 0)) - sys.stdout.buffer.flush() - elif model == "error": - send({"ok": False, "error": "model failed"}) - elif model == "truncated": - send({"ok": True, "count": 2}, b"x") - return - else: - send({"ok": True, "count": 2}, struct.pack("=HH", tokens[-1], len(tokens))) - elif name == "embed": - text = read_exact(command["size"]).decode() - send({"ok": True, "count": 2}, struct.pack("=ff", len(text), command["dimensions"])) - elif name == "exit": - return - elif name == "close": - send({"ok": True}) - return - - -if __name__ == "__main__": - main() diff --git a/src/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp index fa825efcb..4f6869d7c 100644 --- a/src/test/iron_backend/test_worker.cpp +++ b/src/test/iron_backend/test_worker.cpp @@ -4,10 +4,18 @@ #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) { @@ -18,10 +26,81 @@ void SetEnvironment(const char* name, const char* value) { #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")); + send({{"ok", true}}); + 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() { - SetEnvironment("FLM_IRON_WORKER_MODULE", "fake_worker"); +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}); From b25a4edbf3838a40f923ab763d937c090e5715bc Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 30 Oct 2026 11:52:13 -0600 Subject: [PATCH 18/19] fix(iron): serve Llama 3.2 1B Instruct weights --- docs/docs/models/llama.md | 14 ++++++++++++++ src/model_info.json | 2 +- src/model_list.json | 6 +++--- .../test_model_list_platform.cpp | 4 +++- 4 files changed, 21 insertions(+), 5 deletions(-) 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/model_info.json b/src/model_info.json index 5b620b5b2..04642fcc2 100644 --- a/src/model_info.json +++ b/src/model_info.json @@ -2500,7 +2500,7 @@ { "type": "file", "size": 2471645608, - "sha256": "68a2e4be76fa709455a60272fba8e512c02d81c46e6c671cc9449e374fd6809a", + "sha256": "1ff795ff6a07e6a68085d206fb84417da2f083f68391c2843cd2b8ac6df8538f", "path": "model.safetensors" }, { diff --git a/src/model_list.json b/src/model_list.json index 25e35e254..12b5472d9 100644 --- a/src/model_list.json +++ b/src/model_list.json @@ -1485,9 +1485,9 @@ "llama3.2-iron": { "1b": { "supported_platforms": ["aie2p"], - "name": "llama3.2-1b", - "url": "https://huggingface.co/meta-llama/Llama-3.2-1B", - "file_url": "https://huggingface.co/api/models/meta-llama/Llama-3.2-1B/tree/main", + "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", 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 44da73d09..da028aa62 100644 --- a/src/test/model_list_platform/test_model_list_platform.cpp +++ b/src/test/model_list_platform/test_model_list_platform.cpp @@ -143,7 +143,9 @@ void test_bf16_llama_requires_iron() { 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"); + 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"); From c01a0e3c233814d78906a8cb6f7269915b0de9d7 Mon Sep 17 00:00:00 2001 From: andrej Date: Fri, 30 Oct 2026 13:54:31 -0600 Subject: [PATCH 19/19] feat(iron): negotiate protocol version and capabilities --- src/common/iron/worker.cpp | 24 ++++++++++++++++-- src/include/iron/worker.hpp | 10 ++++++++ src/test/iron_backend/test_worker.cpp | 35 ++++++++++++++++++++++++++- 3 files changed, 66 insertions(+), 3 deletions(-) diff --git a/src/common/iron/worker.cpp b/src/common/iron/worker.cpp index 135faaa84..bf77b5aa4 100644 --- a/src/common/iron/worker.cpp +++ b/src/common/iron/worker.cpp @@ -24,6 +24,12 @@ 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"); @@ -84,10 +90,16 @@ Worker::Worker(const std::filesystem::path& model_path, std::string model, 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; - request(command); + 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; @@ -150,6 +162,8 @@ nlohmann::json Worker::request(const nlohmann::json& command, } 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 { @@ -163,11 +177,17 @@ std::vector Worker::logits(const std::vector& tokens) { } } -void Worker::reset() { request({{"command", "reset"}}); } +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}, diff --git a/src/include/iron/worker.hpp b/src/include/iron/worker.hpp index ce078be25..c861023d7 100644 --- a/src/include/iron/worker.hpp +++ b/src/include/iron/worker.hpp @@ -12,6 +12,12 @@ 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, @@ -25,6 +31,9 @@ class Worker { 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: @@ -37,6 +46,7 @@ class Worker { 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/test/iron_backend/test_worker.cpp b/src/test/iron_backend/test_worker.cpp index 4f6869d7c..743802c7f 100644 --- a/src/test/iron_backend/test_worker.cpp +++ b/src/test/iron_backend/test_worker.cpp @@ -55,7 +55,16 @@ int ServeFake() { model = command.at("model").get(); if (model == "llama") TEST_REQUIRE(command.at("context_length") == 32); else TEST_REQUIRE(!command.contains("context_length")); - send({{"ok", true}}); + 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()); @@ -112,6 +121,30 @@ int main(int argc, char** argv) { 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}); });