@@ -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+
375405TEST_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+
25432674TEST_F (CudaBatchedSequenceCacheTest, RewindRespectsTheRingFloor) {
25442675 auto cache_ptr = make (/* initial=*/ kCapacity );
25452676 auto * control = cache_ptr->as <cache::BatchControl>();
0 commit comments