From 3f908ee75ae09bea090b6e429b2d5c79637e902c Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 10:38:06 -0800 Subject: [PATCH] Support sliding-window attention with the native off-graph KV cache --offgraph_cache now exports models with sliding-window layers instead of rejecting them. The window is a property of the cache, not the graph, so the export writes the cache geometry (get_n_caches, get_kv_heads, get_head_dims, get_windows) to the package's metadata, and the golden samples come from a reference cache whose sliding layers are rings. A layer's window follows HF: layer_types when the config has them, else sliding_window on every layer. The runner reads the windows from the metadata, checks the shapes against the graph, and builds ring layers sized for the method's token input. GgmlSequenceCache holds a ring layer at window + max_write - 1 rows from the start. The executor fills one set of row and mask uploads per window from the cache's plan: each row is masked by the logical position it holds, so a ring layer attends all its rows in any physical order, and a query sees only its own window, including within a prefill wider than the window. --- .../native/examples/hf_llm/export_hf_llm.py | 70 +++++++---- backends/native/examples/hf_llm/runtime/BUCK | 1 + .../examples/hf_llm/runtime/LlmRunner.cpp | 96 ++++++++++++--- .../examples/hf_llm/test_export_hf_llm.py | 77 +++++++++++- .../native/runtime/engine/ggml/GgmlCache.cpp | 30 +++-- .../native/runtime/engine/ggml/GgmlCache.h | 17 ++- .../runtime/engine/ggml/GgmlExecutor.cpp | 91 +++++++++----- .../runtime/engine/ggml/GgmlGraphBuilder.h | 31 +++-- .../engine/ggml/GgmlLowerAttention.cpp | 37 ++++-- .../test/runtime/ggml/test_ggml_engine.cpp | 116 ++++++++++++++++-- 10 files changed, 453 insertions(+), 113 deletions(-) diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index e5a53d64f99..260d0f3eda4 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/backends/native/examples/hf_llm/export_hf_llm.py @@ -34,10 +34,12 @@ from executorch.extension.llm.cache.hf_attention import register_offgraph_attention from executorch.extension.llm.cache.reference_cache import ( CacheConfig, + LayerPolicy, SequenceReferenceCache, ) from executorch.extension.llm.cache.update_and_attend import REGISTRY from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter +from executorch.extension.llm.export.model_metadata import write_cache_geometry from torch.export import Dim, export from torchao.quantization import ( Int8DynamicActivationIntxWeightConfig, @@ -212,47 +214,72 @@ def cache_scope(self) -> contextlib.AbstractContextManager: return contextlib.nullcontext() +def _cache_windows(config) -> list[int]: + """Per layer, the sliding window HF attends over, or 0 for full attention. + + Without layer_types (e.g. Mistral), HF windows every layer when + sliding_window is set. + """ + window = getattr(config, "sliding_window", None) + layer_types = getattr(config, "layer_types", None) + if not layer_types: + return [window or 0] * config.num_hidden_layers + unsupported = set(layer_types) - {"full_attention", "sliding_attention"} + if unsupported: + raise ValueError( + f"--offgraph_cache does not support {sorted(unsupported)} layers " + f"({config.model_type})" + ) + if "sliding_attention" in layer_types and not window: + raise ValueError(f"{config.model_type} has sliding layers but no window") + return [window if t == "sliding_attention" else 0 for t in layer_types] + + class OffGraphLastTokenLM(nn.Module): """Last-position logits with the KV cache outside the graph. Each attention layer emits ``kvcache::update_and_attend`` over this step's - K/V. Eager calls (golden samples) run inside ``cache_scope()``, which makes - a reference cache of ``capacity`` positions the active one. + K/V; the cache, not the graph, applies a sliding layer's window, so the + runner learns the windows from ``cache_metadata()``. Eager calls (golden + samples) run inside ``cache_scope()``, which makes a reference cache of + ``capacity`` positions the active one. """ def __init__(self, model: PreTrainedModel, capacity: int) -> None: super().__init__() self.model = model config = model.config.get_text_config() - # The cache op attends every position it holds, so a windowed or - # chunked layer would silently attend past its span. Without - # layer_types (e.g. Mistral), HF windows every layer when - # sliding_window is set. - layer_types = getattr(config, "layer_types", None) or ( - ["sliding_attention"] - if getattr(config, "sliding_window", None) is not None - else [] - ) - unsupported = sorted(set(layer_types) - {"full_attention"}) - if unsupported: - raise ValueError( - f"--offgraph_cache does not support {unsupported} layers " - f"({config.model_type})" - ) - head_dim = ( + self._windows = _cache_windows(config) + self._head_dim = ( getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads ) + self._kv_heads = config.num_key_value_heads self._reference = SequenceReferenceCache( CacheConfig( n_layers=config.num_hidden_layers, - n_kv_heads=config.num_key_value_heads, - head_dim=head_dim, + n_kv_heads=self._kv_heads, + head_dim=self._head_dim, capacity=capacity, + layers=[ + LayerPolicy.ring(w) if w else LayerPolicy.flat() + for w in self._windows + ], ) ) self._key = f"export_hf_llm/{id(self)}" + def cache_metadata(self) -> dict[str, int | list[int]]: + """The cache geometry, under the names LLM runners read it by.""" + n = len(self._windows) + geometry = write_cache_geometry( + [self._kv_heads] * n, [self._head_dim] * n, self._windows + ) + return { + name: value.tolist() if isinstance(value, torch.Tensor) else value + for name, value in geometry.items() + } + def forward( self, input_ids: torch.Tensor, cache_position: torch.Tensor ) -> torch.Tensor: @@ -462,7 +489,8 @@ def main() -> None: report_lowering(ep, "forward") if args.output: - to_native(ep).save(args.output) + metadata = wrapper.cache_metadata() if args.offgraph_cache else None + to_native(ep, metadata=metadata).save(args.output) print(f"saved {args.output}") if prompts: samples = golden_samples(wrapper, prompts, _DECODE_STEPS) diff --git a/backends/native/examples/hf_llm/runtime/BUCK b/backends/native/examples/hf_llm/runtime/BUCK index c2e74421702..d4621cdb274 100644 --- a/backends/native/examples/hf_llm/runtime/BUCK +++ b/backends/native/examples/hf_llm/runtime/BUCK @@ -37,6 +37,7 @@ runtime.cxx_library( "//executorch/backends/native/runtime/graph:scalar_type", "//executorch/examples/models/llama/tokenizer:tiktoken", "//executorch/extension/llm/cache:kv_cache", + "//executorch/extension/llm/runner:constants", ], exported_deps = ["//executorch/backends/native/runtime/engine:engine"], external_deps = ["gflags"], diff --git a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp index dbc1c640639..7e8df20df1a 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp +++ b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp @@ -35,6 +35,7 @@ #include #include #include +#include #include #include #include @@ -126,10 +127,12 @@ size_t kv_cache_length(const ptn::Method& method) { constexpr std::string_view kUpdateAndAttend = "torch.ops.kvcache.update_and_attend.default"; -// The off-graph KV cache a method's update_and_attend calls need: one flat -// layer per cache id, shaped like its K. nullopt if the method has none. +// The off-graph KV cache a method's update_and_attend calls need: one layer per +// cache id, shaped like its K, with the sliding window (or none) the package's +// metadata gives it. nullopt if the method has none. std::optional offgraph_cache_geometry( - const ptn::Method& method) { + const ptn::Method& method, + const ptn::Package& package) { std::map layers; for (const ptn::Node& node : method.graph.nodes) { if (node.target != kUpdateAndAttend) { @@ -184,6 +187,45 @@ std::optional offgraph_cache_geometry( } geometry.layers.push_back(layer); } + + namespace llm = executorch::extension::llm; + const auto per_layer = [&](const char* name) { + std::optional> values = package.metadata_ints(name); + if (values && values->size() != geometry.layers.size()) { + throw std::runtime_error( + std::string("package metadata ") + name + " has " + + std::to_string(values->size()) + " entries for " + + std::to_string(geometry.layers.size()) + " cache layers"); + } + return values; + }; + const auto kv_heads = per_layer(llm::kKVHeads); + const auto head_dims = per_layer(llm::kHeadDims); + const auto windows = per_layer(llm::kWindows); + // Without them a sliding layer would silently attend its whole history. + if (!windows) { + throw std::runtime_error( + std::string("package metadata has no ") + llm::kWindows + + "; export the package with export_hf_llm.py --offgraph_cache"); + } + for (size_t i = 0; i < geometry.layers.size(); ++i) { + cache::LayerGeometry& layer = geometry.layers[i]; + if ((kv_heads && (*kv_heads)[i] != layer.n_kv_heads) || + (head_dims && (*head_dims)[i] != layer.head_dim)) { + throw std::runtime_error( + "package metadata disagrees with the graph on cache layer " + + std::to_string(i) + "'s shape"); + } + const int64_t window = (*windows)[i]; + if (window < 0 || window > std::numeric_limits::max()) { + throw std::runtime_error( + "package metadata gives cache layer " + std::to_string(i) + + " an invalid window"); + } + if (window > 0) { + layer.policy = {cache::LayerPolicy::Kind::Ring, static_cast(window)}; + } + } return geometry; } @@ -195,9 +237,11 @@ struct OffGraphCache { ptn::EngineOptions options; }; +// `max_write` is the most tokens one call writes, which sizes a ring layer. OffGraphCache install_offgraph_cache( const char* backend_id, - const cache::CacheGeometry& geometry) { + const cache::CacheGeometry& geometry, + int max_write) { if (backend_id == nullptr) { throw std::runtime_error( "the package keeps its KV cache off-graph, which this engine's " @@ -211,7 +255,8 @@ OffGraphCache install_offgraph_cache( FLAGS_kv_capacity > 0 ? FLAGS_kv_capacity : FLAGS_max_seq_length, .kv_dtype = static_cast( FLAGS_kv_dtype == "fp16" ? ptn::ScalarType::Half - : ptn::ScalarType::Float)}; + : ptn::ScalarType::Float), + .max_write = max_write}; if (FLAGS_kv_initial_capacity >= 0) { cfg.initial_capacity = FLAGS_kv_initial_capacity; } @@ -237,10 +282,15 @@ OffGraphCache install_offgraph_cache( throw std::runtime_error("failed to publish the KV cache key"); } out.options.emplace(cache::kCacheKeyOption, key); + const auto sliding = std::ranges::count_if( + geometry.layers, [](const cache::LayerGeometry& layer) { + return layer.policy.kind == cache::LayerPolicy::Kind::Ring; + }); std::printf( - "kv cache: off-graph, %zu layers, %d tokens (%d allocated up front), " - "%s\n", + "kv cache: off-graph, %zu layers (%td sliding), %d tokens (%d " + "allocated up front), %s\n", geometry.layers.size(), + sliding, cfg.capacity, std::min(std::max(cfg.initial_capacity, 1), cfg.capacity), FLAGS_kv_dtype.c_str()); @@ -299,7 +349,8 @@ struct Generation { // Prefills `prompt_tokens` from `start_position` in calls of at most // `max_chunk` tokens, then decodes one token per call until `limit` tokens or -// an EOS, streaming the text to stdout. +// an EOS, streaming the text to stdout. With --warmup_prefill it prefills the +// prompt once beforehand, then clears `offgraph_cache` if there is one. Generation generate( ptn::EngineExecutable& executable, size_t logits_output, @@ -310,7 +361,8 @@ Generation generate( int64_t start_position, size_t limit, const std::unordered_set& eos_ids, - const NextToken& next_token_of) { + const NextToken& next_token_of, + cache::CacheControl* offgraph_cache) { std::vector logits; auto set_positions = [&](int64_t start, size_t count) { @@ -421,6 +473,11 @@ Generation generate( "warmup_prefill cannot be combined with logits diagnostics"); } prefill(start_position); + // A sliding layer cannot rewind to the start of a prompt longer than its + // window, so the real prefill starts from an empty cache. + if (offgraph_cache != nullptr) { + offgraph_cache->clear(); + } } Generation generation; @@ -543,8 +600,16 @@ int run_llm( const std::unique_ptr host = make_host(); // Declared before the context, which may hold the cache, so it outlives it. std::optional kv; - if (const auto geometry = offgraph_cache_geometry(method)) { - kv = install_offgraph_cache(cache_backend_id, *geometry); + if (const auto geometry = offgraph_cache_geometry(method, *package)) { + // Serialized input sizes are upper bounds; the tokens are [1, T]. + const std::vector& token_bound = + method.graph.value(method.graph.input_ids.at(0)).tensor_meta().sizes; + if (token_bound.empty() || token_bound.back() <= 0 || + token_bound.back() > std::numeric_limits::max()) { + throw std::runtime_error("token input must have shape [1, sequence]"); + } + kv = install_offgraph_cache( + cache_backend_id, *geometry, static_cast(token_bound.back())); } else if (FLAGS_interactive) { throw std::invalid_argument( "--interactive needs a package exported with --offgraph_cache"); @@ -745,7 +810,8 @@ int run_llm( 0, calls.size(), {}, - check_then_force); + check_then_force, + nullptr); std::printf("%s", report.c_str()); } std::printf( @@ -834,7 +900,8 @@ int run_llm( position, limit, eos_ids, - ptn::greedy_argmax_last_row); + ptn::greedy_argmax_last_row, + nullptr); position += static_cast(tokens.size() + reply.decode_steps); unfed = reply.ended_on_eos ? std::nullopt : std::optional(reply.unfed_token); @@ -859,7 +926,8 @@ int run_llm( 0, limit, eos_ids, - ptn::greedy_argmax_last_row); + ptn::greedy_argmax_last_row, + kv ? kv->control : nullptr); const double prompt_tokens_per_second = generation.prompt_ms > 0.0 ? 1000.0 * static_cast(prompt_tokens.size()) / diff --git a/backends/native/examples/hf_llm/test_export_hf_llm.py b/backends/native/examples/hf_llm/test_export_hf_llm.py index 349fab0144a..d796a35ac66 100644 --- a/backends/native/examples/hf_llm/test_export_hf_llm.py +++ b/backends/native/examples/hf_llm/test_export_hf_llm.py @@ -29,6 +29,7 @@ OffGraphLastTokenLM, ) from executorch.extension.llm.cache.hf_attention import register_offgraph_attention +from executorch.extension.llm.cache.reference_cache import LayerPolicy from executorch.extension.llm.cache.update_and_attend import REGISTRY from transformers import ( AutoModelForCausalLM, @@ -37,11 +38,13 @@ GemmaConfig, GraniteConfig, LlamaConfig, + MistralConfig, PreTrainedModel, Qwen2Config, Qwen3Config, ) +# Architectures gguf_hf maps; the GGUF test covers each. _CONFIGS = { "llama": LlamaConfig, "granite": GraniteConfig, @@ -53,8 +56,10 @@ } -def _tiny(arch: str = "llama", tie_word_embeddings: bool = True) -> PreTrainedModel: - config = _CONFIGS[arch]( +def _tiny( + arch: str = "llama", tie_word_embeddings: bool = True, **config_kwargs +) -> PreTrainedModel: + config = (_CONFIGS | {"mistral": MistralConfig})[arch]( vocab_size=128, hidden_size=64, intermediate_size=128, @@ -64,6 +69,7 @@ def _tiny(arch: str = "llama", tie_word_embeddings: bool = True) -> PreTrainedMo head_dim=16, max_position_embeddings=64, tie_word_embeddings=tie_word_embeddings, + **config_kwargs, ) model = AutoModelForCausalLM.from_config(config).eval().requires_grad_(False) model.config._attn_implementation = "sdpa" @@ -370,9 +376,70 @@ def test_fp16_kv_cache_rejects_sliding_window_layers(self) -> None: with self.assertRaisesRegex(ValueError, "cannot change the KV dtype"): LastTokenLM(_tiny("gemma3"), 32, torch.float16) - def test_offgraph_cache_rejects_sliding_window_layers(self) -> None: - with self.assertRaisesRegex(ValueError, "sliding_attention"): - OffGraphLastTokenLM(_tiny("gemma3"), 32) + def test_offgraph_cache_windows_sliding_layers_like_hf(self) -> None: + # Mistral windows every layer; Gemma3 alternates sliding and full. + register_offgraph_attention("offgraph") + prompts = [[3, 4, 5, 6, 7, 8, 9], [9, 10]] + for arch, kwargs in ( + ("mistral", {"sliding_window": 3}), + ( + "gemma3", + { + "sliding_window": 3, + "layer_types": ["sliding_attention", "full_attention"], + }, + ), + ): + with self.subTest(arch=arch): + torch.manual_seed(7) + model = _tiny(arch, **kwargs) + offgraph_model = copy.deepcopy(model) + offgraph_model.config._attn_implementation = "offgraph" + offgraph = OffGraphLastTokenLM(offgraph_model, 32).eval() + + samples = golden_samples(offgraph, prompts, 3) + + for sample in samples.samples: + # Each call feeds the tokens it appends: the prompt, then one. + sequence = [] + for call in sample.calls: + sequence += call.inputs[0][0].tolist() + with torch.no_grad(): + want = model(torch.tensor([sequence])).logits[:, -1:] + torch.testing.assert_close( + call.outputs[0], want, atol=1e-5, rtol=1e-5 + ) + + def test_offgraph_cache_metadata_gives_each_layer_its_window(self) -> None: + for model, windows in ( + (_tiny(), [0, 0]), + (_tiny("mistral", sliding_window=8), [8, 8]), + (_tiny("mistral", sliding_window=None), [0, 0]), + ( + _tiny( + "gemma3", + sliding_window=8, + layer_types=["sliding_attention", "full_attention"], + ), + [8, 0], + ), + (_tiny("qwen2", use_sliding_window=False), [0, 0]), + ): + with self.subTest(arch=model.config.model_type): + wrapper = OffGraphLastTokenLM(model, 32) + self.assertEqual( + wrapper.cache_metadata(), + { + "get_n_caches": 2, + "get_kv_heads": [2, 2], + "get_head_dims": [16, 16], + "get_windows": windows, + }, + ) + self.assertEqual( + wrapper._reference.config.layers, + [LayerPolicy.ring(w) if w else LayerPolicy.flat() for w in windows], + ) def test_offgraph_cache_rejects_other_attention_layer_types(self) -> None: model = _tiny() diff --git a/backends/native/runtime/engine/ggml/GgmlCache.cpp b/backends/native/runtime/engine/ggml/GgmlCache.cpp index d3aeab3556b..0c79e9d3fc8 100644 --- a/backends/native/runtime/engine/ggml/GgmlCache.cpp +++ b/backends/native/runtime/engine/ggml/GgmlCache.cpp @@ -47,13 +47,25 @@ const cache::CacheGeometry& checked( if (geometry.dflash_block_size) { throw std::invalid_argument("ggml cache: DFlash geometry is unsupported"); } + return geometry; +} + +// SequenceCache's ring: window + max_write - 1 slots, max_write defaulting to +// the window. Positions below capacity never reach a slot past it. +std::vector ring_rows( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) { + std::vector out; + out.reserve(geometry.layers.size()); for (const cache::LayerGeometry& layer : geometry.layers) { - if (layer.policy.kind != cache::LayerPolicy::Kind::Flat) { - throw std::invalid_argument( - "ggml cache: only flat (full-history) layers are supported"); - } + const int window = layer.policy.window; + out.push_back( + layer.policy.kind == cache::LayerPolicy::Kind::Ring + ? std::min( + window + cfg.max_write.value_or(window) - 1, cfg.capacity) + : 0); } - return geometry; + return out; } } // namespace @@ -63,6 +75,7 @@ GgmlSequenceCache::GgmlSequenceCache( const cache::CacheConfig& cfg) : SequenceCache(checked(geometry, cfg), cfg), geometry_(geometry), + ring_rows_(ring_rows(geometry, cfg)), type_(kv_type(cfg.kv_dtype)), initial_rows_(std::clamp(cfg.initial_capacity, 1, cfg.capacity)) {} @@ -101,13 +114,13 @@ void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) { } for (size_t i = 0; i < geometry_.layers.size(); ++i) { const cache::LayerGeometry& layer = geometry_.layers[i]; - Layer out; + Layer out{.window = layer.policy.window}; for (ggml_tensor** t : {&out.k, &out.v}) { *t = ggml_new_tensor_3d( next->context.get(), type_, layer.head_dim, - next->rows, + ring_rows_[i] > 0 ? ring_rows_[i] : next->rows, layer.n_kv_heads); ggml_format_name(*t, "kv_cache.%s.%zu", t == &out.k ? "k" : "v", i); } @@ -122,7 +135,8 @@ void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) { } ggml_backend_buffer_clear(next->buffer.get(), 0); - // Each head's rows are contiguous, but the head stride grows with the rows. + // Each head's rows are contiguous, but a flat layer's head stride grows with + // its rows. A ring layer keeps its size, so its rows land where they were. std::vector staging; for (size_t i = 0; storage_ && i < storage_->layers.size(); ++i) { const Layer& old = storage_->layers[i]; diff --git a/backends/native/runtime/engine/ggml/GgmlCache.h b/backends/native/runtime/engine/ggml/GgmlCache.h index 8f5e46c0e21..dc44c39f3dc 100644 --- a/backends/native/runtime/engine/ggml/GgmlCache.h +++ b/backends/native/runtime/engine/ggml/GgmlCache.h @@ -22,10 +22,12 @@ namespace ptn { // The off-graph KV cache of one sequence for the ggml engine: the neutral // SequenceCache controller plus K/V storage per layer in the buffer type of // the first backend that binds it. Each layer's K and V are ggml -// [head_dim, rows, n_kv_heads] (torch [1, n_kv_heads, rows, head_dim]); row r -// holds logical position r. Storage starts at the config's initial_capacity -// rows and doubles toward capacity as steps need more. Like SequenceCache, it -// is not thread-safe: executables bound to one cache must be serialized. +// [head_dim, rows, n_kv_heads] (torch [1, n_kv_heads, rows, head_dim]). A flat +// layer's row r holds logical position r; its storage starts at the config's +// initial_capacity rows and doubles toward capacity as steps need more. A ring +// layer holds its window + max_write - 1 rows from the start, and the +// SequencePlanner says which rows a step writes and reads. Like SequenceCache, +// it is not thread-safe: executables bound to one cache must be serialized. class GgmlSequenceCache final : public executorch::extension::llm::cache::SequenceCache { public: @@ -34,6 +36,8 @@ class GgmlSequenceCache final struct Layer { ggml_tensor* k = nullptr; ggml_tensor* v = nullptr; + // Sliding window of a ring layer; 0 for a flat one. + int window = 0; }; // One allocation of K/V storage. A regrow moves the rows into a new one; @@ -42,6 +46,7 @@ class GgmlSequenceCache final ggml_context_ptr context; ggml_backend_buffer_ptr buffer; std::vector layers; + // Rows of every flat layer. int rows = 0; }; @@ -53,7 +58,7 @@ class GgmlSequenceCache final GgmlSequenceCache(const GgmlSequenceCache&) = delete; GgmlSequenceCache& operator=(const GgmlSequenceCache&) = delete; - // Ensures storage of at least `rows` rows in `backend`'s default buffer + // Ensures flat layers of at least `rows` rows in `backend`'s default buffer // type, moving the rows it holds into a larger allocation if needed. Throws // if `rows` exceeds capacity or an earlier call used another buffer type. void reserve(ggml_backend_t backend, int rows); @@ -69,6 +74,8 @@ class GgmlSequenceCache final private: executorch::extension::llm::cache::CacheGeometry geometry_; + // Per layer: the fixed rows of a ring layer, 0 for a flat one. + std::vector ring_rows_; ggml_type type_; int initial_rows_; ggml_backend_buffer_type_t buffer_type_ = nullptr; diff --git a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp index 8866f77e75f..22014de402d 100644 --- a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp +++ b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp @@ -50,6 +50,7 @@ namespace ptn { namespace { using namespace ggml_detail; +using executorch::extension::llm::cache::Run; using executorch::extension::llm::cache::SeqStepPlan; // Multiplies each slice of the leading dimension by its scale. @@ -1303,7 +1304,8 @@ class GgmlExecutable final : public EngineExecutable { offgraph_storage_ = kv_cache_->storage(); offgraph_cache_.layers.clear(); for (const GgmlSequenceCache::Layer& layer : offgraph_storage_->layers) { - offgraph_cache_.layers.emplace_back(layer.k, layer.v); + offgraph_cache_.layers.push_back( + {.k = layer.k, .v = layer.v, .window = layer.window}); } } @@ -1354,8 +1356,8 @@ class GgmlExecutable final : public EngineExecutable { return {start, tokens}; } - // Plans the step on the cache and fills the graph's cache rows and mask. - // Returns the plan to commit once the step has run. + // Plans the step on the cache and fills each window's cache rows and mask. + // Returns a plan to commit once the step has run. std::optional prepare_offgraph(GraphInstance& instance) { if (!instance.offgraph) { return std::nullopt; @@ -1366,42 +1368,69 @@ class GgmlExecutable final : public EngineExecutable { "ggml: update_and_attend needs one position per token"); } const int64_t start = offgraph_step_->start; - std::optional plan; - for (size_t layer = 0; layer < offgraph_cache_.layers.size(); ++layer) { - plan = kv_cache_->plan( - static_cast(layer), - static_cast(start), - static_cast(slots.tokens)); - if (!plan || plan->n_write != 1 || plan->n_read != 1 || - plan->read_base_pos != 0 || plan->read[0].start != 0) { + const auto plan_layer = [&](int layer) { + const std::optional plan = kv_cache_->plan( + layer, static_cast(start), static_cast(slots.tokens)); + if (!plan) { throw std::runtime_error( - "ggml: the KV cache planned a non-flat step"); + "ggml: the KV cache cannot plan a " + std::to_string(slots.tokens) + + "-token step at position " + std::to_string(start) + " for layer " + + std::to_string(layer)); } + return *plan; + }; + // Layers outside the method's calls must fit the step too. + for (size_t layer = 0; layer < offgraph_cache_.layers.size(); ++layer) { + plan_layer(static_cast(layer)); } - const int64_t read_end = plan->read[0].len; - if (read_end > slots.keys) { - throw std::logic_error("ggml: KV cache read exceeds the graph's keys"); - } - std::vector& rows = instance.uploads[slots.rows].bytes; - rows.resize(static_cast(slots.tokens) * sizeof(int32_t)); - auto* row = reinterpret_cast(rows.data()); - for (int64_t i = 0; i < slots.tokens; ++i) { - row[i] = plan->write[0].start + static_cast(i); - } - std::vector& mask = instance.uploads[slots.mask].bytes; - mask.resize( - static_cast(slots.tokens * slots.keys) * sizeof(ggml_fp16_t)); - auto* bias = reinterpret_cast(mask.data()); const ggml_fp16_t attend = ggml_fp32_to_fp16(0.0f); const ggml_fp16_t masked = ggml_fp32_to_fp16(-std::numeric_limits::infinity()); - for (int64_t i = 0; i < slots.tokens; ++i) { - for (int64_t j = 0; j < slots.keys; ++j) { - bias[i * slots.keys + j] = - j <= start + i && j < read_end ? attend : masked; + std::optional committed; + for (const OffGraphUploads::Group& group : slots.groups) { + const SeqStepPlan plan = plan_layer(group.layer); + // The logical position each row holds this step, or -1 if it is unread. + std::vector positions(static_cast(group.keys), -1); + int64_t position = plan.read_base_pos; + for (int r = 0; r < plan.n_read; ++r) { + const Run run = plan.read[r]; + if (run.start < 0 || run.start + run.len > group.keys) { + throw std::logic_error( + "ggml: KV cache read exceeds the graph's keys"); + } + for (int i = 0; i < run.len; ++i) { + positions[static_cast(run.start + i)] = position++; + } + } + std::vector& rows = instance.uploads[group.rows].bytes; + rows.resize(static_cast(slots.tokens) * sizeof(int32_t)); + auto* row = reinterpret_cast(rows.data()); + int64_t written = 0; + for (int w = 0; w < plan.n_write; ++w) { + for (int i = 0; i < plan.write[w].len && written < slots.tokens; ++i) { + row[written++] = plan.write[w].start + i; + } + } + if (written != slots.tokens) { + throw std::logic_error("ggml: KV cache planned the wrong rows"); + } + std::vector& mask = instance.uploads[group.mask].bytes; + mask.resize( + static_cast(slots.tokens * group.keys) * sizeof(ggml_fp16_t)); + auto* bias = reinterpret_cast(mask.data()); + for (int64_t i = 0; i < slots.tokens; ++i) { + const int64_t query = start + i; + for (int64_t j = 0; j < group.keys; ++j) { + const int64_t key = positions[static_cast(j)]; + bias[i * group.keys + j] = key >= 0 && key <= query && + (group.window == 0 || query - key < group.window) + ? attend + : masked; + } } + committed = plan; } - return plan; + return committed; } // Small integer constants, e.g. a key offset, read once from the package. diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h index 114bea57e66..84ba6df1dad 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h @@ -90,21 +90,34 @@ struct RowIndexCheck { // The off-graph KV cache as one graph sees it. struct OffGraphCacheBinding { - // Per cache layer: K and V storage, ggml [head_dim, capacity, n_kv_heads]. - std::vector> layers; - // Cache rows the graph's attention reads, a prefix of capacity. + struct Layer { + // K and V storage, ggml [head_dim, rows, n_kv_heads]. + ggml_tensor* k = nullptr; + ggml_tensor* v = nullptr; + // Sliding window of a ring layer, which attends all its rows; 0 for a + // flat layer. + int window = 0; + }; + std::vector layers; + // Rows a flat layer's attention reads, a prefix of its storage. int64_t keys = 0; }; // Uploads of a graph's off-graph attention, which the executor fills from the -// cache's plan before every run. Every layer shares them. +// cache's plan before every run. Layers with the same window share them. struct OffGraphUploads { - // I32 [tokens]: the cache row each token writes. - size_t rows = 0; - // F16 [tokens, keys]: additive mask of the rows each token attends. - size_t mask = 0; + struct Group { + int window = 0; + // A layer with this window, whose plan fills the uploads. + int layer = 0; + // I32 [tokens]: the cache row each token writes. + size_t rows = 0; + // F16 [tokens, keys]: additive mask of the rows each token attends. + size_t mask = 0; + int64_t keys = 0; + }; int64_t tokens = 0; - int64_t keys = 0; + std::vector groups; }; struct GraphInstance { diff --git a/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp b/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp index 3b00c31e966..4f0e00210eb 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp @@ -398,7 +398,8 @@ void GraphBuilder::lower_sdpa(const Node& node) { } // Writes this step's K/V into the cache rows the executor uploads, then -// attends the cache prefix under the uploaded causal mask. Q/K/V are [B,H,T,D]. +// attends the layer's rows under the uploaded mask: a flat layer's prefix, or +// all of a ring layer. Q/K/V are [B,H,T,D]. void GraphBuilder::lower_update_and_attend(const Node& node) { if (offgraph_cache_ == nullptr) { throw std::runtime_error( @@ -416,7 +417,7 @@ void GraphBuilder::lower_update_and_attend(const Node& node) { "ggml: update_and_attend layer_id " + std::to_string(layer) + " is outside the cache's " + std::to_string(num_layers) + " layers"); } - const auto [k_cache, v_cache] = + const auto [k_cache, v_cache, window] = offgraph_cache_->layers[static_cast(layer)]; if (query.sizes.size() != 4 || key.sizes.size() != 4 || val.sizes.size() != 4 || query.sizes[0] != 1 || key.sizes[0] != 1 || @@ -436,28 +437,40 @@ void GraphBuilder::lower_update_and_attend(const Node& node) { "ggml: update_and_attend layer " + std::to_string(layer) + " K/V shape disagrees with the cache geometry"); } - const int64_t keys = offgraph_cache_->keys; + const int64_t keys = window > 0 ? k_cache->ne[1] : offgraph_cache_->keys; ggml_context* ctx = instance_.context.get(); if (!instance_.offgraph) { + instance_.offgraph = OffGraphUploads{.tokens = tokens}; + } else if (instance_.offgraph->tokens != tokens) { + throw std::runtime_error( + "ggml: every update_and_attend in a method must see the same tokens"); + } + std::vector& groups = instance_.offgraph->groups; + auto group = + std::ranges::find(groups, window, &OffGraphUploads::Group::window); + if (group == groups.end()) { ggml_tensor* rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); - ggml_set_name(rows, "kv_cache.rows"); + ggml_format_name(rows, "kv_cache.rows.w%d", window); ggml_set_input(rows); ggml_tensor* mask = new_tensor(ctx, GGML_TYPE_F16, {1, 1, tokens, keys}); - ggml_set_name(mask, "kv_cache.mask"); + ggml_format_name(mask, "kv_cache.mask.w%d", window); ggml_set_input(mask); - instance_.offgraph = OffGraphUploads{ + groups.push_back(OffGraphUploads::Group{ + .window = window, + .layer = static_cast(layer), .rows = instance_.uploads.size(), .mask = instance_.uploads.size() + 1, - .tokens = tokens, - .keys = keys}; + .keys = keys}); instance_.uploads.push_back(Upload{.tensor = rows}); instance_.uploads.push_back(Upload{.tensor = mask}); - } else if (instance_.offgraph->tokens != tokens) { + group = groups.end() - 1; + } else if (group->keys != keys) { throw std::runtime_error( - "ggml: every update_and_attend in a method must see the same tokens"); + "ggml: update_and_attend layers with one window must hold the same " + "rows"); } - ggml_tensor* rows = instance_.uploads[instance_.offgraph->rows].tensor; - ggml_tensor* mask = instance_.uploads[instance_.offgraph->mask].tensor; + ggml_tensor* rows = instance_.uploads[group->rows].tensor; + ggml_tensor* mask = instance_.uploads[group->mask].tensor; const auto write = [&](ggml_tensor* cache, ggml_tensor* step) { // set_rows converts F32 to either cache type, but not F16 to F32 on diff --git a/backends/native/test/runtime/ggml/test_ggml_engine.cpp b/backends/native/test/runtime/ggml/test_ggml_engine.cpp index 147f5b9553d..0c8a9c1e146 100644 --- a/backends/native/test/runtime/ggml/test_ggml_engine.cpp +++ b/backends/native/test/runtime/ggml/test_ggml_engine.cpp @@ -1401,19 +1401,29 @@ struct InstalledCache { } }; +// One layer per entry of `windows`: flat for 0, else a ring of that window. InstalledCache install_cache( int capacity, int initial_capacity, ScalarType kv_dtype = ScalarType::Float, - int head_dim = static_cast(kOffGraphDim)) { + int head_dim = static_cast(kOffGraphDim), + const std::vector& windows = {0}, + std::optional max_write = std::nullopt) { register_ggml_cache_builders(); kv::CacheGeometry geometry; - geometry.layers.push_back(kv::LayerGeometry{ - .n_kv_heads = static_cast(kOffGraphKvHeads), .head_dim = head_dim}); + for (const int window : windows) { + geometry.layers.push_back(kv::LayerGeometry{ + .policy = window > 0 + ? kv::LayerPolicy{kv::LayerPolicy::Kind::Ring, window} + : kv::LayerPolicy{}, + .n_kv_heads = static_cast(kOffGraphKvHeads), + .head_dim = head_dim}); + } const kv::CacheConfig cfg{ .capacity = capacity, .kv_dtype = static_cast(kv_dtype), - .initial_capacity = initial_capacity}; + .initial_capacity = initial_capacity, + .max_write = max_write}; auto built = kv::CacheFactory::global().build( kGgmlCacheBackendId, kv::kind::kSingle, geometry, cfg); if (!built.ok()) { @@ -1434,9 +1444,12 @@ InstalledCache install_cache( } // What the cache should hold: K and V [Hkv, D] per position, written by each -// step, attended causally by naive softmax attention. +// step, attended causally by naive softmax attention, over the last `window` +// positions when it is nonzero. class OffGraphReference { public: + explicit OffGraphReference(int window = 0) : window_(window) {} + // Runs `tokens` tokens from `start` through `executable`, with inputs that // depend on the position and `seed`, and returns the largest difference // from the reference. @@ -1491,8 +1504,10 @@ class OffGraphReference { for (int64_t t = 0; t < tokens; ++t) { const float* query = &q[(h * tokens + t) * kOffGraphDim]; const int64_t keys = start + t + 1; + const int64_t first = + window_ > 0 ? std::max(0, keys - window_) : 0; std::vector weights(keys); - for (int64_t s = 0; s < keys; ++s) { + for (int64_t s = first; s < keys; ++s) { weights[s] = std::exp( std::inner_product( query, @@ -1502,10 +1517,10 @@ class OffGraphReference { kOffGraphScale); } const double total = - std::accumulate(weights.begin(), weights.end(), 0.0); + std::accumulate(weights.begin() + first, weights.end(), 0.0); for (int64_t d = 0; d < kOffGraphDim; ++d) { double sum = 0.0; - for (int64_t s = 0; s < keys; ++s) { + for (int64_t s = first; s < keys; ++s) { sum += weights[s] * v_[s][kv_head * kOffGraphDim + d]; } worst = std::max( @@ -1520,6 +1535,7 @@ class OffGraphReference { } private: + int window_; std::vector> k_; std::vector> v_; }; @@ -1682,6 +1698,90 @@ TEST(GgmlEngineTest, RejectsOffGraphStepsPastCapacityOrNotConsecutive) { EXPECT_THROW(prefill->execute(), std::invalid_argument); } +// A ring layer attends each query's window, within a prefill wider than the +// window too, while decoding wraps the ring several times. +TEST(GgmlEngineTest, OffGraphRingAttentionMatchesReference) { + for (const size_t graphs : {size_t{0}, size_t{4}}) { + for (const int window : {2, 5}) { + for (const auto& [dtype, tol] : + {std::pair{ScalarType::Float, 1e-4f}, + std::pair{ScalarType::Half, 2e-2f}}) { + SCOPED_TRACE( + "graphs " + std::to_string(graphs) + ", window " + + std::to_string(window) + ", " + + std::string(scalar_type_name(dtype))); + InstalledCache cache = install_cache( + 64, 1, dtype, static_cast(kOffGraphDim), {window}, 3); + auto host = GgmlEngineHost::create({.max_cached_graphs = graphs}); + auto prefill_context = make_context( + *host, make_update_and_attend_program(3), cache.options); + auto decode_context = make_context( + *host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + OffGraphReference reference(window); + EXPECT_LT(reference.step(*prefill, 0, 3), tol); + for (int64_t pos = 3; pos < 20; ++pos) { + SCOPED_TRACE(pos); + EXPECT_LT(reference.step(*decode, pos, 1), tol); + } + EXPECT_LT(reference.step(*prefill, 20, 3), tol); + } + } + } +} + +// A flat and a ring layer share one sequence: each layer's executable runs +// every step, the second re-running the first's positions. +TEST(GgmlEngineTest, OffGraphFlatAndRingLayersShareASequence) { + InstalledCache cache = install_cache( + 64, 1, ScalarType::Float, static_cast(kOffGraphDim), {0, 3}, 3); + auto host = GgmlEngineHost::create(); + std::vector> contexts; + std::vector> prefill; + std::vector> decode; + for (int64_t layer = 0; layer < 2; ++layer) { + contexts.push_back(make_context( + *host, make_update_and_attend_program(3, layer), cache.options)); + prefill.push_back(contexts.back()->compile("forward")); + contexts.push_back(make_context( + *host, make_update_and_attend_program(1, layer), cache.options)); + decode.push_back(contexts.back()->compile("forward")); + } + OffGraphReference flat; + OffGraphReference ring(3); + EXPECT_LT(flat.step(*prefill[0], 0, 3), 1e-4f); + EXPECT_LT(ring.step(*prefill[1], 0, 3), 1e-4f); + for (int64_t pos = 3; pos < 12; ++pos) { + SCOPED_TRACE(pos); + EXPECT_LT(flat.step(*decode[0], pos, 1), 1e-4f); + EXPECT_LT(ring.step(*decode[1], pos, 1), 1e-4f); + } +} + +// A ring layer can rewind only as far as the rows it still holds. +TEST(GgmlEngineTest, OffGraphRingRewindsWithinTheRowsItHolds) { + InstalledCache cache = install_cache( + 64, 1, ScalarType::Float, static_cast(kOffGraphDim), {4}, 3); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + OffGraphReference reference(4); + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + for (int64_t pos = 3; pos < 10; ++pos) { + EXPECT_LT(reference.step(*decode, pos, 1), 1e-4f); + } + // The ring holds positions 4..9; the step after a rewind to 7 reads 4..7. + EXPECT_FALSE(cache.control().rewind(6)); + EXPECT_THROW(reference.step(*decode, 6, 1), std::runtime_error); + EXPECT_LT(reference.step(*decode, 7, 1, /*seed=*/4), 1e-4f); + EXPECT_LT(reference.step(*decode, 8, 1, /*seed=*/4), 1e-4f); +} + TEST(GgmlEngineTest, OffGraphAttentionNeedsAMatchingCache) { auto host = GgmlEngineHost::create(); // No cache option.