From 4a9b9446c5eb42d261d4e392ada1e506ddfc3c50 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 2 Mar 2026 19:44:28 -0800 Subject: [PATCH] feat(bench): add PositionalIndexer benchmarks to radix tree benchmark MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .github/workflows/benchmark-radix-tree.yml | 2 +- model_gateway/benches/radix_tree_benchmark.rs | 325 +++++++++++++++++- 2 files changed, 322 insertions(+), 5 deletions(-) diff --git a/.github/workflows/benchmark-radix-tree.yml b/.github/workflows/benchmark-radix-tree.yml index 210cc094fd..e0181fba44 100644 --- a/.github/workflows/benchmark-radix-tree.yml +++ b/.github/workflows/benchmark-radix-tree.yml @@ -51,7 +51,7 @@ jobs: timeout-minutes: 40 run: | source "$HOME/.cargo/env" - cargo bench --bench radix_tree_benchmark -- benchmark_summary --exact 2>&1 | tee benchmark_output.txt + cargo bench --bench radix_tree_benchmark -- benchmark_summary 2>&1 | tee benchmark_output.txt - name: Upload benchmark results if: always() diff --git a/model_gateway/benches/radix_tree_benchmark.rs b/model_gateway/benches/radix_tree_benchmark.rs index 5420c36488..2acb27f566 100644 --- a/model_gateway/benches/radix_tree_benchmark.rs +++ b/model_gateway/benches/radix_tree_benchmark.rs @@ -1,8 +1,9 @@ //! Benchmarks for the radix tree implementations used in cache-aware routing. //! -//! This benchmark tests both implementations: +//! This benchmark tests all three implementations: //! - StringTree: Character-based tree for HTTP router (text input) //! - TokenTree: Token-based tree for gRPC router (pre-tokenized input) +//! - PositionalIndexer: Event-driven indexer for gRPC router (KV cache events) //! //! Run with: cargo bench --bench radix_tree_benchmark //! @@ -25,7 +26,10 @@ use std::{ }; use criterion::{criterion_group, criterion_main, Criterion}; -use kv_index::{StringTree, TokenTree}; +use kv_index::{ + compute_content_hash, compute_request_content_hashes, ContentHash, PositionalIndexer, + SequenceHash, StoredBlock, StringTree, TokenTree, +}; use rand::{ distr::{Alphanumeric, SampleString}, rng as thread_rng, Rng, @@ -393,11 +397,256 @@ macro_rules! bench_token_concurrent { }}; } +// ============================================================================ +// PositionalIndexer Helpers +// ============================================================================ + +/// Generate token chunks of `block_size` tokens each. +fn generate_token_chunks(num_blocks: usize, block_size: usize) -> Vec> { + let mut rng = thread_rng(); + (0..num_blocks) + .map(|_| { + (0..block_size) + .map(|_| rng.random_range(0..50000)) + .collect() + }) + .collect() +} + +fn chunks_to_stored_blocks(chunks: &[Vec]) -> Vec { + chunks + .iter() + .enumerate() + .map(|(i, tokens)| StoredBlock { + seq_hash: SequenceHash(i as u64 + 1), + content_hash: compute_content_hash(tokens), + }) + .collect() +} + +fn flatten_tokens(chunks: &[Vec]) -> Vec { + chunks.iter().flat_map(|c| c.iter().copied()).collect() +} + +/// Build a populated indexer. Workers share `shared_prefix_blocks` initial blocks +/// (simulating common system prompt) then diverge. +fn build_populated_indexer( + workers: &[String], + blocks_per_worker: usize, + block_size: usize, + shared_prefix_blocks: usize, + jump_size: usize, +) -> (Arc, Vec>>) { + let indexer = Arc::new(PositionalIndexer::new(jump_size)); + + let shared_chunks = generate_token_chunks(shared_prefix_blocks, block_size); + let shared_blocks = chunks_to_stored_blocks(&shared_chunks); + + let mut all_worker_chunks = Vec::with_capacity(workers.len()); + + for worker in workers { + indexer.apply_stored(worker, &shared_blocks, None).unwrap(); + + let unique_count = blocks_per_worker.saturating_sub(shared_prefix_blocks); + let unique_chunks = generate_token_chunks(unique_count, block_size); + let unique_blocks: Vec = 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(); + + if !unique_blocks.is_empty() { + let parent = SequenceHash(shared_prefix_blocks as u64); + indexer + .apply_stored(worker, &unique_blocks, Some(parent)) + .unwrap(); + } + + let mut worker_chunks = shared_chunks.clone(); + worker_chunks.extend(unique_chunks); + all_worker_chunks.push(worker_chunks); + } + + (indexer, all_worker_chunks) +} + +// ============================================================================ +// PositionalIndexer Benchmark Macros +// ============================================================================ + +/// Macro for STORE benchmarks with PositionalIndexer +macro_rules! bench_indexer_store { + ($group:expr, $num_workers:expr, $blocks_per_worker:expr, $block_size:expr, $workers:expr) => {{ + let printed = Arc::new(AtomicBool::new(false)); + let bench_name = format!( + "indexer_store_{}w_{}blk_{}bs", + $num_workers, $blocks_per_worker, $block_size + ); + let workers_clone = $workers.clone(); + + $group.bench_function(&bench_name, |b| { + let workers = workers_clone.clone(); + 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); + } + } + let duration = start.elapsed(); + + if !printed.swap(true, Ordering::Relaxed) { + let total_blocks = + iters as f64 * $num_workers as f64 * $blocks_per_worker as f64; + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let latency_us = duration.as_nanos() as f64 / iters as f64 / 1000.0; + let blocks_per_sec = total_blocks / duration.as_secs_f64(); + add_result( + "indexer", + format!( + "{:>3}w | {:>4}blk x {:>2}tok | STORE | {:>8.0} ops/s | {:>7.1} µs | {:>8.0} blk/s", + $num_workers, $blocks_per_worker, $block_size, ops_per_sec, latency_us, blocks_per_sec + ), + ); + } + + duration + }); + }); + }}; +} + +/// Macro for MATCH benchmarks with PositionalIndexer (find_matches — the hot path) +macro_rules! bench_indexer_match { + ($group:expr, $num_workers:expr, $query_blocks:expr, $block_size:expr, + $indexer:expr, $query_hashes:expr) => {{ + let printed = Arc::new(AtomicBool::new(false)); + let bench_name = format!( + "indexer_match_{}w_{}qblk_{}bs", + $num_workers, $query_blocks, $block_size + ); + let indexer_clone = $indexer.clone(); + let hashes_clone = $query_hashes.clone(); + + $group.bench_function(&bench_name, |b| { + let indexer = indexer_clone.clone(); + let hashes = hashes_clone.clone(); + let mut idx = 0; + let printed = printed.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + let result = indexer.find_matches(black_box(&hashes[idx % hashes.len()])); + black_box(result); + idx += 1; + } + let duration = start.elapsed(); + + if !printed.swap(true, Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let latency_us = duration.as_nanos() as f64 / iters as f64 / 1000.0; + let throughput_mblk = + (ops_per_sec * $query_blocks as f64) / 1_000_000.0; + add_result( + "indexer", + format!( + "{:>3}w | {:>4}blk x {:>2}tok | MATCH | {:>8.0} ops/s | {:>7.1} µs | {:>5.2} Mblk/s", + $num_workers, $query_blocks, $block_size, ops_per_sec, latency_us, throughput_mblk + ), + ); + } + + duration + }); + }); + }}; +} + +/// Macro for CONCURRENT benchmarks with PositionalIndexer +macro_rules! bench_indexer_concurrent { + ($group:expr, $num_workers:expr, $block_size:expr, $num_threads:expr, $ops_per_thread:expr) => {{ + let printed = Arc::new(AtomicBool::new(false)); + let bench_name = format!("indexer_concurrent_{}w", $num_workers); + + $group.bench_function(&bench_name, |b| { + let printed = printed.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + let workers = generate_worker_endpoints($num_workers); + let (indexer, worker_chunks) = + build_populated_indexer(&workers, 64, $block_size, 8, 64); + + let handles: Vec<_> = (0..$num_threads) + .map(|t| { + let indexer = Arc::clone(&indexer); + let worker = workers[t % workers.len()].clone(); + let chunks = worker_chunks[t % workers.len()].clone(); + let block_size = $block_size; + + thread::spawn(move || { + 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)); + } + } + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + } + let duration = start.elapsed(); + + if !printed.swap(true, Ordering::Relaxed) { + let total_ops = iters * $num_threads as u64 * $ops_per_thread as u64; + let ops_per_sec = total_ops as f64 / duration.as_secs_f64(); + add_result( + "indexer", + format!( + "{:>3}w | CONCURRENT | {:>7.0} ops/s | {} threads x {} ops", + $num_workers, ops_per_sec, $num_threads, $ops_per_thread + ), + ); + } + + duration + }); + }); + }}; +} + // ============================================================================ // Main Benchmark // ============================================================================ -/// Main benchmark for StringTree and TokenTree +/// Main benchmark for StringTree, TokenTree, and PositionalIndexer fn bench_summary(c: &mut Criterion) { let mut group = c.benchmark_group("benchmark_summary"); @@ -493,6 +742,60 @@ fn bench_summary(c: &mut Criterion) { bench_token_concurrent!(group, num_workers, workers, NUM_THREADS, OPS_PER_THREAD); } + // ======================================================================== + // PositionalIndexer Benchmark + // ======================================================================== + const BLOCK_SIZES: [usize; 2] = [16, 64]; + const BLOCKS_PER_WORKER: [usize; 3] = [64, 256, 1024]; + const QUERY_BLOCK_COUNTS: [usize; 3] = [32, 128, 512]; + const SHARED_PREFIX_BLOCKS: usize = 8; + const JUMP_SIZE: usize = 64; + + for &num_workers in &WORKER_COUNTS { + let workers = generate_worker_endpoints(num_workers); + + for &block_size in &BLOCK_SIZES { + // STORE benchmarks + for &blocks_per_worker in &BLOCKS_PER_WORKER { + bench_indexer_store!(group, num_workers, blocks_per_worker, block_size, workers); + } + + // MATCH benchmarks: build a populated indexer, then query it + let max_blocks = *BLOCKS_PER_WORKER.last().unwrap(); + let (indexer, worker_chunks) = build_populated_indexer( + &workers, + max_blocks, + block_size, + SHARED_PREFIX_BLOCKS, + JUMP_SIZE, + ); + + for &query_blocks in &QUERY_BLOCK_COUNTS { + // Generate query hashes: mix of cached (from worker 0) and novel tokens + let num_queries = 100; + let query_hashes: Vec> = (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) + }) + .collect(); + + bench_indexer_match!( + group, + num_workers, + query_blocks, + block_size, + indexer, + query_hashes + ); + } + } + + bench_indexer_concurrent!(group, num_workers, 16, NUM_THREADS, OPS_PER_THREAD); + } + group.finish(); } @@ -507,12 +810,14 @@ fn print_summary() { // Collect results by category let mut string_results = Vec::new(); let mut token_results = Vec::new(); + let mut indexer_results = Vec::new(); for (key, value) in results.iter() { let category = key.split('_').skip(1).collect::>().join("_"); match category.as_str() { "string" => string_results.push(value.clone()), "token" => token_results.push(value.clone()), + "indexer" => indexer_results.push(value.clone()), _ => {} } } @@ -541,7 +846,19 @@ fn print_summary() { eprintln!("{v}"); } - eprintln!("\n{}", "=".repeat(90)); + eprintln!("\n{}", "=".repeat(95)); + eprintln!("POSITIONALINDEXER (kv_index::PositionalIndexer) — event-driven KV cache routing"); + eprintln!("{}", "=".repeat(95)); + eprintln!( + "{:>4} | {:>15} | {:>6} | {:>10} | {:>9} | {:>12}", + "Work", "Size", "Op", "Throughput", "Latency", "Bandwidth" + ); + eprintln!("{}", "-".repeat(95)); + for v in &indexer_results { + eprintln!("{v}"); + } + + eprintln!("\n{}", "=".repeat(95)); } fn run_benchmarks(c: &mut Criterion) {