[executorch][cuda] Grow sequence-layout ring layers on demand

backends/cuda/runtime/cuda_kv_cache.cpp

@@ -82,7 +82,9 @@ std::vector<CudaKVPool::Layer> sequence_pool_layers(
         ring ? static_cast<int64_t>(layer.policy.window) +
                 cfg.max_write.value_or(1) - 1
              : static_cast<int64_t>(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<std::mutex> 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<int64_t>(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<int64_t>(metrics.logical_length, sequence.plan.length());
     }
-    metrics.flat_capacity = pool_.rows() / columns_;
+    metrics.flat_capacity =
+        std::min<int64_t>(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<size_t>(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<int64_t>(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<int64_t>(keep, held);
       for (const bool value : {false, true}) {
         auto* base = static_cast<uint8_t*>(pool_.layer_data(index, value));
         const cudaError_t error = cudaMemcpy2DAsync(

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<cache::Cache> 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<cache::Cache> make_cuda_batched_sequence_kv_cache(
     const cache::CacheGeometry& geometry,
     const cache::CacheConfig& cfg);

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<int64_t>::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<int64_t>(
         max_rows_, std::max<int64_t>(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<size_t>(live_rows);
+    const size_t live_bytes = row_bytes(layers_[index]) *
+        static_cast<size_t>(std::min(live_rows, current.rows));
     if (live_bytes == 0) {
       continue;
     }

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 {

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<int64_t>(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<int64_t>(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<cu::CudaKVCache>();
+  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<cache::BatchControl>();
+  auto* kv = cache_ptr->as<cu::CudaKVCache>();
+  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<uint8_t> rows(pitch * 4, 0);
+  for (int r = 0; r < 4; ++r) {
+    std::fill_n(rows.begin() + r * pitch, slot, static_cast<uint8_t>(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<uint8_t> 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<cache::BatchControl>();
+  auto* kv = cache_ptr->as<cu::CudaKVCache>();
+  auto fake = container();
+  auto handle = make_handle(fake);
+  ASSERT_TRUE(kv->note_handle(&handle).get());
+  constexpr int64_t kRowBytes =
+      kHeads * kDim * static_cast<int64_t>(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<int32_t>(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<uint8_t> ring(slot * kColumns * kRingRows, 0);
+  for (int r = 0; r < kRingRows; ++r) {
+    std::fill_n(
+        ring.begin() + (r * kColumns + a) * slot,
+        slot,
+        static_cast<uint8_t>(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<uint8_t> 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<cache::BatchControl>();