diff --git a/backends/cuda/runtime/cuda_kv_cache.cpp b/backends/cuda/runtime/cuda_kv_cache.cpp index d0bc284c001..98bc447449c 100644 --- a/backends/cuda/runtime/cuda_kv_cache.cpp +++ b/backends/cuda/runtime/cuda_kv_cache.cpp @@ -82,7 +82,9 @@ std::vector sequence_pool_layers( ring ? static_cast(layer.policy.window) + cfg.max_write.value_or(1) - 1 : static_cast(cfg.capacity), - /*growable=*/!ring}); + // A ring cannot wrap before it holds all its slots, so until then it + // grows as a flat layer does, and a growth copies a plain prefix. + /*growable=*/true}); } return layers; } @@ -195,7 +197,9 @@ class CudaSequenceKVCache final : public cache::SequenceCache, std::lock_guard guard(mutex_); OffGraphKVMetrics metrics; metrics.logical_length = length(); - metrics.flat_capacity = pool_.rows(); + // The pool's shared rows can exceed a small flat layer when a ring + // ceiling dominates, so clamp to what the flat layers hold. + metrics.flat_capacity = std::min(pool_.rows(), capacity()); metrics.growth_count = pool_.growth_count(); metrics.allocated_bytes = pool_.allocated_bytes(); return metrics; @@ -1143,7 +1147,8 @@ class CudaBatchedSequenceCache final : public cache::Cache, metrics.logical_length = std::max(metrics.logical_length, sequence.plan.length()); } - metrics.flat_capacity = pool_.rows() / columns_; + metrics.flat_capacity = + std::min(pool_.rows() / columns_, cfg_.capacity); metrics.growth_count = pool_.growth_count(); metrics.allocated_bytes = pool_.allocated_bytes(); return metrics; @@ -1181,7 +1186,9 @@ class CudaBatchedSequenceCache final : public cache::Cache, layer.n_kv_heads, layer.head_dim, rows_of(layer, cfg) * (*cfg.max_seqs + 1), - /*growable=*/!is_ring(layer)}); + // A ring cannot wrap before it holds all its rows, so until then it + // grows with the flat layers. + /*growable=*/true}); } return layers; } @@ -1206,9 +1213,10 @@ class CudaBatchedSequenceCache final : public cache::Cache, 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_); + const int64_t held = + std::min(layer_rows_[index], pool_.rows() / columns_); + const int64_t rows = + is_ring(layer) ? held : std::min(keep, held); for (const bool value : {false, true}) { auto* base = static_cast(pool_.layer_data(index, value)); const cudaError_t error = cudaMemcpy2DAsync( diff --git a/backends/cuda/runtime/cuda_kv_cache.h b/backends/cuda/runtime/cuda_kv_cache.h index 17f3aac4672..3869f9dc5f1 100644 --- a/backends/cuda/runtime/cuda_kv_cache.h +++ b/backends/cuda/runtime/cuda_kv_cache.h @@ -198,8 +198,8 @@ runtime::Error attach_offgraph_kv_cache( // invalid, so CacheFactory reports the failure instead of the cache asserting // during construction. // -// Flat layers start at cfg.initial_capacity rows and grow geometrically up to -// cfg.capacity; ring layers hold window + max_write - 1 rows from the start. +// Every layer starts at cfg.initial_capacity rows and grows geometrically: a +// flat layer up to cfg.capacity, a ring layer up to window + max_write - 1. // cfg.kv_dtype is an ExecuTorch ScalarType; the slim enum the storage uses // shares its numbering, and unsupported values are rejected here. std::shared_ptr make_cuda_sequence_kv_cache( @@ -230,8 +230,9 @@ inline int batched_sequence_max_segments(int max_seqs) { // 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). +// Every layer starts at cfg.initial_capacity rows and grows geometrically: a +// flat layer to cfg.capacity, a ring layer to min(capacity, window + +// max_write - 1). std::shared_ptr make_cuda_batched_sequence_kv_cache( const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg); diff --git a/backends/cuda/runtime/cuda_kv_pool.cpp b/backends/cuda/runtime/cuda_kv_pool.cpp index de7709c9e7e..74fcf6a2be8 100644 --- a/backends/cuda/runtime/cuda_kv_pool.cpp +++ b/backends/cuda/runtime/cuda_kv_pool.cpp @@ -61,10 +61,9 @@ CudaKVPool::CudaKVPool( side_specs_(std::move(side_buffers)), storage_dtype_(storage_dtype), initial_rows_(initial_rows) { - max_rows_ = std::numeric_limits::max(); for (const Layer& layer : layers_) { if (layer.growable) { - max_rows_ = std::min(max_rows_, layer.declared_rows); + max_rows_ = std::max(max_rows_, layer.declared_rows); } } } @@ -154,7 +153,7 @@ Error CudaKVPool::prepare( if (!allocated_) { ET_CHECK_OK_OR_RETURN_ERROR(allocate_initial(required_rows, stream)); } - if (required_rows > rows_) { + if (required_rows > rows_ && rows_ < max_rows_) { const int64_t next = std::min( max_rows_, std::max(required_rows, rows_ * 2)); ET_CHECK_OK_OR_RETURN_ERROR(grow(next, live_rows, stream)); @@ -295,7 +294,9 @@ Error CudaKVPool::allocate_initial(int64_t required_rows, cudaStream_t stream) { } }; for (const Layer& layer : layers_) { - const int64_t rows = layer.growable ? growable_rows : layer.declared_rows; + const int64_t rows = layer.growable + ? std::min(growable_rows, layer.declared_rows) + : layer.declared_rows; Allocation allocation; const Error error = allocate_layer(layer, rows, stream, allocation); if (error != Error::Ok) { @@ -376,11 +377,12 @@ Error CudaKVPool::grow( }; for (size_t index = 0; index < layers_.size(); ++index) { const Layer& layer = layers_[index]; - if (!layer.growable) { + const int64_t rows = std::min(new_rows, layer.declared_rows); + if (!layer.growable || rows <= allocations_.at(index).rows) { continue; } Allocation replacement; - const Error error = allocate_layer(layer, new_rows, stream, replacement); + const Error error = allocate_layer(layer, rows, stream, replacement); if (error != Error::Ok) { roll_back(); return error; @@ -389,8 +391,8 @@ Error CudaKVPool::grow( } for (const auto& [index, replacement] : replacements) { const Allocation& current = allocations_.at(index); - const size_t live_bytes = - row_bytes(layers_[index]) * static_cast(live_rows); + const size_t live_bytes = row_bytes(layers_[index]) * + static_cast(std::min(live_rows, current.rows)); if (live_bytes == 0) { continue; } diff --git a/backends/cuda/runtime/cuda_kv_pool.h b/backends/cuda/runtime/cuda_kv_pool.h index 02dff799b13..c78f85431bd 100644 --- a/backends/cuda/runtime/cuda_kv_pool.h +++ b/backends/cuda/runtime/cuda_kv_pool.h @@ -29,9 +29,10 @@ namespace executorch::backends::cuda { // the AOTI bindings that point every compiled program at them. Which rows mean // what is the owning cache's business, never this class's. // -// Growable layers share one row count that grows geometrically toward their -// declared rows, moving their storage; fixed layers and side buffers are -// allocated once and never move, so a program may capture their addresses. +// Growable layers share one row count that grows geometrically, each layer +// holding that many rows up to its declared rows, moving their storage; fixed +// layers and side buffers are allocated once and never move, so a program may +// capture their addresses. // // Not thread-safe: the owning cache serializes every call. class CudaKVPool final { diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index bbd4b4b684d..8d753579de0 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -302,7 +302,7 @@ class CudaKVCacheTest : public ::testing::Test { } // namespace -TEST_F(CudaKVCacheTest, FlatLayersGrowGeometricallyAndKeepHistory) { +TEST_F(CudaKVCacheTest, LayersGrowGeometricallyAndKeepHistory) { constexpr int kWindow = 4; constexpr int kMaxWrite = 9; auto cache_ptr = make_cache(flat_and_ring(kWindow), config(32, 4, kMaxWrite)); @@ -314,45 +314,55 @@ TEST_F(CudaKVCacheTest, FlatLayersGrowGeometricallyAndKeepHistory) { // Nothing is allocated until the first step asks for it. EXPECT_EQ(kv.metrics().allocated_bytes, 0); + // A ring starts as small as a flat layer: it cannot wrap before it is full. ASSERT_EQ(kv.prepare_step(3, cudaStreamPerThread), Error::Ok); ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok); constexpr int kRingRows = kWindow + kMaxWrite - 1; + constexpr int64_t kRowBytes = kRow * static_cast(sizeof(uint16_t)); EXPECT_EQ(kv.metrics().flat_capacity, 4); EXPECT_EQ(kv.metrics().growth_count, 0); - EXPECT_EQ( - kv.metrics().allocated_bytes, - 2 * (4 + kRingRows) * kRow * static_cast(sizeof(uint16_t))); + EXPECT_EQ(kv.metrics().allocated_bytes, 2 * (4 + 4) * kRowBytes); - // Stand in for the kernels: the step writes rows [0, 3). + // Stand in for the kernels: the step writes rows [0, 3) of both layers. const auto history = iota_rows(3); - ASSERT_EQ( - cudaMemcpy( - container.bound["flat_k"].data, - history.data(), - history.size() * sizeof(uint16_t), - cudaMemcpyHostToDevice), - cudaSuccess); + for (const char* name : {"flat_k", "ring_k"}) { + ASSERT_EQ( + cudaMemcpy( + container.bound[name].data, + history.data(), + history.size() * sizeof(uint16_t), + cudaMemcpyHostToDevice), + cudaSuccess); + } void* first_flat = container.bound["flat_k"].data; void* first_ring = container.bound["ring_k"].data; ASSERT_EQ(kv.commit_step(3), Error::Ok); - // 3 + 2 > 4: doubles to 8, moves the flat storage, keeps what was written. + // 3 + 2 > 4: doubles to 8, moves both layers, keeps what was written. ASSERT_EQ(kv.prepare_step(2, cudaStreamPerThread), Error::Ok); ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok); EXPECT_EQ(kv.metrics().flat_capacity, 8); EXPECT_EQ(kv.metrics().growth_count, 1); EXPECT_NE(container.bound["flat_k"].data, first_flat); - EXPECT_EQ(container.bound["ring_k"].data, first_ring); + EXPECT_NE(container.bound["ring_k"].data, first_ring); EXPECT_EQ(read_rows(container.bound["flat_k"].data, 3), history); + EXPECT_EQ(read_rows(container.bound["ring_k"].data, 3), history); ASSERT_EQ(kv.commit_step(2), Error::Ok); - // A step wider than doubling asks for: grows straight to what it needs. + // A step wider than doubling asks for: grows straight to what it needs. The + // ring stops at its slots and no longer moves. ASSERT_EQ(kv.prepare_step(9, cudaStreamPerThread), Error::Ok); + ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok); EXPECT_EQ(kv.metrics().flat_capacity, 16); + EXPECT_EQ(kv.metrics().allocated_bytes, 2 * (16 + kRingRows) * kRowBytes); + void* full_ring = container.bound["ring_k"].data; ASSERT_EQ(kv.commit_step(9), Error::Ok); ASSERT_EQ(kv.prepare_step(9, cudaStreamPerThread), Error::Ok); + ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok); EXPECT_EQ(kv.metrics().flat_capacity, 32); EXPECT_EQ(kv.metrics().growth_count, 3); + EXPECT_EQ(kv.metrics().allocated_bytes, 2 * (32 + kRingRows) * kRowBytes); + EXPECT_EQ(container.bound["ring_k"].data, full_ring); ASSERT_EQ(kv.commit_step(9), Error::Ok); // Filling to exactly the declared capacity needs no growth; one more token @@ -372,6 +382,26 @@ TEST_F(CudaKVCacheTest, FlatLayersGrowGeometricallyAndKeepHistory) { kv.forget_handle(&handle); } +TEST_F(CudaKVCacheTest, MetricsClampFlatCapacityAtTheFlatLayers) { + // The ring's ceiling (4 + 8 - 1 = 11) dominates the flat capacity, so + // doubling from 6 rows lands the shared count at 11 while the flat layer + // caps at 8. metrics must report what the flat layers hold. + auto cache_ptr = make_cache(flat_and_ring(4), config(8, 3, 8)); + ASSERT_NE(cache_ptr, nullptr); + auto& kv = *cache_ptr->as(); + auto container = flat_and_ring_container(); + auto handle = make_handle(container); + ASSERT_TRUE(kv.note_handle(&handle).get()); + for (int token = 0; token < 7; ++token) { + ASSERT_EQ(kv.prepare_step(1, cudaStreamPerThread), Error::Ok); + ASSERT_EQ(kv.commit_step(1), Error::Ok); + } + EXPECT_EQ(kv.metrics().logical_length, 7); + EXPECT_EQ(kv.metrics().flat_capacity, 8); + EXPECT_EQ(kv.metrics().growth_count, 2); + kv.forget_handle(&handle); +} + TEST_F(CudaKVCacheTest, BindsTheDeclaredShapeOverTheCurrentAllocation) { constexpr int kWindow = 4; constexpr int kMaxWrite = 4; @@ -2540,6 +2570,107 @@ TEST_F( kv->forget_handle(&handle); } +TEST_F(CudaBatchedSequenceCacheTest, CloneOfAnUnwrappedRingCopiesThePrefix) { + auto cache_ptr = make(/*initial=*/4); + 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); + + // The ring holds kRingRows rows but the pool only 4 per column: the clone + // copies the 4-row prefix, not the whole ring. + const size_t slot = kHeads * kDim * sizeof(uint16_t); + const size_t pitch = slot * kColumns; + std::vector rows(pitch * 4, 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["r_k"].data, + rows.data(), + rows.size(), + cudaMemcpyHostToDevice), + cudaSuccess); + const auto c = control->seq_clone(a, std::nullopt); + ASSERT_TRUE(c.has_value()); + EXPECT_EQ(control->pos(*c), 4); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + std::vector after(rows.size()); + ASSERT_EQ( + cudaMemcpy( + after.data(), + fake.bound["r_k"].data, + after.size(), + cudaMemcpyDeviceToHost), + cudaSuccess); + for (int r = 0; r < 4; ++r) { + EXPECT_EQ(after[r * pitch + *c * slot], r + 1) << r; + } + kv->forget_handle(&handle); +} + +TEST_F(CudaBatchedSequenceCacheTest, RingLayersGrowUntilTheyHoldTheirSlots) { + auto cache_ptr = make(/*initial=*/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()); + constexpr int64_t kRowBytes = + kHeads * kDim * static_cast(sizeof(uint16_t)) * kColumns; + const int32_t a = *control->seq_new(); + + // Two rows per column for both layers, not the ring's kRingRows up front. + ASSERT_EQ(step(*cache_ptr, handle, {a, a}), Error::Ok); + EXPECT_EQ(kv->metrics().allocated_bytes, 2 * (2 + 2) * kRowBytes); + + // Past kRingRows the ring stops growing while the flat layer keeps going. + ASSERT_EQ( + step(*cache_ptr, handle, std::vector(kMaxWrite, a)), Error::Ok); + constexpr int kLength = 2 + kMaxWrite; + static_assert(kLength > kRingRows); + EXPECT_EQ(control->pos(a), kLength); + EXPECT_EQ(kv->metrics().flat_capacity, kLength); + EXPECT_EQ( + kv->metrics().allocated_bytes, 2 * (kLength + kRingRows) * kRowBytes); + + // A wrapped ring is cloned whole. + const size_t slot = kHeads * kDim * sizeof(uint16_t); + std::vector ring(slot * kColumns * kRingRows, 0); + for (int r = 0; r < kRingRows; ++r) { + std::fill_n( + ring.begin() + (r * kColumns + a) * slot, + slot, + static_cast(r + 1)); + } + ASSERT_EQ( + cudaMemcpy( + fake.bound["r_k"].data, + ring.data(), + ring.size(), + cudaMemcpyHostToDevice), + cudaSuccess); + const auto c = control->seq_clone(a, std::nullopt); + ASSERT_TRUE(c.has_value()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + std::vector after(ring.size()); + ASSERT_EQ( + cudaMemcpy( + after.data(), + fake.bound["r_k"].data, + after.size(), + cudaMemcpyDeviceToHost), + cudaSuccess); + for (int r = 0; r < kRingRows; ++r) { + EXPECT_EQ(after[(r * kColumns + *c) * slot], r + 1) << r; + } + kv->forget_handle(&handle); +} + TEST_F(CudaBatchedSequenceCacheTest, RewindRespectsTheRingFloor) { auto cache_ptr = make(/*initial=*/kCapacity); auto* control = cache_ptr->as();