From 77fee2d637a148dd888b4e1b316ddef57518ce4b Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 16:46:35 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/runtime/cuda_kv_cache.cpp | 539 ++++++++++++++++++ backends/cuda/runtime/cuda_kv_cache.h | 19 + backends/cuda/runtime/cuda_kv_pool.cpp | 7 + backends/cuda/runtime/cuda_kv_pool.h | 3 + .../cuda/runtime/test/test_cuda_kv_cache.cpp | 308 ++++++++++ extension/llm/cache/cache.h | 7 +- 6 files changed, 882 insertions(+), 1 deletion(-) diff --git a/backends/cuda/runtime/cuda_kv_cache.cpp b/backends/cuda/runtime/cuda_kv_cache.cpp index 87800103363..70755db7ed4 100644 --- a/backends/cuda/runtime/cuda_kv_cache.cpp +++ b/backends/cuda/runtime/cuda_kv_cache.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -790,6 +791,517 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { Error error_{Error::Ok}; }; +// Batched-sequence layout: every sequence keeps a private history in its own +// column of each layer's storage, as MLX's batched-sequence cache keeps one +// buffer per sequence. A layer's storage is [rows, columns, H, D] (declared +// [1, rows * columns, H, D]): row = the position, or position % rows on a ring +// layer, so each sequence is contiguous by position and attention reads only +// its own column, bounded by its own length. The last column is scratch, +// written by static-width padding and never read. +// +// The neutral SequenceCache plans each sequence: admission against its +// capacity, ring limits, and rewind floors. The pool owns the bytes; rows grow +// together for every column. Before each forward, prepare_step writes each +// token's column (seq_cols) and one (q_start, q_len, column, length) row per +// run of one sequence's query rows (seq_info) into buffers the program reads +// at fixed addresses, so a captured CUDA graph serves any mix of sequences. +// +// A DFlash draft step declares, per owner, its new context then its block; +// only the block is queried, it sees all of that owner's history, and it is +// dropped again after the forward, as on the cell layout. +constexpr char kSeqColsFqn[] = "__et_offgraph_kv_seq_cols"; +constexpr char kSeqInfoFqn[] = "__et_offgraph_kv_seq_info"; +constexpr size_t kSeqColsBuffer = 0; +constexpr size_t kSeqInfoBuffer = 1; + +class CudaBatchedSequenceCache final : public cache::Cache, + public cache::BatchControl, + public CudaKVCache { + public: + CudaBatchedSequenceCache( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg, + slimc10::ScalarType storage_dtype) + : geometry_(geometry), + cfg_(cfg), + max_seqs_(*cfg.max_seqs), + columns_(*cfg.max_seqs + 1), + max_write_(*cfg.max_write), + max_segments_(batched_sequence_max_segments(*cfg.max_seqs)), + dflash_block_size_(geometry.dflash_block_size), + pool_( + pool_layers(geometry, cfg), + {{kSeqColsFqn, slimc10::ScalarType::Long, {*cfg.max_write}}, + {kSeqInfoFqn, + slimc10::ScalarType::Int, + {batched_sequence_max_segments(*cfg.max_seqs), 4}}}, + storage_dtype, + static_cast(cfg.initial_capacity) * (*cfg.max_seqs + 1)) { + for (const cache::LayerGeometry& layer : geometry.layers) { + layer_rows_.push_back(rows_of(layer, cfg)); + } + } + + CudaBatchedSequenceCache(const CudaBatchedSequenceCache&) = delete; + CudaBatchedSequenceCache& operator=(const CudaBatchedSequenceCache&) = delete; + + ~CudaBatchedSequenceCache() override { + if (staged_ != nullptr) { + (void)cudaEventSynchronize(staged_); + (void)cudaEventDestroy(staged_); + } + if (staging_ != nullptr) { + (void)cudaFreeHost(staging_); + } + } + + // -- CacheControl / BatchControl. + + // Every sequence's capacity, summed: the most the cache holds at once. + int capacity() const override { + return cfg_.capacity * max_seqs_; + } + + // Drops every sequence and keeps the storage, so a reset neither regrows + // nor invalidates captured graphs. + void clear() override { + std::lock_guard guard(mutex_); + seqs_.clear(); + step_seq_ids_.clear(); + } + + bool declare_step(const std::vector& seq_ids) override { + std::lock_guard guard(mutex_); + if (seq_ids.empty() || seq_ids.size() > static_cast(max_write_)) { + return false; + } + std::map counts; + int runs = 0; + for (size_t i = 0; i < seq_ids.size(); ++i) { + if (seqs_.count(seq_ids[i]) == 0) { + return false; + } + ++counts[seq_ids[i]]; + runs += (i == 0 || seq_ids[i] != seq_ids[i - 1]) ? 1 : 0; + } + if (runs > max_segments_ || + (dflash_block_size_ && runs != static_cast(counts.size()))) { + return false; + } + for (const auto& [seq_id, count] : counts) { + const cache::SequenceCache& plan = seqs_.at(seq_id).plan; + if (dflash_block_size_ && count < *dflash_block_size_) { + return false; + } + for (size_t layer = 0; layer < geometry_.layers.size(); ++layer) { + if (!plan.plan(static_cast(layer), plan.length(), count)) { + return false; + } + } + } + step_seq_ids_ = seq_ids; + return true; + } + + std::optional max_seqs() const override { + return max_seqs_; + } + + std::optional seq_new() override { + std::lock_guard guard(mutex_); + const auto column = free_column(); + if (!column) { + return std::nullopt; + } + seqs_.emplace( + *column, Sequence{*column, cache::SequenceCache(geometry_, cfg_)}); + return *column; + } + + // Copies the retained prefix into a free column: a sequence's history is + // private, so a fork cannot share it. + std::optional seq_clone(int32_t src, std::optional upto) + override { + std::lock_guard guard(mutex_); + const auto source = seqs_.find(src); + const auto column = free_column(); + if (source == seqs_.end() || !column || !pool_.allocated()) { + return std::nullopt; + } + const int length = source->second.plan.length(); + const int keep = std::min(upto.value_or(length), length); + cache::SequenceCache plan(source->second.plan); + if (keep <= 0 || !plan.rewind(keep) || + copy_column(source->second.column, *column, keep) != Error::Ok) { + return std::nullopt; + } + seqs_.emplace(*column, Sequence{*column, std::move(plan)}); + return *column; + } + + bool seq_rm(int32_t seq_id) override { + std::lock_guard guard(mutex_); + return seqs_.erase(seq_id) > 0; + } + + bool rewind(int32_t seq_id, int position) override { + std::lock_guard guard(mutex_); + const auto it = seqs_.find(seq_id); + return it != seqs_.end() && it->second.plan.rewind(position); + } + + int pos(int32_t seq_id) const override { + std::lock_guard guard(mutex_); + const auto it = seqs_.find(seq_id); + return it == seqs_.end() ? 0 : it->second.plan.length(); + } + + // -- CudaKVCache. + + runtime::Result note_handle(CudaDelegateHandle* handle) override { + std::lock_guard guard(mutex_); + auto serves = pool_.note_handle(handle); + if (!serves.ok()) { + error_ = serves.error(); + } + return serves; + } + + void forget_handle(CudaDelegateHandle* handle) override { + std::lock_guard guard(mutex_); + pool_.forget_handle(handle); + } + + Error rebind_for_execute(CudaDelegateHandle* handle) override { + std::lock_guard guard(mutex_); + if (!pool_.serves(handle)) { + return Error::Ok; + } + ET_CHECK_OK_OR_RETURN_ERROR(error_); + ET_CHECK_OR_RETURN_ERROR( + pool_.allocated(), + InvalidState, + "offgraph_kv: prepare_step must run before execute"); + return pool_.bind(handle); + } + + 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( + !dflash_block_size_ || pending_owners_.empty(), + InvalidState, + "offgraph_kv: an uncommitted draft step; clear required"); + if (!validated_) { + const Error valid = pool_.validate(); + if (valid != Error::Ok) { + error_ = valid; + return valid; + } + validated_ = true; + } + const int tokens = static_cast(step_seq_ids_.size()); + ET_CHECK_OR_RETURN_ERROR( + write_length > 0 && write_length <= max_write_ && tokens > 0 && + tokens <= write_length, + InvalidArgument, + "offgraph_kv: a %lld-token program cannot run %d declared tokens", + static_cast(write_length), + tokens); + + // One run per segment, each continuing its sequence. + struct Run { + int32_t seq_id; + int start; + int length; + }; + std::vector runs; + std::map next; + for (int i = 0; i < tokens; ++i) { + const int32_t seq_id = step_seq_ids_[i]; + if (runs.empty() || runs.back().seq_id != seq_id) { + runs.push_back({seq_id, i, 0}); + } + ++runs.back().length; + next.emplace(seq_id, seqs_.at(seq_id).plan.length()); + } + // Rows every column must hold once this step lands, and those it holds now. + int64_t live = 0; + int64_t required = 0; + for (const auto& [seq_id, sequence] : seqs_) { + live = std::max(live, sequence.plan.length()); + } + std::map ends = next; + for (const Run& run : runs) { + ends[run.seq_id] += run.length; + } + for (const auto& [seq_id, end] : ends) { + required = std::max(required, end); + } + ET_CHECK_OK_OR_RETURN_ERROR(pool_.prepare( + std::min(required, cfg_.capacity) * columns_, + live * columns_, + stream)); + + ET_CHECK_OK_OR_RETURN_ERROR(reserve_staging( + static_cast(max_write_) * sizeof(int64_t) + + static_cast(max_segments_) * 4 * sizeof(int32_t))); + auto* cols = reinterpret_cast(staging_); + auto* info = reinterpret_cast(cols + max_write_); + std::fill(cols, cols + write_length, static_cast(max_seqs_)); + std::fill(info, info + max_segments_ * 4, 0); + int q_start = 0; + for (size_t index = 0; index < runs.size(); ++index) { + const Run& run = runs[index]; + const int column = seqs_.at(run.seq_id).column; + std::fill(cols + run.start, cols + run.start + run.length, column); + int32_t* row = info + index * 4; + const int end = next[run.seq_id] + run.length; + next[run.seq_id] = end; + // A draft queries only each owner's block, packed owner by owner. + const int queries = dflash_block_size_ ? *dflash_block_size_ : run.length; + row[0] = dflash_block_size_ ? q_start : run.start; + row[1] = queries; + row[2] = column; + row[3] = end; + q_start += queries; + } + const Error enqueued = upload(static_cast(write_length), stream); + const Error fenced = fence_staging(stream); + ET_CHECK_OK_OR_RETURN_ERROR(enqueued); + ET_CHECK_OK_OR_RETURN_ERROR(fenced); + + // The rows are claimed now, as the cell layout claims its cells; the + // forward only fills them. + for (const Run& run : runs) { + cache::SequenceCache& plan = seqs_.at(run.seq_id).plan; + const auto step = plan.plan(0, plan.length(), run.length); + ET_CHECK_OR_RETURN_ERROR( + step.has_value(), InvalidArgument, "offgraph_kv: uncommittable step"); + plan.commit(*step); + if (dflash_block_size_) { + pending_owners_.push_back(run.seq_id); + } + } + step_seq_ids_.clear(); + return Error::Ok; + } + + Error commit_step(int64_t write_length) override { + std::lock_guard guard(mutex_); + ET_CHECK_OR_RETURN_ERROR( + write_length > 0, InvalidArgument, "write length must be positive"); + ET_CHECK_OK_OR_RETURN_ERROR(error_); + ET_CHECK_OK_OR_RETURN_ERROR(pool_.mark_step_done()); + // A draft's block is scratch: the next draft rewrites it. + for (const int32_t owner : pending_owners_) { + const auto it = seqs_.find(owner); + if (it != seqs_.end() && + !it->second.plan.rewind( + it->second.plan.length() - *dflash_block_size_)) { + pending_owners_.clear(); + return Error::InvalidState; + } + } + pending_owners_.clear(); + return Error::Ok; + } + + // logical_length is the longest sequence: every column holds that many rows. + OffGraphKVMetrics metrics() const override { + std::lock_guard guard(mutex_); + OffGraphKVMetrics metrics; + for (const auto& [seq_id, sequence] : seqs_) { + metrics.logical_length = + std::max(metrics.logical_length, sequence.plan.length()); + } + metrics.flat_capacity = pool_.rows() / columns_; + metrics.growth_count = pool_.growth_count(); + metrics.allocated_bytes = pool_.allocated_bytes(); + return metrics; + } + + protected: + void* face(cache::FaceId id) override { + return cache::expose(this, id); + } + + private: + struct Sequence { + int column; + cache::SequenceCache plan; + }; + + // A ring never needs more rows than a sequence can hold; it then never wraps. + static int64_t rows_of( + const cache::LayerGeometry& layer, + const cache::CacheConfig& cfg) { + if (!is_ring(layer)) { + return cfg.capacity; + } + return std::min( + cfg.capacity, + static_cast(layer.policy.window) + *cfg.max_write - 1); + } + + static std::vector pool_layers( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) { + std::vector layers; + for (const cache::LayerGeometry& layer : geometry.layers) { + layers.push_back(CudaKVPool::Layer{ + layer.n_kv_heads, + layer.head_dim, + rows_of(layer, cfg) * (*cfg.max_seqs + 1), + /*growable=*/!is_ring(layer)}); + } + return layers; + } + + std::optional free_column() const { + for (int column = 0; column < max_seqs_; ++column) { + if (seqs_.count(column) == 0) { + return column; + } + } + return std::nullopt; + } + + // Copies rows [0, keep) of every flat layer, and a ring layer's whole ring, + // from one column to another, ordered on the latest step's stream. + Error copy_column(int from, int to, int keep) { + for (size_t index = 0; index < geometry_.layers.size(); ++index) { + const cache::LayerGeometry& layer = geometry_.layers[index]; + const size_t slot = static_cast(layer.n_kv_heads) * + layer.head_dim * slimc10::elementSize(storage_dtype()); + const size_t pitch = slot * columns_; + const int64_t rows = is_ring(layer) + ? layer_rows_[index] + : std::min(keep, pool_.rows() / columns_); + for (const bool value : {false, true}) { + auto* base = static_cast(pool_.layer_data(index, value)); + const cudaError_t error = cudaMemcpy2DAsync( + base + slot * to, + pitch, + base + slot * from, + pitch, + slot, + static_cast(rows), + cudaMemcpyDeviceToDevice, + pool_.stream()); + ET_CHECK_OR_RETURN_ERROR( + error == cudaSuccess, + Internal, + "offgraph_kv: cannot copy a sequence's history: %s", + cudaGetErrorString(error)); + } + } + return Error::Ok; + } + + slimc10::ScalarType storage_dtype() const { + slimc10::ScalarType dtype = slimc10::ScalarType::BFloat16; + (void)storage_dtype_of(cfg_.kv_dtype, dtype); + return dtype; + } + + Error upload(size_t width, cudaStream_t stream) { + const size_t cols_bytes = static_cast(max_write_) * sizeof(int64_t); + for (const auto& [index, offset, bytes] : + {std::tuple{kSeqColsBuffer, size_t{0}, width * sizeof(int64_t)}, + std::tuple{ + kSeqInfoBuffer, + cols_bytes, + static_cast(max_segments_) * 4 * sizeof(int32_t)}}) { + const cudaError_t error = cudaMemcpyAsync( + pool_.side_buffer(index), + staging_ + offset, + bytes, + cudaMemcpyHostToDevice, + stream); + ET_CHECK_OR_RETURN_ERROR( + error == cudaSuccess, + Internal, + "offgraph_kv: cannot write the step's sequences: %s", + cudaGetErrorString(error)); + } + return Error::Ok; + } + + // Waits for the previous step's copies to finish reading the staging, then + // makes it hold at least `bytes`. + Error reserve_staging(size_t bytes) { + if (staged_ != nullptr) { + const cudaError_t error = cudaEventSynchronize(staged_); + ET_CHECK_OR_RETURN_ERROR( + error == cudaSuccess, + Internal, + "offgraph_kv: cannot wait for the previous step's uploads: %s", + cudaGetErrorString(error)); + } + if (bytes <= staging_bytes_) { + return Error::Ok; + } + if (staging_ != nullptr) { + (void)cudaFreeHost(staging_); + staging_ = nullptr; + staging_bytes_ = 0; + } + void* data = nullptr; + const cudaError_t error = cudaMallocHost(&data, bytes); + ET_CHECK_OR_RETURN_ERROR( + error == cudaSuccess, + MemoryAllocationFailed, + "offgraph_kv: cannot allocate %zu bytes of pinned staging: %s", + bytes, + cudaGetErrorString(error)); + staging_ = static_cast(data); + staging_bytes_ = bytes; + return Error::Ok; + } + + Error fence_staging(cudaStream_t stream) { + cudaError_t error = cudaSuccess; + if (staged_ == nullptr) { + error = cudaEventCreateWithFlags(&staged_, cudaEventDisableTiming); + } + if (error == cudaSuccess) { + error = cudaEventRecord(staged_, stream); + } + if (error != cudaSuccess) { + (void)cudaStreamSynchronize(stream); + ET_LOG( + Error, + "offgraph_kv: cannot fence the step uploads: %s", + cudaGetErrorString(error)); + return Error::Internal; + } + return Error::Ok; + } + + // Guards everything below. The runner drives the verbs and the delegate the + // steps, from whichever threads they run on. + mutable std::mutex mutex_; + + const cache::CacheGeometry geometry_; + const cache::CacheConfig cfg_; + const int max_seqs_; + const int columns_; + const int max_write_; + const int max_segments_; + const std::optional dflash_block_size_; + std::vector layer_rows_; + CudaKVPool pool_; + std::map seqs_; + std::vector step_seq_ids_; + std::vector pending_owners_; + uint8_t* staging_{nullptr}; + size_t staging_bytes_{0}; + cudaEvent_t staged_{nullptr}; + bool validated_{false}; + Error error_{Error::Ok}; +}; + } // namespace std::shared_ptr make_cuda_sequence_kv_cache( @@ -857,6 +1369,32 @@ std::shared_ptr make_cuda_cell_kv_cache( std::vector(windows.begin(), windows.end())); } +std::shared_ptr make_cuda_batched_sequence_kv_cache( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) { + slimc10::ScalarType storage_dtype = slimc10::ScalarType::BFloat16; + if (!cache::valid(geometry, cfg) || 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; + } + // The storage is one column per sequence and the step buffers are declared + // [max_write] and [segments, 4]; all are fixed by the program. + if (!cfg.max_seqs || !cfg.max_write || *cfg.max_write > cfg.capacity || + (geometry.dflash_block_size && + (*geometry.dflash_block_size <= 0 || + *geometry.dflash_block_size > *cfg.max_write))) { + ET_LOG( + Error, + "offgraph_kv: the batched-sequence layout requires max_seqs and " + "max_write"); + return nullptr; + } + return std::make_shared( + geometry, cfg, storage_dtype); +} + runtime::Error attach_offgraph_kv_cache( CudaDelegateHandle& handle, const char* cache_key, @@ -904,6 +1442,7 @@ const bool cuda_cache_builders_registered = [] { } builders[] = { {cache::kind::kSingle, make_cuda_sequence_kv_cache}, {cache::kind::kBatchedCell, make_cuda_cell_kv_cache}, + {cache::kind::kBatchedSequence, make_cuda_batched_sequence_kv_cache}, // The cell layout is the batch layout this backend serves. {cache::kind::kBatched, make_cuda_cell_kv_cache}, }; diff --git a/backends/cuda/runtime/cuda_kv_cache.h b/backends/cuda/runtime/cuda_kv_cache.h index e4da35f53c9..17f3aac4672 100644 --- a/backends/cuda/runtime/cuda_kv_cache.h +++ b/backends/cuda/runtime/cuda_kv_cache.h @@ -217,4 +217,23 @@ std::shared_ptr make_cuda_cell_kv_cache( const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg); +// Segments of a batched-sequence step: runs of one sequence's query rows. The +// program declares its seq_info buffer [segments, 4] from the same rule. +inline int batched_sequence_max_segments(int max_seqs) { + return 2 * max_seqs; +} + +// Builder for cache::kind::kBatchedSequence: many sequences, each with a +// private history in its own column, for a program lowered in the +// batched-sequence layout. +// +// cfg.capacity is one sequence's capacity (the program's maximum_capacity), +// cfg.max_seqs its sequence columns and cfg.max_write its widest step; all +// fix shapes the program declared, so they must match the export exactly. +// Flat layers start at cfg.initial_capacity rows and grow geometrically to +// cfg.capacity; ring layers hold min(capacity, window + max_write - 1). +std::shared_ptr make_cuda_batched_sequence_kv_cache( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg); + } // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/cuda_kv_pool.cpp b/backends/cuda/runtime/cuda_kv_pool.cpp index 556ba52b34a..c52a0031d5d 100644 --- a/backends/cuda/runtime/cuda_kv_pool.cpp +++ b/backends/cuda/runtime/cuda_kv_pool.cpp @@ -199,6 +199,13 @@ Error CudaKVPool::synchronize(bool all_streams) { return Error::Ok; } +void* CudaKVPool::layer_data(size_t index, bool value) const { + if (index >= allocations_.size()) { + return nullptr; + } + return value ? allocations_[index].v : allocations_[index].k; +} + 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 96b927c7d61..02d9267cc7a 100644 --- a/backends/cuda/runtime/cuda_kv_pool.h +++ b/backends/cuda/runtime/cuda_kv_pool.h @@ -99,6 +99,9 @@ class CudaKVPool final { // allocated(). void* side_buffer(size_t index) const; size_t side_buffer_bytes(size_t index) const; + // Device pointer of layer `index`'s K or V storage, which moves when a + // growable layer grows. Precondition: allocated(). + void* layer_data(size_t index, bool value) const; // The stream of the latest step, on which side-buffer writes must be issued. cudaStream_t stream() const { return stream_; diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index 0c32da2079f..3cd435fe02f 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -2133,3 +2134,310 @@ TEST_F(CudaCellCacheTest, BuilderValidatesAndIsRegisteredForBatchedKinds) { EXPECT_NE(built.get()->as(), nullptr) << kind; } } + +namespace { + +class CudaBatchedSequenceCacheTest : public CudaKVCacheTest { + protected: + static constexpr int kCapacity = 16; + static constexpr int kMaxWrite = 8; + static constexpr int kMaxSeqs = 3; + static constexpr int kColumns = kMaxSeqs + 1; + static constexpr int kWindow = 2; + // min(capacity, window + max_write - 1) + static constexpr int kRingRows = kWindow + kMaxWrite - 1; + static constexpr int kSegments = 2 * kMaxSeqs; + + static cache::CacheConfig sequence_config(int initial = 4) { + auto cfg = config(kCapacity, initial, kMaxWrite); + cfg.max_seqs = kMaxSeqs; + return cfg; + } + + 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_batched_sequence_kv_cache( + geometry, sequence_config(initial)); + } + + // What a program lowered in the batched-sequence layout declares for make(). + static FakeContainer container() { + FakeContainer container{ + {"f_k", "f_v", "r_k", "r_v", "seq_cols", "seq_info"}, + {"__et_offgraph_kv_layer_0_k", + "__et_offgraph_kv_layer_0_v", + "__et_offgraph_kv_layer_1_k", + "__et_offgraph_kv_layer_1_v", + "__et_offgraph_kv_seq_cols", + "__et_offgraph_kv_seq_info"}, + {}, + 0}; + for (size_t index = 0; index < 4; ++index) { + declare( + container, + container.fqns[index], + slimc10::ScalarType::BFloat16, + {1, (index < 2 ? kCapacity : kRingRows) * kColumns, kHeads, kDim}); + } + declare( + container, + "__et_offgraph_kv_seq_cols", + slimc10::ScalarType::Long, + {kMaxWrite}); + declare( + container, + "__et_offgraph_kv_seq_info", + slimc10::ScalarType::Int, + {kSegments, 4}); + return container; + } + + static std::vector read_cols(void* device, int count) { + std::vector values(count); + EXPECT_EQ( + cudaMemcpy( + values.data(), + device, + count * sizeof(int64_t), + cudaMemcpyDeviceToHost), + cudaSuccess); + return values; + } + + using Info = std::vector>; + static Info read_info(void* device, int rows) { + std::vector flat(kSegments * 4); + EXPECT_EQ( + cudaMemcpy( + flat.data(), + device, + flat.size() * sizeof(int32_t), + cudaMemcpyDeviceToHost), + cudaSuccess); + Info out(rows); + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < 4; ++j) { + out[i][j] = flat[i * 4 + j]; + } + } + return out; + } + + static Error step( + cache::Cache& cache, + cu::CudaDelegateHandle& handle, + const std::vector& seq_ids, + int64_t width = 0) { + auto* control = cache.as(); + auto* kv = cache.as(); + if (!control->declare_step(seq_ids)) { + return Error::InvalidArgument; + } + if (width == 0) { + width = static_cast(seq_ids.size()); + } + ET_CHECK_OK_OR_RETURN_ERROR(kv->prepare_step(width, cudaStreamPerThread)); + ET_CHECK_OK_OR_RETURN_ERROR(kv->rebind_for_execute(&handle)); + return kv->commit_step(width); + } +}; + +} // namespace + +TEST_F(CudaBatchedSequenceCacheTest, EachSequenceWritesItsOwnColumn) { + auto cache_ptr = make(); + ASSERT_NE(cache_ptr, nullptr); + auto* control = cache_ptr->as(); + auto* kv = cache_ptr->as(); + ASSERT_NE(control, nullptr); + ASSERT_NE(kv, nullptr); + EXPECT_EQ(control->max_seqs(), kMaxSeqs); + auto fake = container(); + auto handle = make_handle(fake); + ASSERT_TRUE(kv->note_handle(&handle).get()); + const int32_t a = *control->seq_new(); + const int32_t b = *control->seq_new(); + + // Prefilled together in a program two wider: padding names the scratch. + ASSERT_EQ(step(*cache_ptr, handle, {a, a, a, b, b}, 7), Error::Ok); + EXPECT_EQ( + read_cols(fake.bound["seq_cols"].data, 7), + std::vector({0, 0, 0, 1, 1, kMaxSeqs, kMaxSeqs})); + EXPECT_EQ( + read_info(fake.bound["seq_info"].data, 3), + Info({{0, 3, 0, 3}, {3, 2, 1, 2}, {0, 0, 0, 0}})); + EXPECT_EQ( + fake.bound["f_k"].sizes, + std::vector({1, kCapacity * kColumns, kHeads, kDim})); + EXPECT_EQ(fake.bound["seq_info"].dtype, slimc10::ScalarType::Int); + EXPECT_EQ(kv->metrics().flat_capacity, 4); + + // Decoding in the other order keeps each one's column and own length. + ASSERT_EQ(step(*cache_ptr, handle, {b, a}), Error::Ok); + EXPECT_EQ( + read_cols(fake.bound["seq_cols"].data, 2), std::vector({1, 0})); + EXPECT_EQ( + read_info(fake.bound["seq_info"].data, 2), + Info({{0, 1, 1, 3}, {1, 1, 0, 4}})); + EXPECT_EQ(control->pos(a), 4); + EXPECT_EQ(control->pos(b), 3); + EXPECT_EQ(kv->metrics().logical_length, 4); + + // The longest sequence decides the rows; growth doubles them. + ASSERT_EQ(step(*cache_ptr, handle, {a}), Error::Ok); + const auto metrics = kv->metrics(); + EXPECT_EQ(metrics.flat_capacity, 8); + EXPECT_EQ(metrics.growth_count, 1); + kv->forget_handle(&handle); +} + +TEST_F(CudaBatchedSequenceCacheTest, AdmissionIsPerSequenceAndBounded) { + auto cache_ptr = make(/*initial=*/kCapacity); + 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()); + std::vector ids; + for (int i = 0; i < kMaxSeqs; ++i) { + ids.push_back(*control->seq_new()); + } + EXPECT_FALSE(control->seq_new().has_value()); + // Unknown ids, and steps past max_write, are refused. + EXPECT_FALSE(control->declare_step({ids[0], 99})); + EXPECT_FALSE( + control->declare_step(std::vector(kMaxWrite + 1, ids[0]))); + // A sequence fills its own capacity without touching another's room. + for (int i = 0; i < kCapacity / kMaxWrite; ++i) { + ASSERT_EQ( + step(*cache_ptr, handle, std::vector(kMaxWrite, ids[0])), + Error::Ok); + } + EXPECT_EQ(control->pos(ids[0]), kCapacity); + EXPECT_FALSE(control->declare_step({ids[0]})); + EXPECT_TRUE(control->declare_step({ids[1]})); + ASSERT_TRUE(control->seq_rm(ids[0])); + EXPECT_EQ(control->seq_new(), ids[0]); + EXPECT_EQ(control->pos(ids[0]), 0); + kv->forget_handle(&handle); +} + +TEST_F(CudaBatchedSequenceCacheTest, CloneCopiesThePrefixIntoAFreeColumn) { + auto cache_ptr = make(/*initial=*/8); + 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_EQ(step(*cache_ptr, handle, {a, a, a, a}), Error::Ok); + + // Fill a's column with a recognizable pattern, then clone its first 3 rows. + const size_t slot = kHeads * kDim * sizeof(uint16_t); + const size_t pitch = slot * kColumns; + std::vector rows(pitch * 8, 0); + for (int r = 0; r < 4; ++r) { + std::fill_n(rows.begin() + r * pitch, slot, static_cast(r + 1)); + } + ASSERT_EQ( + cudaMemcpy( + fake.bound["f_k"].data, + rows.data(), + rows.size(), + cudaMemcpyHostToDevice), + cudaSuccess); + const auto c = control->seq_clone(a, 3); + ASSERT_TRUE(c.has_value()); + EXPECT_NE(*c, a); + EXPECT_EQ(control->pos(*c), 3); + EXPECT_EQ(control->pos(a), 4); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + std::vector after(rows.size()); + ASSERT_EQ( + cudaMemcpy( + after.data(), + fake.bound["f_k"].data, + after.size(), + cudaMemcpyDeviceToHost), + cudaSuccess); + for (int r = 0; r < 4; ++r) { + EXPECT_EQ(after[r * pitch + *c * slot], r < 3 ? r + 1 : 0) << r; + } + // The clone continues at its own length. + ASSERT_EQ(step(*cache_ptr, handle, {*c}), Error::Ok); + EXPECT_EQ(read_info(fake.bound["seq_info"].data, 1), Info({{0, 1, *c, 4}})); + EXPECT_FALSE(control->seq_clone(99, std::nullopt).has_value()); + kv->forget_handle(&handle); +} + +TEST_F(CudaBatchedSequenceCacheTest, RewindRespectsTheRingFloor) { + auto cache_ptr = make(/*initial=*/kCapacity); + 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_EQ( + step(*cache_ptr, handle, std::vector(kMaxWrite, a)), Error::Ok); + ASSERT_EQ( + step(*cache_ptr, handle, std::vector(kMaxWrite, a)), Error::Ok); + // A ring of kRingRows rows retains back to 16 - window = 14 only. + EXPECT_FALSE(control->rewind(a, 2)); + EXPECT_TRUE(control->rewind(a, 15)); + EXPECT_EQ(control->pos(a), 15); + EXPECT_FALSE(control->rewind(a, 16)); + kv->forget_handle(&handle); +} + +TEST_F(CudaBatchedSequenceCacheTest, DFlashQueriesOnlyEachOwnersBlock) { + constexpr int kBlock = 2; + auto cache_ptr = make(/*initial=*/kCapacity, kBlock); + 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(); + const int32_t b = *control->seq_new(); + // a: 3 new context + block; b: 1 new context + block; padded to 8. + ASSERT_EQ(step(*cache_ptr, handle, {a, a, a, a, a, b, b, b}), Error::Ok); + EXPECT_EQ( + read_cols(fake.bound["seq_cols"].data, 8), + std::vector({0, 0, 0, 0, 0, 1, 1, 1})); + // Queries pack block by block; each sees its context and its block. + EXPECT_EQ( + read_info(fake.bound["seq_info"].data, 2), + Info({{0, kBlock, 0, 5}, {kBlock, kBlock, 1, 3}})); + // The blocks are dropped after the forward. + EXPECT_EQ(control->pos(a), 3); + EXPECT_EQ(control->pos(b), 1); + // An owner split into two runs, or one short of its block, is refused. + EXPECT_FALSE(control->declare_step({a, a, b, b, a})); + EXPECT_FALSE(control->declare_step({a, b, b})); + kv->forget_handle(&handle); +} + +TEST_F(CudaBatchedSequenceCacheTest, BuilderValidatesAndIsRegistered) { + auto no_seqs = sequence_config(); + no_seqs.max_seqs.reset(); + EXPECT_EQ( + cu::make_cuda_batched_sequence_kv_cache(flat_and_ring(kWindow), no_seqs), + nullptr); + auto no_write = sequence_config(); + no_write.max_write.reset(); + EXPECT_EQ( + cu::make_cuda_batched_sequence_kv_cache(flat_and_ring(kWindow), no_write), + nullptr); + auto built = cache::CacheFactory::global().build( + cu::kCudaBackendId, + cache::kind::kBatchedSequence, + flat_and_ring(kWindow), + sequence_config()); + ASSERT_TRUE(built.ok()); + EXPECT_NE(built.get()->as(), nullptr); + EXPECT_NE(built.get()->as(), nullptr); +} diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index 09fab2fc0c6..89c7f5d9530 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -174,6 +174,10 @@ struct ET_EXPERIMENTAL CacheConfig { // Max tokens per step; a ring layer sizes slots to window + max_write - 1. // Unset = each ring layer uses its own window. std::optional max_write; + // Sequences a batch layout must hold at once, for a byte layer whose + // storage width is fixed when the program is compiled. Unset = no fixed + // bound; the capacity they share is the only limit. + std::optional max_seqs; }; ET_EXPERIMENTAL inline bool valid(const CacheGeometry& geometry) { @@ -201,7 +205,8 @@ ET_EXPERIMENTAL inline bool valid(const CacheConfig& cfg) { // initial_capacity may be 0 but not negative, and may exceed capacity -- the // byte layer clamps it. return cfg.capacity > 0 && cfg.initial_capacity >= 0 && - (!cfg.max_write || *cfg.max_write > 0); + (!cfg.max_write || *cfg.max_write > 0) && + (!cfg.max_seqs || *cfg.max_seqs > 0); } ET_EXPERIMENTAL inline bool valid(