Skip to content

perf: build KV slot mappings by page ranges #577

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
T4t4KAU wants to merge 2 commits into
base: InfiniLM-v0.2.9c
Choose a base branch
from
Open

perf: build KV slot mappings by page ranges #577

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
622c468
perf: build KV slot mappings by page ranges
T4t4KAU Sep 17, 2026
99c2254
chore: remove added unit tests
T4t4KAU Sep 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 17 additions & 17 deletions python/infinilm/llm/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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