fix(kv-cache): grow sequences by block delta, not append — prevents block leak + device block-table overrun (RTX 5090 verified)

runtime/src/kv_cache.cpp

@@ -63,15 +63,26 @@ KVCacheManager::~KVCacheManager() {
 
 bool KVCacheManager::allocate(uint64_t seq_id, int num_tokens) {
     const int need = (num_tokens + impl_->cfg.block_size - 1) / impl_->cfg.block_size;
-    if ((int)impl_->free_list.size() < need || impl_->free_slots.empty()) return false;
     if (need > kMaxBlocksPerSeq) return false;
 
+    // `need` is the TOTAL blocks the sequence must own for num_tokens. Grow only by the
+    // delta over what it already holds -- the seq_slot reuse branch below is designed to
+    // re-allocate an existing seq_id, and appending `need` blocks unconditionally would
+    // (a) leak the blocks it already held and (b) push seq_blocks past kMaxBlocksPerSeq,
+    // making the cudaMemcpy of blocks.data() overrun the fixed-width device table row into
+    // the next sequence. A fresh seq_id has have==0, so this is identical to the old path.
     auto& blocks = impl_->seq_blocks[seq_id];
-    for (int i = 0; i < need; i++) { blocks.push_back(impl_->free_list.back()); impl_->free_list.pop_back(); }
+    const int have = (int)blocks.size();
+    const int add  = need - have;
+    auto it = impl_->seq_slot.find(seq_id);
+    const bool need_new_slot = (it == impl_->seq_slot.end());
+    // Check every precondition before mutating any free-list / slot state (fail clean).
+    if (add > 0 && (int)impl_->free_list.size() < add) return false;
+    if (need_new_slot && impl_->free_slots.empty()) return false;
+    for (int i = 0; i < add; i++) { blocks.push_back(impl_->free_list.back()); impl_->free_list.pop_back(); }
 
     int slot;
-    auto it = impl_->seq_slot.find(seq_id);
-    if (it != impl_->seq_slot.end()) slot = it->second;
+    if (!need_new_slot) slot = it->second;
     else { slot = impl_->free_slots.back(); impl_->free_slots.pop_back(); impl_->seq_slot[seq_id] = slot; }
 
     cu(cudaMemcpy(impl_->d_block_tables + (size_t)slot * kMaxBlocksPerSeq, blocks.data(),

runtime/tests/CMakeLists.txt

@@ -21,3 +21,10 @@ add_test(NAME decode_runner_gpu_test COMMAND decode_runner_gpu_test)
 add_executable(qwen35_gpu_test qwen35_gpu_test.cpp)
 target_link_libraries(qwen35_gpu_test PRIVATE sparkinfer_runtime CUDA::cudart)
 add_test(NAME qwen35_gpu_test COMMAND qwen35_gpu_test)
+
+add_executable(kv_cache_gpu_test kv_cache_gpu_test.cpp)
+target_link_libraries(kv_cache_gpu_test PRIVATE sparkinfer_runtime CUDA::cudart)
+add_test(NAME kv_cache_gpu_test COMMAND kv_cache_gpu_test)
+
+add_executable(kv_allocate_repro kv_allocate_repro.cpp)
+target_link_libraries(kv_allocate_repro PRIVATE sparkinfer_runtime CUDA::cudart)

runtime/tests/kv_allocate_repro.cpp

@@ -0,0 +1,58 @@
+// Standalone repro: allocate(1,100) -> grow allocate(1,200) on live seq_id.
+// Matches the bug report scenario (block_size=16 -> 7 blocks, grow to 13).
+
+#include "sparkinfer/kv_cache.h"
+
+#include <cuda_runtime.h>
+#include <cstdio>
+
+using sparkinfer::KVCacheConfig;
+using sparkinfer::KVCacheManager;
+
+static int used(KVCacheManager& kv) {
+    return kv.num_total_blocks() - kv.num_free_blocks();
+}
+
+int main() {
+    int ndev = 0;
+    if (cudaGetDeviceCount(&ndev) != cudaSuccess || ndev == 0) {
+        printf("[SKIP] no CUDA device\n");
+        return 0;
+    }
+
+    cudaDeviceProp prop{};
+    cudaGetDeviceProperties(&prop, 0);
+
+    KVCacheConfig cfg{};
+    cfg.num_layers = 2;
+    cfg.num_kv_heads = 4;
+    cfg.head_dim = 128;
+    cfg.block_size = 16;
+
+    KVCacheManager kv(cfg, 32ull * 1024 * 1024);
+    const int total = kv.num_total_blocks();
+
+    if (!kv.allocate(1, 100)) { printf("RESULT: FAIL (allocate 100)\n"); return 1; }
+    const int u100 = used(kv);
+    printf("after allocate(1,100):  used=%d  (expect 7)\n", u100);
+    if (u100 != 7) { printf("RESULT: FAIL\n"); return 1; }
+
+    if (!kv.allocate(1, 200)) { printf("RESULT: FAIL (allocate 200)\n"); return 1; }
+    const int u200 = used(kv);
+    printf("after allocate(1,200):  used=%d  (correct=13)\n", u200);
+    if (u200 != 13) { printf("RESULT: FAIL (expected 13, got %d)\n", u200); return 1; }
+
+    if (!kv.allocate(1, 200)) { printf("RESULT: FAIL (idempotent)\n"); return 1; }
+    const int u200b = used(kv);
+    printf("re-allocate same size:  used=%d  (expect 13)\n", u200b);
+    if (u200b != 13) { printf("RESULT: FAIL\n"); return 1; }
+
+    kv.free(1);
+    const int free_after = kv.num_free_blocks();
+    printf("free(1):                free=%d  (total=%d)\n", free_after, total);
+    if (free_after != total) { printf("RESULT: FAIL (leak: %d blocks missing)\n", total - free_after); return 1; }
+
+    printf("GPU: %s (sm_%d.%d)\n", prop.name, prop.major, prop.minor);
+    printf("RESULT: PASS\n");
+    return 0;
+}

runtime/tests/kv_cache_gpu_test.cpp

@@ -0,0 +1,84 @@
+// GPU test for KVCacheManager block growth — exercises seq_id re-allocation.
+
+#include "sparkinfer/kv_cache.h"
+
+#include <cuda_runtime.h>
+#include <cstdio>
+#include <vector>
+
+using sparkinfer::KVCacheConfig;
+using sparkinfer::KVCacheManager;
+
+int main() {
+    int ndev = 0;
+    if (cudaGetDeviceCount(&ndev) != cudaSuccess || ndev == 0) {
+        printf("[SKIP] no CUDA device — kv_cache_gpu_test requires a GPU\n");
+        return 0;
+    }
+
+    cudaDeviceProp prop{};
+    cudaGetDeviceProperties(&prop, 0);
+
+    constexpr int block_size = 16;
+    KVCacheConfig cfg{};
+    cfg.num_layers = 2;
+    cfg.num_kv_heads = 4;
+    cfg.head_dim = 128;
+    cfg.block_size = block_size;
+
+    KVCacheManager kv(cfg, 32ull * 1024 * 1024);
+    const int total = kv.num_total_blocks();
+    const int free0 = kv.num_free_blocks();
+    if (free0 != total) {
+        printf("[FAIL] expected all blocks free at start (got %d, total %d)\n", free0, total);
+        return 1;
+    }
+
+    if (!kv.allocate(0, 32)) { printf("[FAIL] first allocate(0, 32)\n"); return 1; }
+    const int free1 = kv.num_free_blocks();
+    if (free1 != free0 - 2) {
+        printf("[FAIL] first allocate should take 2 blocks (free %d -> %d)\n", free0, free1);
+        return 1;
+    }
+
+    std::vector<int> seq0_blocks(2);
+    cudaMemcpy(seq0_blocks.data(), kv.block_table(0), 2 * sizeof(int), cudaMemcpyDeviceToHost);
+
+    if (!kv.allocate(0, 64)) { printf("[FAIL] re-allocate(0, 64)\n"); return 1; }
+    const int free2 = kv.num_free_blocks();
+    if (free2 != free1 - 2) {
+        printf("[FAIL] re-allocate should add 2 blocks (free %d -> %d, want %d)\n",
+               free1, free2, free1 - 2);
+        return 1;
+    }
+
+    std::vector<int> seq0_grown(4);
+    cudaMemcpy(seq0_grown.data(), kv.block_table(0), 4 * sizeof(int), cudaMemcpyDeviceToHost);
+    for (int i = 0; i < 2; i++) {
+        if (seq0_grown[i] != seq0_blocks[i]) {
+            printf("[FAIL] re-allocate changed existing block id at %d (%d != %d)\n",
+                   i, seq0_grown[i], seq0_blocks[i]);
+            return 1;
+        }
+    }
+
+    if (!kv.allocate(1, 32)) { printf("[FAIL] allocate(1, 32)\n"); return 1; }
+    std::vector<int> seq1_blocks(2);
+    cudaMemcpy(seq1_blocks.data(), kv.block_table(1), 2 * sizeof(int), cudaMemcpyDeviceToHost);
+    if (seq1_blocks[0] == seq0_grown[0] && seq1_blocks[1] == seq0_grown[1]) {
+        printf("[FAIL] seq 1 block table aliases seq 0 (table row corruption)\n");
+        return 1;
+    }
+
+    const int free3 = kv.num_free_blocks();
+    if (!kv.allocate(0, 64)) { printf("[FAIL] idempotent re-allocate(0, 64)\n"); return 1; }
+    if (kv.num_free_blocks() != free3) {
+        printf("[FAIL] idempotent re-allocate consumed blocks (%d -> %d)\n",
+               free3, kv.num_free_blocks());
+        return 1;
+    }
+
+    printf("[PASS] kv_cache_gpu_test on %s: fresh allocate, delta re-grow, no leak, "
+           "seq isolation (%d blocks total)\n", prop.name, total);
+    return 0;
+}