diff --git a/runtime/src/kv_cache.cpp b/runtime/src/kv_cache.cpp index 0aefec4..4ef5914 100644 --- a/runtime/src/kv_cache.cpp +++ b/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(), diff --git a/runtime/tests/CMakeLists.txt b/runtime/tests/CMakeLists.txt index 4802f2e..0745fc3 100644 --- a/runtime/tests/CMakeLists.txt +++ b/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) diff --git a/runtime/tests/kv_allocate_repro.cpp b/runtime/tests/kv_allocate_repro.cpp new file mode 100644 index 0000000..20d6045 --- /dev/null +++ b/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 +#include + +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; +} diff --git a/runtime/tests/kv_cache_gpu_test.cpp b/runtime/tests/kv_cache_gpu_test.cpp new file mode 100644 index 0000000..9e78ec5 --- /dev/null +++ b/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 +#include +#include + +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 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 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 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; +}