feat: add shared-memory radix-tree prefix cache reuse for host paged KV by lausannel · Pull Request #85 · batchgen-project/batchgen

lausannel · 2026-02-27T18:14:31Z

Added AllocatePagesForSequencesWithPrefix method to HostPagedKVWorkerView for handling sequences with prefix tokens.
Introduced PrefixSequenceState struct to manage state related to prefix reuse.
Implemented methods for updating, committing, and removing prefix states.
Enhanced the HostKVPrefixCache with a harness for testing prefix cache behavior.
Updated bindings in Python to expose new functionality for prefix reuse.
Added tests for prefix reuse scenarios, including cache hits and eviction behavior.
Modified configuration structure to include parameters for prefix reuse settings.

Description

Briefly describe your changes.

Motivation

Explain why this change is needed and what problem it solves.
If it fixes an issue, link it (e.g., close #123).

Type of Change

Bug fix
New feature
Breaking change
Documentation update

Checklist

I have read the CONTRIBUTION guide.
I have updated the tests (if applicable).
I have updated the documentation (if applicable).

…ated components - Added AllocatePagesForSequencesWithPrefix method to HostPagedKVWorkerView for handling sequences with prefix tokens. - Introduced PrefixSequenceState struct to manage state related to prefix reuse. - Implemented methods for updating, committing, and removing prefix states. - Enhanced the HostKVPrefixCache with a harness for testing prefix cache behavior. - Updated bindings in Python to expose new functionality for prefix reuse. - Added tests for prefix reuse scenarios, including cache hits and eviction behavior. - Modified configuration structure to include parameters for prefix reuse settings.

…aults

…ce processing

Copilot

Pull request overview

This PR adds a shared-memory radix-tree prefix KV cache for host paged KV, enabling prefix tokens that were previously computed and offloaded to host memory to be reused across requests. When a new sequence shares a common prefix with a previously-offloaded sequence, those KV pages are borrowed directly rather than re-copied, saving memory and I/O bandwidth.

Changes:

New HostKVPrefixCache class implementing a radix-tree + LRU-eviction prefix store, fully backed by shared memory so multiple processes sharing the same SHM region benefit uniformly.
New AllocatePagesForSequencesWithPrefix API on HostPagedKVWorkerView and HostPagedKVBackend for prefix-aware batch page allocation with transactional rollback on failure.
Python-side plumbing: new config fields (enable_prefix_reuse, capacity knobs), updated bindings, worker integration, and tests.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
`core/KV_Storage/host_paged_kv_prefix_cache.h/.cpp`	New radix-tree prefix cache with LRU eviction
`core/KV_Storage/host_paged_kv_backend.h/.cpp`	Backend extended with prefix-aware allocation, page-refcount tracking, and seq-page linked-list nodes
`core/KV_Storage/host_paged_kv_worker_view.h`	Worker-level prefix state management, skip-token logic for offload, and decode-token accumulation
`core/batchgen_Binding.cpp`	Exposes `HostKVPrefixCacheHarness` and new worker/config fields to Python
`core/data_structures.h` / `core/utils.cpp`	Python-facing config struct updated with prefix reuse knobs
`core/KV_Storage/host_paged_kv_config_utils.h`	Propagates new config fields from Python config to C++ config
`batchgen/config/config.py`	Python `HostPagedKVConfig` dataclass gains prefix reuse fields and auto-default logic
`batchgen/config/engine_config_parser.py`	Calls `apply_runtime_defaults()` after parsing YAML config
`batchgen/kv_cache/host_kv_mananger_config.py`	`build_host_kv_config*` helpers initialize prefix reuse config fields
`batchgen/batchgen_worker.py`	Routes prefill allocation through prefix-aware path; accumulates decode token IDs
`op_builder/core_engine.py`	Adds `host_paged_kv_prefix_cache.cpp` to build sources; minor style cleanup
`test/paged_kv/test_prefix_cache_binding.py`	Unit tests for `HostKVPrefixCacheHarness`
`test/paged_kv/test_host_prefix_reuse.py`	Integration tests for end-to-end prefix reuse through the worker view

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-03T18:21:17Z

@@ -987,6 +989,31 @@ def _append_decode_kv_to_host_async(
 			# Write position is current position (0-indexed)
 			write_pos = seq.current_context_length - 1
 			sequence_lengths.append(write_pos)
+			decode_pos = write_pos - seq.prompt_length
+			if decode_pos >= 0:
+				query_entry = self.query_book.get(local_idx)
+				decoded_tokens = (
+					None if query_entry is None else query_entry.decoded_tokens
+				)
+				if (
+					decode_pos >= seq.decoded_length
+					or decoded_tokens is None
+					or decoded_tokens.dim() < 2
+					or decode_pos >= decoded_tokens.shape[1]
+				):
+					has_decode_tokens = False
+				else:
+					decode_token_ids.append(
+						int(decoded_tokens[0, decode_pos].item())
+					)
+			elif (
+				seq.input_ids is None
+				or write_pos < 0
+				or write_pos >= seq.input_ids.shape[1]
+			):
+				has_decode_tokens = False
+			else:
+				decode_token_ids.append(int(seq.input_ids[0, write_pos].item()))

 			# Track for debugging (only first few sequences)
 			if len(append_diag) < 3 and seq.decoded_length > 1:
@@ -996,6 +1023,11 @@ def _append_decode_kv_to_host_async(
 					'decoded_len': seq.decoded_length,
 					'write_pos': write_pos,
 				})
+		if not has_decode_tokens:
+			logging.debug(
+				f"Rank {self.rank}: PrefixCache decode token capture unavailable "
+				f"(layer={layer_idx}, batch_size={len(batch)}), skip decode token cache update"
+			)


In the _append_decode_kv_to_host_async method (both at line ~976 and ~6809), when has_decode_tokens becomes False for any one sequence in the batch, the entire decode_token_ids list is discarded (passed as None) for the whole batch. This means if a single sequence in a batch lacks token-ID information, no sequences in that batch get their decode tokens tracked for prefix reuse. The partially-constructed decode_token_ids list also has an inconsistent length relative to sequence_ids at the point has_decode_tokens is set to False, which could cause subtle count mismatches if this logic ever changes to not discard the list. Consider accumulating None entries for missing sequences to preserve per-sequence token tracking.

Copilot · 2026-03-03T18:21:18Z

+    const std::int32_t terminal_node = UpsertRadixPathLocked(tokens, token_count);
+    if (terminal_node == kHostKVInvalidIndex) {
+        return false;
+    }
+
+    if (radix_nodes_[terminal_node].terminal_entry != kHostKVInvalidIndex) {
+        TouchPrefixEntryLocked(radix_nodes_[terminal_node].terminal_entry);
+        return true;
+    }
+
+    if (!EnsurePrefixCapacityLocked(0, 0, 1, pages.size(), pages.size())) {
+        return false;
+    }


In CommitPrefixLocked, UpsertRadixPathLocked at line 589 may allocate new radix nodes/edges. If the subsequent second call to EnsurePrefixCapacityLocked at line 599 returns false (capacity cannot be freed), the function returns false at line 601 but the nodes/edges newly allocated by UpsertRadixPathLocked are now orphaned: they have no terminal entry, are not tracked by LRU, and can never be evicted. This leaks radix nodes and edges from the pool, reducing available capacity permanently until a reset.

To fix this, the second EnsurePrefixCapacityLocked call should be done before UpsertRadixPathLocked, or the function should undo the partial path (e.g., by re-searching and pruning the newly-created chain if the terminal node has no entry).

Copilot · 2026-03-03T18:21:18Z

+        config.sequence_page_node_capacity = std::max(
+            config.num_pages, static_cast<std::size_t>(config.num_pages * 4));
+    }
+    if (config.radix_node_capacity == 0) {
+        config.radix_node_capacity =
+            std::max<std::size_t>(4096, config.num_pages / 2);
+    }


