From 21c1444e3c449284d3532c03ae7a400ffe2aad48 Mon Sep 17 00:00:00 2001 From: T4t4KAU Date: Thu, 17 Sep 2026 20:11:34 +0800 Subject: [PATCH] perf: avoid scanning KV cache pages for capacity checks --- python/infinilm/llm/cache_manager.py | 16 +- test/llm/test_cache_manager.py | 310 +++++++++++++++++++++++++++ 2 files changed, 320 insertions(+), 6 deletions(-) create mode 100644 test/llm/test_cache_manager.py diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index 6c045e4ee..d743d083b 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -81,6 +81,8 @@ def __init__(self, num_blocks: int, block_size: int): self.hash_to_block_ids: Dict[BlockHash, Set[int]] = {} self.free_block_ids: deque[int] = deque(range(num_blocks)) self.used_block_ids: Set[int] = set() + # Count physical pages with references, including shared prefix pages only once. + self._num_referenced_blocks = 0 def __repr__(self) -> str: return ( @@ -94,6 +96,7 @@ def _allocate_block(self) -> Block: assert block.ref_count == 0, f"Block {block_id} ref_count not zero" block.reset() self.used_block_ids.add(block_id) + self._num_referenced_blocks += 1 return block def _remove_block_hash(self, block: Block) -> None: @@ -126,12 +129,8 @@ def get_num_free_blocks(self) -> int: return len(self.free_block_ids) def get_total_usable_blocks(self) -> int: - freeable_used_blocks = sum( - 1 - for block_id in self.used_block_ids - if self.blocks[block_id].ref_count == 0 - ) - return len(self.free_block_ids) + freeable_used_blocks + # Both never-allocated and unreferenced cached pages are usable. + return self.num_blocks - self._num_referenced_blocks def get_computed_blocks( self, @@ -151,6 +150,8 @@ def get_computed_blocks( block_id = next(iter(block_ids)) block = self.blocks[block_id] assert block.hash == block_hash and block_id in self.used_block_ids + if block.ref_count == 0: + self._num_referenced_blocks += 1 block.ref_count += 1 cached_block_table.append(block_id) return cached_block_table, len(cached_block_table) * self.block_size @@ -277,6 +278,7 @@ def truncate_blocks( for block_id in discarded_block_ids: block = self.blocks[block_id] block.ref_count = 0 + self._num_referenced_blocks -= 1 self._deallocate_block(block_id) return block_table[:keep_blocks] @@ -333,6 +335,8 @@ def free_blocks(self, block_table: Sequence[int]) -> None: block = self.blocks[block_id] assert block.ref_count > 0, "block ref_count must be greater than 0" block.ref_count -= 1 + if block.ref_count == 0: + self._num_referenced_blocks -= 1 def try_free_blocks(self, num_required: int) -> bool: """Evict unreferenced blocks until the requested capacity is available.""" diff --git a/test/llm/test_cache_manager.py b/test/llm/test_cache_manager.py new file mode 100644 index 000000000..72bc137e0 --- /dev/null +++ b/test/llm/test_cache_manager.py @@ -0,0 +1,310 @@ +"""CPU-only cache lifecycle and scheduler regression tests. + +Run with: python -m unittest discover -s test/llm -v +""" + +import importlib.util +import random +import sys +import unittest +from collections import Counter +from pathlib import Path +from types import SimpleNamespace + + +def load_control_plane(): + """Load real control-plane modules without the GPU package initializer.""" + directory = Path(__file__).resolve().parents[2] / "python/infinilm/llm" + modules = {} + missing = object() + originals = {} + try: + for name in ( + "prefix_cache", + "sampling_params", + "request", + "cache_manager", + "scheduler", + ): + qualified_name = f"infinilm.llm.{name}" + spec = importlib.util.spec_from_file_location( + qualified_name, directory / f"{name}.py" + ) + module = importlib.util.module_from_spec(spec) + originals[qualified_name] = sys.modules.get(qualified_name, missing) + sys.modules[qualified_name] = module + spec.loader.exec_module(module) + modules[name] = module + finally: + # Restore only our modules; dependencies may contain non-reloadable extensions. + for name, original in originals.items(): + if original is missing: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + return SimpleNamespace(**modules) + + +modules = load_control_plane() +BlockManager = modules.cache_manager.BlockManager + + +class ScanningBlockManager(BlockManager): + """The previous capacity query, used as a scheduler decision oracle.""" + + def get_total_usable_blocks(self): + return len(self.free_block_ids) + sum( + self.blocks[block_id].ref_count == 0 for block_id in self.used_block_ids + ) + + +class BlockManagerTest(unittest.TestCase): + def assert_state(self, manager, owners): + references = Counter(block for table in owners for block in table) + self.assertEqual( + manager.get_total_usable_blocks(), manager.num_blocks - len(references) + ) + self.assertEqual( + [block.ref_count for block in manager.blocks], + [references[i] for i in range(manager.num_blocks)], + ) + free = set(manager.free_block_ids) + self.assertEqual(len(free), len(manager.free_block_ids)) + self.assertFalse(free & manager.used_block_ids) + self.assertEqual(free | manager.used_block_ids, set(range(manager.num_blocks))) + indexed = {} + for block in manager.blocks: + if block.hash: + self.assertIn(block.block_id, manager.used_block_ids) + indexed.setdefault(block.hash, set()).add(block.block_id) + self.assertEqual(manager.hash_to_block_ids, indexed) + + def test_shared_prefix_counts_physical_pages_once(self): + manager = BlockManager(4, 4) + table, _ = manager.allocate_slots(8) + hashes = [b"a" * 16, b"b" * 16] + manager.publish_computed_blocks(table, hashes, 0, 8) + shared, hit = manager.get_computed_blocks(hashes, 8) + self.assertEqual(hit, 8) + self.assert_state(manager, [table, shared]) + manager.free_blocks(table) + self.assert_state(manager, [shared]) + manager.free_blocks(shared) + self.assert_state(manager, []) + repinned, hit = manager.get_computed_blocks(hashes, 4) + self.assertEqual(hit, 4) + self.assert_state(manager, [repinned]) + manager.free_blocks(repinned) + self.assertTrue(manager.try_free_blocks(4)) + self.assert_state(manager, []) + + def test_failed_allocation_after_reclaim_and_retry(self): + manager = BlockManager(3, 4) + live, _ = manager.allocate_slots(4) + released, _ = manager.allocate_slots(8) + manager.free_blocks(released) + self.assertIsNone(manager.allocate_slots(12)) + self.assert_state(manager, [live]) + replacement, _ = manager.allocate_slots(8) + self.assert_state(manager, [live, replacement]) + with self.assertRaisesRegex(RuntimeError, "No available"): + manager.append_slots(replacement, 9, 1) + self.assert_state(manager, [live, replacement]) + + def test_prefix_reuse_then_decode_crosses_block_boundary(self): + manager = BlockManager(4, 4) + prefix, _ = manager.allocate_slots(4) + hashes = [b"a" * 16] + manager.publish_computed_blocks(prefix, hashes, 0, 4) + manager.free_blocks(prefix) + reused, hit = manager.get_computed_blocks(hashes, 4) + table, _ = manager.allocate_slots(8, hit, reused) + self.assert_state(manager, [table]) + self.assertEqual(manager.get_total_usable_blocks(), 1) + table, slot = manager.append_slot(table, 13) + self.assertEqual(slot, table[-1] * 4) + self.assert_state(manager, [table]) + self.assertEqual(manager.get_total_usable_blocks(), 0) + manager.free_blocks(table) + self.assert_state(manager, []) + + def test_speculative_rollback_and_rejected_truncation(self): + manager = BlockManager(6, 4) + table, _ = manager.allocate_slots(5) + table, _ = manager.append_slots(table, 6, 8) + self.assert_state(manager, [table]) + table = manager.truncate_blocks(table, 6) + self.assertEqual(len(table), 2) + self.assert_state(manager, [table]) + manager.publish_computed_blocks(table, [b"a" * 16, b"b" * 16], 0, 8) + with self.assertRaisesRegex(RuntimeError, "must not be published"): + manager.truncate_blocks(table, 4) + self.assert_state(manager, [table]) + manager.free_blocks(table) + self.assert_state(manager, []) + + def test_duplicate_hash_survives_other_copy_eviction(self): + manager = BlockManager(2, 4) + tables = [manager.allocate_slots(4)[0] for _ in range(2)] + hashes = [b"a" * 16] + for table in tables: + manager.publish_computed_blocks(table, hashes, 0, 4) + manager.free_blocks(tables[0]) + self.assertTrue(manager.try_free_blocks(1)) + self.assertEqual(manager.hash_to_block_ids[hashes[0]], set(tables[1])) + shared, _ = manager.get_computed_blocks(hashes, 4) + self.assert_state(manager, [tables[1], shared]) + + def test_randomized_lifecycles(self): + for seed in range(10): + with self.subTest(seed=seed): + rng = random.Random(seed) + manager = BlockManager(23, 4) + owners = [] + for _ in range(1000): + action = rng.randrange(7) + if action == 0 or not owners: + tokens = rng.randint(1, 20) + result = manager.allocate_slots(tokens) + if result is not None: + table, _ = result + owners.append([table, tokens, False]) + elif action == 1: + table, _, _ = owners.pop(rng.randrange(len(owners))) + manager.free_blocks(table) + elif action == 2: + owner = rng.choice(owners) + table, tokens, published = owner + if not published and tokens >= 4: + hashes = [ + bytes([i % 4 + 1]) * 16 for i in range(tokens // 4) + ] + manager.publish_computed_blocks(table, hashes, 0, tokens) + owner[2] = True + elif action == 3: + hashes = [ + bytes([i % 4 + 1]) * 16 for i in range(rng.randint(1, 6)) + ] + table, tokens = manager.get_computed_blocks( + hashes, rng.randint(0, 24) + ) + if table: + owners.append([table, tokens, True]) + elif action == 4: + owner = rng.choice(owners) + if not owner[2]: + extra = rng.randint(0, 9) + try: + table, _ = manager.append_slots( + owner[0], owner[1] + 1, extra + ) + except RuntimeError: + pass + else: + owner[:2] = [table, owner[1] + extra] + elif action == 5: + owner = rng.choice(owners) + if not owner[2]: + tokens = rng.randint(1, owner[1]) + owner[:2] = [ + manager.truncate_blocks(owner[0], tokens), + tokens, + ] + else: + manager.try_free_blocks(rng.randint(1, 23)) + self.assert_state(manager, [owner[0] for owner in owners]) + + +class SchedulerCapacityTest(unittest.TestCase): + def run_lifecycle(self, manager_type): + scheduler = modules.scheduler.Scheduler( + max_batch_size=4, num_blocks=16, block_size=4 + ) + scheduler.cache_manager = manager_type(16, 4) + scheduler.speculative_cache_ops = modules.scheduler.SpeculativeCacheOps( + scheduler.cache_manager + ) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + trace = [] + hits = 0 + for step in range(40): + if step in (0, 5, 10, 15): + for i in range(3): + req = modules.request.InferenceRequest( + request_id=f"{step}-{i}", + prompt_token_ids=[(i + step // 10) % 2] * 9, + sampling_params=modules.sampling_params.SamplingParams( + max_tokens=3 + ), + ) + scheduler.add_request(req) + output = scheduler.schedule() + if output is None: + continue + trace.append( + ( + output.is_prefill, + [ + ( + req.request_id, + tuple(req.block_table), + tuple(req.slot_mapping), + req.num_local_cached_tokens, + ) + for req in output.scheduled_requests + ], + ) + ) + for req in output.scheduled_requests: + if output.is_prefill: + hits += req.num_local_cached_tokens > 0 + scheduler.commit_computed_tokens(req, req.get_total_length()) + req.append_generated_token_id(7) + if req.get_num_generated_tokens() == 3: + req.mark_finished(modules.request.FinishReason.LENGTH) + scheduler.complete_requests(output.scheduled_requests) + manager = scheduler.cache_manager + self.assertEqual( + manager.get_total_usable_blocks(), + ScanningBlockManager.get_total_usable_blocks(manager), + ) + trace.append( + (manager.get_total_usable_blocks(), tuple(manager.free_block_ids)) + ) + self.assertTrue(scheduler.waiting_queue.sync_q.empty()) + self.assertTrue(scheduler.running_queue.sync_q.empty()) + self.assertGreater(hits, 0) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 16) + return trace + + def test_scheduler_decisions_match_scanning_capacity(self): + self.assertEqual( + self.run_lifecycle(BlockManager), self.run_lifecycle(ScanningBlockManager) + ) + + def test_remote_transfer_delays_release(self): + scheduler = modules.scheduler.Scheduler(num_blocks=4, block_size=4) + self.addCleanup(scheduler.waiting_queue.close) + self.addCleanup(scheduler.running_queue.close) + scheduler.connector = SimpleNamespace( + request_finished=lambda *args: (True, None) + ) + request = modules.request.InferenceRequest( + request_id="remote", prompt_token_ids=[1] * 4 + ) + request.block_table, _ = scheduler.cache_manager.allocate_slots(4) + request.mark_canceled() + scheduler.complete_requests([request]) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 3) + output = SimpleNamespace( + kv_connector_output=SimpleNamespace(finished_sending={"remote"}) + ) + scheduler.update_from_output(output) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 4) + scheduler.update_from_output(output) + self.assertEqual(scheduler.cache_manager.get_total_usable_blocks(), 4) + + +if __name__ == "__main__": + unittest.main()