From 8abeae2e5ce2b718e7ba85bd4e89114cd5e9039a Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 20 Oct 2026 15:52:42 -0700 Subject: [PATCH] [executorch][cuda] Run packed noncausal DFlash draft steps on the cell cache Pull Request resolved: https://github.com/pytorch/executorch/pull/23978 MLX batches DFlash drafts by packing every owner's new target context and query block into one noncausal draft step. Let the CUDA cell cache serve the same step so the CUDA DFlash executor can reuse the off-graph KV path. - CacheGeometry::dflash_block_size (already on main) now selects a packed draft layout on CUDA: context rows commit, query rows are scratch. - CudaCellCache prepares packed draft steps (declared context + block rows per owner) and lowers the noncausal mask the draft attends with. - lower_offgraph_kv / cuda_backend learn the draft step width spec. - Drop the "DFlash is unsupported on CUDA" builder guard. ghstack-source-id: 453099239 @exported-using-ghexport Differential Revision: [D125705710](https://our.internmc.facebook.com/intern/diff/D125705710/) --- backends/cuda/CMakeLists.txt | 5 + backends/cuda/cuda_backend.py | 11 +- backends/cuda/passes/lower_offgraph_kv.py | 78 ++++- backends/cuda/runtime/cuda_backend.cpp | 74 ++++- backends/cuda/runtime/cuda_kv_cache.cpp | 136 ++++++-- backends/cuda/runtime/cuda_kv_pool.cpp | 16 + backends/cuda/runtime/cuda_kv_pool.h | 3 + .../runtime/test/test_cuda_cache_routing.cpp | 267 +++++++++++++++ .../cuda/runtime/test/test_cuda_kv_cache.cpp | 210 +++++++++++- .../cuda/tests/test_dflash_offgraph_kv.py | 306 ++++++++++++++++++ extension/llm/cache/cache.h | 4 +- extension/llm/cache/cell_cache.cpp | 43 ++- extension/llm/cache/cell_cache.h | 6 +- extension/llm/cache/test/cache_test.cpp | 100 ++++++ 14 files changed, 1209 insertions(+), 50 deletions(-) create mode 100644 backends/cuda/runtime/test/test_cuda_cache_routing.cpp create mode 100644 backends/cuda/tests/test_dflash_offgraph_kv.py diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 25261c30487..50354e308b5 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -509,5 +509,10 @@ if(BUILD_TESTING) aoti_cuda_backend ) target_compile_definitions(test_cuda_kv_cache PRIVATE CUDA_AVAILABLE=1) + et_cxx_test( + test_cuda_cache_routing SOURCES runtime/test/test_cuda_cache_routing.cpp + EXTRA_LIBS aoti_cuda_backend + ) + target_compile_definitions(test_cuda_cache_routing PRIVATE CUDA_AVAILABLE=1) endif() endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index aad054fa288..787d2502b63 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -802,9 +802,18 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] ) triton_kernel_mode = mode passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()] + cell_layout = any( + spec.key == OFFGRAPH_KV_COMPILE_SPEC + and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell" + for spec in compile_specs + ) for spec in compile_specs: if spec.key == OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC: - passes.append(CheckOffGraphKVStepWidthPass(spec.value)) + passes.append( + CheckOffGraphKVStepWidthPass( + spec.value, layout="cell" if cell_layout else "sequence" + ) + ) for spec in compile_specs: if spec.key == OFFGRAPH_KV_COMPILE_SPEC: passes.append( diff --git a/backends/cuda/passes/lower_offgraph_kv.py b/backends/cuda/passes/lower_offgraph_kv.py index 60bf6821a73..0567e809139 100644 --- a/backends/cuda/passes/lower_offgraph_kv.py +++ b/backends/cuda/passes/lower_offgraph_kv.py @@ -215,19 +215,20 @@ def cell_step( so no causal alignment describes visibility and the mask is always explicit. ``read_len`` bounds the sweep to the occupied extent. - ``cells`` and ``mask`` are declared at the widest step; this one uses - their first T rows. ``k``/``v`` are BHSD; the pools are BSHD. sdpa returns + ``cells`` and ``mask`` are declared at the widest step; writes use the + first K/V-width cells and attention uses the first query-width mask rows. + ``k``/``v`` are BHSD; the pools are BSHD. sdpa returns the query's dtype; ``out_dtype`` is the neutral op's requested output dtype. """ - width = q.shape[2] - slots = cells[:width] + query_width = q.shape[2] + slots = cells[: k.shape[2]] _write(k_pool, slots, k) _write(v_pool, slots, v) out = torch.ops.triton.sdpa( q, k_pool.transpose(1, 2), v_pool.transpose(1, 2), - mask[:, :, :width, :], + mask[:, :, :query_width, :], 0.0, False, scale, @@ -319,6 +320,53 @@ def _source_input(node): return node +def _cell_position_source(node): + # Edge conversion emits copy variants of the value-preserving view ops. + # Match overloads exactly rather than admitting arbitrary name substrings. + preserving = tuple( + op + for name in ( + *_VALUE_PRESERVING_OPS, + "view_copy", + "unsqueeze_copy", + "squeeze_copy", + "expand_copy", + "alias_copy", + ) + for op in ( + getattr(torch.ops.aten, name).default, + getattr(exir_ops.edge.aten, name).default, + ) + ) + selects = ( + torch.ops.aten.index_select.default, + exir_ops.edge.aten.index_select.default, + ) + width = node.meta["val"].shape[0] + selected = False + while node.op == "call_function": + if node.target in preserving: + node = node.args[0] + elif node.target in selects and not selected: + source, dim, index = node.args + shape = source.meta["val"].shape + if ( + dim != 0 + or shape[0] != width + or tuple(node.meta["val"].shape) != tuple(shape) + or index.meta["val"].dim() != 1 + or index.meta["val"].shape[0] != shape[0] + ): + raise ValueError( + "off-graph cell position index_select must preserve full width on dim 0" + ) + selected = True + node = source + else: + break + return node + + class CheckOffGraphKVStepWidthPass: """Check the declared step width against the delegated program. @@ -331,7 +379,10 @@ class CheckOffGraphKVStepWidthPass: requires_exported_program = True - def __init__(self, value: bytes) -> None: + def __init__(self, value: bytes, layout: str = "sequence") -> None: + if layout not in ("sequence", "cell"): + raise ValueError(f"invalid off-graph KV layout {layout!r}") + self._cell_layout = layout == "cell" text = value.decode("utf-8") fields = text.split(":") if len(fields) != 2 or not all(f.isdigit() for f in fields): @@ -358,11 +409,16 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram: ) declared = inputs[self._input] position = node.args[3] - if _source_input(position) is not declared: + source = ( + _cell_position_source(position) + if self._cell_layout + else _source_input(position) + ) + if source is not declared: raise ValueError( f"off-graph KV step width names input {self._input} " f"({declared.name}), but the cache op's position comes " - f"from {_source_input(position).name}" + f"from {source.name}" ) shape = declared.meta["val"].shape width = position.meta["val"].shape[0] @@ -372,6 +428,12 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram: f"{declared.name} {tuple(shape)}, which is not the step's " f"{width} positions" ) + k, v = (arg.meta["val"] for arg in node.args[1:3]) + if k.shape[2] != width or v.shape[2] != width: + raise ValueError( + "off-graph KV step width requires matching K/V extent and " + f"positions, got {k.shape[2]}, {v.shape[2]} and {width}" + ) return exported_program diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index c8fea3e206c..1f989ab1b5b 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -109,6 +109,61 @@ Error commit_offgraph_kv_step(CudaDelegateHandle* handle, int64_t width) { return handle->kv_cache->commit_step(width); } +// A per-method actual key overrides the artifact's logical role. Without that +// override, the logical role names a load-time option containing the actual +// key; it is never itself used as a process-global registry ID. +Result resolve_cuda_cache_key( + const BackendInitContext& context, + ArrayRef compile_specs) { + constexpr const char* option = + ::executorch::extension::llm::cache::kCacheKeyOption; + std::optional logical; + for (const CompileSpec& spec : compile_specs) { + if (std::strcmp(spec.key, option) != 0) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + !logical && spec.value.buffer != nullptr && spec.value.nbytes > 0, + InvalidArgument, + "init: duplicate or empty logical CUDA cache key"); + std::string_view value( + static_cast(spec.value.buffer), spec.value.nbytes); + if (value.back() == '\0') { + value.remove_suffix(1); + } + ET_CHECK_OR_RETURN_ERROR( + !value.empty() && value.size() < runtime::kMaxOptionKeyLength && + value.find('\0') == std::string_view::npos && value != option, + InvalidArgument, + "init: invalid logical CUDA cache key"); + logical = std::string(value); + } + const auto actual = context.get_runtime_spec(option); + ET_CHECK_OR_RETURN_ERROR( + actual.ok() || actual.error() == Error::NotFound || !logical, + InvalidArgument, + "init: CUDA cache runtime override must be a string"); + if (actual.ok() && actual.get() != nullptr && *actual.get() != '\0') { + return std::string(actual.get()); + } + if (!logical) { + return std::string{}; + } + const auto mapped = context.get_runtime_spec(logical->c_str()); + ET_CHECK_OK_OR_RETURN_ERROR( + mapped.error(), + "init: logical CUDA cache '%s' requires a string runtime mapping to an " + "installed cache key, or a per-method '%s' override", + logical->c_str(), + option); + ET_CHECK_OR_RETURN_ERROR( + mapped.get() != nullptr && *mapped.get() != '\0', + InvalidArgument, + "init: logical CUDA cache '%s' has an empty runtime mapping", + logical->c_str()); + return std::string(mapped.get()); +} + class ET_EXPERIMENTAL CudaBackend final : public ::executorch::runtime::BackendInterface { private: @@ -350,6 +405,10 @@ class ET_EXPERIMENTAL CudaBackend final } } + const auto resolved_cache_key = + resolve_cuda_cache_key(context, compile_specs); + ET_CHECK_OK_OR_RETURN_ERROR(resolved_cache_key.error()); + std::string so_blob_key; std::string weights_blob_key; CudaWeightCache::Metadata fqn_weights; @@ -486,21 +545,16 @@ class ET_EXPERIMENTAL CudaBackend final // Runtime-owned off-graph buffers must capture their AOTI names before // the serialized constants update installs the ordinary weight set. // - // The cache is found by the key the runner published as a runtime spec. - // No key means an in-graph model, whose KV state is ordinary mutable - // buffers -- that is the whole of the off-graph/in-graph detection, and it - // needs nothing from the user. + // Resolve each delegate's logical cache role through this method's load + // options. The explicit actual-key override retains the ordinary runner + // contract; a logical name is never used directly as a registry ID. // // The key alone does not decide it, though: a program whose storage was // lowered off-graph carries no bytes for it, so running one without a cache // would hand its kernels unbound memory. That requirement is read from the // program's own constants and checked whether or not a key was passed. - const char* cache_key = nullptr; - if (auto spec = context.get_runtime_spec( - ::executorch::extension::llm::cache::kCacheKeyOption); - spec.ok() && spec.get() != nullptr && *spec.get() != '\0') { - cache_key = spec.get(); - } + const char* cache_key = + resolved_cache_key->empty() ? nullptr : resolved_cache_key->c_str(); const auto needs_offgraph_kv = requires_offgraph_kv_storage(*handle); ET_CHECK_OK_OR_RETURN_ERROR(needs_offgraph_kv.error()); if (cache_key != nullptr) { diff --git a/backends/cuda/runtime/cuda_kv_cache.cpp b/backends/cuda/runtime/cuda_kv_cache.cpp index ed4638868b2..91a63798c7e 100644 --- a/backends/cuda/runtime/cuda_kv_cache.cpp +++ b/backends/cuda/runtime/cuda_kv_cache.cpp @@ -288,6 +288,7 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { : cache::CellCache(geometry, table_config(cfg)), pool_cells_(cfg.capacity), max_write_(*cfg.max_write), + dflash_block_size_(geometry.dflash_block_size), windows_(std::move(windows)), pool_( pool_layers(geometry, cfg), @@ -327,15 +328,32 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { void clear() override { std::lock_guard guard(mutex_); + if (dflash_block_size_ && pool_.allocated()) { + // A failed forward may be abandoned on another host thread, where + // cudaStreamPerThread no longer names the originating stream. + if (pool_.synchronize(/*all_streams=*/true) != Error::Ok) { + // Recycling cells a forward may still write is unsafe; a later clear() + // retries the fence. + poisoned_ = true; + return; + } + } cache::CellCache::clear(); step_seq_ids_.clear(); + active_owners_.clear(); + pending_width_ = 0; + poisoned_ = false; } bool declare_step(const std::vector& seq_ids) override { std::lock_guard guard(mutex_); // A rejected declaration leaves an earlier accepted one standing, as the // base does. - if (!cache::CellCache::declare_step(seq_ids)) { + if (poisoned_ || + (dflash_block_size_ && + (pending_width_ != 0 || + seq_ids.size() > static_cast(max_write_))) || + !cache::CellCache::declare_step(seq_ids)) { return false; } step_seq_ids_ = seq_ids; @@ -350,16 +368,25 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { std::optional seq_clone(int32_t src, std::optional upto) override { std::lock_guard guard(mutex_); + if (dflash_block_size_ && (pending_width_ != 0 || poisoned_)) { + return std::nullopt; + } return cache::CellCache::seq_clone(src, upto); } bool seq_rm(int32_t seq_id) override { std::lock_guard guard(mutex_); + if (dflash_block_size_ && pending_width_ != 0) { + return false; + } return cache::CellCache::seq_rm(seq_id); } bool rewind(int32_t seq_id, int position) override { std::lock_guard guard(mutex_); + if (dflash_block_size_ && (pending_width_ != 0 || poisoned_)) { + return false; + } return cache::CellCache::rewind(seq_id, position); } @@ -389,7 +416,15 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { pool_.allocated(), InvalidState, "offgraph_kv: prepare_step must run before execute"); - return pool_.bind(handle); + ET_CHECK_OR_RETURN_ERROR( + !dflash_block_size_ || (!poisoned_ && pending_width_ > 0), + InvalidState, + "offgraph_kv: execute requires a prepared, healthy draft step"); + const Error bound = pool_.bind(handle); + if (bound != Error::Ok && dflash_block_size_) { + poisoned_ = true; + } + return bound; } // Places the declared step and writes it where the program reads it. Runs on @@ -397,6 +432,10 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { Error prepare_step(int64_t write_length, cudaStream_t stream) override { std::lock_guard guard(mutex_); ET_CHECK_OK_OR_RETURN_ERROR(error_); + ET_CHECK_OR_RETURN_ERROR( + !poisoned_ && (!dflash_block_size_ || pending_width_ == 0), + InvalidState, + "offgraph_kv: uncommitted or failed draft co-batch; clear required"); if (!validated_) { const Error valid = pool_.validate(); if (valid != Error::Ok) { @@ -414,7 +453,8 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { max_write_); const int tokens = static_cast(step_seq_ids_.size()); ET_CHECK_OR_RETURN_ERROR( - tokens > 0 && tokens <= width, + tokens > 0 && tokens <= width && + (!dflash_block_size_ || tokens == width), InvalidState, "offgraph_kv: a %d-token program cannot run %d declared tokens", width, @@ -447,18 +487,31 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { stream)); const int64_t scratch = pool_.rows() - 1; + if (dflash_block_size_) { + active_owners_.clear(); + for (const int32_t owner : step_seq_ids_) { + if (active_owners_.empty() || active_owners_.back() != owner) { + active_owners_.push_back(owner); + } + } + pending_width_ = width; + } std::vector steps; steps.reserve(windows_.size()); for (size_t index = 0; index < windows_.size(); ++index) { const cache::CellStep* step = place_step(window_layers_[index], positions_.data(), tokens); - ET_CHECK_OR_RETURN_ERROR( - step != nullptr, - InvalidArgument, - "offgraph_kv: the declared step does not place"); + if (step == nullptr) { + poisoned_ = dflash_block_size_.has_value(); + return Error::InvalidArgument; + } steps.push_back(step); } - ET_CHECK_OK_OR_RETURN_ERROR(upload_step(steps, width, scratch, stream)); + const Error uploaded = upload_step(steps, width, scratch, stream); + if (uploaded != Error::Ok) { + poisoned_ = dflash_block_size_.has_value(); + return uploaded; + } step_seq_ids_.clear(); return Error::Ok; } @@ -470,7 +523,29 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { ET_CHECK_OR_RETURN_ERROR( write_length > 0, InvalidArgument, "write length must be positive"); ET_CHECK_OK_OR_RETURN_ERROR(error_); - return pool_.mark_step_done(); + if (!dflash_block_size_) { + return pool_.mark_step_done(); + } + if (poisoned_ || pending_width_ == 0 || write_length != pending_width_) { + poisoned_ = true; + return Error::InvalidState; + } + const Error marked = pool_.mark_step_done(); + const Error fenced = marked == Error::Ok ? pool_.synchronize() : marked; + if (fenced != Error::Ok) { + poisoned_ = true; + return fenced; + } + for (const int32_t owner : active_owners_) { + if (!cache::CellCache::rewind( + owner, cache::CellCache::pos(owner) - *dflash_block_size_)) { + poisoned_ = true; + return Error::InvalidState; + } + } + active_owners_.clear(); + pending_width_ = 0; + return Error::Ok; } // logical_length is the read extent: every occupied cell is below it. @@ -551,9 +626,12 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { const cache::CellStep& first = *steps.front(); const size_t rows = static_cast(width); const size_t tokens = first.cells.size(); + const size_t queries = + first.query_rows.empty() ? tokens : first.query_rows.size(); + const size_t mask_rows = dflash_block_size_ ? queries : rows; const size_t read_len = static_cast(first.read_len); const size_t cells_bytes = rows * sizeof(int64_t); - const size_t mask_bytes = rows * read_len; + const size_t mask_bytes = mask_rows * read_len; ET_CHECK_OK_OR_RETURN_ERROR(reserve_staging( cells_bytes + sizeof(int64_t) + steps.size() * mask_bytes)); auto* cells = reinterpret_cast(staging_); @@ -563,14 +641,22 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { uint8_t* masks = staging_ + cells_bytes + sizeof(int64_t); for (size_t index = 0; index < steps.size(); ++index) { uint8_t* mask = masks + index * mask_bytes; - std::memcpy(mask, steps[index]->mask_bits.data(), tokens * read_len); - const uint8_t* last = mask + (tokens - 1) * read_len; - for (size_t row = tokens; row < rows; ++row) { + const cache::CellStep& step = *steps[index]; + for (size_t row = 0; row < queries; ++row) { + const size_t source = + step.query_rows.empty() ? row : step.query_rows[row]; + std::memcpy( + mask + row * read_len, + step.mask_bits.data() + source * read_len, + read_len); + } + const uint8_t* last = mask + (queries - 1) * read_len; + for (size_t row = queries; row < mask_rows; ++row) { std::memcpy(mask + row * read_len, last, read_len); } } const Error enqueued = - enqueue_uploads(steps.size(), rows, read_len, stream); + enqueue_uploads(steps.size(), rows, mask_rows, read_len, stream); // Fenced even when a copy failed: those issued before it may still read // the staging. const Error fenced = fence_staging(stream); @@ -584,6 +670,7 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { Error enqueue_uploads( size_t windows, size_t width, + size_t mask_rows, size_t read_len, cudaStream_t stream) { const size_t cells_bytes = width * sizeof(int64_t); @@ -607,10 +694,10 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { const cudaError_t error = cudaMemcpy2DAsync( pool_.side_buffer(kFirstMaskBuffer + index), static_cast(pool_cells_), - masks + index * width * read_len, + masks + index * mask_rows * read_len, read_len, read_len, - width, + mask_rows, cudaMemcpyHostToDevice, stream); ET_CHECK_OR_RETURN_ERROR( @@ -686,11 +773,15 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { // pitch. const int pool_cells_; const int max_write_; + const std::optional dflash_block_size_; const std::vector windows_; std::vector window_layers_; CudaKVPool pool_; std::vector step_seq_ids_; std::vector positions_; + std::vector active_owners_; + int pending_width_{0}; + bool poisoned_{false}; // Pinned host staging for one step's uploads, and the event marking where // the last step's copies from it end. uint8_t* staging_{nullptr}; @@ -712,8 +803,8 @@ std::shared_ptr make_cuda_sequence_kv_cache( slimc10::ScalarType storage_dtype = slimc10::ScalarType::BFloat16; // Checked before constructing: SequenceCache asserts on an invalid geometry, // and returning null lets CacheFactory report the failure instead. - if (!cache::valid(geometry, cfg) || cfg.initial_capacity <= 0 || - cfg.initial_capacity > cfg.capacity || + if (!cache::valid(geometry, cfg) || geometry.dflash_block_size || + cfg.initial_capacity <= 0 || cfg.initial_capacity > cfg.capacity || !storage_dtype_of(cfg.kv_dtype, storage_dtype)) { ET_LOG(Error, "offgraph_kv: invalid cache geometry or config"); return nullptr; @@ -735,10 +826,6 @@ std::shared_ptr make_cuda_sequence_kv_cache( std::shared_ptr make_cuda_cell_kv_cache( const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg) { - if (geometry.dflash_block_size.has_value()) { - ET_LOG(Error, "offgraph_kv: DFlash is unsupported on CUDA"); - return nullptr; - } slimc10::ScalarType storage_dtype = slimc10::ScalarType::BFloat16; if (!cache::valid(geometry, cfg) || cfg.initial_capacity <= 0 || cfg.initial_capacity > cfg.capacity || @@ -753,7 +840,10 @@ std::shared_ptr make_cuda_cell_kv_cache( return nullptr; } // One pool row is the padding scratch; the table needs at least one cell. - if (cfg.capacity < 2) { + if (cfg.capacity < 2 || + (geometry.dflash_block_size && + (*geometry.dflash_block_size >= cfg.capacity || + *geometry.dflash_block_size > *cfg.max_write))) { ET_LOG(Error, "offgraph_kv: the cell layout needs at least two cells"); return nullptr; } diff --git a/backends/cuda/runtime/cuda_kv_pool.cpp b/backends/cuda/runtime/cuda_kv_pool.cpp index ccc62ebc88f..556ba52b34a 100644 --- a/backends/cuda/runtime/cuda_kv_pool.cpp +++ b/backends/cuda/runtime/cuda_kv_pool.cpp @@ -183,6 +183,22 @@ Error CudaKVPool::mark_step_done() { return Error::Ok; } +Error CudaKVPool::synchronize(bool all_streams) { + if (!device_known_) { + return Error::Ok; + } + const auto device = CUDAGuard::create(device_); + ET_CHECK_OK_OR_RETURN_ERROR(device.error()); + const cudaError_t error = + all_streams ? cudaDeviceSynchronize() : cudaStreamSynchronize(stream_); + ET_CHECK_OR_RETURN_ERROR( + error == cudaSuccess, + Internal, + "offgraph_kv: cannot synchronize the step: %s", + cudaGetErrorString(error)); + return Error::Ok; +} + void* CudaKVPool::side_buffer(size_t index) const { return index < side_buffers_.size() ? side_buffers_[index] : nullptr; } diff --git a/backends/cuda/runtime/cuda_kv_pool.h b/backends/cuda/runtime/cuda_kv_pool.h index 3f1f24ec793..96b927c7d61 100644 --- a/backends/cuda/runtime/cuda_kv_pool.h +++ b/backends/cuda/runtime/cuda_kv_pool.h @@ -85,6 +85,9 @@ class CudaKVPool final { // follow if it runs on another. runtime::Error mark_step_done(); + // Host fence on the pool's device before cells can be reclaimed or reset. + runtime::Error synchronize(bool all_streams = false); + // Points `handle`'s constants at the current storage. Cached until storage // moves. Precondition: allocated(). runtime::Error bind(CudaDelegateHandle* handle); diff --git a/backends/cuda/runtime/test/test_cuda_cache_routing.cpp b/backends/cuda/runtime/test/test_cuda_cache_routing.cpp new file mode 100644 index 00000000000..f7c230e0904 --- /dev/null +++ b/backends/cuda/runtime/test/test_cuda_cache_routing.cpp @@ -0,0 +1,267 @@ +/* + * 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. + */ + +#include +#include +#include + +#include +#include + +namespace executorch::backends::cuda { +// Internal load-time resolver implemented by the CUDA delegate. +runtime::Result resolve_cuda_cache_key( + const runtime::BackendInitContext& context, + runtime::ArrayRef compile_specs); +} // namespace executorch::backends::cuda + +namespace { +namespace cache = executorch::extension::llm::cache; +namespace rt = executorch::runtime; +using executorch::backends::cuda::resolve_cuda_cache_key; + +class MarkerCache final : public cache::Cache { + protected: + void* face(cache::FaceId) override { + return nullptr; + } +}; + +template +rt::BackendInitContext context(rt::BackendOptions& options) { + auto view = options.view(); + return rt::BackendInitContext( + nullptr, + nullptr, + nullptr, + nullptr, + rt::Span(view.data(), view.size())); +} + +rt::CompileSpec role(const std::string& logical) { + return { + cache::kCacheKeyOption, + {const_cast(logical.data()), logical.size()}}; +} + +std::string installed_key(const cache::InstallGuard& guard) { + rt::BackendOptions<1> options; + EXPECT_EQ(guard.set_option(options), rt::Error::Ok); + const char* key = nullptr; + EXPECT_EQ(options.get_option(cache::kCacheKeyOption, key), rt::Error::Ok); + return key ? key : ""; +} + +TEST(CudaCacheRoutingTest, LegacyNoKeyKeepsInGraphBehavior) { + rt::BackendOptions<1> options; + auto resolved = resolve_cuda_cache_key(context(options), {}); + ASSERT_TRUE(resolved.ok()); + EXPECT_TRUE(resolved->empty()); +} + +TEST(CudaCacheRoutingTest, LegacyGuardOptionResolvesActualRegistryEntry) { + auto cache_ptr = std::make_shared(); + cache::InstallGuard guard(cache_ptr); + rt::BackendOptions<1> options; + ASSERT_EQ(guard.set_option(options), rt::Error::Ok); + auto resolved = resolve_cuda_cache_key(context(options), {}); + ASSERT_TRUE(resolved.ok()); + EXPECT_EQ(cache::CacheRegistry::global().get(*resolved), cache_ptr); +} + +TEST(CudaCacheRoutingTest, TargetAndDraftRolesMapToDistinctGuardKeys) { + auto target = std::make_shared(); + auto draft = std::make_shared(); + cache::InstallGuard target_guard(target), draft_guard(draft); + rt::BackendOptions<2> options; + ASSERT_EQ( + options.set_option( + "muse_glimmer_target", installed_key(target_guard).c_str()), + rt::Error::Ok); + ASSERT_EQ( + options.set_option( + "muse_glimmer_draft", installed_key(draft_guard).c_str()), + rt::Error::Ok); + const std::string names[] = {"muse_glimmer_target", "muse_glimmer_draft"}; + const std::shared_ptr expected[] = {target, draft}; + for (size_t i = 0; i < 2; ++i) { + auto spec = role(names[i]); + auto resolved = resolve_cuda_cache_key(context(options), {&spec, 1}); + ASSERT_TRUE(resolved.ok()); + EXPECT_NE(*resolved, names[i]); + EXPECT_EQ(cache::CacheRegistry::global().get(*resolved), expected[i]); + EXPECT_EQ(cache::CacheRegistry::global().get(names[i]), nullptr); + } +} + +TEST(CudaCacheRoutingTest, PerMethodActualOverrideWinsOverLogicalMapping) { + auto target = std::make_shared(); + auto draft = std::make_shared(); + cache::InstallGuard target_guard(target), draft_guard(draft); + rt::BackendOptions<2> options; + ASSERT_EQ(draft_guard.set_option(options), rt::Error::Ok); + ASSERT_EQ( + options.set_option( + "muse_glimmer_draft", installed_key(target_guard).c_str()), + rt::Error::Ok); + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + auto resolved = resolve_cuda_cache_key(context(options), {&spec, 1}); + ASSERT_TRUE(resolved.ok()); + EXPECT_EQ(cache::CacheRegistry::global().get(*resolved), draft); +} + +TEST(CudaCacheRoutingTest, PerMethodActualOverrideNeedsNoLogicalMapping) { + auto draft = std::make_shared(); + cache::InstallGuard guard(draft); + rt::BackendOptions<1> options; + ASSERT_EQ(guard.set_option(options), rt::Error::Ok); + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + auto resolved = resolve_cuda_cache_key(context(options), {&spec, 1}); + ASSERT_TRUE(resolved.ok()); + EXPECT_EQ(cache::CacheRegistry::global().get(*resolved), draft); +} + +TEST(CudaCacheRoutingTest, SameLogicalRoleInTwoModelsDoesNotCollide) { + auto a = std::make_shared(), b = std::make_shared(); + cache::InstallGuard ga(a), gb(b); + rt::BackendOptions<1> oa, ob; + ASSERT_EQ( + oa.set_option("muse_glimmer_draft", installed_key(ga).c_str()), + rt::Error::Ok); + ASSERT_EQ( + ob.set_option("muse_glimmer_draft", installed_key(gb).c_str()), + rt::Error::Ok); + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + auto ra = resolve_cuda_cache_key(context(oa), {&spec, 1}); + auto rb = resolve_cuda_cache_key(context(ob), {&spec, 1}); + ASSERT_TRUE(ra.ok()); + ASSERT_TRUE(rb.ok()); + EXPECT_NE(*ra, *rb); + EXPECT_EQ(cache::CacheRegistry::global().get(*ra), a); + EXPECT_EQ(cache::CacheRegistry::global().get(*rb), b); +} + +TEST(CudaCacheRoutingTest, MissingLogicalMappingIsNotARawRegistryLookup) { + rt::BackendOptions<1> options; + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {&spec, 1}).error(), + rt::Error::NotFound); +} + +TEST(CudaCacheRoutingTest, EmptyGenericOverrideFallsBackToRoleMapping) { + rt::BackendOptions<2> options; + ASSERT_EQ(options.set_option(cache::kCacheKeyOption, ""), rt::Error::Ok); + ASSERT_EQ( + options.set_option("muse_glimmer_draft", "cache-test"), rt::Error::Ok); + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + auto resolved = resolve_cuda_cache_key(context(options), {&spec, 1}); + ASSERT_TRUE(resolved.ok()); + EXPECT_EQ(*resolved, "cache-test"); +} + +TEST(CudaCacheRoutingTest, EmptyOrWrongTypeMappingIsRejected) { + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + rt::BackendOptions<1> options; + ASSERT_EQ(options.set_option("muse_glimmer_draft", ""), rt::Error::Ok); + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {&spec, 1}).error(), + rt::Error::InvalidArgument); + ASSERT_EQ(options.set_option("muse_glimmer_draft", 42), rt::Error::Ok); + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {&spec, 1}).error(), + rt::Error::InvalidArgument); +} + +TEST(CudaCacheRoutingTest, WrongTypeOverrideIsRejectedForLogicalRole) { + rt::BackendOptions<1> options; + ASSERT_EQ(options.set_option(cache::kCacheKeyOption, true), rt::Error::Ok); + const std::string logical = "muse_glimmer_draft"; + auto spec = role(logical); + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {&spec, 1}).error(), + rt::Error::InvalidArgument); +} + +TEST(CudaCacheRoutingTest, LegacyWrongTypeOverrideKeepsNoKeyBehavior) { + rt::BackendOptions<1> options; + ASSERT_EQ(options.set_option(cache::kCacheKeyOption, true), rt::Error::Ok); + auto resolved = resolve_cuda_cache_key(context(options), {}); + ASSERT_TRUE(resolved.ok()); + EXPECT_TRUE(resolved->empty()); +} + +TEST(CudaCacheRoutingTest, MalformedCompileKeysAreRejectedEvenWithOverride) { + rt::BackendOptions<1> options; + ASSERT_EQ( + options.set_option(cache::kCacheKeyOption, "cache-test"), rt::Error::Ok); + const std::string invalid[] = { + "", + std::string("bad\0role", 8), + std::string(rt::kMaxOptionKeyLength, 'x'), + cache::kCacheKeyOption}; + for (const auto& logical : invalid) { + auto spec = role(logical); + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {&spec, 1}).error(), + rt::Error::InvalidArgument); + } + rt::CompileSpec null_spec{cache::kCacheKeyOption, {nullptr, 4}}; + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {&null_spec, 1}).error(), + rt::Error::InvalidArgument); +} + +TEST(CudaCacheRoutingTest, DuplicateCompileKeysAreRejected) { + rt::BackendOptions<1> options; + const std::string logical = "muse_glimmer_draft"; + rt::CompileSpec specs[] = {role(logical), role(logical)}; + EXPECT_EQ( + resolve_cuda_cache_key(context(options), {specs, 2}).error(), + rt::Error::InvalidArgument); +} + +TEST(CudaCacheRoutingTest, OptionalTrailingNulIsAcceptedButNotEmbeddedNuls) { + rt::BackendOptions<1> options; + ASSERT_EQ(options.set_option("draft", "cache-test"), rt::Error::Ok); + const std::string logical("draft\0", 6); + auto spec = role(logical); + auto resolved = resolve_cuda_cache_key(context(options), {&spec, 1}); + ASSERT_TRUE(resolved.ok()); + EXPECT_EQ(*resolved, "cache-test"); +} + +TEST( + CudaCacheRoutingTest, + ResolvedKeyOwnsItsBytesAndAcquiredCacheOutlivesGuard) { + auto cache_ptr = std::make_shared(); + std::string key; + std::shared_ptr acquired; + { + cache::InstallGuard guard(cache_ptr); + rt::BackendOptions<1> options; + ASSERT_EQ(guard.set_option(options), rt::Error::Ok); + auto resolved = resolve_cuda_cache_key(context(options), {}); + ASSERT_TRUE(resolved.ok()); + key = *resolved; + ASSERT_EQ( + options.set_option(cache::kCacheKeyOption, "changed"), rt::Error::Ok); + EXPECT_EQ(*resolved, key); + acquired = cache::CacheRegistry::global().get(*resolved); + EXPECT_EQ(acquired, cache_ptr); + } + EXPECT_EQ(cache::CacheRegistry::global().get(key), nullptr); + EXPECT_EQ(acquired, cache_ptr); +} +} // namespace diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index c6461dcdaac..2d996a1e535 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -16,6 +16,7 @@ #include #include +#include #include #include #include @@ -1410,9 +1411,13 @@ class CudaCellCacheTest : public CudaKVCacheTest { static constexpr int kMaxWrite = 8; static constexpr int kWindow = 2; - static std::shared_ptr make(int initial = 4) { + static std::shared_ptr make( + int initial = 4, + std::optional block_size = std::nullopt) { + auto geometry = flat_and_ring(kWindow); + geometry.dflash_block_size = block_size; return cu::make_cuda_cell_kv_cache( - flat_and_ring(kWindow), config(kCells, initial, kMaxWrite)); + geometry, config(kCells, initial, kMaxWrite)); } // What a program lowered in the cell layout for make() declares: every @@ -1868,6 +1873,207 @@ TEST_F(CudaCellCacheTest, GrowthGivesAHandleAboutToCaptureOneEagerRunFirst) { kv->forget_handle(&early); } +TEST_F(CudaCellCacheTest, DFlashUploadsKVWideCellsAndQueryWideOwnerMasks) { + auto cache_ptr = make(4, 2); + ASSERT_NE(cache_ptr, nullptr); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(), b = *control->seq_new(); + ASSERT_EQ(step(*cache_ptr, handle, {b, b, b, b, b, a, a, a}), Error::Ok); + EXPECT_EQ( + read_longs(fake.bound["cells"].data, 8), + std::vector({0, 1, 2, 3, 4, 5, 6, 7})); + EXPECT_EQ(read_longs(fake.bound["read_len"].data, 1)[0], 8); + using Mask = std::vector>; + EXPECT_EQ( + read_mask(fake.bound["mask0"].data, 4, 8), + Mask( + {{1, 1, 1, 1, 1, 0, 0, 0}, + {1, 1, 1, 1, 1, 0, 0, 0}, + {0, 0, 0, 0, 0, 1, 1, 1}, + {0, 0, 0, 0, 0, 1, 1, 1}})); + EXPECT_EQ( + read_mask(fake.bound["mask2"].data, 4, 8), + Mask( + {{0, 1, 1, 1, 1, 0, 0, 0}, + {0, 1, 1, 1, 1, 0, 0, 0}, + {0, 0, 0, 0, 0, 1, 1, 1}, + {0, 0, 0, 0, 0, 1, 1, 1}})); + EXPECT_EQ(control->pos(b), 3); + EXPECT_EQ(control->pos(a), 1); + // Block removal leaves holes; the next owner order is reversed. + handle.cuda_graph_state.phase = cu::CudaGraphPhase::Replay; + void* cells = fake.bound["cells"].data; + ASSERT_EQ(step(*cache_ptr, handle, {a, a, b, b}), Error::Ok); + EXPECT_EQ( + read_longs(fake.bound["cells"].data, 4), + std::vector({3, 4, 6, 7})); + EXPECT_EQ( + read_mask(fake.bound["mask0"].data, 4, 8), + Mask( + {{0, 0, 0, 1, 1, 1, 0, 0}, + {0, 0, 0, 1, 1, 1, 0, 0}, + {1, 1, 1, 0, 0, 0, 1, 1}, + {1, 1, 1, 0, 0, 0, 1, 1}})); + EXPECT_EQ(handle.cuda_graph_state.phase, cu::CudaGraphPhase::Replay); + EXPECT_EQ(fake.bound["cells"].data, cells); + EXPECT_EQ(control->pos(b), 3); + EXPECT_EQ(control->pos(a), 1); + kv->forget_handle(&handle); +} + +TEST_F( + CudaCellCacheTest, + DFlashRejectsGroupingAndPaddedTokenCountsBeforeAllocation) { + auto cache_ptr = make(4, 2); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(), b = *control->seq_new(); + EXPECT_FALSE(control->declare_step({a, a, b, b, a, a})); + EXPECT_FALSE(control->declare_step({a, a, b})); + ASSERT_TRUE(control->declare_step({b, b, a, a})); + EXPECT_EQ(kv->prepare_step(5, cudaStreamPerThread), Error::InvalidState); + EXPECT_EQ(kv->prepare_step(3, cudaStreamPerThread), Error::InvalidState); + EXPECT_EQ(kv->metrics().allocated_bytes, 0); + EXPECT_EQ(control->pos(a), 0); + EXPECT_EQ(control->pos(b), 0); + ASSERT_EQ(kv->prepare_step(4, cudaStreamPerThread), Error::Ok); + ASSERT_EQ(kv->commit_step(4), Error::Ok); + EXPECT_EQ(control->pos(a), 0); + EXPECT_EQ(control->pos(b), 0); + kv->forget_handle(&handle); +} + +TEST_F( + CudaCellCacheTest, + DFlashAbandonedCoBatchCannotContinueUntilWholeCacheCleared) { + auto cache_ptr = make(4, 2); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(), b = *control->seq_new(); + ASSERT_TRUE(control->declare_step({a, a, a, b, b, b})); + ASSERT_EQ(kv->prepare_step(6, cudaStreamPerThread), Error::Ok); + // A failed forward does not commit: neither owner may continue or fork. + EXPECT_FALSE(control->declare_step({a, a})); + EXPECT_FALSE(control->rewind(a, 0)); + EXPECT_FALSE(control->seq_clone(b, std::nullopt)); + EXPECT_EQ(kv->prepare_step(6, cudaStreamPerThread), Error::InvalidState); + EXPECT_EQ(kv->commit_step(5), Error::InvalidState); + EXPECT_EQ(kv->commit_step(6), Error::InvalidState); + control->clear(); + EXPECT_EQ(kv->metrics().logical_length, 0); + const int32_t reused = *control->seq_new(); + ASSERT_EQ(step(*cache_ptr, handle, {reused, reused, reused}), Error::Ok); + EXPECT_EQ(control->pos(reused), 1); + kv->forget_handle(&handle); +} + +TEST_F(CudaCellCacheTest, DFlashWarmupCommitsAllContextWithOneQueryPerOwner) { + auto cache_ptr = make(4, 0); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(), b = *control->seq_new(); + ASSERT_EQ(step(*cache_ptr, handle, {b, b, b, a}), Error::Ok); + EXPECT_EQ( + read_mask(fake.bound["mask2"].data, 2, 4), + std::vector>({{0, 1, 1, 0}, {0, 0, 0, 1}})); + EXPECT_EQ(control->pos(b), 3); + EXPECT_EQ(control->pos(a), 1); + kv->forget_handle(&handle); +} + +TEST_F( + CudaCellCacheTest, + DFlashCommitFencesGPUWorkBeforeDroppingTemporaryCells) { + auto cache_ptr = make(4, 2); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(); + ASSERT_TRUE(control->declare_step({a, a, a})); + ASSERT_EQ(kv->prepare_step(3, cudaStreamPerThread), Error::Ok); + ASSERT_EQ(kv->rebind_for_execute(&handle), Error::Ok); + std::atomic finished{false}; + ASSERT_EQ( + cudaLaunchHostFunc( + cudaStreamPerThread, + [](void* state) { + std::this_thread::sleep_for(std::chrono::milliseconds(30)); + static_cast*>(state)->store(true); + }, + &finished), + cudaSuccess); + ASSERT_EQ(kv->commit_step(3), Error::Ok); + EXPECT_TRUE(finished.load()); + EXPECT_EQ(control->pos(a), 1); + kv->forget_handle(&handle); +} + +TEST_F(CudaCellCacheTest, DFlashClearFromAnotherThreadFencesAbandonedForward) { + auto cache_ptr = make(4, 2); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(); + ASSERT_TRUE(control->declare_step({a, a, a})); + ASSERT_EQ(kv->prepare_step(3, cudaStreamPerThread), Error::Ok); + ASSERT_EQ(kv->rebind_for_execute(&handle), Error::Ok); + std::atomic finished{false}; + ASSERT_EQ( + cudaLaunchHostFunc( + cudaStreamPerThread, + [](void* state) { + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + static_cast*>(state)->store(true); + }, + &finished), + cudaSuccess); + bool finished_when_cleared = false; + std::thread cleanup([&] { + control->clear(); + finished_when_cleared = finished.load(); + }); + cleanup.join(); + EXPECT_TRUE(finished_when_cleared); + EXPECT_EQ(kv->metrics().logical_length, 0); + const int32_t reused = *control->seq_new(); + ASSERT_EQ(step(*cache_ptr, handle, {reused, reused, reused}), Error::Ok); + EXPECT_EQ(control->pos(reused), 1); + kv->forget_handle(&handle); +} + +TEST_F(CudaCellCacheTest, DFlashBuilderRejectsUnsupportedSequenceLayout) { + auto geometry = flat_and_ring(kWindow); + for (const int block : {0, 2}) { + geometry.dflash_block_size = block; + EXPECT_EQ(make_cache(geometry, config(kCells, 4, kMaxWrite)), nullptr); + } + geometry.dflash_block_size = -1; + EXPECT_EQ( + cu::make_cuda_cell_kv_cache(geometry, config(kCells, 4, kMaxWrite)), + nullptr); + geometry.dflash_block_size = kMaxWrite + 1; + EXPECT_EQ( + cu::make_cuda_cell_kv_cache(geometry, config(kCells, 4, kMaxWrite)), + nullptr); +} + TEST_F(CudaCellCacheTest, BuilderValidatesAndIsRegisteredForBatchedKinds) { auto no_max_write = config(kCells, 4, kMaxWrite); no_max_write.max_write.reset(); diff --git a/backends/cuda/tests/test_dflash_offgraph_kv.py b/backends/cuda/tests/test_dflash_offgraph_kv.py new file mode 100644 index 00000000000..2ccb93b154b --- /dev/null +++ b/backends/cuda/tests/test_dflash_offgraph_kv.py @@ -0,0 +1,306 @@ +# 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. + +import json +from unittest.mock import patch + +import pytest +import torch +import torch.nn.functional as F +from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.passes.lower_offgraph_kv import ( + cell_step, + CheckOffGraphKVStepWidthPass, + LowerOffGraphKVPass, + OFFGRAPH_KV_CELLS_FQN, + OFFGRAPH_KV_COMPILE_SPEC, + offgraph_kv_mask_fqn, + OFFGRAPH_KV_READ_LEN_FQN, + OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, + parse_offgraph_kv_manifest, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.extension.llm.cache.update_and_attend import ( # noqa: F401 + update_and_attend, +) + + +SCALE = 0.125 + + +class _DraftAttention(torch.nn.Module): + def __init__(self, transform="select"): + super().__init__() + self.transform = transform + + def forward(self, q, k, v, position, order): + if self.transform == "select": + position = position.index_select(0, order) + elif self.transform == "short_select": + position = position.index_select(0, order[:-1]) + elif self.transform == "wrong_dim_select": + position = position.reshape(-1, 1).index_select(1, order[:1] * 0) + elif self.transform == "add": + position = position + 1 + elif self.transform == "flip": + position = position.flip(0) + elif self.transform == "double_select": + position = position.index_select(0, order).index_select(0, order) + return torch.ops.kvcache.update_and_attend( + q, k, v, position.reshape(-1, 1), 0, SCALE, torch.bfloat16 + ) + + +def _args(device="cpu", q_width=4, kv_width=8): + q = torch.ones(1, 4, q_width, 64, device=device, dtype=torch.bfloat16) + k = torch.zeros(1, 2, kv_width, 64, device=device, dtype=torch.bfloat16) + v = torch.arange(kv_width, device=device, dtype=torch.bfloat16) + v = v.reshape(1, 1, kv_width, 1).expand_as(k).contiguous() + return ( + q, + k, + v, + torch.arange(kv_width, device=device), + torch.arange(kv_width - 1, -1, -1, device=device), + ) + + +def _edge(transform="select", args=None, dynamic=False): + args = _args() if args is None else args + shapes = None + if dynamic: + queries = torch.export.Dim("queries", min=2, max=16) + writes = torch.export.Dim("writes", min=2, max=32) + shapes = ({2: queries}, {2: writes}, {2: writes}, {0: writes}, {0: writes}) + program = torch.export.export( + _DraftAttention(transform), args, dynamic_shapes=shapes, strict=False + ) + return to_edge( + program, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ).exported_program() + + +def _manifest(max_cells=512): + return json.dumps( + { + "version": 1, + "layout": "cell", + "maximum_capacity": max_cells, + "max_cells": max_cells, + "max_write": 16, + "layers": [{"layer_id": 0, "policy": "flat"}], + } + ).encode() + + +@pytest.mark.parametrize("dynamic", [False, True]) +def test_cell_checker_accepts_full_width_position_select(dynamic): + edge = _edge(dynamic=dynamic) + assert CheckOffGraphKVStepWidthPass(b"3:0", layout="cell")(edge) is edge + + +@pytest.mark.parametrize("transform", ["select", "add", "flip", "double_select"]) +def test_ordinary_checker_remains_strict(transform): + with pytest.raises(ValueError, match="position comes from"): + CheckOffGraphKVStepWidthPass(b"3:0")(_edge(transform)) + + +@pytest.mark.parametrize("transform", ["add", "flip", "double_select"]) +def test_cell_checker_rejects_unvalidated_transforms(transform): + with pytest.raises(ValueError, match="position comes from"): + CheckOffGraphKVStepWidthPass(b"3:0", layout="cell")(_edge(transform)) + + +@pytest.mark.parametrize("transform", ["short_select", "wrong_dim_select"]) +def test_cell_checker_rejects_non_full_token_axis_select(transform): + with pytest.raises(ValueError, match="preserve full width"): + CheckOffGraphKVStepWidthPass(b"3:0", layout="cell")(_edge(transform)) + + +@pytest.mark.parametrize( + "spec,error", + [(b"1:2", "position comes from"), (b"3:1", "not the step"), (b"5:0", "takes 5")], +) +def test_cell_checker_validates_delegate_input_and_dimension(spec, error): + with pytest.raises(ValueError, match=error): + CheckOffGraphKVStepWidthPass(spec, layout="cell")(_edge()) + + +@pytest.mark.parametrize("which", [1, 2]) +def test_cell_checker_rejects_mismatched_kv_extent(which): + args = list(_args()) + args[which] = args[which][:, :, :-1].contiguous() + with pytest.raises(ValueError, match="matching K/V extent"): + CheckOffGraphKVStepWidthPass(b"3:0", layout="cell")(_edge(args=tuple(args))) + + +def test_ordinary_checker_accepts_direct_position(): + edge = _edge("direct", args=_args(q_width=8)) + assert CheckOffGraphKVStepWidthPass(b"3:0")(edge) is edge + + +def test_backend_uses_manifest_layout_for_checker_independent_of_spec_order(): + for layout in ("cell", "sequence"): + manifest = dict(json.loads(_manifest()), layout=layout) + specs = [ + CompileSpec(OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC, b"3:0"), + CompileSpec(OFFGRAPH_KV_COMPILE_SPEC, json.dumps(manifest).encode()), + ] + for ordered in (specs, list(reversed(specs))): + passes = CudaBackend.get_custom_passes(ordered) + checker = next( + p for p in passes if isinstance(p, CheckOffGraphKVStepWidthPass) + ) + assert passes.index(checker) < next( + i for i, p in enumerate(passes) if isinstance(p, LowerOffGraphKVPass) + ) + if layout == "cell": + assert checker(_edge()) is not None + else: + with pytest.raises(ValueError, match="position comes from"): + checker(_edge()) + + +def test_cpu_lowering_uses_kv_cells_and_query_mask_rows(): + edge = _edge() + CheckOffGraphKVStepWidthPass(b"3:0", layout="cell")(edge) + lowered = LowerOffGraphKVPass(parse_offgraph_kv_manifest(_manifest()))(edge) + nodes = list(lowered.graph.nodes) + assert not any( + n.target == exir_ops.edge.kvcache.update_and_attend.default for n in nodes + ) + copies = [n for n in nodes if n.target == torch.ops.aten.index_copy_.default] + assert len(copies) == 2 + assert all(tuple(n.args[2].meta["val"].shape) == (8,) for n in copies) + (attention,) = [n for n in nodes if n.target == torch.ops.triton.sdpa.default] + assert tuple(attention.args[3].meta["val"].shape) == (1, 1, 4, 512) + assert attention.args[5] is False + + +def _runtime(device, read_len): + args = _args(device) + q, k, v = args[:3] + # Owner 0: retained cell 0, three context rows, two block rows. + # Owner 1: retained cell 1, one context row, two block rows. + slots = torch.tensor([3, 6, 2, 5, 4, 9, 7, 8], device=device) + cells = torch.full((16,), 511, dtype=torch.long, device=device) + cells[:8] = slots + k_pool = torch.full((1, 512, 2, 64), 10000.0, dtype=q.dtype, device=device) + v_pool = torch.full_like(k_pool, -10000.0) + k_pool[:, :2].zero_() + v_pool[:, 0].fill_(10) + v_pool[:, 1].fill_(20) + mask = torch.zeros(1, 1, 16, 512, dtype=torch.bool, device=device) + mask[0, 0, :2, [0, 3, 6, 2, 5, 4]] = True + mask[0, 0, 2:4, [1, 9, 7, 8]] = True + # Stale visible tail must be excluded by the device scalar, not the mask. + mask[..., read_len:] = True + expected = torch.cat( + [ + F.scaled_dot_product_attention( + q[:, :, :2].float(), + torch.cat([k_pool[:, :1].transpose(1, 2), k[:, :, :5]], 2).float(), + torch.cat([v_pool[:, :1].transpose(1, 2), v[:, :, :5]], 2).float(), + scale=SCALE, + enable_gqa=True, + ), + F.scaled_dot_product_attention( + q[:, :, 2:].float(), + torch.cat([k_pool[:, 1:2].transpose(1, 2), k[:, :, 5:]], 2).float(), + torch.cat([v_pool[:, 1:2].transpose(1, 2), v[:, :, 5:]], 2).float(), + scale=SCALE, + enable_gqa=True, + ), + ], + 2, + ) + return ( + args, + (k_pool, v_pool, cells, torch.tensor([read_len], device=device), mask), + expected, + ) + + +def _cpu_sdpa(q, k, v, mask, dropout, causal, scale, gqa, read_len): + extent = int(read_len.item()) + return F.scaled_dot_product_attention( + q.float(), + k[:, :, :extent].float(), + v[:, :, :extent].float(), + attn_mask=mask[..., :extent], + scale=scale, + enable_gqa=gqa, + ).to(q.dtype) + + +def test_cpu_cell_step_writes_all_kv_and_attends_only_block_queries(): + args, runtime, expected = _runtime("cpu", 10) + with patch.object(torch.ops.triton, "sdpa", _cpu_sdpa): + actual = cell_step(*args[:3], *runtime, SCALE, torch.float32) + torch.testing.assert_close(actual, expected, atol=0.03, rtol=0) + k_pool, v_pool, cells = runtime[:3] + torch.testing.assert_close(k_pool[:, cells[:8]], args[1].transpose(1, 2)) + torch.testing.assert_close(v_pool[:, cells[:8]], args[2].transpose(1, 2)) + assert torch.all(k_pool[:, 10:] == 10000) + assert torch.all(v_pool[:, 10:] == -10000) + + +@pytest.mark.parametrize("mode", ["eager", "compile", "graph"]) +@pytest.mark.parametrize("read_len", [10, 300]) +def test_gpu_unequal_width_owner_visibility_and_poisoned_tail(mode, read_len): + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("CUDA BF16 required") + args, runtime, expected = _runtime("cuda", read_len) + fn = cell_step + if mode != "eager": + fn = torch.compile(cell_step, fullgraph=True) + inputs = (*args[:3], *runtime, SCALE) + actual = fn(*inputs) + if mode == "graph": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + fn(*inputs) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = fn(*inputs) + # Replay must consume new K/V and runtime side-buffer data at fixed addresses. + args[2].add_(2) + expected[:, :, :2].add_(2 * 5 / 6) + expected[:, :, 2:].add_(2 * 3 / 4) + runtime[3].fill_(10) + runtime[4][..., 10:] = True + graph.replay() + torch.testing.assert_close(actual.float(), expected, atol=0.04, rtol=0) + k_pool, v_pool, cells = runtime[:3] + torch.testing.assert_close(k_pool[:, cells[:8]], args[1].transpose(1, 2)) + torch.testing.assert_close(v_pool[:, cells[:8]], args[2].transpose(1, 2)) + assert torch.all(k_pool[:, read_len:] == 10000) + assert torch.all(v_pool[:, read_len:] == -10000) + + +def test_gpu_lowered_draft_program_matches_owner_reference(): + if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported(): + pytest.skip("CUDA BF16 required") + args, buffers, expected = _runtime("cuda", 10) + edge = _edge(args=args) + CheckOffGraphKVStepWidthPass(b"3:0", layout="cell")(edge) + lowered = LowerOffGraphKVPass(parse_offgraph_kv_manifest(_manifest()))(edge) + names = ( + "__et_offgraph_kv_layer_0_k", + "__et_offgraph_kv_layer_0_v", + OFFGRAPH_KV_CELLS_FQN, + OFFGRAPH_KV_READ_LEN_FQN, + offgraph_kv_mask_fqn(0), + ) + for name, value in zip(names, buffers): + lowered.constants[name] = value + actual = lowered.module()(*args) + torch.testing.assert_close(actual.float(), expected, atol=0.04, rtol=0) diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 6342c234667..09fab2fc0c6 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -185,7 +185,9 @@ ET_EXPERIMENTAL inline bool valid(const CacheGeometry& geometry) { if (layer.n_kv_heads <= 0 || layer.head_dim <= 0) { return false; } - if ((layer.policy.kind == LayerPolicy::Kind::Flat && + if ((layer.policy.kind != LayerPolicy::Kind::Flat && + layer.policy.kind != LayerPolicy::Kind::Ring) || + (layer.policy.kind == LayerPolicy::Kind::Flat && layer.policy.window != 0) || (layer.policy.kind == LayerPolicy::Kind::Ring && layer.policy.window <= 0)) { diff --git a/extension/llm/cache/cell_cache.cpp b/extension/llm/cache/cell_cache.cpp index 3ad85eedbf6..907bdc000ef 100644 --- a/extension/llm/cache/cell_cache.cpp +++ b/extension/llm/cache/cell_cache.cpp @@ -18,6 +18,7 @@ namespace cache { CellCache::CellCache(const CacheGeometry& geometry, const CacheConfig& cfg) : capacity_(cfg.capacity), + dflash_block_size_(geometry.dflash_block_size), pos_(cfg.capacity, -1), owners_(cfg.capacity, 0) { assert(valid(geometry, cfg)); @@ -59,6 +60,22 @@ bool CellCache::declare_step(const std::vector& seq_ids) { return false; } } + if (dflash_block_size_) { + uint64_t seen = 0; + for (size_t begin = 0; begin < seq_ids.size();) { + const int32_t owner = seq_ids[begin]; + size_t end = begin + 1; + while (end < seq_ids.size() && seq_ids[end] == owner) { + ++end; + } + if ((seen & bit(owner)) || + end - begin < static_cast(*dflash_block_size_)) { + return false; + } + seen |= bit(owner); + begin = end; + } + } step_seq_ids_ = seq_ids; declared_ = true; invalidate_steps(); @@ -161,7 +178,8 @@ int CellCache::used_end() const { const CellStep* CellCache::place_step(int layer, const int32_t* positions, int length) { - if (layer < 0 || layer >= static_cast(windows_.size())) { + if (positions == nullptr || length <= 0 || layer < 0 || + layer >= static_cast(windows_.size())) { return nullptr; // layer out of range } if (placed_) { @@ -281,6 +299,19 @@ const CellStep& CellCache::step_for(int window) { step.read_len = used_end_; step.cells = cells_; step.mask_bits = build_mask(window); + if (dflash_block_size_) { + const int queries = std::max(1, *dflash_block_size_); + for (int begin = 0; begin < step.length;) { + int end = begin + 1; + while (end < step.length && step_seq_ids_[end] == step_seq_ids_[begin]) { + ++end; + } + for (int row = end - queries; row < end; ++row) { + step.query_rows.push_back(row); + } + begin = end; + } + } return steps_.emplace(window, std::move(step)).first->second; } @@ -291,14 +322,18 @@ std::vector CellCache::build_mask(int window) const { for (int i = 0; i < length; ++i) { const uint64_t tok_bit = bit(step_seq_ids_[i]); const int32_t tok_pos = step_pos_[i]; - // A flat layer reaches back to the start; a windowed one to its window. - const int32_t oldest = window > 0 ? tok_pos - window + 1 : 0; + // Draft history is anchored at context end, shared by all block queries. + const int32_t context_end = + info_[step_seq_ids_[i]].max_pos + 1 - dflash_block_size_.value_or(0); + const int32_t oldest = window > 0 + ? (dflash_block_size_ ? context_end - window : tok_pos - window + 1) + : 0; uint8_t* row = bits.data() + static_cast(i) * used_end_; for (int j = 0; j < used_end_; ++j) { row[j] = static_cast( pos_[j] >= 0 && // occupied: a freed cell holds nothing (owners_[j] & tok_bit) && // one of this query's sequences - pos_[j] <= tok_pos && // not the future; <= so a query sees itself + (dflash_block_size_ || pos_[j] <= tok_pos) && // full draft block pos_[j] >= oldest); // within the window } } diff --git a/extension/llm/cache/cell_cache.h b/extension/llm/cache/cell_cache.h index 5062dbe08ea..4ab8365dff9 100644 --- a/extension/llm/cache/cell_cache.h +++ b/extension/llm/cache/cell_cache.h @@ -32,8 +32,11 @@ namespace cache { struct ET_EXPERIMENTAL CellStep { int length; int read_len; // the window is cells [0, read_len) - std::vector cells; // cell per query token + std::vector cells; // cell per written K/V token std::vector mask_bits; // [length, read_len], 1 = attend + // Full-mask rows to use as queries, in first-owner-appearance block order. + // Empty for ordinary attention, which uses every row. + std::vector query_rows; }; // Backend face of the cell layout. The step's first layer places its tokens and @@ -150,6 +153,7 @@ class ET_EXPERIMENTAL CellCache : public Cache, std::vector build_mask(int window) const; int capacity_; + const std::optional dflash_block_size_; std::vector pos_; // per cell; -1 = free std::vector owners_; // per cell; owning-sequence bitset int used_count_ = 0; // occupied cells, so admission stays O(1) diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index e7d19c80326..c073428d1ed 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -846,6 +846,106 @@ TEST_F(CacheTest, KvSharedLayerReservesIdempotently) { EXPECT_EQ(c.place(0, {7, 8, 9}), nullptr); // different tokens, never declared } +TEST_F(CacheTest, DFlashCellRowsUseOwnerOrderAndContextAnchoredWindow) { + CacheGeometry geometry{{flat_layer(), ring_layer(2), ring_layer(2)}}; + geometry.dflash_block_size = 2; + CellCache cache(geometry, config(32)); + const int32_t a = *cache.seq_new(), b = *cache.seq_new(); + // b is first despite its larger id; context lengths are unequal. + const auto args = flatten_step({{b, 0, 5}, {a, 0, 3}}); + ASSERT_TRUE(cache.declare_step(args.seq_ids)); + const auto* flat = cache.place_step(0, args.positions.data(), 8); + const auto* windowed = cache.place_step(1, args.positions.data(), 8); + ASSERT_NE(flat, nullptr); + ASSERT_NE(windowed, nullptr); + EXPECT_EQ(flat->query_rows, (std::vector{3, 4, 6, 7})); + EXPECT_EQ(windowed->query_rows, flat->query_rows); + EXPECT_EQ(flat->cells, (std::vector{0, 1, 2, 3, 4, 5, 6, 7})); + EXPECT_EQ(row(*flat, 3), "11111..."); + EXPECT_EQ(row(*flat, 4), "11111..."); + EXPECT_EQ(row(*flat, 6), ".....111"); + // b's context ends at 3: retain positions 1,2 and the full block 3,4. + EXPECT_EQ(row(*windowed, 3), ".1111..."); + EXPECT_EQ(row(*windowed, 4), ".1111..."); + EXPECT_EQ(row(*windowed, 6), ".....111"); + EXPECT_EQ(cache.place_step(2, args.positions.data(), 8), windowed); + EXPECT_EQ(cache.free_cells(), 24); +} + +TEST_F(CacheTest, DFlashCellRejectsMalformedGroupingWithoutPlacement) { + CacheGeometry geometry{{flat_layer()}}; + geometry.dflash_block_size = 2; + CellCache cache(geometry, config(16)); + const int32_t a = *cache.seq_new(), b = *cache.seq_new(); + ASSERT_TRUE(cache.declare_step({b, b, a, a})); + EXPECT_FALSE(cache.declare_step({a, a, b, b, a, a})); + EXPECT_FALSE(cache.declare_step({a, a, b})); + EXPECT_FALSE(cache.declare_step({a, -1})); + const int32_t wrong[] = {0, 1, 0, 2}; + EXPECT_EQ(cache.place_step(0, wrong, 4), nullptr); + EXPECT_EQ(cache.place_step(0, wrong, 3), nullptr); + EXPECT_EQ(cache.place_step(0, nullptr, 4), nullptr); + EXPECT_EQ(cache.free_cells(), 16); + EXPECT_EQ(cache.pos(a), 0); + EXPECT_EQ(cache.pos(b), 0); + const int32_t positions[] = {0, 1, 0, 1}; + ASSERT_NE(cache.place_step(0, positions, 4), nullptr); + EXPECT_EQ(cache.free_cells(), 12); +} + +TEST_F(CacheTest, DFlashCellRewindRetainsContextAndReusesTemporaryHoles) { + CacheGeometry geometry{{flat_layer()}}; + geometry.dflash_block_size = 2; + CellCache cache(geometry, config(16)); + const int32_t a = *cache.seq_new(), b = *cache.seq_new(); + auto args = flatten_step({{b, 0, 5}, {a, 0, 3}}); + ASSERT_TRUE(cache.declare_step(args.seq_ids)); + ASSERT_NE(cache.place_step(0, args.positions.data(), 8), nullptr); + ASSERT_TRUE(cache.rewind(b, cache.pos(b) - 2)); + ASSERT_TRUE(cache.rewind(a, cache.pos(a) - 2)); + EXPECT_EQ(cache.pos(b), 3); + EXPECT_EQ(cache.pos(a), 1); + EXPECT_EQ(cache.free_cells(), 12); + const auto fork = cache.seq_clone(b, std::nullopt); + ASSERT_TRUE(fork); + args = flatten_step({{a, 1, 2}, {*fork, 3, 2}}); + ASSERT_TRUE(cache.declare_step(args.seq_ids)); + const auto* step = cache.place_step(0, args.positions.data(), 4); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->cells, (std::vector{3, 4, 6, 7})); + EXPECT_EQ(step->query_rows, (std::vector{0, 1, 2, 3})); + EXPECT_EQ(row(*step, 0), "...111.."); + EXPECT_EQ(row(*step, 2), "111...11"); + ASSERT_TRUE(cache.seq_rm(a)); + EXPECT_EQ(*cache.seq_new(), a); +} + +TEST_F(CacheTest, DFlashZeroWarmupSelectsOneQueryPerOwnerAndPersistsContext) { + CacheGeometry geometry{{flat_layer(), ring_layer(2)}}; + geometry.dflash_block_size = 0; + CellCache cache(geometry, config(16)); + const int32_t a = *cache.seq_new(), b = *cache.seq_new(); + const auto args = flatten_step({{b, 0, 4}, {a, 0, 1}}); + ASSERT_TRUE(cache.declare_step(args.seq_ids)); + const auto* step = cache.place_step(1, args.positions.data(), 5); + ASSERT_NE(step, nullptr); + EXPECT_EQ(step->query_rows, (std::vector{3, 4})); + EXPECT_EQ(row(*step, 3), "..11."); + EXPECT_EQ(row(*step, 4), "....1"); + EXPECT_EQ(cache.pos(b), 4); + EXPECT_EQ(cache.pos(a), 1); +} + +TEST_F(CacheTest, DFlashGeometryRejectsNegativeBlocksAndUnknownPolicy) { + CacheGeometry geometry{{flat_layer()}}; + geometry.dflash_block_size = -1; + EXPECT_FALSE(executorch::extension::llm::cache::valid(geometry)); + geometry.dflash_block_size = 0; + EXPECT_TRUE(executorch::extension::llm::cache::valid(geometry)); + geometry.layers[0].policy.kind = static_cast(42); + EXPECT_FALSE(executorch::extension::llm::cache::valid(geometry)); +} + TEST_F(CacheTest, CellRejectedDeclarationLeavesTheLastOneStanding) { Cells c(16); const int32_t s0 = c.seq_new();