In SanitizeConfig (in host_paged_kv_backend.h, line 115-118), sequence_page_node_capacity is auto-computed as std::max(config.num_pages, config.num_pages * 4). Since config.num_pages * 4 will always be ≥ config.num_pages (for positive num_pages), the std::max here is redundant and the expression simply yields num_pages * 4. The same pattern is replicated in batchgen/config/config.py (line 201) and batchgen/kv_cache/host_kv_mananger_config.py (line 187). This is a minor maintenance concern but it also exists in host_paged_kv_backend.h which was newly added in this PR.

Suggested change

config.sequence_page_node_capacity = std::max(

config.num_pages, static_cast<std::size_t>(config.num_pages * 4));

}

if (config.radix_node_capacity == 0) {

config.radix_node_capacity =

std::max<std::size_t>(4096, config.num_pages / 2);

}

config.sequence_page_node_capacity = config.num_pages * 4;

}

if (config.radix_node_capacity == 0) {

config.radix_node_capacity =

std::max<std::size_t>(4096, config.num_pages / 2);

}

}

Copilot · 2026-03-03T18:21:19Z

+    void RemovePrefixStates(const std::vector<std::int64_t>& sequence_ids) {
+        std::lock_guard<std::mutex> guard(prefix_state_mutex_);
+        for (std::int64_t sequence_id : sequence_ids) {
+            prefix_states_.erase(sequence_id);
+        }
+    }
+


