Skip to content

Commit a1da392

Browse files
committed
[executorch][cuda] Grow sequence-layout ring layers on demand
Pull Request resolved: #24076 A ring layer cannot wrap before it holds all its slots, so until then it grows with the flat layers instead of allocating window + max_write - 1 rows per sequence up front. Growable layers now cap at their own declared rows. ghstack-source-id: 468443991 @exported-using-ghexport Differential Revision: [D126851225](https://our.internmc.facebook.com/intern/diff/D126851225/)
1 parent 5309e75 commit a1da392

5 files changed

Lines changed: 180 additions & 37 deletions

File tree

‎backends/cuda/runtime/cuda_kv_cache.cpp‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,9 @@ std::vector<CudaKVPool::Layer> sequence_pool_layers(
8282
ring ? static_cast<int64_t>(layer.policy.window) +
8383
cfg.max_write.value_or(1) - 1
8484
: static_cast<int64_t>(cfg.capacity),
85-
/*growable=*/!ring});
85+
// A ring cannot wrap before it holds all its slots, so until then it
86+
// grows as a flat layer does, and a growth copies a plain prefix.
87+
/*growable=*/true});
8688
}
8789
return layers;
8890
}
@@ -195,7 +197,9 @@ class CudaSequenceKVCache final : public cache::SequenceCache,
195197
std::lock_guard<std::mutex> guard(mutex_);
196198
OffGraphKVMetrics metrics;
197199
metrics.logical_length = length();
198-
metrics.flat_capacity = pool_.rows();
200+
// The pool's shared rows can exceed a small flat layer when a ring
201+
// ceiling dominates, so clamp to what the flat layers hold.
202+
metrics.flat_capacity = std::min<int64_t>(pool_.rows(), capacity());
199203
metrics.growth_count = pool_.growth_count();
200204
metrics.allocated_bytes = pool_.allocated_bytes();
201205
return metrics;
@@ -1143,7 +1147,8 @@ class CudaBatchedSequenceCache final : public cache::Cache,
11431147
metrics.logical_length =
11441148
std::max<int64_t>(metrics.logical_length, sequence.plan.length());
11451149
}
1146-
metrics.flat_capacity = pool_.rows() / columns_;
1150+
metrics.flat_capacity =
1151+
std::min<int64_t>(pool_.rows() / columns_, cfg_.capacity);
11471152
metrics.growth_count = pool_.growth_count();
11481153
metrics.allocated_bytes = pool_.allocated_bytes();
11491154
return metrics;
@@ -1181,7 +1186,9 @@ class CudaBatchedSequenceCache final : public cache::Cache,
11811186
layer.n_kv_heads,
11821187
layer.head_dim,
11831188
rows_of(layer, cfg) * (*cfg.max_seqs + 1),
1184-
/*growable=*/!is_ring(layer)});
1189+
// A ring cannot wrap before it holds all its rows, so until then it
1190+
// grows with the flat layers.
1191+
/*growable=*/true});
11851192
}
11861193
return layers;
11871194
}
@@ -1206,9 +1213,10 @@ class CudaBatchedSequenceCache final : public cache::Cache,
12061213
const size_t slot = static_cast<size_t>(layer.n_kv_heads) *
12071214
layer.head_dim * slimc10::elementSize(storage_dtype());
12081215
const size_t pitch = slot * columns_;
1209-
const int64_t rows = is_ring(layer)
1210-
? layer_rows_[index]
1211-
: std::min<int64_t>(keep, pool_.rows() / columns_);
1216+
const int64_t held =
1217+
std::min(layer_rows_[index], pool_.rows() / columns_);
1218+
const int64_t rows =
1219+
is_ring(layer) ? held : std::min<int64_t>(keep, held);
12121220
for (const bool value : {false, true}) {
12131221
auto* base = static_cast<uint8_t*>(pool_.layer_data(index, value));
12141222
const cudaError_t error = cudaMemcpy2DAsync(

‎backends/cuda/runtime/cuda_kv_cache.h‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -198,8 +198,8 @@ runtime::Error attach_offgraph_kv_cache(
198198
// invalid, so CacheFactory reports the failure instead of the cache asserting
199199
// during construction.
200200
//
201-
// Flat layers start at cfg.initial_capacity rows and grow geometrically up to
202-
// cfg.capacity; ring layers hold window + max_write - 1 rows from the start.
201+
// Every layer starts at cfg.initial_capacity rows and grows geometrically: a
202+
// flat layer up to cfg.capacity, a ring layer up to window + max_write - 1.
203203
// cfg.kv_dtype is an ExecuTorch ScalarType; the slim enum the storage uses
204204
// shares its numbering, and unsupported values are rejected here.
205205
std::shared_ptr<cache::Cache> make_cuda_sequence_kv_cache(
@@ -230,8 +230,9 @@ inline int batched_sequence_max_segments(int max_seqs) {
230230
// cfg.capacity is one sequence's capacity (the program's maximum_capacity),
231231
// cfg.max_seqs its sequence columns and cfg.max_write its widest step; all
232232
// fix shapes the program declared, so they must match the export exactly.
233-
// Flat layers start at cfg.initial_capacity rows and grow geometrically to
234-
// cfg.capacity; ring layers hold min(capacity, window + max_write - 1).
233+
// Every layer starts at cfg.initial_capacity rows and grows geometrically: a
234+
// flat layer to cfg.capacity, a ring layer to min(capacity, window +
235+
// max_write - 1).
235236
std::shared_ptr<cache::Cache> make_cuda_batched_sequence_kv_cache(
236237
const cache::CacheGeometry& geometry,
237238
const cache::CacheConfig& cfg);

‎backends/cuda/runtime/cuda_kv_pool.cpp‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,9 @@ CudaKVPool::CudaKVPool(
6161
side_specs_(std::move(side_buffers)),
6262
storage_dtype_(storage_dtype),
6363
initial_rows_(initial_rows) {
64-
max_rows_ = std::numeric_limits<int64_t>::max();
6564
for (const Layer& layer : layers_) {
6665
if (layer.growable) {
67-
max_rows_ = std::min(max_rows_, layer.declared_rows);
66+
max_rows_ = std::max(max_rows_, layer.declared_rows);
6867
}
6968
}
7069
}
@@ -154,7 +153,7 @@ Error CudaKVPool::prepare(
154153
if (!allocated_) {
155154
ET_CHECK_OK_OR_RETURN_ERROR(allocate_initial(required_rows, stream));
156155
}
157-
if (required_rows > rows_) {
156+
if (required_rows > rows_ && rows_ < max_rows_) {
158157
const int64_t next = std::min<int64_t>(
159158
max_rows_, std::max<int64_t>(required_rows, rows_ * 2));
160159
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) {
295294
}
296295
};
297296
for (const Layer& layer : layers_) {
298-
const int64_t rows = layer.growable ? growable_rows : layer.declared_rows;
297+
const int64_t rows = layer.growable
298+
? std::min(growable_rows, layer.declared_rows)
299+
: layer.declared_rows;
299300
Allocation allocation;
300301
const Error error = allocate_layer(layer, rows, stream, allocation);
301302
if (error != Error::Ok) {
@@ -376,11 +377,12 @@ Error CudaKVPool::grow(
376377
};
377378
for (size_t index = 0; index < layers_.size(); ++index) {
378379
const Layer& layer = layers_[index];
379-
if (!layer.growable) {
380+
const int64_t rows = std::min(new_rows, layer.declared_rows);
381+
if (!layer.growable || rows <= allocations_.at(index).rows) {
380382
continue;
381383
}
382384
Allocation replacement;
383-
const Error error = allocate_layer(layer, new_rows, stream, replacement);
385+
const Error error = allocate_layer(layer, rows, stream, replacement);
384386
if (error != Error::Ok) {
385387
roll_back();
386388
return error;
@@ -389,8 +391,8 @@ Error CudaKVPool::grow(
389391
}
390392
for (const auto& [index, replacement] : replacements) {
391393
const Allocation& current = allocations_.at(index);
392-
const size_t live_bytes =
393-
row_bytes(layers_[index]) * static_cast<size_t>(live_rows);
394+
const size_t live_bytes = row_bytes(layers_[index]) *
395+
static_cast<size_t>(std::min(live_rows, current.rows));
394396
if (live_bytes == 0) {
395397
continue;
396398
}

‎backends/cuda/runtime/cuda_kv_pool.h‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,10 @@ namespace executorch::backends::cuda {
2929
// the AOTI bindings that point every compiled program at them. Which rows mean
3030
// what is the owning cache's business, never this class's.
3131
//
32-
// Growable layers share one row count that grows geometrically toward their
33-
// declared rows, moving their storage; fixed layers and side buffers are
34-
// allocated once and never move, so a program may capture their addresses.
32+
// Growable layers share one row count that grows geometrically, each layer
33+
// holding that many rows up to its declared rows, moving their storage; fixed
34+
// layers and side buffers are allocated once and never move, so a program may
35+
// capture their addresses.
3536
//
3637
// Not thread-safe: the owning cache serializes every call.
3738
class CudaKVPool final {

‎backends/cuda/runtime/test/test_cuda_kv_cache.cpp‎

Lines changed: 146 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,7 @@ class CudaKVCacheTest : public ::testing::Test {
302302

303303
} // namespace
304304

305-
TEST_F(CudaKVCacheTest, FlatLayersGrowGeometricallyAndKeepHistory) {
305+
TEST_F(CudaKVCacheTest, LayersGrowGeometricallyAndKeepHistory) {
306306
constexpr int kWindow = 4;
307307
constexpr int kMaxWrite = 9;
308308
auto cache_ptr = make_cache(flat_and_ring(kWindow), config(32, 4, kMaxWrite));
@@ -314,45 +314,55 @@ TEST_F(CudaKVCacheTest, FlatLayersGrowGeometricallyAndKeepHistory) {
314314
// Nothing is allocated until the first step asks for it.
315315
EXPECT_EQ(kv.metrics().allocated_bytes, 0);
316316

317+
// A ring starts as small as a flat layer: it cannot wrap before it is full.
317318
ASSERT_EQ(kv.prepare_step(3, cudaStreamPerThread), Error::Ok);
318319
ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok);
319320
constexpr int kRingRows = kWindow + kMaxWrite - 1;
321+
constexpr int64_t kRowBytes = kRow * static_cast<int64_t>(sizeof(uint16_t));
320322
EXPECT_EQ(kv.metrics().flat_capacity, 4);
321323
EXPECT_EQ(kv.metrics().growth_count, 0);
322-
EXPECT_EQ(
323-
kv.metrics().allocated_bytes,
324-
2 * (4 + kRingRows) * kRow * static_cast<int64_t>(sizeof(uint16_t)));
324+
EXPECT_EQ(kv.metrics().allocated_bytes, 2 * (4 + 4) * kRowBytes);
325325

326-
// Stand in for the kernels: the step writes rows [0, 3).
326+
// Stand in for the kernels: the step writes rows [0, 3) of both layers.
327327
const auto history = iota_rows(3);
328-
ASSERT_EQ(
329-
cudaMemcpy(
330-
container.bound["flat_k"].data,
331-
history.data(),
332-
history.size() * sizeof(uint16_t),
333-
cudaMemcpyHostToDevice),
334-
cudaSuccess);
328+
for (const char* name : {"flat_k", "ring_k"}) {
329+
ASSERT_EQ(
330+
cudaMemcpy(
331+
container.bound[name].data,
332+
history.data(),
333+
history.size() * sizeof(uint16_t),
334+
cudaMemcpyHostToDevice),
335+
cudaSuccess);
336+
}
335337
void* first_flat = container.bound["flat_k"].data;
336338
void* first_ring = container.bound["ring_k"].data;
337339
ASSERT_EQ(kv.commit_step(3), Error::Ok);
338340

339-
// 3 + 2 > 4: doubles to 8, moves the flat storage, keeps what was written.
341+
// 3 + 2 > 4: doubles to 8, moves both layers, keeps what was written.
340342
ASSERT_EQ(kv.prepare_step(2, cudaStreamPerThread), Error::Ok);
341343
ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok);
342344
EXPECT_EQ(kv.metrics().flat_capacity, 8);
343345
EXPECT_EQ(kv.metrics().growth_count, 1);
344346
EXPECT_NE(container.bound["flat_k"].data, first_flat);
345-
EXPECT_EQ(container.bound["ring_k"].data, first_ring);
347+
EXPECT_NE(container.bound["ring_k"].data, first_ring);
346348
EXPECT_EQ(read_rows(container.bound["flat_k"].data, 3), history);
349+
EXPECT_EQ(read_rows(container.bound["ring_k"].data, 3), history);
347350
ASSERT_EQ(kv.commit_step(2), Error::Ok);
348351

349-
// A step wider than doubling asks for: grows straight to what it needs.
352+
// A step wider than doubling asks for: grows straight to what it needs. The
353+
// ring stops at its slots and no longer moves.
350354
ASSERT_EQ(kv.prepare_step(9, cudaStreamPerThread), Error::Ok);
355+
ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok);
351356
EXPECT_EQ(kv.metrics().flat_capacity, 16);
357+
EXPECT_EQ(kv.metrics().allocated_bytes, 2 * (16 + kRingRows) * kRowBytes);
358+
void* full_ring = container.bound["ring_k"].data;
352359
ASSERT_EQ(kv.commit_step(9), Error::Ok);
353360
ASSERT_EQ(kv.prepare_step(9, cudaStreamPerThread), Error::Ok);
361+
ASSERT_EQ(kv.rebind_for_execute(&handle), Error::Ok);
354362
EXPECT_EQ(kv.metrics().flat_capacity, 32);
355363
EXPECT_EQ(kv.metrics().growth_count, 3);
364+
EXPECT_EQ(kv.metrics().allocated_bytes, 2 * (32 + kRingRows) * kRowBytes);
365+
EXPECT_EQ(container.bound["ring_k"].data, full_ring);
356366
ASSERT_EQ(kv.commit_step(9), Error::Ok);
357367

358368
// Filling to exactly the declared capacity needs no growth; one more token
@@ -372,6 +382,26 @@ TEST_F(CudaKVCacheTest, FlatLayersGrowGeometricallyAndKeepHistory) {
372382
kv.forget_handle(&handle);
373383
}
374384

385+
TEST_F(CudaKVCacheTest, MetricsClampFlatCapacityAtTheFlatLayers) {
386+
// The ring's ceiling (4 + 8 - 1 = 11) dominates the flat capacity, so
387+
// doubling from 6 rows lands the shared count at 11 while the flat layer
388+
// caps at 8. metrics must report what the flat layers hold.
389+
auto cache_ptr = make_cache(flat_and_ring(4), config(8, 3, 8));
390+
ASSERT_NE(cache_ptr, nullptr);
391+
auto& kv = *cache_ptr->as<cu::CudaKVCache>();
392+
auto container = flat_and_ring_container();
393+
auto handle = make_handle(container);
394+
ASSERT_TRUE(kv.note_handle(&handle).get());
395+
for (int token = 0; token < 7; ++token) {
396+
ASSERT_EQ(kv.prepare_step(1, cudaStreamPerThread), Error::Ok);
397+
ASSERT_EQ(kv.commit_step(1), Error::Ok);
398+
}
399+
EXPECT_EQ(kv.metrics().logical_length, 7);
400+
EXPECT_EQ(kv.metrics().flat_capacity, 8);
401+
EXPECT_EQ(kv.metrics().growth_count, 2);
402+
kv.forget_handle(&handle);
403+
}
404+
375405
TEST_F(CudaKVCacheTest, BindsTheDeclaredShapeOverTheCurrentAllocation) {
376406
constexpr int kWindow = 4;
377407
constexpr int kMaxWrite = 4;
@@ -2540,6 +2570,107 @@ TEST_F(
25402570
kv->forget_handle(&handle);
25412571
}
25422572

2573+
TEST_F(CudaBatchedSequenceCacheTest, CloneOfAnUnwrappedRingCopiesThePrefix) {
2574+
auto cache_ptr = make(/*initial=*/4);
2575+
auto* control = cache_ptr->as<cache::BatchControl>();
2576+
auto* kv = cache_ptr->as<cu::CudaKVCache>();
2577+
auto fake = container();
2578+
auto handle = make_handle(fake);
2579+
ASSERT_TRUE(kv->note_handle(&handle).get());
2580+
const int32_t a = *control->seq_new();
2581+
ASSERT_EQ(step(*cache_ptr, handle, {a, a, a, a}), Error::Ok);
2582+
2583+
// The ring holds kRingRows rows but the pool only 4 per column: the clone
2584+
// copies the 4-row prefix, not the whole ring.
2585+
const size_t slot = kHeads * kDim * sizeof(uint16_t);
2586+
const size_t pitch = slot * kColumns;
2587+
std::vector<uint8_t> rows(pitch * 4, 0);
2588+
for (int r = 0; r < 4; ++r) {
2589+
std::fill_n(rows.begin() + r * pitch, slot, static_cast<uint8_t>(r + 1));
2590+
}
2591+
ASSERT_EQ(
2592+
cudaMemcpy(
2593+
fake.bound["r_k"].data,
2594+
rows.data(),
2595+
rows.size(),
2596+
cudaMemcpyHostToDevice),
2597+
cudaSuccess);
2598+
const auto c = control->seq_clone(a, std::nullopt);
2599+
ASSERT_TRUE(c.has_value());
2600+
EXPECT_EQ(control->pos(*c), 4);
2601+
ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
2602+
std::vector<uint8_t> after(rows.size());
2603+
ASSERT_EQ(
2604+
cudaMemcpy(
2605+
after.data(),
2606+
fake.bound["r_k"].data,
2607+
after.size(),
2608+
cudaMemcpyDeviceToHost),
2609+
cudaSuccess);
2610+
for (int r = 0; r < 4; ++r) {
2611+
EXPECT_EQ(after[r * pitch + *c * slot], r + 1) << r;
2612+
}
2613+
kv->forget_handle(&handle);
2614+
}
2615+
2616+
TEST_F(CudaBatchedSequenceCacheTest, RingLayersGrowUntilTheyHoldTheirSlots) {
2617+
auto cache_ptr = make(/*initial=*/2);
2618+
auto* control = cache_ptr->as<cache::BatchControl>();
2619+
auto* kv = cache_ptr->as<cu::CudaKVCache>();
2620+
auto fake = container();
2621+
auto handle = make_handle(fake);
2622+
ASSERT_TRUE(kv->note_handle(&handle).get());
2623+
constexpr int64_t kRowBytes =
2624+
kHeads * kDim * static_cast<int64_t>(sizeof(uint16_t)) * kColumns;
2625+
const int32_t a = *control->seq_new();
2626+
2627+
// Two rows per column for both layers, not the ring's kRingRows up front.
2628+
ASSERT_EQ(step(*cache_ptr, handle, {a, a}), Error::Ok);
2629+
EXPECT_EQ(kv->metrics().allocated_bytes, 2 * (2 + 2) * kRowBytes);
2630+
2631+
// Past kRingRows the ring stops growing while the flat layer keeps going.
2632+
ASSERT_EQ(
2633+
step(*cache_ptr, handle, std::vector<int32_t>(kMaxWrite, a)), Error::Ok);
2634+
constexpr int kLength = 2 + kMaxWrite;
2635+
static_assert(kLength > kRingRows);
2636+
EXPECT_EQ(control->pos(a), kLength);
2637+
EXPECT_EQ(kv->metrics().flat_capacity, kLength);
2638+
EXPECT_EQ(
2639+
kv->metrics().allocated_bytes, 2 * (kLength + kRingRows) * kRowBytes);
2640+
2641+
// A wrapped ring is cloned whole.
2642+
const size_t slot = kHeads * kDim * sizeof(uint16_t);
2643+
std::vector<uint8_t> ring(slot * kColumns * kRingRows, 0);
2644+
for (int r = 0; r < kRingRows; ++r) {
2645+
std::fill_n(
2646+
ring.begin() + (r * kColumns + a) * slot,
2647+
slot,
2648+
static_cast<uint8_t>(r + 1));
2649+
}
2650+
ASSERT_EQ(
2651+
cudaMemcpy(
2652+
fake.bound["r_k"].data,
2653+
ring.data(),
2654+
ring.size(),
2655+
cudaMemcpyHostToDevice),
2656+
cudaSuccess);
2657+
const auto c = control->seq_clone(a, std::nullopt);
2658+
ASSERT_TRUE(c.has_value());
2659+
ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess);
2660+
std::vector<uint8_t> after(ring.size());
2661+
ASSERT_EQ(
2662+
cudaMemcpy(
2663+
after.data(),
2664+
fake.bound["r_k"].data,
2665+
after.size(),
2666+
cudaMemcpyDeviceToHost),
2667+
cudaSuccess);
2668+
for (int r = 0; r < kRingRows; ++r) {
2669+
EXPECT_EQ(after[(r * kColumns + *c) * slot], r + 1) << r;
2670+
}
2671+
kv->forget_handle(&handle);
2672+
}
2673+
25432674
TEST_F(CudaBatchedSequenceCacheTest, RewindRespectsTheRingFloor) {
25442675
auto cache_ptr = make(/*initial=*/kCapacity);
25452676
auto* control = cache_ptr->as<cache::BatchControl>();

0 commit comments

Comments
 (0)