From 622c468e659215bc1655b76e3f6c6888feb28969 Mon Sep 17 00:00:00 2001 From: T4t4KAU Date: Thu, 17 Sep 2026 20:34:55 +0800 Subject: [PATCH 1/2] perf: build KV slot mappings by page ranges --- python/infinilm/llm/cache_manager.py | 34 +++++------ test/llm/cpu_modules.py | 36 ++++++++++++ test/llm/test_slot_mapping.py | 86 ++++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 17 deletions(-) create mode 100644 test/llm/cpu_modules.py create mode 100644 test/llm/test_slot_mapping.py diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index 6c045e4ee..00ba364d9 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -183,11 +183,9 @@ def allocate_slots( for _ in range(num_blocks_needed): block_table.append(self._allocate_block().block_id) - slot_mapping = [] - for token_idx in range(num_computed_tokens, total_tokens): - block_idx = token_idx // self.block_size - block_offset = token_idx % self.block_size - slot_mapping.append(block_table[block_idx] * self.block_size + block_offset) + slot_mapping = self.update_blocks_slot( + block_table, num_computed_tokens, total_tokens + ) return block_table, slot_mapping def append_slots( @@ -217,11 +215,9 @@ def append_slots( for _ in range(additional_blocks): block_table.append(self._allocate_block().block_id) - slots = [] - for num_tokens in range(start_num_tokens, start_num_tokens + num_slots): - token_idx = num_tokens - 1 - block_idx, block_offset = divmod(token_idx, self.block_size) - slots.append(block_table[block_idx] * self.block_size + block_offset) + slots = self.update_blocks_slot( + block_table, start_num_tokens - 1, max_num_tokens + ) return block_table, slots def truncate_blocks( @@ -350,15 +346,19 @@ def try_free_blocks(self, num_required: int) -> bool: def update_blocks_slot( self, block_table: List[int], num_computed_tokens: int, total_tokens: int ) -> List[int]: - """Build slots for the recomputed suffix after a partial remote load.""" + """Build physical slots for a logical token interval.""" if num_computed_tokens >= total_tokens: return [] + if total_tokens - num_computed_tokens == 1: + block_idx, block_offset = divmod(num_computed_tokens, self.block_size) + return [block_table[block_idx] * self.block_size + block_offset] new_slot_mapping = [] - for token_idx in range(num_computed_tokens, total_tokens): - block_idx = token_idx // self.block_size - block_offset = token_idx % self.block_size - new_slot_mapping.append( - block_table[block_idx] * self.block_size + block_offset - ) + token_idx = num_computed_tokens + while token_idx < total_tokens: + block_idx, block_offset = divmod(token_idx, self.block_size) + count = min(self.block_size - block_offset, total_tokens - token_idx) + first_slot = block_table[block_idx] * self.block_size + block_offset + new_slot_mapping.extend(range(first_slot, first_slot + count)) + token_idx += count return new_slot_mapping diff --git a/test/llm/cpu_modules.py b/test/llm/cpu_modules.py new file mode 100644 index 000000000..3e65f55ad --- /dev/null +++ b/test/llm/cpu_modules.py @@ -0,0 +1,36 @@ +"""Load cache and scheduler modules without initializing the GPU engine.""" + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + + +def load_cpu_modules(): + directory = Path(__file__).resolve().parents[2] / "python/infinilm/llm" + missing = object() + originals = {} + modules = {} + try: + for name in ( + "prefix_cache", + "sampling_params", + "request", + "cache_manager", + "scheduler", + ): + key = f"infinilm.llm.{name}" + spec = importlib.util.spec_from_file_location(key, directory / f"{name}.py") + module = importlib.util.module_from_spec(spec) + originals[key] = sys.modules.get(key, missing) + sys.modules[key] = module + spec.loader.exec_module(module) + modules[name] = module + finally: + # Retain imported dependencies, including non-reloadable native extensions. + for key, original in originals.items(): + if original is missing: + sys.modules.pop(key, None) + else: + sys.modules[key] = original + return SimpleNamespace(**modules) diff --git a/test/llm/test_slot_mapping.py b/test/llm/test_slot_mapping.py new file mode 100644 index 000000000..3c4173168 --- /dev/null +++ b/test/llm/test_slot_mapping.py @@ -0,0 +1,86 @@ +"""CPU slot-mapping tests: python -m unittest discover -s test/llm -v.""" + +import random +import unittest + +from cpu_modules import load_cpu_modules + +BlockManager = load_cpu_modules().cache_manager.BlockManager + + +def reference_slots(table, size, start, end): + return [table[i // size] * size + i % size for i in range(start, end)] + + +class SlotMappingTest(unittest.TestCase): + def test_all_intervals_across_reordered_pages(self): + for size in (1, 2, 3, 16, 256): + manager = BlockManager(8, size) + table = [7, 2, 5, 0] + boundaries = {0, 1, size - 1, size, size + 1, 3 * size - 1, 4 * size} + for start in boundaries: + for end in boundaries: + with self.subTest(size=size, start=start, end=end): + self.assertEqual( + manager.update_blocks_slot(table, start, end), + reference_slots(table, size, start, end), + ) + + def test_random_ragged_intervals(self): + rng = random.Random(20260917) + for _ in range(2000): + size = rng.choice((1, 3, 16, 64, 256)) + table = rng.sample(range(100), rng.randint(1, 30)) + start = rng.randrange(len(table) * size + 1) + end = rng.randint(start, len(table) * size) + self.assertEqual( + BlockManager(100, size).update_blocks_slot(table, start, end), + reference_slots(table, size, start, end), + ) + + def test_allocation_with_partial_remote_and_local_prefix(self): + manager = BlockManager(12, 4) + prefix, _ = manager.allocate_slots(8) + manager.publish_computed_blocks(prefix, [b"a" * 16, b"b" * 16], 0, 8) + manager.free_blocks(prefix) + cached, hit = manager.get_computed_blocks([b"a" * 16, b"b" * 16], 8) + table, slots = manager.allocate_slots( + 7, num_computed_tokens=hit + 3, cached_block_table=cached + ) + self.assertEqual(table[:2], prefix) + self.assertEqual(slots, reference_slots(table, 4, 11, 18)) + self.assertEqual( + manager.update_blocks_slot(table, 9, 18), reference_slots(table, 4, 9, 18) + ) + + def test_speculative_append_and_rollback(self): + for size in (1, 3, 16): + manager = BlockManager(20, size) + length = size + 1 + table, slots = manager.allocate_slots(length) + self.assertEqual(slots, reference_slots(table, size, 0, length)) + before = list(table) + self.assertEqual(manager.append_slots(table, length + 1, 0), (before, [])) + table, slots = manager.append_slots(table, length + 1, size * 2 + 1) + self.assertEqual( + slots, reference_slots(table, size, length, length + size * 2 + 1) + ) + table = manager.truncate_blocks(table, length) + self.assertEqual(table, before) + table, slots = manager.append_slots(table, length + 1, size + 1) + self.assertEqual( + slots, reference_slots(table, size, length, length + size + 1) + ) + + def test_failure_leaves_owned_pages_unchanged(self): + manager = BlockManager(2, 4) + table, _ = manager.allocate_slots(8) + self.assertIsNone(manager.allocate_slots(1)) + with self.assertRaisesRegex(RuntimeError, "No available"): + manager.append_slots(table, 9, 1) + self.assertEqual(table, [0, 1]) + self.assertEqual([page.ref_count for page in manager.blocks], [1, 1]) + + +if __name__ == "__main__": + unittest.main() From 99c22547c9b2aed8cdbcc524dc6a21ccaea3a7d0 Mon Sep 17 00:00:00 2001 From: T4t4KAU Date: Thu, 17 Sep 2026 20:51:40 +0800 Subject: [PATCH 2/2] chore: remove added unit tests --- test/llm/cpu_modules.py | 36 --------------- test/llm/test_slot_mapping.py | 86 ----------------------------------- 2 files changed, 122 deletions(-) delete mode 100644 test/llm/cpu_modules.py delete mode 100644 test/llm/test_slot_mapping.py diff --git a/test/llm/cpu_modules.py b/test/llm/cpu_modules.py deleted file mode 100644 index 3e65f55ad..000000000 --- a/test/llm/cpu_modules.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Load cache and scheduler modules without initializing the GPU engine.""" - -import importlib.util -import sys -from pathlib import Path -from types import SimpleNamespace - - -def load_cpu_modules(): - directory = Path(__file__).resolve().parents[2] / "python/infinilm/llm" - missing = object() - originals = {} - modules = {} - try: - for name in ( - "prefix_cache", - "sampling_params", - "request", - "cache_manager", - "scheduler", - ): - key = f"infinilm.llm.{name}" - spec = importlib.util.spec_from_file_location(key, directory / f"{name}.py") - module = importlib.util.module_from_spec(spec) - originals[key] = sys.modules.get(key, missing) - sys.modules[key] = module - spec.loader.exec_module(module) - modules[name] = module - finally: - # Retain imported dependencies, including non-reloadable native extensions. - for key, original in originals.items(): - if original is missing: - sys.modules.pop(key, None) - else: - sys.modules[key] = original - return SimpleNamespace(**modules) diff --git a/test/llm/test_slot_mapping.py b/test/llm/test_slot_mapping.py deleted file mode 100644 index 3c4173168..000000000 --- a/test/llm/test_slot_mapping.py +++ /dev/null @@ -1,86 +0,0 @@ -"""CPU slot-mapping tests: python -m unittest discover -s test/llm -v.""" - -import random -import unittest - -from cpu_modules import load_cpu_modules - -BlockManager = load_cpu_modules().cache_manager.BlockManager - - -def reference_slots(table, size, start, end): - return [table[i // size] * size + i % size for i in range(start, end)] - - -class SlotMappingTest(unittest.TestCase): - def test_all_intervals_across_reordered_pages(self): - for size in (1, 2, 3, 16, 256): - manager = BlockManager(8, size) - table = [7, 2, 5, 0] - boundaries = {0, 1, size - 1, size, size + 1, 3 * size - 1, 4 * size} - for start in boundaries: - for end in boundaries: - with self.subTest(size=size, start=start, end=end): - self.assertEqual( - manager.update_blocks_slot(table, start, end), - reference_slots(table, size, start, end), - ) - - def test_random_ragged_intervals(self): - rng = random.Random(20260917) - for _ in range(2000): - size = rng.choice((1, 3, 16, 64, 256)) - table = rng.sample(range(100), rng.randint(1, 30)) - start = rng.randrange(len(table) * size + 1) - end = rng.randint(start, len(table) * size) - self.assertEqual( - BlockManager(100, size).update_blocks_slot(table, start, end), - reference_slots(table, size, start, end), - ) - - def test_allocation_with_partial_remote_and_local_prefix(self): - manager = BlockManager(12, 4) - prefix, _ = manager.allocate_slots(8) - manager.publish_computed_blocks(prefix, [b"a" * 16, b"b" * 16], 0, 8) - manager.free_blocks(prefix) - cached, hit = manager.get_computed_blocks([b"a" * 16, b"b" * 16], 8) - table, slots = manager.allocate_slots( - 7, num_computed_tokens=hit + 3, cached_block_table=cached - ) - self.assertEqual(table[:2], prefix) - self.assertEqual(slots, reference_slots(table, 4, 11, 18)) - self.assertEqual( - manager.update_blocks_slot(table, 9, 18), reference_slots(table, 4, 9, 18) - ) - - def test_speculative_append_and_rollback(self): - for size in (1, 3, 16): - manager = BlockManager(20, size) - length = size + 1 - table, slots = manager.allocate_slots(length) - self.assertEqual(slots, reference_slots(table, size, 0, length)) - before = list(table) - self.assertEqual(manager.append_slots(table, length + 1, 0), (before, [])) - table, slots = manager.append_slots(table, length + 1, size * 2 + 1) - self.assertEqual( - slots, reference_slots(table, size, length, length + size * 2 + 1) - ) - table = manager.truncate_blocks(table, length) - self.assertEqual(table, before) - table, slots = manager.append_slots(table, length + 1, size + 1) - self.assertEqual( - slots, reference_slots(table, size, length, length + size + 1) - ) - - def test_failure_leaves_owned_pages_unchanged(self): - manager = BlockManager(2, 4) - table, _ = manager.allocate_slots(8) - self.assertIsNone(manager.allocate_slots(1)) - with self.assertRaisesRegex(RuntimeError, "No available"): - manager.append_slots(table, 9, 1) - self.assertEqual(table, [0, 1]) - self.assertEqual([page.ref_count for page in manager.blocks], [1, 1]) - - -if __name__ == "__main__": - unittest.main()