Summary

This draft fixes pending KV cache transfer tracking so primary and secondary pools no longer collide when they use the same dense memory-pool slot index.

KVCacheIndex::get() strips the secondary-pool bit, but the transfer manager used that dense value as the key for pending read/write CUDA events. That aliases physical locations such as primary slot 0 and secondary slot 0. The change adds a pool-qualified pending-transfer key and routes onboard/offload read/write tracking through it.

Test Coverage

Added a deterministic regression test in kvCacheManagerTest that schedules:

• primary slot 0 -> secondary slot 1
• secondary slot 0 -> primary slot 1

The test verifies the transfer manager keeps two pending reads and two pending writes, one per physical (pool, slot) location. This would collapse to one read/write entry with dense-only keys.

Validation

• git diff --check
• pre-commit on commit, with clang-format passing

Not yet run: GPU C++ gtest. Suggested command:

./cpp/build/tests/unit_tests/batch_manager/kvCacheManagerTest --gtest_filter=KVCacheManagerTest.KVCacheTransferManager*

Pre-commit note: skipped waive list check, validate-test-lists, and type-check on the successful commit because scripts/check_test_list.py failed in this local hook environment with TypeError: unsupported operand type(s) for |: 'type' and 'NoneType', unrelated to this C++ change.

Summary by CodeRabbit

• Refactor

  • Improved internal synchronization logic for KV cache transfer operations.

• Tests

  • Added test infrastructure and unit test for KV cache transfer management, including validation of pending transfer tracking across primary and secondary memory slots.