The method RemovePrefixStates (taking a vector of sequence IDs) is defined but never called anywhere in the codebase. The single-sequence variant RemovePrefixState is called from UnregisterSequence, and ClearPrefixStates is called from ReleaseAllSequences. The vector overload is dead code and should either be used (e.g., from UnregisterSequences) or removed.

Suggested change

void RemovePrefixStates(const std::vector<std::int64_t>& sequence_ids) {

std::lock_guard<std::mutex> guard(prefix_state_mutex_);

for (std::int64_t sequence_id : sequence_ids) {

prefix_states_.erase(sequence_id);

}

}

…performance

lausannel added 4 commits February 27, 2026 18:12

feat: enhance prefix cache with decode token support and capacity def…

b4a2ae3

…aults

feat: enhance prefix cache logging for reused and skipped tokens

23debbc

feat: add decode token handling in BatchGenWorker for improved sequen…

9490730

…ce processing

lausannel requested a review from Copilot March 3, 2026 18:15

Copilot started reviewing on behalf of lausannel March 3, 2026 18:16 View session

Copilot AI reviewed Mar 3, 2026

View reviewed changes

lausannel added 19 commits March 4, 2026 18:04

Merge remote-tracking branch 'origin' into feature/prefix-cache

01bbfec

feat: update sequence handling in HostPagedKVWorkerView for improved …

32f8dcb

…performance

Merge origin/main into feature/prefix-cache

ed65a7c

fix: restore namespace scope in host paged kv backend

9955bd8

feat: allow prefix reuse via env

8311d06

feat: add prefix cache server flag

db58e50

fix: use venv pip in install script

680762a

revert: restore install script

a50ac5c

fix: lazily import server worker entrypoint

42c997e

Merge remote-tracking branch 'origin/main' into feature/prefix-cache

05d33f4

fix: initialize empty gpu page table on idle ranks

fde3312

fix: keep prefix cache runtime config in sync

927d892

fix: persist worker host kv config for runtime sync

2d58cd8

test: expand prefix cache coverage

1a51939

test: add multiprocess prefix reuse coverage

7adf834

test: fix multiprocess prefix stats assertions

551d05a

Merge remote-tracking branch 'origin' into feature/prefix-cache

67f730f

merge

865a726

merge

d0163ab

Andrewxu313 mentioned this pull request Apr 26, 2026

feat: add staged page-level prefix reuse #138

Draft

16 tasks

github-actions Bot added the ci:run Trigger build + GPU regression on H20 label May 27, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add shared-memory radix-tree prefix cache reuse for host paged KV - #85

feat: add shared-memory radix-tree prefix cache reuse for host paged KV#85
lausannel wants to merge 23 commits into
mainfrom
feature/prefix-cache

lausannel commented Feb 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Mar 3, 2026

Uh oh!

Copilot AI Mar 3, 2026

Uh oh!

Copilot AI Mar 3, 2026

Uh oh!

Copilot AI Mar 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

lausannel commented Feb 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation

Type of Change

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

lausannel commented Feb 27, 2026 •

edited

Loading