From c45d68e2d3229e679dbb53c92706d7429d2dcba2 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 7 Aug 2026 16:48:28 -0700 Subject: [PATCH 1/4] [MLX] C++ runner for HF LLMs with the off-graph KV cache --- .github/workflows/mlx.yml | 110 ++++ backends/mlx/examples/llm/CMakeLists.txt | 57 ++ backends/mlx/examples/llm/run_llm_hf.cpp | 662 +++++++++++++++++++++++ 3 files changed, 829 insertions(+) create mode 100644 backends/mlx/examples/llm/CMakeLists.txt create mode 100644 backends/mlx/examples/llm/run_llm_hf.cpp diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 65c599d042d..44bc8b03bf3 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -678,3 +678,113 @@ jobs: exit 1 fi echo "::endgroup::" + + # Off-graph KV cache: the cache is a runtime object, so this path can only be + # exercised by the C++ runner (pybindings cannot bind a cache_key). Also the + # only coverage that the layout the export publishes is actually consumable. + test-mlx-llm-offgraph: + # Requires HuggingFace secrets — skip on fork PRs. + needs: run-decision + if: | + (github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request') && + (github.event_name == 'pull_request' || needs.run-decision.outputs.is-full-run == 'true') + strategy: + fail-fast: false + matrix: + model: + - id: "unsloth/Llama-3.2-1B-Instruct" + name: "llama-1b" + chat: "llama3" + runner: "macos-14-xlarge" + - id: "unsloth/gemma-3-1b-it" + name: "gemma3-1b" + chat: "gemma" + runner: "macos-14-xlarge" + # Only model with KV sharing: 15 caches for 35 layers, and sliding and + # full-attention caches with different head shapes. + - id: "google/gemma-4-E2B-it" + name: "gemma4-e2b" + chat: "gemma4" + runner: "macos-15-xlarge" + uses: pytorch/test-infra/.github/workflows/macos_job.yml@main + secrets: inherit + with: + default-packages: "" + job-name: test-mlx-llm-offgraph-${{ matrix.model.name }} + runner: ${{ matrix.model.runner }} + python-version: "3.12" + submodules: recursive + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + secrets-env: EXECUTORCH_HF_TOKEN + # Higher than the sibling LLM job: this one also builds and installs + # ExecuTorch so the standalone runner project can link against it. + timeout: 120 + script: | + set -eux + export HF_HUB_DISABLE_XET=1 + + MODEL_ID="${{ matrix.model.id }}" + MODEL_NAME="${{ matrix.model.name }}" + CHAT="${{ matrix.model.chat }}" + + echo "::group::Install ExecuTorch and build the MLX runtime" + ${CONDA_RUN} python install_executorch.py > /dev/null + ${CONDA_RUN} cmake --preset mlx-release + ${CONDA_RUN} cmake --build cmake-out --target install -j$(( $(sysctl -n hw.ncpu) - 1 )) + echo "::endgroup::" + + echo "::group::Build the C++ runner" + # Standalone find_package(executorch) project, so it needs the install above. + ${CONDA_RUN} cmake -S backends/mlx/examples/llm \ + -B cmake-out/backends/mlx/examples/llm -DCMAKE_BUILD_TYPE=Release + ${CONDA_RUN} cmake --build cmake-out/backends/mlx/examples/llm \ + -j$(( $(sysctl -n hw.ncpu) - 1 )) + RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf + if [ ! -x "${RUNNER}" ]; then + echo "Failed: runner not found at ${RUNNER}" + exit 1 + fi + echo "::endgroup::" + + echo "::group::Install LLM requirements" + ${CONDA_RUN} pip install -U "huggingface_hub[cli]<1.0" + ${CONDA_RUN} huggingface-cli login --token $SECRET_EXECUTORCH_HF_TOKEN + OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt) + ${CONDA_RUN} pip install transformers "optimum-executorch @ git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION}" + if [ "${MODEL_ID}" = "google/gemma-4-E2B-it" ]; then + # Gemma 4 needs a newer Transformers than the CI-wide pin. Keep this + # on the same commit test-mlx-llm validated against. + GEMMA4_TRANSFORMERS_COMMIT=61461a7bcb458db7cf6eeea49678b9ab776a7821 + ${CONDA_RUN} pip install -U "transformers @ git+https://github.com/huggingface/transformers.git@${GEMMA4_TRANSFORMERS_COMMIT}" + fi + echo "::endgroup::" + + echo "::group::Export ${MODEL_NAME} off-graph" + ${CONDA_RUN} python -m executorch.backends.mlx.examples.llm.export_llm_hf \ + --model-id "${MODEL_ID}" \ + --output /tmp/${MODEL_NAME}_offgraph.pte \ + --use-offgraph-cache \ + --max-seq-len 1024 \ + --dtype bf16 \ + --qlinear 4w + echo "::endgroup::" + + echo "::group::Run ${MODEL_NAME} off-graph inference" + # The cache geometry comes from the .pte; only capacity is given here. + TOKENIZER=$(${CONDA_RUN} python -c \ + "from huggingface_hub import hf_hub_download; print(hf_hub_download('${MODEL_ID}', 'tokenizer.json'))") + OUTPUT=$(${RUNNER} \ + --pte /tmp/${MODEL_NAME}_offgraph.pte \ + --tokenizer "${TOKENIZER}" \ + --chat "${CHAT}" \ + --kv-max-capacity 1024 \ + --prompt "What is the capital of France?" \ + --max-new-tokens 50 2>&1) + echo "$OUTPUT" + if echo "$OUTPUT" | grep -iq "Paris"; then + echo "Success: 'Paris' found in output" + else + echo "Failed: Expected 'Paris' not found in output" + exit 1 + fi + echo "::endgroup::" diff --git a/backends/mlx/examples/llm/CMakeLists.txt b/backends/mlx/examples/llm/CMakeLists.txt new file mode 100644 index 00000000000..c36f042fd6e --- /dev/null +++ b/backends/mlx/examples/llm/CMakeLists.txt @@ -0,0 +1,57 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# C++ runner for off-graph-cache HF models (export_llm_hf --use-offgraph-cache). +# Unlike run_llm_hf.py (pybindings), this binds the cache via cache_key, so it +# is the run path for off-graph .pte files. + +cmake_minimum_required(VERSION 3.24) +project(mlx_run_llm_hf) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) +include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) + +set(_common_include_directories ${EXECUTORCH_ROOT}/..) +set(_json_include + ${EXECUTORCH_ROOT}/extension/llm/tokenizers/third-party/json/single_include +) +# MLXExecutor.h reaches schema_generated.h, which needs flatbuffers. The MLX +# backend exposes that include under BUILD_INTERFACE, which does not reach a +# separate find_package() project like this one. +set(_flatbuffers_include ${EXECUTORCH_ROOT}/third-party/flatbuffers/include) + +list(APPEND CMAKE_FIND_ROOT_PATH ${CMAKE_CURRENT_BINARY_DIR}/../../../..) +find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH) +executorch_target_link_options_shared_lib(executorch) + +set(link_libraries executorch extension_module extension_tensor + extension_llm_cache +) + +if(NOT TARGET mlxdelegate) + message(FATAL_ERROR "mlx_run_llm_hf requires the MLX backend (mlxdelegate)") +endif() +list(APPEND link_libraries mlxdelegate mlx) +executorch_target_link_options_shared_lib(mlxdelegate) + +# CPU kernels for the ops that stay outside the delegate (e.g. the cache +# bookkeeping copy_ the HF export wrappers emit). +if(TARGET optimized_native_cpu_ops_lib) + list(APPEND link_libraries optimized_native_cpu_ops_lib) + executorch_target_link_options_shared_lib(optimized_native_cpu_ops_lib) +endif() + +list(APPEND link_libraries tokenizers::tokenizers) + +add_executable(mlx_run_llm_hf run_llm_hf.cpp) +target_include_directories( + mlx_run_llm_hf PUBLIC ${_common_include_directories} ${_json_include} + ${_flatbuffers_include} +) +target_link_libraries(mlx_run_llm_hf PUBLIC ${link_libraries}) diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp new file mode 100644 index 00000000000..13a6c995de2 --- /dev/null +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -0,0 +1,662 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// C++ runner for HuggingFace LLMs on the MLX backend. Unlike the pybindings +// run_llm_hf.py, it can bind the off-graph KV cache: with --kv-cache it builds +// an MLXSequenceCache, installs it in the process-global registry, and passes +// its cache_key as a load-time backend option (the rendezvous init() reads). +// Its shape comes from constant methods the export publishes; the flags below +// only choose policy. Without --kv-max-capacity it runs an in-graph model +// unchanged -- so the same binary compares both cache paths. Greedy decode. +// +// Usage: +// run_llm_hf --pte --tokenizer \ +// [--kv-max-capacity N] [--kv-storage-dtype bf16|fp16|fp32] \ +// [--kv-initial-capacity N] [--kv-max-write N] \ +// [--kv-windows ] \ +// [--prompt "..."] [--max-new-tokens N] [--chat llama3|gemma|gemma4|0] \ +// [--warmup N] [--iters N] [--interactive 1] + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wconversion" +#pragma clang diagnostic ignored "-Wsign-conversion" +#include +#include +#pragma clang diagnostic pop + +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +using ::executorch::extension::make_tensor_ptr; +using ::executorch::extension::Module; +using ::executorch::extension::TensorPtr; +using ::executorch::runtime::Error; +using ::executorch::runtime::EValue; + +namespace cache = ::executorch::extension::llm::cache; + +namespace { + +std::vector split(const std::string& s, char delim) { + std::vector out; + size_t pos = 0; + while (pos <= s.size()) { + const size_t d = s.find(delim, pos); + out.push_back( + s.substr(pos, d == std::string::npos ? std::string::npos : d - pos)); + if (d == std::string::npos) { + break; + } + pos = d + 1; + } + return out; +} + +bool parse_int_list( + const std::string& spec, + char delim, + std::vector& out) { + for (const std::string& field : split(spec, delim)) { + if (field.empty()) { + return false; + } + try { + out.push_back(std::stoi(field)); + } catch (const std::exception&) { + return false; + } + } + return true; +} + +int storage_dtype(const std::string& name) { + using S = ::executorch::runtime::etensor::ScalarType; + if (name == "bf16") { + return static_cast(S::BFloat16); + } + if (name == "fp16") { + return static_cast(S::Half); + } + if (name == "fp32") { + return static_cast(S::Float); + } + return -1; +} + +// Constant methods the export publishes (get_n_caches and friends). They carry +// no delegate, so reading them only needs the program loaded -- which is what +// lets the cache be built before forward's backend init consumes its key. +std::optional const_int(Module& module, const char* name) { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isInt()) { + return std::nullopt; + } + return r->at(0).toInt(); +} + +std::optional> const_ints(Module& module, const char* name) { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isTensor()) { + return std::nullopt; + } + const auto t = r->at(0).toTensor(); + if (t.scalar_type() != ::executorch::aten::ScalarType::Int) { + return std::nullopt; + } + const int32_t* p = t.const_data_ptr(); + return std::vector(p, p + t.numel()); +} + +// Replace the model's own attention pattern with `spec`, a comma-separated list +// of windows repeating over the caches (0 = flat). One entry makes every layer +// sliding. Only the policy changes; each cache keeps the geometry the .pte +// declared, so this cannot desync from the graph. +bool apply_window_override(const std::string& spec, cache::CacheConfig& cfg) { + std::vector pattern; + if (!parse_int_list(spec, ',', pattern) || pattern.empty()) { + return false; + } + for (size_t l = 0; l < cfg.layers.size(); ++l) { + const int w = pattern[l % pattern.size()]; + cfg.layers[l].policy = w > 0 + ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, w} + : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; + } + return cache::valid(cfg); +} + +// argmax over the last position's vocab row of a [1, T, vocab] logits tensor, +// reading whatever float dtype the op emitted. +int64_t argmax_last(const ::executorch::aten::Tensor& logits) { + const auto dim = logits.dim(); + const int64_t vocab = logits.size(dim - 1); + const int64_t t = dim >= 2 ? logits.size(dim - 2) : 1; + const int64_t offset = (t - 1) * vocab; // start of the last row + auto scan = [&](auto* data) { + int64_t best = 0; + float best_v = -1e30f; + for (int64_t i = 0; i < vocab; ++i) { + const float v = static_cast(data[offset + i]); + if (v > best_v) { + best_v = v; + best = i; + } + } + return best; + }; + switch (logits.scalar_type()) { + case ::executorch::aten::ScalarType::Float: + return scan(logits.const_data_ptr()); + case ::executorch::aten::ScalarType::Half: + return scan(logits.const_data_ptr<::executorch::aten::Half>()); + case ::executorch::aten::ScalarType::BFloat16: + return scan(logits.const_data_ptr<::executorch::aten::BFloat16>()); + default: + throw std::runtime_error("argmax_last: unsupported logits dtype"); + } +} + +// Human-readable name for a kv_dtype (an ET ScalarType int). Only the +// storage dtypes the pool uses are named; anything else prints its raw value. +std::string dtype_name(int st) { + using S = ::executorch::runtime::etensor::ScalarType; + switch (static_cast(st)) { + case S::Half: + return "Half(fp16)"; + case S::Float: + return "Float(fp32)"; + case S::BFloat16: + return "BFloat16"; + default: + return "scalar_type_" + std::to_string(st); + } +} + +// One user turn wrapped in the model's instruct template. Returns false for an +// unknown template name. The leading BOS belongs to the first turn only, so a +// continuing conversation passes with_bos=false. +bool wrap_turn( + const std::string& chat, + const std::string& prompt, + bool with_bos, + std::string& out) { + if (chat == "0") { + out = prompt; + } else if (chat == "llama3") { + out = std::string(with_bos ? "<|begin_of_text|>" : "") + + "<|start_header_id|>user<|end_header_id|>\n\n" + prompt + + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; + } else if (chat == "gemma") { + out = std::string(with_bos ? "" : "") + "user\n" + + prompt + "\nmodel\n"; + } else if (chat == "gemma4") { + // Gemma 4 renamed the turn markers; its own is also the eos. + out = std::string(with_bos ? "" : "") + "<|turn>user\n" + prompt + + "\n<|turn>model\n"; + } else { + return false; + } + return true; +} + +std::string +arg(int argc, char** argv, const std::string& key, const std::string& def) { + for (int i = 1; i + 1 < argc; ++i) { + if (key == argv[i]) { + return argv[i + 1]; + } + } + return def; +} + +} // namespace + +int main(int argc, char** argv) { + const std::string pte = arg(argc, argv, "--pte", ""); + const std::string tok_path = arg(argc, argv, "--tokenizer", ""); + // Off-graph: how much history the cache may hold. Its presence selects the + // off-graph path; the cache's shape comes from the .pte, so only policy is + // set here. + const int kv_capacity = std::stoi(arg(argc, argv, "--kv-max-capacity", "0")); + const std::string kv_dtype = arg(argc, argv, "--kv-storage-dtype", "bf16"); + // Optional: impose an attention pattern other than the model's own, e.g. + // "512" to make every layer sliding. + const std::string kv_windows = arg(argc, argv, "--kv-windows", ""); + const std::string prompt = arg(argc, argv, "--prompt", "The quick brown fox"); + const int max_new = std::stoi(arg(argc, argv, "--max-new-tokens", "50")); + // Instruct chat template to wrap the prompt in: llama3 (default) or gemma. + // Raw text confuses an instruct model into emitting turn markers; 0 disables. + const std::string chat = arg(argc, argv, "--chat", "llama3"); + // Off-graph only: the cache pool's starting size; it grows (doubling) up to + // capacity. -1 keeps the CacheConfig default. Small values force growth. + const int initial_capacity = + std::stoi(arg(argc, argv, "--kv-initial-capacity", "-1")); + // Off-graph only: the largest single step, which sizes a ring layer's slots + // to window + max_write - 1. -1 lets each ring layer use its own window. + // Prefill goes in one step, so this must be at least the prompt length. + const int max_write = std::stoi(arg(argc, argv, "--kv-max-write", "-1")); + // Benchmarking: `warmup` throwaway iters (absorb JIT + pool growth) then + // `iters` measured; per-iter tok/s plus a mean +/- stddev summary. + // Multi-turn chat on stdin instead of a single prompt; off-graph only. + const bool interactive = arg(argc, argv, "--interactive", "0") != "0"; + const int warmup = std::stoi(arg(argc, argv, "--warmup", "0")); + const int iters = std::stoi(arg(argc, argv, "--iters", "1")); + if (pte.empty() || tok_path.empty()) { + std::cerr << "Required: --pte --tokenizer " + "[--kv-max-capacity N for off-graph models]\n"; + return 1; + } + + try { + // Tokenizer. + ::tokenizers::HFTokenizer tokenizer; + if (tokenizer.load(tok_path) != ::tokenizers::Error::Ok) { + std::cerr << "Failed to load tokenizer: " << tok_path << std::endl; + return 1; + } + + // Off-graph models (update_and_attend) need a cache bound via cache_key; + // in-graph models (mlx::kv_cache_update) don't -- omit --kv-max-capacity + // for those. session/options are outer-scoped: session must outlive the + // Module (it keeps the cache in the registry) and mlx_opts must outlive + // load_method() (the map holds a view into it). + std::optional session; + ::executorch::runtime::BackendOptions<1> mlx_opts; + ::executorch::runtime::LoadBackendOptionsMap options_map; + const bool off_graph = kv_capacity > 0; + + // Load the program but not forward: the cache must exist before forward's + // backend init reads its key, and the layout it needs is published by + // constant methods in the same file. + Module module(pte); + if (module.load() != Error::Ok) { + std::cerr << "Failed to load " << pte << std::endl; + return 1; + } + + if (off_graph) { + cache::CacheConfig cfg{}; + cfg.capacity = kv_capacity; + cfg.kv_dtype = storage_dtype(kv_dtype); + if (cfg.kv_dtype < 0) { + std::cerr << "Invalid --kv-storage-dtype: " << kv_dtype + << " (bf16|fp16|fp32)" << std::endl; + return 1; + } + const auto n_caches = const_int(module, "get_n_caches"); + const auto kv_heads = const_ints(module, "get_kv_heads"); + const auto head_dims = const_ints(module, "get_head_dims"); + const auto windows = const_ints(module, "get_windows"); + if (!n_caches || !kv_heads || !head_dims || !windows) { + std::cerr << "No KV cache layout in " << pte + << "; re-export with --use-offgraph-cache" << std::endl; + return 1; + } + cfg.n_layers = static_cast(*n_caches); + cfg.layers.clear(); + cfg.layers.reserve(static_cast(cfg.n_layers)); + for (size_t l = 0; l < static_cast(cfg.n_layers); ++l) { + cache::LayerConfig lc{}; + lc.n_kv_heads = (*kv_heads)[l]; + lc.head_dim = (*head_dims)[l]; + lc.policy = (*windows)[l] > 0 + ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, (*windows)[l]} + : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; + cfg.layers.push_back(lc); + } + if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) { + std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl; + return 1; + } + if (!cache::valid(cfg)) { + std::cerr << "Invalid cache config" << std::endl; + return 1; + } + if (initial_capacity >= 0) { + cfg.initial_capacity = initial_capacity; + } + if (max_write > 0) { + cfg.max_write = max_write; + } + auto built = cache::CacheBuilderRegistry::global().build( + ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); + if (!built.ok()) { + std::cerr << "Failed to build cache: " + << static_cast(built.error()) << std::endl; + return 1; + } + session.emplace(cache::make_unique_key(), built.get()); + + // Announce the cache shape: the same .pte runs under whatever config this + // invocation asks for -- capacity, storage dtype, flat/ring layers -- + // with no re-export. The footprint lines below then show it growing at + // runtime. + int flat = 0, ring = 0, ring_window = 0; + for (int l = 0; l < cfg.n_layers; ++l) { + const cache::LayerConfig& lc = + cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; + if (lc.policy.kind == cache::LayerPolicy::Kind::Ring) { + ++ring; + ring_window = lc.policy.window; + } else { + ++flat; + } + } + std::cout << "\n[cache] off-graph seq | capacity=" << cfg.capacity + << " initial=" << cfg.initial_capacity + << " kv_dtype=" << dtype_name(cfg.kv_dtype); + if (cfg.max_write) { + std::cout << " max_write=" << *cfg.max_write; + } + std::cout << "\n " << cfg.n_layers << " layers: " << flat + << " flat"; + if (ring > 0) { + std::cout << " + " << ring << " ring(window " << ring_window << ")"; + } + std::cout << std::endl; + if (mlx_opts.set_option( + ::executorch::backends::mlx::kCacheKeyKey, + session->key().c_str()) != Error::Ok || + options_map.set_options( + ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != + Error::Ok) { + std::cerr << "Failed to set cache_key option" << std::endl; + return 1; + } + } + + if (module.load_method( + "forward", + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + off_graph ? &options_map : nullptr) != Error::Ok) { + std::cerr << "Failed to load forward" << std::endl; + return 1; + } + // Weights-only baseline, so the deltas below isolate the cache. + const double mem_at_load = ::mlx::core::get_active_memory() / 1048576.0; + std::cout << "[mem] after load : " << mem_at_load << " MiB" << std::endl; + + // Encode. HFTokenizer maps special-token markers in the string to their + // ids, so the template's <|...|> tokens encode correctly; it already + // carries <|begin_of_text|>, so pass bos=0 to avoid a doubled BOS. + std::string enc_input; + if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { + std::cerr << "Unknown --chat template: " << chat + << " (expected llama3, gemma, gemma4, or 0)" << std::endl; + return 1; + } + // The template carries its own BOS, so only a raw prompt asks for one. + const int8_t bos = chat == "0" ? 1 : 0; + auto enc = tokenizer.encode(enc_input, bos, /*eos=*/0); + if (!enc.ok()) { + std::cerr << "Encode failed" << std::endl; + return 1; + } + std::vector tokens = std::move(*enc); + const int prompt_len = static_cast(tokens.size()); + + // Stop on end-of-text and (for chat) the turn-end token <|eot_id|>. + std::vector stop_ids = {tokenizer.eos_tok()}; + if (chat != "0") { + const char* turn_end = chat == "llama3" ? "<|eot_id|>" + : chat == "gemma4" ? "" + : ""; + if (auto eot = tokenizer.piece_to_id(turn_end); eot.ok()) { + stop_ids.push_back(*eot); + } + } + auto is_stop = [&](int64_t t) { + for (uint64_t s : stop_ids) { + if (t == static_cast(s)) { + return true; + } + } + return false; + }; + + auto step = [&](const std::vector& ids, + const std::vector& pos) { + auto in = + make_tensor_ptr({1, (int)ids.size()}, std::vector(ids)); + auto cp = make_tensor_ptr({(int)pos.size()}, std::vector(pos)); + auto out = module.execute("forward", {in, cp}); + if (!out.ok()) { + throw std::runtime_error("execute failed"); + } + return argmax_last(out->at(0).toTensor()); + }; + + // Multi-turn: history stays in the cache, so each turn only prefills its + // own tokens at the running position. /reset and /undo drive the cache's + // control face directly -- off-graph only, since an in-graph cache gives + // the runner no handle to its state. + if (interactive) { + if (!off_graph) { + std::cerr << "--interactive requires --kv-max-capacity\n"; + return 1; + } + auto* control = session->control(); + std::cout << "Multi-turn chat. /reset clears, /undo drops the last turn, " + "/undo N drops N tokens, /quit exits.\n"; + int64_t position = 0; + int64_t turn_start = 0; // position this turn began at, for /undo + std::string line; + while (std::cout << "\n> " && std::getline(std::cin, line)) { + if (line == "/quit") { + break; + } + if (line == "/reset") { + control->clear(); + position = turn_start = 0; + std::cout << "[cleared]\n"; + continue; + } + if (line == "/undo" || line.rfind("/undo ", 0) == 0) { + // Bare /undo drops the last turn; /undo N drops N tokens. + int64_t target = turn_start; + if (line.size() > 6) { + try { + const int64_t n = std::stoll(line.substr(6)); + target = n >= position ? 0 : position - n; + } catch (const std::exception&) { + std::cout << "[usage: /undo [n_tokens]]\n"; + continue; + } + } + if (control->rewind(static_cast(target))) { + position = target; + turn_start = std::min(turn_start, position); + std::cout << "[rewound to " << position << "]\n"; + } else { + // A sliding-window layer has physically dropped those cells. + std::cout << "[cannot rewind to " << target << "]\n"; + } + continue; + } + if (line.empty()) { + continue; + } + + std::string turn; + wrap_turn(chat, line, /*with_bos=*/position == 0, turn); + auto te = tokenizer.encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); + if (!te.ok()) { + std::cerr << "Encode failed\n"; + continue; + } + const int n = static_cast(te->size()); + // Admit the turn if its prompt plus one token fits; reserving the whole + // max_new budget up front would report "full" with most of the cache + // still free. Generation is then clamped to the room that remains. + if (!control->can_extend(n + 1)) { + std::cout << "[cache full: " << position << "/" << control->capacity() + << ", turn " << n << " tokens" + << (control->can_extend(1) ? "" : ", length at capacity") + << ", use /reset]\n"; + continue; + } + const int budget = std::min( + max_new, control->capacity() - static_cast(position) - n); + + turn_start = position; + std::vector tin(te->begin(), te->end()), tpos; + for (int i = 0; i < n; ++i) { + tpos.push_back(position + i); + } + int64_t next = step(tin, tpos); + position += n; + + uint64_t prev = te->back(); + for (int i = 0; i < budget && !is_stop(next); ++i) { + if (auto piece = tokenizer.decode(prev, static_cast(next)); + piece.ok()) { + std::cout << *piece << std::flush; + } + prev = static_cast(next); + next = step({next}, {position}); + ++position; + } + std::cout << "\n[" << position << "/" << control->capacity() + << " tokens" + << (budget < max_new ? ", generation capped by capacity" : "") + << "]\n"; + } + return 0; + } + + std::vector ids(tokens.begin(), tokens.end()), prefill_pos; + for (int i = 0; i < prompt_len; ++i) { + prefill_pos.push_back(i); + } + auto ms = [](auto a, auto b) { + return std::chrono::duration(b - a).count(); + }; + // Sequence length against the configured ceiling, with what MLX actually + // holds for it. Pools start at initial_capacity and grow by doubling, so + // the bytes lag the token count in steps; bf16 storage (kv_dtype 15) halves + // them vs fp32 (6). + auto print_footprint = [&](const char* when, int len) { + if (!session) { + return; + } + const int cap = session->control()->capacity(); + const double pct = cap > 0 ? 100.0 * len / cap : 0.0; + std::cout << "[cache] " << when << ": " << len << " / " << cap + << " tokens (" << pct << "%)" << std::endl; + const double mem = ::mlx::core::get_active_memory() / 1048576.0; + std::cout << "[mem] " << when << ": " << mem << " MiB (+" + << (mem - mem_at_load) << " MiB since load)" << std::endl; + }; + + // warmup iters absorb JIT + pool growth; measured iters record tok/s. The + // off-graph cache is cleared between iters so each prefill starts at length + // 0 (in-graph overwrites by position, so no reset needed there). + const int total_iters = warmup + std::max(1, iters); + std::vector pf_tps, dc_tps; + for (int iter = 0; iter < total_iters; ++iter) { + if (iter > 0 && off_graph) { + session->control()->clear(); + } + const bool measured = iter >= warmup; + const bool print_text = (iter == 0); + // Report the footprint from the last iteration: a measured one, and in + // steady state once any pool growth has settled. + const bool print_mem = (iter == total_iters - 1); + + const auto t0 = std::chrono::steady_clock::now(); + int64_t next = step(ids, prefill_pos); + const auto t1 = std::chrono::steady_clock::now(); + if (print_text || print_mem) { + std::cout << "\n"; // blank line separating this section from the banner + } + if (print_mem) { + print_footprint("after prefill", prompt_len); + } + if (print_text) { + std::cout << "\n"; // blank line before the streamed generation + } + + uint64_t prev = tokens.back(); + int generated = 0; + for (int i = 0; i < max_new; ++i) { + if (is_stop(next)) { + break; + } + if (print_text) { + if (auto piece = tokenizer.decode(prev, static_cast(next)); + piece.ok()) { + std::cout << *piece << std::flush; + } + } + prev = static_cast(next); + ++generated; + next = step({next}, {prompt_len + i}); + } + const auto t2 = std::chrono::steady_clock::now(); + if (print_text) { + std::cout << "\n\n"; // close the generation line + blank separator + } + if (print_mem) { + // trailing space aligns the colon with the "after prefill" line above + print_footprint("after decode ", prompt_len + generated); + } + + const double pf = ms(t0, t1), dc = ms(t1, t2); + const double pf_t = prompt_len / (pf / 1000.0); + const double dc_t = dc > 0 ? generated / (dc / 1000.0) : 0.0; + std::cout << "\n[iter " << iter << (measured ? "" : " warmup") + << "] prefill " << pf_t << " tok/s (" << prompt_len << " tok, " + << pf << " ms) | decode " << dc_t << " tok/s (" << generated + << " tok, " << dc << " ms)\n"; + if (measured) { + pf_tps.push_back(pf_t); + dc_tps.push_back(dc_t); + } + } + + auto summarize = [](const char* label, const std::vector& v) { + double m = 0.0; + for (double x : v) { + m += x; + } + m /= static_cast(v.size()); + double var = 0.0; + for (double x : v) { + var += (x - m) * (x - m); + } + const double sd = v.size() > 1 + ? std::sqrt(var / static_cast(v.size() - 1)) + : 0.0; + std::cout << label << ": " << m << " +/- " << sd + << " tok/s (n=" << v.size() << ")\n"; + }; + std::cout << "\n"; + summarize("prefill", pf_tps); + summarize("decode ", dc_tps); + return 0; + } catch (const std::exception& e) { + std::cerr << "Error: " << e.what() << std::endl; + return 1; + } +} From 977de443e40d604da9f399cb188ec0c9da0cc690 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Tue, 11 Aug 2026 09:40:16 -0700 Subject: [PATCH 2/4] Address review on the MLX off-graph runner --- backends/mlx/examples/llm/CMakeLists.txt | 12 +- backends/mlx/examples/llm/README.md | 116 ++++++++++ backends/mlx/examples/llm/run_llm_hf.cpp | 267 ++++++++++++++--------- 3 files changed, 295 insertions(+), 100 deletions(-) diff --git a/backends/mlx/examples/llm/CMakeLists.txt b/backends/mlx/examples/llm/CMakeLists.txt index c36f042fd6e..74944f143ae 100644 --- a/backends/mlx/examples/llm/CMakeLists.txt +++ b/backends/mlx/examples/llm/CMakeLists.txt @@ -30,8 +30,13 @@ list(APPEND CMAKE_FIND_ROOT_PATH ${CMAKE_CURRENT_BINARY_DIR}/../../../..) find_package(executorch CONFIG REQUIRED FIND_ROOT_PATH_BOTH) executorch_target_link_options_shared_lib(executorch) +# As in examples/models/llama: point gflags_DIR at the copy the executorch build +# configured, since a separate find_package() project cannot see it. +set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../../third-party/gflags) +find_package(gflags REQUIRED) + set(link_libraries executorch extension_module extension_tensor - extension_llm_cache + extension_llm_cache gflags ) if(NOT TARGET mlxdelegate) @@ -55,3 +60,8 @@ target_include_directories( ${_flatbuffers_include} ) target_link_libraries(mlx_run_llm_hf PUBLIC ${link_libraries}) + +# The copy helper is gated on EXECUTORCH_BUILD_MLX, which the installed config +# does not set; reaching here means mlxdelegate exists. +set(EXECUTORCH_BUILD_MLX ON) +executorch_target_copy_mlx_metallib(mlx_run_llm_hf) diff --git a/backends/mlx/examples/llm/README.md b/backends/mlx/examples/llm/README.md index 738bbfb8c14..ea3dc498b1c 100644 --- a/backends/mlx/examples/llm/README.md +++ b/backends/mlx/examples/llm/README.md @@ -7,6 +7,7 @@ This example demonstrates how to export and run LLMs using the MLX delegate for - **Export**: Convert HuggingFace LLMs to ExecuTorch format with MLX delegate - **Quantization**: Optional INT4/INT8 weight quantization via TorchAO - **KV Cache**: Efficient KV cache implementation for autoregressive generation +- **Off-graph KV Cache**: Optional runtime-owned cache, sized and configured per run instead of at export - **Custom Ops**: Uses `mlx::custom_sdpa` and `mlx::kv_cache_update` for optimal execution on MLX - **Pybindings**: Run inference using ExecuTorch Python bindings - **Gemma 4**: Text-only export and run flow supports processor-backed checkpoints such as `google/gemma-4-E2B-it` @@ -23,6 +24,7 @@ pip install transformers optimum-executorch |--------|-------------| | `export_llm_hf` | Export LLMs using optimum-executorch pipeline, with optional custom MLX SDPA/KV cache | | `run_llm_hf` | Run exported models with token-by-token generation | +| `run_llm_hf.cpp` | C++ runner; the run path for off-graph-cache exports, and runs in-graph ones too | For exporting via the ExecuTorch LLM pipeline (e.g. `examples/models/llama`), use `--mlx` to enable the MLX delegate. @@ -54,6 +56,13 @@ python -m executorch.backends.mlx.examples.llm.export_llm_hf \ --qlinear 4w \ --qembedding 4w +# Off-graph KV cache: the cache is owned by the runtime, not baked into the .pte +python -m executorch.backends.mlx.examples.llm.export_llm_hf \ + --model-id "unsloth/gemma-3-1b-it" \ + --output gemma3_offgraph.pte \ + --use-offgraph-cache \ + --max-seq-len 1024 + # Gemma 4 text-only export python -m executorch.backends.mlx.examples.llm.export_llm_hf \ --model-id "google/gemma-4-E2B-it" \ @@ -86,6 +95,11 @@ pip install -U "transformers @ git+https://github.com/huggingface/transformers.g | `--no-tie-word-embeddings` | `False` | Disable re-tying lm_head to embedding after quantization | | `--use-custom-sdpa` | `False` | Use MLX custom SDPA (`mlx::custom_sdpa`) | | `--use-custom-kv-cache` | `False` | Use MLX custom KV cache (`mlx::kv_cache_update`) | +| `--use-offgraph-cache` | `False` | Use the off-graph KV cache (`kvcache::update_and_attend`); replaces the two flags above | + +Off-graph exports keep no cache in the `.pte`, so the pybindings `run_llm_hf` +cannot run them — use [`mlx_run_llm_hf`](#mlx_run_llm_hf-c) below, which builds +the cache and binds it at load time. --- @@ -123,6 +137,108 @@ python -m executorch.backends.mlx.examples.llm.run_llm_hf \ --- +## `mlx_run_llm_hf` (C++) + +Native runner for models exported with `--use-offgraph-cache`. That cache lives +outside the graph and is handed to the backend by key at load time, which the +pybindings `run_llm_hf` above cannot do. It also runs in-graph `.pte` files +unchanged — omit `--kv-max-capacity` — so the same binary compares both cache +paths. Greedy decode. + +### Build + +A standalone `find_package(executorch)` project, so ExecuTorch must be installed +first: + +```bash +cmake --preset mlx-release +cmake --build cmake-out --target install -j$(( $(sysctl -n hw.ncpu) - 1 )) + +cmake -S backends/mlx/examples/llm -B cmake-out/backends/mlx/examples/llm \ + -DCMAKE_BUILD_TYPE=Release +cmake --build cmake-out/backends/mlx/examples/llm -j$(( $(sysctl -n hw.ncpu) - 1 )) +``` + +The binary lands at `cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf`. + +### Run + +```bash +python -m executorch.backends.mlx.examples.llm.export_llm_hf \ + --model-id unsloth/gemma-3-1b-it \ + --output gemma3_offgraph.pte \ + --use-offgraph-cache \ + --max-seq-len 1024 + +cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf \ + --pte gemma3_offgraph.pte \ + --tokenizer ~/.cache/huggingface/hub/models--unsloth--gemma-3-1b-it/snapshots/*/tokenizer.json \ + --chat gemma \ + --kv-max-capacity 1024 \ + --prompt "What is the capital of France?" \ + --max-new-tokens 50 +``` + +`--chat` selects the instruct template — `llama3`, `gemma`, `gemma4`, or `0` for +raw text. It matters: raw text confuses an instruct model into emitting turn +markers, and using the wrong template invalidates a comparison between two +`.pte` files. + +### Configuring the cache at run time + +Only the cache *geometry* is fixed at export — how many caches, their KV heads, +head dims and windows, which the export publishes as constant methods the runner +reads before building the cache. Everything else is chosen per run, with no +re-export. The runner reports what it built: + +``` +[cache] off-graph seq | capacity=1024 initial=512 kv_dtype=BFloat16 + 26 layers: 4 flat + 22 ring(window 512) +``` + +`--kv-windows` overrides the attention pattern, repeating a comma-separated list +over the caches (`0` = flat). The geometry each cache declared is untouched, so +this cannot desync from the graph. For gemma-3-1b, whose 26 layers are 22 +sliding at 512 and 4 full: + +```bash +# (omitted) 26 layers: 4 flat + 22 ring(window 512) +--kv-windows 512,512,512,512,512,0 # the same, spelling out gemma-3's 5:1 period +--kv-windows 0 # 26 layers: 26 flat +--kv-windows 512 # 26 layers: 0 flat + 26 ring(window 512) +--kv-windows 512,256 # 26 layers: 0 flat + 13 ring(256) + 13 ring(512) +``` + +A ring layer allocates its whole `window + chunk - 1` slots up front, while a +flat layer starts at `--kv-initial-capacity` and doubles as the sequence grows, +so an all-flat cache can look smaller than a sliding one early in a run and +larger later. Prefill runs in `--prefill-chunk-size` steps precisely so that +the ring is bounded by the chunk rather than by the prompt: prefilling 24k +tokens in one step would need 541 MiB of ring here, against 24 MiB at the +512 default. + +### Options + +`--help` lists every flag with its default. + +| Option | Default | Description | +|--------|---------|-------------| +| `--pte` | *(required)* | Path to .pte file | +| `--tokenizer` | *(required)* | Path to `tokenizer.json` | +| `--prompt` | `The quick brown fox` | Input prompt | +| `--max-new-tokens` | `50` | Tokens to generate, excluding the prompt | +| `--chat` | `llama3` | Chat template: `llama3`, `gemma`, `gemma4`, or `0` to disable | +| `--kv-max-capacity` | `0` | Off-graph: history the cache may hold. Setting it selects the off-graph path | +| `--kv-storage-dtype` | `bf16` | Off-graph: KV storage dtype (`bf16`, `fp16`, `fp32`) | +| `--kv-initial-capacity` | `-1` | Off-graph: starting pool size; grows by doubling up to capacity | +| `--prefill-chunk-size` | `512` | Tokens per prefill step; also the largest single write, so it sizes a ring layer to `window + chunk - 1`. Must not exceed the sequence length the `.pte` was exported with | +| `--kv-windows` | *(model's own)* | Off-graph: attention pattern override, e.g. `512` | +| `--interactive` | `false` | Multi-turn chat on stdin; off-graph only | +| `--warmup` | `0` | Throwaway iterations before measuring | +| `--iters` | `1` | Measured iterations; reports tok/s and a mean +/- stddev | + +--- + ## Architecture The `export_llm_hf` script uses optimum-executorch's `CausalLMExportableModule` by default. When custom flags are enabled, it uses `TorchExportableModuleWithStaticCache` from HuggingFace transformers, with optional MLX-specific replacements: diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp index 13a6c995de2..6fc8d1b5ff1 100644 --- a/backends/mlx/examples/llm/run_llm_hf.cpp +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -7,20 +7,17 @@ */ // C++ runner for HuggingFace LLMs on the MLX backend. Unlike the pybindings -// run_llm_hf.py, it can bind the off-graph KV cache: with --kv-cache it builds -// an MLXSequenceCache, installs it in the process-global registry, and passes -// its cache_key as a load-time backend option (the rendezvous init() reads). -// Its shape comes from constant methods the export publishes; the flags below -// only choose policy. Without --kv-max-capacity it runs an in-graph model +// run_llm_hf.py, it can bind the off-graph KV cache: with --kv-max-capacity it +// builds an MLXSequenceCache, installs it in the process-global registry, and +// passes its cache_key as a load-time backend option (the rendezvous init() +// reads). Its shape comes from constant methods the export publishes; the flags +// below only choose policy. Without --kv-max-capacity it runs an in-graph model // unchanged -- so the same binary compares both cache paths. Greedy decode. // // Usage: -// run_llm_hf --pte --tokenizer \ -// [--kv-max-capacity N] [--kv-storage-dtype bf16|fp16|fp32] \ -// [--kv-initial-capacity N] [--kv-max-write N] \ -// [--kv-windows ] \ -// [--prompt "..."] [--max-new-tokens N] [--chat llama3|gemma|gemma4|0] \ -// [--warmup N] [--iters N] [--interactive 1] +// run_llm_hf --pte --tokenizer [flags] +// +// --help lists every flag with its default. #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wconversion" @@ -37,16 +34,69 @@ #include +#include #include #include #include #include #include +#include #include #include #include +DEFINE_string(pte, "", "Model .pte file."); +DEFINE_string(tokenizer, "", "tokenizer.json for the model."); +DEFINE_string(prompt, "The quick brown fox", "Prompt to generate from."); +DEFINE_int32(max_new_tokens, 50, "Tokens to generate, excluding the prompt."); +DEFINE_string( + chat, + "llama3", + "Instruct chat template to wrap the prompt in: llama3, gemma, gemma4, or 0 " + "to disable. Raw text confuses an instruct model into emitting turn " + "markers."); +DEFINE_int32( + kv_max_capacity, + 0, + "Off-graph: how much history the cache may hold. Setting it selects the " + "off-graph path; the cache's shape comes from the .pte, so the kv_ flags " + "only choose policy."); +DEFINE_string( + kv_storage_dtype, + "bf16", + "Off-graph: KV storage dtype, bf16|fp16|fp32."); +DEFINE_int32( + kv_initial_capacity, + -1, + "Off-graph: the cache pool's starting size; it grows (doubling) up to " + "capacity. -1 keeps the CacheConfig default. Small values force growth."); +DEFINE_int32( + prefill_chunk_size, + 512, + "Tokens per prefill step. This is the largest single write, so it also " + "sizes a ring layer to window + chunk - 1 -- without it a long prompt " + "would need a ring as large as itself, defeating the window. Must not " + "exceed the sequence length the .pte was exported with."); +DEFINE_string( + kv_windows, + "", + "Off-graph: impose an attention pattern other than the model's own, e.g. " + "\"512\" to make every layer sliding."); +DEFINE_bool( + interactive, + false, + "Multi-turn chat on stdin instead of a single prompt; off-graph only."); +DEFINE_int32( + warmup, + 0, + "Throwaway iterations before measuring, to absorb JIT and pool growth."); +DEFINE_int32( + iters, + 1, + "Measured iterations; reports per-iter tok/s and a mean +/- stddev " + "summary."); + using ::executorch::extension::make_tensor_ptr; using ::executorch::extension::Module; using ::executorch::extension::TensorPtr; @@ -127,6 +177,37 @@ std::optional> const_ints(Module& module, const char* name) { return std::vector(p, p + t.numel()); } +// Fill in the cache geometry the export published: get_n_caches, then one +// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat). +// Only the shape comes from the .pte -- capacity, dtype and any override stay +// with the flags. False means this is not an off-graph model. +bool read_kv_layout(Module& module, cache::CacheConfig& cfg) { + const auto n_caches = const_int(module, "get_n_caches"); + const auto kv_heads = const_ints(module, "get_kv_heads"); + const auto head_dims = const_ints(module, "get_head_dims"); + const auto windows = const_ints(module, "get_windows"); + if (!n_caches || !kv_heads || !head_dims || !windows) { + return false; + } + const size_t n = static_cast(*n_caches); + if (kv_heads->size() != n || head_dims->size() != n || windows->size() != n) { + return false; + } + cfg.n_layers = static_cast(n); + cfg.layers.clear(); + cfg.layers.reserve(n); + for (size_t l = 0; l < n; ++l) { + cache::LayerConfig lc{}; + lc.n_kv_heads = (*kv_heads)[l]; + lc.head_dim = (*head_dims)[l]; + lc.policy = (*windows)[l] > 0 + ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, (*windows)[l]} + : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; + cfg.layers.push_back(lc); + } + return true; +} + // Replace the model's own attention pattern with `spec`, a comma-separated list // of windows repeating over the caches (0 = flat). One entry makes every layer // sliding. Only the policy changes; each cache keeps the geometry the .pte @@ -192,6 +273,36 @@ std::string dtype_name(int st) { } } +// Announce the cache shape: the same .pte runs under whatever config this +// invocation asks for -- capacity, storage dtype, flat/ring layers -- with no +// re-export. The footprint lines printed later then show it growing at runtime. +void print_cache_summary(const cache::CacheConfig& cfg) { + // Ring layers grouped by window: --kv-windows can give each layer its own, + // and the pools are sized per layer, so a single number would misreport them. + std::map ring; + int flat = 0; + for (int l = 0; l < cfg.n_layers; ++l) { + const cache::LayerConfig& lc = + cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; + if (lc.policy.kind == cache::LayerPolicy::Kind::Ring) { + ++ring[lc.policy.window]; + } else { + ++flat; + } + } + std::cout << "\n[cache] off-graph seq | capacity=" << cfg.capacity + << " initial=" << cfg.initial_capacity + << " kv_dtype=" << dtype_name(cfg.kv_dtype); + if (cfg.max_write) { + std::cout << " max_write=" << *cfg.max_write; + } + std::cout << "\n " << cfg.n_layers << " layers: " << flat << " flat"; + for (const auto& [window, n] : ring) { + std::cout << " + " << n << " ring(window " << window << ")"; + } + std::cout << std::endl; +} + // One user turn wrapped in the model's instruct template. Returns false for an // unknown template name. The leading BOS belongs to the first turn only, so a // continuing conversation passes with_bos=false. @@ -219,48 +330,24 @@ bool wrap_turn( return true; } -std::string -arg(int argc, char** argv, const std::string& key, const std::string& def) { - for (int i = 1; i + 1 < argc; ++i) { - if (key == argv[i]) { - return argv[i + 1]; - } - } - return def; -} - } // namespace int main(int argc, char** argv) { - const std::string pte = arg(argc, argv, "--pte", ""); - const std::string tok_path = arg(argc, argv, "--tokenizer", ""); - // Off-graph: how much history the cache may hold. Its presence selects the - // off-graph path; the cache's shape comes from the .pte, so only policy is - // set here. - const int kv_capacity = std::stoi(arg(argc, argv, "--kv-max-capacity", "0")); - const std::string kv_dtype = arg(argc, argv, "--kv-storage-dtype", "bf16"); - // Optional: impose an attention pattern other than the model's own, e.g. - // "512" to make every layer sliding. - const std::string kv_windows = arg(argc, argv, "--kv-windows", ""); - const std::string prompt = arg(argc, argv, "--prompt", "The quick brown fox"); - const int max_new = std::stoi(arg(argc, argv, "--max-new-tokens", "50")); - // Instruct chat template to wrap the prompt in: llama3 (default) or gemma. - // Raw text confuses an instruct model into emitting turn markers; 0 disables. - const std::string chat = arg(argc, argv, "--chat", "llama3"); - // Off-graph only: the cache pool's starting size; it grows (doubling) up to - // capacity. -1 keeps the CacheConfig default. Small values force growth. - const int initial_capacity = - std::stoi(arg(argc, argv, "--kv-initial-capacity", "-1")); - // Off-graph only: the largest single step, which sizes a ring layer's slots - // to window + max_write - 1. -1 lets each ring layer use its own window. - // Prefill goes in one step, so this must be at least the prompt length. - const int max_write = std::stoi(arg(argc, argv, "--kv-max-write", "-1")); - // Benchmarking: `warmup` throwaway iters (absorb JIT + pool growth) then - // `iters` measured; per-iter tok/s plus a mean +/- stddev summary. - // Multi-turn chat on stdin instead of a single prompt; off-graph only. - const bool interactive = arg(argc, argv, "--interactive", "0") != "0"; - const int warmup = std::stoi(arg(argc, argv, "--warmup", "0")); - const int iters = std::stoi(arg(argc, argv, "--iters", "1")); + gflags::ParseCommandLineFlags(&argc, &argv, true); + + const std::string& pte = FLAGS_pte; + const std::string& tok_path = FLAGS_tokenizer; + const std::string& kv_dtype = FLAGS_kv_storage_dtype; + const std::string& kv_windows = FLAGS_kv_windows; + const std::string& prompt = FLAGS_prompt; + const std::string& chat = FLAGS_chat; + const int kv_capacity = FLAGS_kv_max_capacity; + const int max_new = FLAGS_max_new_tokens; + const int initial_capacity = FLAGS_kv_initial_capacity; + const int chunk = FLAGS_prefill_chunk_size; + const bool interactive = FLAGS_interactive; + const int warmup = FLAGS_warmup; + const int iters = FLAGS_iters; if (pte.empty() || tok_path.empty()) { std::cerr << "Required: --pte --tokenizer " "[--kv-max-capacity N for off-graph models]\n"; @@ -303,27 +390,11 @@ int main(int argc, char** argv) { << " (bf16|fp16|fp32)" << std::endl; return 1; } - const auto n_caches = const_int(module, "get_n_caches"); - const auto kv_heads = const_ints(module, "get_kv_heads"); - const auto head_dims = const_ints(module, "get_head_dims"); - const auto windows = const_ints(module, "get_windows"); - if (!n_caches || !kv_heads || !head_dims || !windows) { + if (!read_kv_layout(module, cfg)) { std::cerr << "No KV cache layout in " << pte << "; re-export with --use-offgraph-cache" << std::endl; return 1; } - cfg.n_layers = static_cast(*n_caches); - cfg.layers.clear(); - cfg.layers.reserve(static_cast(cfg.n_layers)); - for (size_t l = 0; l < static_cast(cfg.n_layers); ++l) { - cache::LayerConfig lc{}; - lc.n_kv_heads = (*kv_heads)[l]; - lc.head_dim = (*head_dims)[l]; - lc.policy = (*windows)[l] > 0 - ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, (*windows)[l]} - : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; - cfg.layers.push_back(lc); - } if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) { std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl; return 1; @@ -335,9 +406,8 @@ int main(int argc, char** argv) { if (initial_capacity >= 0) { cfg.initial_capacity = initial_capacity; } - if (max_write > 0) { - cfg.max_write = max_write; - } + // Prefill is chunked, so the chunk is the largest step the cache sees. + cfg.max_write = chunk; auto built = cache::CacheBuilderRegistry::global().build( ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); if (!built.ok()) { @@ -347,33 +417,7 @@ int main(int argc, char** argv) { } session.emplace(cache::make_unique_key(), built.get()); - // Announce the cache shape: the same .pte runs under whatever config this - // invocation asks for -- capacity, storage dtype, flat/ring layers -- - // with no re-export. The footprint lines below then show it growing at - // runtime. - int flat = 0, ring = 0, ring_window = 0; - for (int l = 0; l < cfg.n_layers; ++l) { - const cache::LayerConfig& lc = - cfg.layers.size() == 1 ? cfg.layers.front() : cfg.layers[l]; - if (lc.policy.kind == cache::LayerPolicy::Kind::Ring) { - ++ring; - ring_window = lc.policy.window; - } else { - ++flat; - } - } - std::cout << "\n[cache] off-graph seq | capacity=" << cfg.capacity - << " initial=" << cfg.initial_capacity - << " kv_dtype=" << dtype_name(cfg.kv_dtype); - if (cfg.max_write) { - std::cout << " max_write=" << *cfg.max_write; - } - std::cout << "\n " << cfg.n_layers << " layers: " << flat - << " flat"; - if (ring > 0) { - std::cout << " + " << ring << " ring(window " << ring_window << ")"; - } - std::cout << std::endl; + print_cache_summary(cfg); if (mlx_opts.set_option( ::executorch::backends::mlx::kCacheKeyKey, session->key().c_str()) != Error::Ok || @@ -418,11 +462,13 @@ int main(int argc, char** argv) { // Stop on end-of-text and (for chat) the turn-end token <|eot_id|>. std::vector stop_ids = {tokenizer.eos_tok()}; + std::optional turn_end_id; if (chat != "0") { const char* turn_end = chat == "llama3" ? "<|eot_id|>" : chat == "gemma4" ? "" : ""; if (auto eot = tokenizer.piece_to_id(turn_end); eot.ok()) { + turn_end_id = static_cast(*eot); stop_ids.push_back(*eot); } } @@ -447,6 +493,21 @@ int main(int argc, char** argv) { return argmax_last(out->at(0).toTensor()); }; + // Prefill in chunks, so a ring layer holds window + chunk - 1 slots rather + // than growing with the prompt. Only the last chunk's token is kept; the + // earlier ones exist to place their K/V in the cache. + auto prefill = [&](const std::vector& ids, + const std::vector& pos) { + int64_t next = 0; + for (size_t off = 0; off < ids.size(); off += chunk) { + const size_t n = std::min(static_cast(chunk), ids.size() - off); + next = step( + {ids.begin() + off, ids.begin() + off + n}, + {pos.begin() + off, pos.begin() + off + n}); + } + return next; + }; + // Multi-turn: history stays in the cache, so each turn only prefills its // own tokens at the running position. /reset and /undo drive the cache's // control face directly -- off-graph only, since an in-graph cache gives @@ -524,7 +585,7 @@ int main(int argc, char** argv) { for (int i = 0; i < n; ++i) { tpos.push_back(position + i); } - int64_t next = step(tin, tpos); + int64_t next = prefill(tin, tpos); position += n; uint64_t prev = te->back(); @@ -537,6 +598,14 @@ int main(int argc, char** argv) { next = step({next}, {position}); ++position; } + // The turn-end token stops generation, so it is neither printed nor + // fed back -- but the next turn opens without closing this one, and an + // unterminated assistant turn compounds over a session. Commit it, at + // the cost of one extra step per turn. + if (turn_end_id && next == *turn_end_id && control->can_extend(1)) { + step({next}, {position}); + ++position; + } std::cout << "\n[" << position << "/" << control->capacity() << " tokens" << (budget < max_new ? ", generation capped by capacity" : "") @@ -585,7 +654,7 @@ int main(int argc, char** argv) { const bool print_mem = (iter == total_iters - 1); const auto t0 = std::chrono::steady_clock::now(); - int64_t next = step(ids, prefill_pos); + int64_t next = prefill(ids, prefill_pos); const auto t1 = std::chrono::steady_clock::now(); if (print_text || print_mem) { std::cout << "\n"; // blank line separating this section from the banner From 1aa2e73bde29fe4be6ec2debd88a0ea571d14b54 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Tue, 11 Aug 2026 10:42:46 -0700 Subject: [PATCH 3/4] Address review: reuse the shared LLM helpers, take the chunk from metadata --- backends/mlx/examples/llm/CMakeLists.txt | 10 +- backends/mlx/examples/llm/README.md | 12 +- backends/mlx/examples/llm/export_llm_hf.py | 31 +++ backends/mlx/examples/llm/run_llm_hf.cpp | 240 ++++++++++----------- 4 files changed, 155 insertions(+), 138 deletions(-) diff --git a/backends/mlx/examples/llm/CMakeLists.txt b/backends/mlx/examples/llm/CMakeLists.txt index 74944f143ae..fa041b7129a 100644 --- a/backends/mlx/examples/llm/CMakeLists.txt +++ b/backends/mlx/examples/llm/CMakeLists.txt @@ -35,8 +35,14 @@ executorch_target_link_options_shared_lib(executorch) set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../../third-party/gflags) find_package(gflags REQUIRED) -set(link_libraries executorch extension_module extension_tensor - extension_llm_cache gflags +set(link_libraries + executorch + extension_module + extension_tensor + extension_llm_cache + extension_llm_runner + extension_llm_sampler + gflags ) if(NOT TARGET mlxdelegate) diff --git a/backends/mlx/examples/llm/README.md b/backends/mlx/examples/llm/README.md index ea3dc498b1c..a3fe12ca6b3 100644 --- a/backends/mlx/examples/llm/README.md +++ b/backends/mlx/examples/llm/README.md @@ -96,6 +96,7 @@ pip install -U "transformers @ git+https://github.com/huggingface/transformers.g | `--use-custom-sdpa` | `False` | Use MLX custom SDPA (`mlx::custom_sdpa`) | | `--use-custom-kv-cache` | `False` | Use MLX custom KV cache (`mlx::kv_cache_update`) | | `--use-offgraph-cache` | `False` | Use the off-graph KV cache (`kvcache::update_and_attend`); replaces the two flags above | +| `--prefill-chunk-size` | `512` | Off-graph: tokens per prefill step, published as `get_prefill_chunk_size` for the runner. It is the largest single write, so a ring layer is sized `window + chunk - 1`; it may not exceed the sliding window or the exported sequence length | Off-graph exports keep no cache in the `.pte`, so the pybindings `run_llm_hf` cannot run them — use [`mlx_run_llm_hf`](#mlx_run_llm_hf-c) below, which builds @@ -212,10 +213,10 @@ sliding at 512 and 4 full: A ring layer allocates its whole `window + chunk - 1` slots up front, while a flat layer starts at `--kv-initial-capacity` and doubles as the sequence grows, so an all-flat cache can look smaller than a sliding one early in a run and -larger later. Prefill runs in `--prefill-chunk-size` steps precisely so that +larger later. Prefill runs in steps of the chunk size the export published, so the ring is bounded by the chunk rather than by the prompt: prefilling 24k -tokens in one step would need 541 MiB of ring here, against 24 MiB at the -512 default. +tokens in one step would need 541 MiB of ring here, against 24 MiB at the 512 +default. ### Options @@ -227,15 +228,14 @@ tokens in one step would need 541 MiB of ring here, against 24 MiB at the | `--tokenizer` | *(required)* | Path to `tokenizer.json` | | `--prompt` | `The quick brown fox` | Input prompt | | `--max-new-tokens` | `50` | Tokens to generate, excluding the prompt | +| `--temperature` | `0` | Sampling temperature; 0 is greedy argmax, which is what makes two `.pte` files comparable | | `--chat` | `llama3` | Chat template: `llama3`, `gemma`, `gemma4`, or `0` to disable | | `--kv-max-capacity` | `0` | Off-graph: history the cache may hold. Setting it selects the off-graph path | | `--kv-storage-dtype` | `bf16` | Off-graph: KV storage dtype (`bf16`, `fp16`, `fp32`) | | `--kv-initial-capacity` | `-1` | Off-graph: starting pool size; grows by doubling up to capacity | -| `--prefill-chunk-size` | `512` | Tokens per prefill step; also the largest single write, so it sizes a ring layer to `window + chunk - 1`. Must not exceed the sequence length the `.pte` was exported with | | `--kv-windows` | *(model's own)* | Off-graph: attention pattern override, e.g. `512` | | `--interactive` | `false` | Multi-turn chat on stdin; off-graph only | -| `--warmup` | `0` | Throwaway iterations before measuring | -| `--iters` | `1` | Measured iterations; reports tok/s and a mean +/- stddev | +| `--warmup` | `false` | Run once before measuring, to absorb JIT and pool growth | --- diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py index adebc121e51..d1800bf91df 100644 --- a/backends/mlx/examples/llm/export_llm_hf.py +++ b/backends/mlx/examples/llm/export_llm_hf.py @@ -357,6 +357,7 @@ def _export_with_offgraph_cache( no_tie_word_embeddings: bool = False, qlinear_group_size: Optional[int] = None, qembedding_group_size: Optional[int] = None, + prefill_chunk_size: int = 512, ) -> None: """ Export using the off-graph KV cache op (kvcache::update_and_attend). @@ -385,6 +386,14 @@ def _export_with_offgraph_cache( } torch_dtype = torch_dtype_map.get(dtype, torch.bfloat16) + # The chunk is the largest single step, so it cannot exceed the sequence + # length traced below. Checked before the model loads, which is slow. + if prefill_chunk_size < 1 or prefill_chunk_size > max_seq_len - 1: + raise ValueError( + f"--prefill-chunk-size {prefill_chunk_size} must be in " + f"[1, {max_seq_len - 1}], the largest step this export traces" + ) + register_mlx_offgraph_attention() logger.info("Registered MLX off-graph attention (update_and_attend)") @@ -429,6 +438,16 @@ def _export_with_offgraph_cache( cache_windows = [ sliding_window if t == "sliding_attention" else 0 for t in layer_types ] + # Beyond the window the chunk, not the window, decides how much a sliding + # layer holds, which is the opposite of the point. + sliding = [w for w in cache_windows if w > 0] + if sliding and prefill_chunk_size > min(sliding): + raise ValueError( + f"--prefill-chunk-size {prefill_chunk_size} exceeds the sliding " + f"window {min(sliding)}: a ring layer holds window + chunk - 1 " + "slots, so the chunk would dominate the window" + ) + kv_metadata = { # The cache count, not num_hidden_layers: gemma-4 E2B shares KV across # its tail, so 15 caches for 35 layers. @@ -436,6 +455,7 @@ def _export_with_offgraph_cache( "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32), "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32), "get_windows": torch.tensor(cache_windows, dtype=torch.int32), + "get_prefill_chunk_size": prefill_chunk_size, } logger.info( f"KV cache layout: {len(layer_types)} caches, " @@ -512,6 +532,7 @@ def export_llama_hf( no_tie_word_embeddings: bool = False, qlinear_group_size: Optional[int] = None, qembedding_group_size: Optional[int] = None, + prefill_chunk_size: int = 512, ) -> None: """ Export a HuggingFace Llama model to ExecuTorch with MLX backend. @@ -544,6 +565,7 @@ def export_llama_hf( no_tie_word_embeddings=no_tie_word_embeddings, qlinear_group_size=qlinear_group_size, qembedding_group_size=qembedding_group_size, + prefill_chunk_size=prefill_chunk_size, ) elif use_custom_sdpa or use_custom_kv_cache: logger.info( @@ -637,6 +659,14 @@ def main(): help="Use the off-graph KV cache op (kvcache::update_and_attend); " "replaces --use-custom-sdpa/--use-custom-kv-cache", ) + parser.add_argument( + "--prefill-chunk-size", + type=int, + default=512, + help="Off-graph: tokens per prefill step, published for the runner. " + "It is the largest single write, so a ring layer is sized " + "window + chunk - 1; it may not exceed the sliding window", + ) args = parser.parse_args() @@ -654,6 +684,7 @@ def main(): no_tie_word_embeddings=args.no_tie_word_embeddings, qlinear_group_size=args.qlinear_group_size, qembedding_group_size=args.qembedding_group_size, + prefill_chunk_size=args.prefill_chunk_size, ) diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp index 6fc8d1b5ff1..37ae7bc8921 100644 --- a/backends/mlx/examples/llm/run_llm_hf.cpp +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -12,10 +12,11 @@ // passes its cache_key as a load-time backend option (the rendezvous init() // reads). Its shape comes from constant methods the export publishes; the flags // below only choose policy. Without --kv-max-capacity it runs an in-graph model -// unchanged -- so the same binary compares both cache paths. Greedy decode. +// unchanged -- so the same binary compares both cache paths. Greedy decode +// unless --temperature is set. // // Usage: -// run_llm_hf --pte --tokenizer [flags] +// run_llm_hf --pte --tokenizer [flags] // // --help lists every flag with its default. @@ -29,10 +30,14 @@ #include #include #include +#include +#include +#include +#include #include #include -#include +#include #include #include @@ -44,12 +49,22 @@ #include #include #include +#include #include DEFINE_string(pte, "", "Model .pte file."); -DEFINE_string(tokenizer, "", "tokenizer.json for the model."); +DEFINE_string( + tokenizer, + "", + "Tokenizer file; any format the shared loader accepts (tokenizer.json, " + "tiktoken, sentencepiece)."); DEFINE_string(prompt, "The quick brown fox", "Prompt to generate from."); DEFINE_int32(max_new_tokens, 50, "Tokens to generate, excluding the prompt."); +DEFINE_double( + temperature, + 0.0, + "Sampling temperature. 0 is greedy argmax, which is what makes two .pte " + "files comparable; above 0 samples and the run stops being reproducible."); DEFINE_string( chat, "llama3", @@ -71,13 +86,6 @@ DEFINE_int32( -1, "Off-graph: the cache pool's starting size; it grows (doubling) up to " "capacity. -1 keeps the CacheConfig default. Small values force growth."); -DEFINE_int32( - prefill_chunk_size, - 512, - "Tokens per prefill step. This is the largest single write, so it also " - "sizes a ring layer to window + chunk - 1 -- without it a long prompt " - "would need a ring as large as itself, defeating the window. Must not " - "exceed the sequence length the .pte was exported with."); DEFINE_string( kv_windows, "", @@ -87,15 +95,10 @@ DEFINE_bool( interactive, false, "Multi-turn chat on stdin instead of a single prompt; off-graph only."); -DEFINE_int32( +DEFINE_bool( warmup, - 0, - "Throwaway iterations before measuring, to absorb JIT and pool growth."); -DEFINE_int32( - iters, - 1, - "Measured iterations; reports per-iter tok/s and a mean +/- stddev " - "summary."); + false, + "Run once before measuring, to absorb JIT and pool growth."); using ::executorch::extension::make_tensor_ptr; using ::executorch::extension::Module; @@ -178,17 +181,21 @@ std::optional> const_ints(Module& module, const char* name) { } // Fill in the cache geometry the export published: get_n_caches, then one -// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat). -// Only the shape comes from the .pte -- capacity, dtype and any override stay -// with the flags. False means this is not an off-graph model. +// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat), +// plus get_prefill_chunk_size, which the export validates against the sliding +// window and which becomes max_write -- the largest step the cache may see. +// Capacity and dtype stay with the flags. False means this is not an off-graph +// model. bool read_kv_layout(Module& module, cache::CacheConfig& cfg) { const auto n_caches = const_int(module, "get_n_caches"); const auto kv_heads = const_ints(module, "get_kv_heads"); const auto head_dims = const_ints(module, "get_head_dims"); const auto windows = const_ints(module, "get_windows"); - if (!n_caches || !kv_heads || !head_dims || !windows) { + const auto chunk = const_int(module, "get_prefill_chunk_size"); + if (!n_caches || !kv_heads || !head_dims || !windows || !chunk) { return false; } + cfg.max_write = static_cast(*chunk); const size_t n = static_cast(*n_caches); if (kv_heads->size() != n || head_dims->size() != n || windows->size() != n) { return false; @@ -212,6 +219,10 @@ bool read_kv_layout(Module& module, cache::CacheConfig& cfg) { // of windows repeating over the caches (0 = flat). One entry makes every layer // sliding. Only the policy changes; each cache keeps the geometry the .pte // declared, so this cannot desync from the graph. +// +// The export sizes the chunk to the model's own window; narrowing the window +// here would leave the ring (window + max_write - 1) sized by the chunk +// instead, so the chunk follows the window down. bool apply_window_override(const std::string& spec, cache::CacheConfig& cfg) { std::vector pattern; if (!parse_int_list(spec, ',', pattern) || pattern.empty()) { @@ -223,38 +234,16 @@ bool apply_window_override(const std::string& spec, cache::CacheConfig& cfg) { ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, w} : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; } - return cache::valid(cfg); -} - -// argmax over the last position's vocab row of a [1, T, vocab] logits tensor, -// reading whatever float dtype the op emitted. -int64_t argmax_last(const ::executorch::aten::Tensor& logits) { - const auto dim = logits.dim(); - const int64_t vocab = logits.size(dim - 1); - const int64_t t = dim >= 2 ? logits.size(dim - 2) : 1; - const int64_t offset = (t - 1) * vocab; // start of the last row - auto scan = [&](auto* data) { - int64_t best = 0; - float best_v = -1e30f; - for (int64_t i = 0; i < vocab; ++i) { - const float v = static_cast(data[offset + i]); - if (v > best_v) { - best_v = v; - best = i; - } + int narrowest = 0; // smallest ring window in the pattern; 0 if all flat + for (int w : pattern) { + if (w > 0 && (narrowest == 0 || w < narrowest)) { + narrowest = w; } - return best; - }; - switch (logits.scalar_type()) { - case ::executorch::aten::ScalarType::Float: - return scan(logits.const_data_ptr()); - case ::executorch::aten::ScalarType::Half: - return scan(logits.const_data_ptr<::executorch::aten::Half>()); - case ::executorch::aten::ScalarType::BFloat16: - return scan(logits.const_data_ptr<::executorch::aten::BFloat16>()); - default: - throw std::runtime_error("argmax_last: unsupported logits dtype"); } + if (cfg.max_write && narrowest > 0 && narrowest < *cfg.max_write) { + cfg.max_write = narrowest; + } + return cache::valid(cfg); } // Human-readable name for a kv_dtype (an ET ScalarType int). Only the @@ -343,11 +332,10 @@ int main(int argc, char** argv) { const std::string& chat = FLAGS_chat; const int kv_capacity = FLAGS_kv_max_capacity; const int max_new = FLAGS_max_new_tokens; + const float temperature = static_cast(FLAGS_temperature); const int initial_capacity = FLAGS_kv_initial_capacity; - const int chunk = FLAGS_prefill_chunk_size; const bool interactive = FLAGS_interactive; - const int warmup = FLAGS_warmup; - const int iters = FLAGS_iters; + const bool warmup = FLAGS_warmup; if (pte.empty() || tok_path.empty()) { std::cerr << "Required: --pte --tokenizer " "[--kv-max-capacity N for off-graph models]\n"; @@ -355,9 +343,10 @@ int main(int argc, char** argv) { } try { - // Tokenizer. - ::tokenizers::HFTokenizer tokenizer; - if (tokenizer.load(tok_path) != ::tokenizers::Error::Ok) { + // The shared loader sniffs the format, so --tokenizer takes any of the + // files the other runners accept, not just tokenizer.json. + auto tokenizer = ::executorch::extension::llm::load_tokenizer(tok_path); + if (!tokenizer) { std::cerr << "Failed to load tokenizer: " << tok_path << std::endl; return 1; } @@ -371,11 +360,15 @@ int main(int argc, char** argv) { ::executorch::runtime::BackendOptions<1> mlx_opts; ::executorch::runtime::LoadBackendOptionsMap options_map; const bool off_graph = kv_capacity > 0; + // Tokens per prefill step, from the .pte. 0 means one step: an + // in-graph model publishes no chunk and has no ring to bound. + int prefill_chunk = 0; // Load the program but not forward: the cache must exist before forward's // backend init reads its key, and the layout it needs is published by // constant methods in the same file. Module module(pte); + const long load_start_ms = ::executorch::extension::llm::time_in_ms(); if (module.load() != Error::Ok) { std::cerr << "Failed to load " << pte << std::endl; return 1; @@ -406,8 +399,6 @@ int main(int argc, char** argv) { if (initial_capacity >= 0) { cfg.initial_capacity = initial_capacity; } - // Prefill is chunked, so the chunk is the largest step the cache sees. - cfg.max_write = chunk; auto built = cache::CacheBuilderRegistry::global().build( ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); if (!built.ok()) { @@ -415,6 +406,7 @@ int main(int argc, char** argv) { << static_cast(built.error()) << std::endl; return 1; } + prefill_chunk = cfg.max_write ? *cfg.max_write : 0; session.emplace(cache::make_unique_key(), built.get()); print_cache_summary(cfg); @@ -437,6 +429,11 @@ int main(int argc, char** argv) { std::cerr << "Failed to load forward" << std::endl; return 1; } + // Timings reported at the end, in the shared runner's format. + ::executorch::extension::llm::Stats stats; + stats.model_load_start_ms = load_start_ms; + stats.model_load_end_ms = ::executorch::extension::llm::time_in_ms(); + // Weights-only baseline, so the deltas below isolate the cache. const double mem_at_load = ::mlx::core::get_active_memory() / 1048576.0; std::cout << "[mem] after load : " << mem_at_load << " MiB" << std::endl; @@ -452,7 +449,7 @@ int main(int argc, char** argv) { } // The template carries its own BOS, so only a raw prompt asks for one. const int8_t bos = chat == "0" ? 1 : 0; - auto enc = tokenizer.encode(enc_input, bos, /*eos=*/0); + auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); if (!enc.ok()) { std::cerr << "Encode failed" << std::endl; return 1; @@ -460,16 +457,19 @@ int main(int argc, char** argv) { std::vector tokens = std::move(*enc); const int prompt_len = static_cast(tokens.size()); - // Stop on end-of-text and (for chat) the turn-end token <|eot_id|>. - std::vector stop_ids = {tokenizer.eos_tok()}; + // End-of-text from the model's metadata when it publishes any, else the + // tokenizer's. The turn-end token is ours: it depends on --chat, which the + // .pte knows nothing about. + std::unordered_set stop_ids = + ::executorch::extension::llm::get_eos_ids(tokenizer.get(), &module); std::optional turn_end_id; if (chat != "0") { const char* turn_end = chat == "llama3" ? "<|eot_id|>" : chat == "gemma4" ? "" : ""; - if (auto eot = tokenizer.piece_to_id(turn_end); eot.ok()) { + if (auto eot = tokenizer->piece_to_id(turn_end); eot.ok()) { turn_end_id = static_cast(*eot); - stop_ids.push_back(*eot); + stop_ids.insert(*eot); } } auto is_stop = [&](int64_t t) { @@ -481,6 +481,12 @@ int main(int argc, char** argv) { return false; }; + // One Sampler for the whole run, as the shared runner does: constructing + // one per token would reseed its RNG from the wall clock every time. Built + // on first use because the vocab size comes from the logits -- this export + // publishes no get_vocab_size. + std::optional<::executorch::extension::llm::Sampler> sampler; + auto step = [&](const std::vector& ids, const std::vector& pos) { auto in = @@ -490,7 +496,16 @@ int main(int argc, char** argv) { if (!out.ok()) { throw std::runtime_error("execute failed"); } - return argmax_last(out->at(0).toTensor()); + const auto& logits = out->at(0).toTensor(); + if (!sampler) { + sampler.emplace( + static_cast(logits.size(logits.dim() - 1)), temperature); + } + stats.on_sampling_begin(); + const int32_t tok = + ::executorch::extension::llm::sample_from_logits(logits, *sampler); + stats.on_sampling_end(); + return static_cast(tok); }; // Prefill in chunks, so a ring layer holds window + chunk - 1 slots rather @@ -498,9 +513,11 @@ int main(int argc, char** argv) { // earlier ones exist to place their K/V in the cache. auto prefill = [&](const std::vector& ids, const std::vector& pos) { + const size_t step_size = + prefill_chunk > 0 ? static_cast(prefill_chunk) : ids.size(); int64_t next = 0; - for (size_t off = 0; off < ids.size(); off += chunk) { - const size_t n = std::min(static_cast(chunk), ids.size() - off); + for (size_t off = 0; off < ids.size(); off += step_size) { + const size_t n = std::min(step_size, ids.size() - off); next = step( {ids.begin() + off, ids.begin() + off + n}, {pos.begin() + off, pos.begin() + off + n}); @@ -561,7 +578,7 @@ int main(int argc, char** argv) { std::string turn; wrap_turn(chat, line, /*with_bos=*/position == 0, turn); - auto te = tokenizer.encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); + auto te = tokenizer->encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); if (!te.ok()) { std::cerr << "Encode failed\n"; continue; @@ -590,7 +607,7 @@ int main(int argc, char** argv) { uint64_t prev = te->back(); for (int i = 0; i < budget && !is_stop(next); ++i) { - if (auto piece = tokenizer.decode(prev, static_cast(next)); + if (auto piece = tokenizer->decode(prev, static_cast(next)); piece.ok()) { std::cout << *piece << std::flush; } @@ -638,31 +655,23 @@ int main(int argc, char** argv) { << (mem - mem_at_load) << " MiB since load)" << std::endl; }; - // warmup iters absorb JIT + pool growth; measured iters record tok/s. The - // off-graph cache is cleared between iters so each prefill starts at length - // 0 (in-graph overwrites by position, so no reset needed there). - const int total_iters = warmup + std::max(1, iters); - std::vector pf_tps, dc_tps; - for (int iter = 0; iter < total_iters; ++iter) { + // One optional warmup run to absorb JIT and pool growth, then one measured + // run, as the shared LLM runners do. Repeats belong in a harness that + // restarts the process: clear() rewinds the sequence but leaves the pools + // at their grown size, so an in-process repeat cannot see reallocation. + for (int iter = 0; iter < (warmup ? 2 : 1); ++iter) { + const bool measured = !warmup || iter == 1; if (iter > 0 && off_graph) { session->control()->clear(); } - const bool measured = iter >= warmup; - const bool print_text = (iter == 0); - // Report the footprint from the last iteration: a measured one, and in - // steady state once any pool growth has settled. - const bool print_mem = (iter == total_iters - 1); - - const auto t0 = std::chrono::steady_clock::now(); + stats.inference_start_ms = ::executorch::extension::llm::time_in_ms(); int64_t next = prefill(ids, prefill_pos); - const auto t1 = std::chrono::steady_clock::now(); - if (print_text || print_mem) { - std::cout << "\n"; // blank line separating this section from the banner - } - if (print_mem) { + stats.prompt_eval_end_ms = ::executorch::extension::llm::time_in_ms(); + // prefill returns the first generated token, so TTFT ends with prefill + stats.first_token_ms = stats.prompt_eval_end_ms; + if (measured) { + std::cout << "\n"; print_footprint("after prefill", prompt_len); - } - if (print_text) { std::cout << "\n"; // blank line before the streamed generation } @@ -672,57 +681,28 @@ int main(int argc, char** argv) { if (is_stop(next)) { break; } - if (print_text) { - if (auto piece = tokenizer.decode(prev, static_cast(next)); + if (measured) { + if (auto piece = tokenizer->decode(prev, static_cast(next)); piece.ok()) { - std::cout << *piece << std::flush; + ::executorch::extension::llm::safe_printf(piece->c_str()); + fflush(stdout); } } prev = static_cast(next); ++generated; next = step({next}, {prompt_len + i}); } - const auto t2 = std::chrono::steady_clock::now(); - if (print_text) { + stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); + if (measured) { std::cout << "\n\n"; // close the generation line + blank separator - } - if (print_mem) { // trailing space aligns the colon with the "after prefill" line above print_footprint("after decode ", prompt_len + generated); - } - - const double pf = ms(t0, t1), dc = ms(t1, t2); - const double pf_t = prompt_len / (pf / 1000.0); - const double dc_t = dc > 0 ? generated / (dc / 1000.0) : 0.0; - std::cout << "\n[iter " << iter << (measured ? "" : " warmup") - << "] prefill " << pf_t << " tok/s (" << prompt_len << " tok, " - << pf << " ms) | decode " << dc_t << " tok/s (" << generated - << " tok, " << dc << " ms)\n"; - if (measured) { - pf_tps.push_back(pf_t); - dc_tps.push_back(dc_t); + stats.num_prompt_tokens = prompt_len; + stats.num_generated_tokens = generated; } } - - auto summarize = [](const char* label, const std::vector& v) { - double m = 0.0; - for (double x : v) { - m += x; - } - m /= static_cast(v.size()); - double var = 0.0; - for (double x : v) { - var += (x - m) * (x - m); - } - const double sd = v.size() > 1 - ? std::sqrt(var / static_cast(v.size() - 1)) - : 0.0; - std::cout << label << ": " << m << " +/- " << sd - << " tok/s (n=" << v.size() << ")\n"; - }; - std::cout << "\n"; - summarize("prefill", pf_tps); - summarize("decode ", dc_tps); + std::cout << std::endl; + ::executorch::extension::llm::print_report(stats); return 0; } catch (const std::exception& e) { std::cerr << "Error: " << e.what() << std::endl; From 589d63ca1838ad07ffa594a1ead754a37103754c Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Tue, 11 Aug 2026 12:51:35 -0700 Subject: [PATCH 4/4] Stop rejecting unknown attention kwargs --- backends/mlx/llm/hf_attention.py | 25 +++++++------------------ 1 file changed, 7 insertions(+), 18 deletions(-) diff --git a/backends/mlx/llm/hf_attention.py b/backends/mlx/llm/hf_attention.py index 1977ef3ce7b..dcb20af04ec 100644 --- a/backends/mlx/llm/hf_attention.py +++ b/backends/mlx/llm/hf_attention.py @@ -65,17 +65,12 @@ def mlx_sdpa_with_start_pos_forward( Returns (output, None) where output is [B, seq_len, num_heads, head_dim] (BSHD). """ - # HuggingFace calls every registered implementation with the same argument - # list. Drop the ones handled elsewhere, and refuse the rest rather than - # ignoring them silently. softcap and head_mask stay named because models - # pass them unconditionally, often as None -- the value is what matters. - kwargs.pop("is_causal", None) - kwargs.pop("use_cache", None) - kwargs.pop("sliding_window", None) # the export routes sliding layers away - if kwargs.pop("dropout", 0.0): + # Refuse only what would silently change the result. Anything else in + # kwargs is ignored: HuggingFace forwards model-level arguments down into + # attention -- gemma 4 sends `labels` -- so the set is open-ended and + # rejecting the unknown breaks models that pass something harmless. + if kwargs.get("dropout"): raise ValueError("mlx attention does not support dropout") - if kwargs: - raise ValueError(f"mlx attention got unsupported args: {sorted(kwargs)}") if softcap is not None: raise ValueError("mlx attention does not support softcap") if head_mask is not None: @@ -152,15 +147,9 @@ def mlx_offgraph_attention_forward( Returns (output, None) where output is [B, q_len, num_heads, head_dim] (BSHD). """ - kwargs.pop("is_causal", None) - kwargs.pop("use_cache", None) - kwargs.pop("sliding_window", None) # the cache applies the window - if kwargs.pop("dropout", 0.0): + # As above: reject only what changes the result, ignore the rest. + if kwargs.get("dropout"): raise ValueError("mlx_offgraph attention does not support dropout") - if kwargs: - raise ValueError( - f"mlx_offgraph attention got unsupported args: {sorted(kwargs)}" - ) if softcap is not None: raise ValueError("mlx_offgraph attention does not support softcap") if head_mask is not None: