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..fa041b7129a --- /dev/null +++ b/backends/mlx/examples/llm/CMakeLists.txt @@ -0,0 +1,73 @@ +# 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) + +# 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_runner + extension_llm_sampler + gflags +) + +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}) + +# 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..a3fe12ca6b3 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,12 @@ 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 | +| `--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 +the cache and binds it at load time. --- @@ -123,6 +138,107 @@ 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 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. + +### 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 | +| `--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 | +| `--kv-windows` | *(model's own)* | Off-graph: attention pattern override, e.g. `512` | +| `--interactive` | `false` | Multi-turn chat on stdin; off-graph only | +| `--warmup` | `false` | Run once before measuring, to absorb JIT and pool growth | + +--- + ## 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/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 new file mode 100644 index 00000000000..37ae7bc8921 --- /dev/null +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -0,0 +1,711 @@ +/* + * 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-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 +// unless --temperature is set. +// +// Usage: +// run_llm_hf --pte --tokenizer [flags] +// +// --help lists every flag with its default. + +#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 +#include +#include +#include +#include +#include +#include +#include + +DEFINE_string(pte, "", "Model .pte file."); +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", + "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_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_bool( + warmup, + false, + "Run once before measuring, to absorb JIT and pool growth."); + +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()); +} + +// 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), +// 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"); + 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; + } + 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 +// 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()) { + 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}; + } + 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; + } + } + 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 +// 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); + } +} + +// 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. +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; +} + +} // namespace + +int main(int argc, char** argv) { + 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 float temperature = static_cast(FLAGS_temperature); + const int initial_capacity = FLAGS_kv_initial_capacity; + const bool interactive = FLAGS_interactive; + 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"; + return 1; + } + + try { + // 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; + } + + // 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; + // 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; + } + + 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; + } + if (!read_kv_layout(module, cfg)) { + std::cerr << "No KV cache layout in " << pte + << "; re-export with --use-offgraph-cache" << std::endl; + return 1; + } + 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; + } + 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; + } + prefill_chunk = cfg.max_write ? *cfg.max_write : 0; + session.emplace(cache::make_unique_key(), built.get()); + + print_cache_summary(cfg); + 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; + } + // 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; + + // 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()); + + // 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()) { + turn_end_id = static_cast(*eot); + stop_ids.insert(*eot); + } + } + auto is_stop = [&](int64_t t) { + for (uint64_t s : stop_ids) { + if (t == static_cast(s)) { + return true; + } + } + 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 = + 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"); + } + 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 + // 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) { + 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 += 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}); + } + 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 + // 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 = prefill(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; + } + // 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" : "") + << "]\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; + }; + + // 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(); + } + stats.inference_start_ms = ::executorch::extension::llm::time_in_ms(); + int64_t next = prefill(ids, prefill_pos); + 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); + 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 (measured) { + if (auto piece = tokenizer->decode(prev, static_cast(next)); + piece.ok()) { + ::executorch::extension::llm::safe_printf(piece->c_str()); + fflush(stdout); + } + } + prev = static_cast(next); + ++generated; + next = step({next}, {prompt_len + i}); + } + stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); + if (measured) { + std::cout << "\n\n"; // close the generation line + blank separator + // trailing space aligns the colon with the "after prefill" line above + print_footprint("after decode ", prompt_len + generated); + stats.num_prompt_tokens = prompt_len; + stats.num_generated_tokens = generated; + } + } + std::cout << std::endl; + ::executorch::extension::llm::print_report(stats); + return 0; + } catch (const std::exception& e) { + std::cerr << "Error: " << e.what() << std::endl; + return 1; + } +} 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: