Repository navigation
feat(bench): add PositionalIndexer benchmarks to radix tree benchmark - #579
Conversation
Add comprehensive benchmarks for the event-driven PositionalIndexer alongside the existing StringTree and TokenTree benchmarks in radix_tree_benchmark.rs. What changed: - model_gateway/benches/radix_tree_benchmark.rs: Added PositionalIndexer helper functions (generate_token_chunks, chunks_to_stored_blocks, flatten_tokens, build_populated_indexer), three benchmark macros (bench_indexer_store!, bench_indexer_match!, bench_indexer_concurrent!), benchmark invocations in bench_summary() covering STORE, MATCH, and CONCURRENT operations across worker counts [10, 100, 500], block sizes [16, 64], blocks per worker [64, 256, 1024], and query sizes [32, 128, 512], plus summary table output in print_summary(). - .github/workflows/benchmark-radix-tree.yml: Remove --exact flag from cargo bench invocation which prevented individual benchmarks within the benchmark_summary group from running (substring match is needed). Why: The PositionalIndexer (event-driven KV cache routing) was added in PR #571 but had no benchmarks. This covers the hot path (find_matches at 5-170µs depending on worker count) and background operations (apply_stored), matching the existing benchmark report format. How: Follows the same macro-based pattern as StringTree/TokenTree benchmarks with iter_custom, BTreeMap result collection, and formatted summary tables. The build_populated_indexer helper simulates realistic cache state with shared system prompt prefix blocks and unique per-worker blocks. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
github-actions
Bot
added
ci
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the benchmarking capabilities for the Highlights
Changelog
Ignored Files
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
📝 WalkthroughWalkthroughBenchmark configuration updated to remove flag restriction, and radix tree benchmark suite expanded to comprehensively test PositionalIndexer across STORE, MATCH, and CONCURRENT scenarios alongside existing StringTree and TokenTree implementations, including new data preparation helpers and results aggregation. Changes
Sequence DiagramsequenceDiagram
participant Harness as Benchmark Harness
participant DataGen as Data Generators
participant Indexer as PositionalIndexer
participant Runner as Benchmark Runner
participant Results as Results Collector
Harness->>DataGen: generate_token_chunks()
DataGen->>DataGen: chunks_to_stored_blocks()
DataGen->>DataGen: flatten_tokens()
DataGen-->>Harness: prepared data
Harness->>Indexer: build_populated_indexer()
Indexer-->>Harness: initialized indexer
loop for each scenario (STORE/MATCH/CONCURRENT)
Harness->>Runner: execute benchmark
Runner->>Indexer: perform operations
Indexer-->>Runner: results
Runner->>Results: aggregate metrics
end
Results->>Results: format output (throughput, latency, bandwidth)
Results-->>Harness: summary report
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a comprehensive benchmark suite for PositionalIndexer, which is a great addition for performance tracking. The changes are well-structured and follow the existing patterns in the benchmark file. I have a couple of suggestions to improve the implementation of the new benchmarks. One is to refactor some duplicated code for better maintainability, and the other is to adjust the concurrent benchmark to more accurately measure the performance of the find_matches operation under contention.
| let unique_blocks: Vec<StoredBlock> = unique_chunks | ||
| .iter() | ||
| .enumerate() | ||
| .map(|(i, tokens)| StoredBlock { | ||
| seq_hash: SequenceHash((shared_prefix_blocks + i) as u64 + 1), | ||
| content_hash: compute_content_hash(tokens), | ||
| }) | ||
| .collect(); |
There was a problem hiding this comment.
The logic for creating unique_blocks is very similar to the chunks_to_stored_blocks helper function. To improve code reuse and maintainability, you could consider refactoring chunks_to_stored_blocks to accept a starting index for the seq_hash calculation. This would allow you to replace these lines with a single call to the modified helper, for example: let unique_blocks = chunks_to_stored_blocks(&unique_chunks, shared_prefix_blocks); (after updating the helper and other call sites).
References
- Extract duplicated logic into a shared helper function to improve maintainability and reduce redundancy.
| let mut rng = thread_rng(); | ||
| for i in 0..$ops_per_thread { | ||
| if i % 3 == 0 { | ||
| // Read: find_matches | ||
| let query_tokens = flatten_tokens(&chunks); | ||
| let content_hashes = compute_request_content_hashes( | ||
| &query_tokens, | ||
| block_size, | ||
| ); | ||
| black_box(indexer.find_matches(&content_hashes)); | ||
| } else { | ||
| // Write: apply_stored | ||
| let new_chunks = generate_token_chunks(4, block_size); | ||
| let blocks = chunks_to_stored_blocks(&new_chunks); | ||
| let parent = SequenceHash(rng.random_range(1u64..65)); | ||
| let _ = | ||
| indexer.apply_stored(&worker, &blocks, Some(parent)); | ||
| } | ||
| } |
There was a problem hiding this comment.
In the concurrent benchmark, the data preparation for the read path (flatten_tokens and compute_request_content_hashes) is performed inside the hot loop. Since the chunks data doesn't change within the thread's execution, these computations can be hoisted out of the for loop. This change will ensure the benchmark more accurately measures the performance of find_matches under contention, rather than including the data preparation overhead in every iteration.
let mut rng = thread_rng();
let query_tokens = flatten_tokens(&chunks);
let content_hashes =
compute_request_content_hashes(&query_tokens, block_size);
for i in 0..$ops_per_thread {
if i % 3 == 0 {
// Read: find_matches
black_box(indexer.find_matches(&content_hashes));
} else {
// Write: apply_stored
let new_chunks = generate_token_chunks(4, block_size);
let blocks = chunks_to_stored_blocks(&new_chunks);
let parent = SequenceHash(rng.random_range(1u64..65));
let _ =
indexer.apply_stored(&worker, &blocks, Some(parent));
}
}![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/benches/radix_tree_benchmark.rs`:
- Around line 586-589: The benchmark currently rebuilds worker endpoints and
repopulates the indexer inside the timed section (calls to
generate_worker_endpoints and build_populated_indexer), which mixes setup cost
with measured workload; move the calls to
generate_worker_endpoints($num_workers) and build_populated_indexer(&workers,
64, $block_size, 8, 64) out of the timed benchmark block so that the timed
section only runs the workload (reads/writes) against the already-populated
indexer and worker_chunks, and ensure any per-iteration state that must be reset
is done outside timing or via cloned lightweight copies before starting the
timer.
- Around line 774-782: The comment points out query_hashes is built only from
indexed worker_chunks rather than a mixture; modify how query_hashes is created
to combine cached queries taken from worker_chunks and novel queries generated
on the fly. In the block that builds query_hashes (variable names: num_queries,
query_hashes, worker_chunks, query_blocks, block_size), produce a mix (e.g.,
half cached, half novel) by mapping over indices: for cached indices reuse
flatten_tokens(&chunks[..take]) + compute_request_content_hashes, and for novel
indices synthesize new token sequences (or call the existing token generator
used elsewhere) then call compute_request_content_hashes on those; ensure the
final query_hashes Vec preserves order/length and still uses block_size and
query_blocks variables.
- Line 501: The benchmark is silently ignoring failures from
indexer.apply_stored which can make failed writes count as successful ops;
change the measured calls to check the Result instead of discarding it (e.g.,
use expect/unwrap or match and panic/log on Err) for the apply_stored invocation
used with black_box(worker) and black_box(&blocks) so any failure aborts the
benchmark and only true successes are measured; apply the same fix to the other
occurrences around the later calls (the ones at the spots corresponding to lines
613–614).
- Around line 499-500: The benchmark currently measures token generation and
hashing because generate_token_chunks and chunks_to_stored_blocks are called
inside the timed loop; move the generation and conversion out of the measured
section so only apply_stored is timed: precompute chunks with
generate_token_chunks($blocks_per_worker, $block_size) and convert them to
blocks via chunks_to_stored_blocks(...) before entering the benchmark timing
loop (or store a Vec of blocks and reuse/clone as needed inside each iteration),
then call apply_stored on those precomputed blocks inside the timed loop to
isolate throughput measurement.
ℹ️ Review info
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
📒 Files selected for processing (2)
.github/workflows/benchmark-radix-tree.ymlmodel_gateway/benches/radix_tree_benchmark.rs
| let chunks = generate_token_chunks($blocks_per_worker, $block_size); | ||
| let blocks = chunks_to_stored_blocks(&chunks); |
There was a problem hiding this comment.
STORE benchmark timing includes input-generation overhead.
Line 499 and Line 500 run token generation + block hashing inside the measured loop, so this benchmark is not isolating apply_stored throughput.
Proposed fix (pre-generate blocks outside timed loop)
$group.bench_function(&bench_name, |b| {
let workers = workers_clone.clone();
+ let worker_blocks: Vec<Vec<StoredBlock>> = (0..workers.len())
+ .map(|_| {
+ let chunks = generate_token_chunks($blocks_per_worker, $block_size);
+ chunks_to_stored_blocks(&chunks)
+ })
+ .collect();
let printed = printed.clone();
b.iter_custom(|iters| {
let start = Instant::now();
for _ in 0..iters {
let indexer = PositionalIndexer::new(64);
- for worker in &workers {
- let chunks = generate_token_chunks($blocks_per_worker, $block_size);
- let blocks = chunks_to_stored_blocks(&chunks);
- let _ = indexer.apply_stored(black_box(worker), black_box(&blocks), None);
+ for (i, worker) in workers.iter().enumerate() {
+ let blocks = &worker_blocks[i];
+ let _ = indexer.apply_stored(black_box(worker), black_box(blocks), None);
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| let chunks = generate_token_chunks($blocks_per_worker, $block_size); | |
| let blocks = chunks_to_stored_blocks(&chunks); | |
| $group.bench_function(&bench_name, |b| { | |
| let workers = workers_clone.clone(); | |
| let worker_blocks: Vec<Vec<StoredBlock>> = (0..workers.len()) | |
| .map(|_| { | |
| let chunks = generate_token_chunks($blocks_per_worker, $block_size); | |
| chunks_to_stored_blocks(&chunks) | |
| }) | |
| .collect(); | |
| let printed = printed.clone(); | |
| b.iter_custom(|iters| { | |
| let start = Instant::now(); | |
| for _ in 0..iters { | |
| let indexer = PositionalIndexer::new(64); | |
| for (i, worker) in workers.iter().enumerate() { | |
| let blocks = &worker_blocks[i]; | |
| let _ = indexer.apply_stored(black_box(worker), black_box(blocks), None); | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/benches/radix_tree_benchmark.rs` around lines 499 - 500, The
benchmark currently measures token generation and hashing because
generate_token_chunks and chunks_to_stored_blocks are called inside the timed
loop; move the generation and conversion out of the measured section so only
apply_stored is timed: precompute chunks with
generate_token_chunks($blocks_per_worker, $block_size) and convert them to
blocks via chunks_to_stored_blocks(...) before entering the benchmark timing
loop (or store a Vec of blocks and reuse/clone as needed inside each iteration),
then call apply_stored on those precomputed blocks inside the timed loop to
isolate throughput measurement.
| for worker in &workers { | ||
| let chunks = generate_token_chunks($blocks_per_worker, $block_size); | ||
| let blocks = chunks_to_stored_blocks(&chunks); | ||
| let _ = indexer.apply_stored(black_box(worker), black_box(&blocks), None); |
There was a problem hiding this comment.
apply_stored failures are silently ignored in measured paths.
Line 501 and Line 613-614 discard the Result, so failed writes can be counted as successful ops and skew throughput numbers.
Proposed fix (enforce success-path measurement)
- let _ = indexer.apply_stored(black_box(worker), black_box(&blocks), None);
+ indexer
+ .apply_stored(black_box(worker), black_box(&blocks), None)
+ .expect("STORE benchmark input must be valid");
...
- let _ =
- indexer.apply_stored(&worker, &blocks, Some(parent));
+ indexer
+ .apply_stored(&worker, &blocks, Some(parent))
+ .expect("CONCURRENT benchmark write must be valid");Also applies to: 613-614
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/benches/radix_tree_benchmark.rs` at line 501, The benchmark is
silently ignoring failures from indexer.apply_stored which can make failed
writes count as successful ops; change the measured calls to check the Result
instead of discarding it (e.g., use expect/unwrap or match and panic/log on Err)
for the apply_stored invocation used with black_box(worker) and
black_box(&blocks) so any failure aborts the benchmark and only true successes
are measured; apply the same fix to the other occurrences around the later calls
(the ones at the spots corresponding to lines 613–614).
| let workers = generate_worker_endpoints($num_workers); | ||
| let (indexer, worker_chunks) = | ||
| build_populated_indexer(&workers, 64, $block_size, 8, 64); | ||
|
|
There was a problem hiding this comment.
Concurrent benchmark measures setup + workload together.
Line 586-589 rebuilds worker endpoints and repopulates the indexer inside the timed section. That inflates/obscures the mixed read/write throughput signal.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/benches/radix_tree_benchmark.rs` around lines 586 - 589, The
benchmark currently rebuilds worker endpoints and repopulates the indexer inside
the timed section (calls to generate_worker_endpoints and
build_populated_indexer), which mixes setup cost with measured workload; move
the calls to generate_worker_endpoints($num_workers) and
build_populated_indexer(&workers, 64, $block_size, 8, 64) out of the timed
benchmark block so that the timed section only runs the workload (reads/writes)
against the already-populated indexer and worker_chunks, and ensure any
per-iteration state that must be reset is done outside timing or via cloned
lightweight copies before starting the timer.
| // Generate query hashes: mix of cached (from worker 0) and novel tokens | ||
| let num_queries = 100; | ||
| let query_hashes: Vec<Vec<ContentHash>> = (0..num_queries) | ||
| .map(|i| { | ||
| let chunks = &worker_chunks[i % worker_chunks.len()]; | ||
| let take = query_blocks.min(chunks.len()); | ||
| let tokens = flatten_tokens(&chunks[..take]); | ||
| compute_request_content_hashes(&tokens, block_size) | ||
| }) |
There was a problem hiding this comment.
Query corpus is not mixed despite the inline comment.
Line 774 says “mix of cached and novel tokens,” but Line 776-782 builds all queries from already-indexed worker_chunks.
Proposed fix (actually mix cached + novel queries)
let query_hashes: Vec<Vec<ContentHash>> = (0..num_queries)
.map(|i| {
- let chunks = &worker_chunks[i % worker_chunks.len()];
- let take = query_blocks.min(chunks.len());
- let tokens = flatten_tokens(&chunks[..take]);
- compute_request_content_hashes(&tokens, block_size)
+ if i % 2 == 0 {
+ let chunks = &worker_chunks[i % worker_chunks.len()];
+ let take = query_blocks.min(chunks.len());
+ let tokens = flatten_tokens(&chunks[..take]);
+ compute_request_content_hashes(&tokens, block_size)
+ } else {
+ let novel_chunks = generate_token_chunks(query_blocks, block_size);
+ let novel_tokens = flatten_tokens(&novel_chunks);
+ compute_request_content_hashes(&novel_tokens, block_size)
+ }
})
.collect();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@model_gateway/benches/radix_tree_benchmark.rs` around lines 774 - 782, The
comment points out query_hashes is built only from indexed worker_chunks rather
than a mixture; modify how query_hashes is created to combine cached queries
taken from worker_chunks and novel queries generated on the fly. In the block
that builds query_hashes (variable names: num_queries, query_hashes,
worker_chunks, query_blocks, block_size), produce a mix (e.g., half cached, half
novel) by mapping over indices: for cached indices reuse
flatten_tokens(&chunks[..take]) + compute_request_content_hashes, and for novel
indices synthesize new token sequences (or call the existing token generator
used elsewhere) then call compute_request_content_hashes on those; ensure the
final query_hashes Vec preserves order/length and still uses block_size and
query_blocks variables.
1 participant
Summary
PositionalIndexer(event-driven KV cache routing) to the existing radix tree benchmark suite--exactflag in CI workflow that prevented benchmarks from runningRefs: #571
What changed
model_gateway/benches/radix_tree_benchmark.rs(+317 lines):generate_token_chunks,chunks_to_stored_blocks,flatten_tokens,build_populated_indexer(builds realistic cache state with shared system prompt prefix + unique per-worker blocks)bench_indexer_store!(apply_stored throughput),bench_indexer_match!(find_matches — the routing hot path),bench_indexer_concurrent!(mixed read/write with 32 threads)print_summary()for the new POSITIONALINDEXER section.github/workflows/benchmark-radix-tree.yml:--exactflag fromcargo benchinvocation — individual benchmarks within thebenchmark_summarygroup have names likebenchmark_summary/string_insert_10w_4096c, so--exactmatches nothing. Substring match is needed.Why
The
PositionalIndexerwas added in #571 (event-driven cache-aware routing) but had no benchmarks. The hot path (find_matches) runs on every gRPC routing decision when KV cache events are available, so its latency profile matters for production sizing.Test plan
cargo bench --bench radix_tree_benchmark --no-run— compiles cleanlycargo clippy --bench radix_tree_benchmark -- -D warnings— no warningscargo bench --bench radix_tree_benchmark -- benchmark_summary— all three sections produce results:Summary by CodeRabbit
Tests
Chores