From ff850e76ed6d8ad74e794f46fabad4724bac20f1 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 18 Mar 2026 21:02:51 -0700 Subject: [PATCH 1/7] [None][fix] Fix KvCacheAwareRouter tokenization and add aiperf/router support to disagg benchmark - Fix KvCacheAwareRouter._tokenize() to handle ChatCompletionRequest with models whose apply_chat_template returns a string (e.g. Kimi-K2.5). Use two-step tokenize: apply_chat_template(tokenize=False) then encode(). - Extract _get_tokenizer() helper to avoid duplication. - Add use_aiperf benchmark mode to submit.py for mooncake_trace datasets. - Add router_config support to convert_allocations_to_server_config() so kv_cache_aware routing is included in the generated server_config.yaml. - Add server_config_extra merge for perf_metrics and other server options. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../disaggregated/slurm/benchmark/submit.py | 21 +++++++++++++--- tensorrt_llm/serve/router.py | 25 ++++++++++++++++--- 2 files changed, 39 insertions(+), 7 deletions(-) diff --git a/examples/disaggregated/slurm/benchmark/submit.py b/examples/disaggregated/slurm/benchmark/submit.py index a1bcedaf5087..68248539f443 100644 --- a/examples/disaggregated/slurm/benchmark/submit.py +++ b/examples/disaggregated/slurm/benchmark/submit.py @@ -116,7 +116,8 @@ def assign_servers( return allocations -def convert_allocations_to_server_config(allocations, server_port=8333): +def convert_allocations_to_server_config(allocations, server_port=8333, + router_config=None): generation_servers = {} context_servers = {} server_hostname = None @@ -130,6 +131,8 @@ def convert_allocations_to_server_config(allocations, server_port=8333): f"{list(instance['nodes'].keys())[0]}:{instance['port']}") server_config_entry = {'num_instances': num_servers, 'urls': urls} + if router_config: + server_config_entry['router'] = router_config.copy() if server_type == "GEN": generation_servers = server_config_entry @@ -484,7 +487,12 @@ def submit_job(config, log_dir, dry_run): json.dump(allocations, f, indent=2) # Generate disagg server config - server_config = convert_allocations_to_server_config(allocations) + router_config = config.get('router_config', None) + server_config = convert_allocations_to_server_config( + allocations, router_config=router_config) + # Merge server_config_extra into disagg server config + if 'server_config_extra' in config: + server_config.update(config['server_config_extra']) with open(os.path.join(log_dir, "server_config_base.yaml"), "w") as f: yaml.dump(server_config, f) disagg_server_hostname = server_config['hostname'] @@ -606,7 +614,14 @@ def submit_job(config, log_dir, dry_run): benchmark_prefix = client_slurm_prefix + [ f"--export \"{convert_envs_to_str(env_var)}\"" ] - if benchmark_config['use_nv_sa_benchmark']: + if benchmark_config.get('use_aiperf', False): + benchmark_cmd = [ + f"bash {os.path.join(script_dir, 'run_benchmark_aiperf.sh')}", + f"'{env_config['model_path']}' '{benchmark_config['dataset_file']}' {benchmark_config['multi_round']} {gen_num} '{benchmark_config['concurrency_list']}' {benchmark_config['streaming']} '{log_dir}' {disagg_server_hostname} {disagg_server_port} {ucx_warmup_requests}", + f"&> {log_dir}/6_bench.log" + ] + client_cmds.append(" ".join(benchmark_prefix + benchmark_cmd)) + elif benchmark_config['use_nv_sa_benchmark']: if benchmark_config['mode'] == "gen_only": print( f"[ERROR] SA benchmark client script is not supported for gen_only mode" diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index c64eae2e3f6f..94280234b710 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -644,7 +644,27 @@ def __init__(self, self._max_batch_size = max_batch_size self._tokens_per_block = tokens_per_block + def _get_tokenizer(self, model: str): + if model not in self._tokenizers: + self._tokenizers[model] = AutoTokenizer.from_pretrained(model) + return self._tokenizers[model] + def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: + # Handle ChatCompletionRequest (has messages, not prompt) + if isinstance(request, ChatCompletionRequest): + if request.prompt_token_ids is not None: + return [request.prompt_token_ids] + # TODO: send tokenize-only request instead of tokenizing locally + tokenizer = self._get_tokenizer(request.model) + messages = [{"role": m["role"], "content": m.get("content", "")} + for m in request.messages + if "role" in m] + text = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False) + token_ids = tokenizer.encode(text, add_special_tokens=False) + return [token_ids] + + # Handle CompletionRequest (has prompt) prompts = request.prompt if isinstance(prompts, list) and isinstance(prompts[0], list): return prompts @@ -656,10 +676,7 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: assert isinstance(prompts, list) and isinstance(prompts[0], str) # TODO: send tokenize-only request instead of tokenizing locally - if request.model not in self._tokenizers: - self._tokenizers[request.model] = AutoTokenizer.from_pretrained( - request.model) - tokenizer = self._tokenizers[request.model] + tokenizer = self._get_tokenizer(request.model) return [tokenizer(prompt)["input_ids"] for prompt in prompts] async def get_next_server( From 3f947684cc02189359ae9ce5a56d2c7924906ca2 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 18 Mar 2026 21:06:46 -0700 Subject: [PATCH 2/7] [None][fix] Add run_benchmark_aiperf.sh for aiperf-based disagg benchmarking Script used by submit.py when use_aiperf=true. Installs the correct aiperf version with trust_remote_code support and runs mooncake_trace dataset benchmarks against the disagg serving endpoint. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../slurm/benchmark/run_benchmark_aiperf.sh | 102 ++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100755 examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh diff --git a/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh b/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh new file mode 100755 index 000000000000..1f2d19b18417 --- /dev/null +++ b/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh @@ -0,0 +1,102 @@ +#!/bin/bash + +# aiperf-based benchmark script for disaggregated serving +# Args: model_name dataset_file multi_round num_gen_servers concurrency_list streaming log_path hostname port ucx_warmup_requests + +set -e +set -u +trap 'echo "Error occurred at line $LINENO"; exit 1' ERR + +if [ "$#" -lt 10 ]; then + echo "Error: Missing required arguments, got $# arguments, args: $@" + echo "Usage: $0 model_name dataset_file multi_round num_gen_servers concurrency_list streaming log_path hostname port ucx_warmup_requests" + exit 1 +fi + +model_name=$1 +dataset_file=$2 +multi_round=$3 +num_gen_servers=$4 +concurrency_list=$5 +streaming=$6 +log_path=$7 +hostname=$8 +port=$9 +ucx_warmup_requests=${10} + +# check process id is not 0 +if [[ ${SLURM_PROCID} != "0" ]]; then + echo "Process id is ${SLURM_PROCID} for loadgen, exiting" + exit 0 +fi + +# Always install/upgrade aiperf to ensure we have the version with trust_remote_code fix +# (container may have an older version with parallel_decode.py that lacks trust_remote_code) +echo "Installing aiperf..." +pip install --force-reinstall --no-deps 'aiperf @ git+https://github.com/ai-dynamo/aiperf.git@ac3d91652e5e024bfb4ac38d48603423aad666bc' + +# warmup requests for ucx connections +if [ "${ucx_warmup_requests}" -gt 0 ]; then + echo "warming up ucx connections with small requests... ${ucx_warmup_requests}" + python -m tensorrt_llm.serve.scripts.benchmark_serving \ + --model ${model_name} \ + --dataset-name random \ + --random-ids \ + --random-input-len 100 \ + --random-output-len 10 \ + --num-prompts ${ucx_warmup_requests} \ + --host ${hostname} \ + --port ${port} \ + --ignore-eos \ + --trust-remote-code \ + --non-streaming + echo "UCX warmup done" +fi + +# Trust remote code globally for custom tokenizers in parallel workers +export HF_HUB_TRUST_REMOTE_CODE=1 + +echo "Hostname: ${hostname}, Port: ${port}" +echo "Starting aiperf benchmark..." + +concurrency_list=$(echo "${concurrency_list}" | tr ',' ' ') +for concurrency in ${concurrency_list}; do + concurrency=$((concurrency)) + request_count=$((concurrency * multi_round)) + # benchmark_duration: 20min per round + benchmark_duration=$((multi_round * 1200)) + echo "Benchmarking with concurrency ${concurrency} ... ${request_count} requests, duration ${benchmark_duration}s" + mkdir -p ${log_path}/concurrency_${concurrency} + + aiperf profile \ + -m ${model_name} \ + --tokenizer ${model_name} \ + --tokenizer-trust-remote-code \ + --url http://${hostname}:${port} \ + --streaming \ + --ui simple \ + --input-file ${dataset_file} \ + --artifact-dir ${log_path}/concurrency_${concurrency} \ + --concurrency ${concurrency} \ + --concurrency-ramp-duration 60 \ + --custom-dataset-type mooncake_trace \ + --benchmark-duration ${benchmark_duration} \ + --benchmark-grace-period 60 \ + --workers-max 200 \ + --request-timeout-seconds 1200 \ + --profile-export-level records \ + --extra-inputs ignore_eos:true \ + --request-count ${request_count} \ + --record-processors 8 + + echo "Benchmark with concurrency ${concurrency} done" +done + +# Fetch perf metrics from disagg server +echo "Fetching perf metrics from http://${hostname}:${port}/perf_metrics ..." +curl -s "http://${hostname}:${port}/perf_metrics" > ${log_path}/perf_metrics.json 2>&1 || true +if [ -s "${log_path}/perf_metrics.json" ]; then + echo "Perf metrics saved to ${log_path}/perf_metrics.json" +else + echo "Warning: perf_metrics response was empty or endpoint not available" +fi From eb45a0c3211fdb14f114c7daebbca5353bf6815a Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Wed, 18 Mar 2026 21:35:04 -0700 Subject: [PATCH 3/7] [None][fix] Set token IDs on request after router tokenization to avoid re-tokenization KvCacheAwareRouter now sets prompt_token_ids (ChatCompletionRequest) or replaces prompt with token IDs (CompletionRequest) after tokenizing, so the downstream worker server skips redundant tokenization. Also adds proper ChatCompletionRequest handling via apply_chat_template. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/router.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 94280234b710..f7225d291d2b 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -654,14 +654,17 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: if isinstance(request, ChatCompletionRequest): if request.prompt_token_ids is not None: return [request.prompt_token_ids] - # TODO: send tokenize-only request instead of tokenizing locally tokenizer = self._get_tokenizer(request.model) - messages = [{"role": m["role"], "content": m.get("content", "")} - for m in request.messages - if "role" in m] - text = tokenizer.apply_chat_template( - messages, add_generation_prompt=True, tokenize=False) - token_ids = tokenizer.encode(text, add_special_tokens=False) + token_ids = tokenizer.apply_chat_template( + [ + msg if isinstance(msg, dict) else dict(msg) + for msg in request.messages + ], + add_generation_prompt=request.add_generation_prompt, + tokenize=True, + ) + # Set prompt_token_ids so the worker server skips re-tokenization + request.prompt_token_ids = token_ids return [token_ids] # Handle CompletionRequest (has prompt) @@ -675,9 +678,12 @@ def _tokenize(self, request: OpenAIRequest) -> list[list[int]]: else: assert isinstance(prompts, list) and isinstance(prompts[0], str) - # TODO: send tokenize-only request instead of tokenizing locally tokenizer = self._get_tokenizer(request.model) - return [tokenizer(prompt)["input_ids"] for prompt in prompts] + token_lists = [tokenizer(prompt)["input_ids"] for prompt in prompts] + # Replace string prompts with token IDs so the worker server + # skips re-tokenization + request.prompt = token_lists if len(token_lists) > 1 else token_lists[0] + return token_lists async def get_next_server( self, From 4ed6c89101cc315f4566efc01af1f767816063f6 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 19 Mar 2026 19:43:36 -0700 Subject: [PATCH 4/7] [None][test] Add multi-turn conversation routing test for KvCacheAwareRouter Test that consecutive turns of a multi-turn conversation are routed to the same server due to KV cache prefix hits. Simulates two concurrent sessions with disjoint token sequences inspired by agentic_data traces, verifying that accumulated context prefixes correctly match cached blocks across 3 turns. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/disaggregated/test_router.py | 146 ++++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 73310f5f03ab..c1fd4b18f88f 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -319,6 +319,152 @@ async def test_kv_cache_aware_router(servers): assert servers == ("server3", "server1", "server2") +@pytest.mark.asyncio +async def test_kv_cache_aware_router_multi_turn_conversation(): + """Test that consecutive turns of a multi-turn conversation route to the + same server due to KV cache prefix hits. + + Simulates two concurrent sessions inspired by + agentic_data/dataset_sample2000.jsonl session sess-fca58a1f44cd: + Turn 0: 68 hash_ids (system prompt + first user input) + Turn 1: 9 hash_ids (second user input, accumulated with turn 0) + Turn 2: 6 hash_ids (third user input, accumulated with turn 1) + + Scaled down to 10, 3, 2 blocks for test manageability. Each hash_id + maps to a deterministic block of tokens (mirroring aiperf's + HashIdRandomGenerator). The router should prefer the server that + already caches the conversation's prefix. + """ + server_list = ["server1", "server2", "server3"] + tokens_per_block = 32 + + router = KvCacheAwareRouter( + server_role=None, + servers=server_list, + use_tokens=False, + max_batch_size=64, + tokens_per_block=tokens_per_block, + ) + + # -- helpers ---------------------------------------------------------- + def hash_id_to_block(hash_id: int) -> list[int]: + """Deterministic token block per hash_id (mirrors aiperf corpus sampling).""" + return [(hash_id * 7 + i) % 50000 for i in range(tokens_per_block)] + + def build_tokens(hash_ids: list[int]) -> list[int]: + tokens = [] + for hid in hash_ids: + tokens.extend(hash_id_to_block(hid)) + # Append one extra token so the last full block is included in hashing. + # (KvCacheManager excludes the very last token from block keys.) + tokens.append(0) + return tokens + + # -- dataset-inspired hash_ids per turn (new blocks only) ------------- + # Session A (the conversation under test) + sess_a_turn0_hids = list(range(10)) # 10 blocks + sess_a_turn1_hids = list(range(100, 103)) # 3 new blocks + sess_a_turn2_hids = list(range(200, 202)) # 2 new blocks + + # Session B (competing traffic on a different server) + sess_b_turn0_hids = list(range(500, 510)) # 10 completely different blocks + + # -- build accumulated token sequences -------------------------------- + # Turn 0: just the first turn's tokens + sess_a_turn0_tokens = build_tokens(sess_a_turn0_hids) + + # Turn 1 accumulated: turn 0 tokens + simulated assistant reply + new user tokens + assistant_reply_tokens = [9999] * (tokens_per_block * 2) # 2 blocks of reply + sess_a_turn1_tokens = build_tokens( + sess_a_turn0_hids + [9990, 9991] + sess_a_turn1_hids + ) + # (hash_ids 9990/9991 stand in for the assistant-reply blocks) + + # Turn 2 accumulated: extends turn 1 further + sess_a_turn2_tokens = build_tokens( + sess_a_turn0_hids + [9990, 9991] + sess_a_turn1_hids + [9992, 9993] + + sess_a_turn2_hids + ) + + sess_b_tokens = build_tokens(sess_b_turn0_hids) + + # -- Round 1: initial routing (empty caches) -------------------------- + # Route both sessions concurrently so load-balancing spreads them to + # different servers (with equal KV cache misses, ties are broken by load). + req_a0 = CompletionRequest(model="TinyLlama", prompt=[sess_a_turn0_tokens]) + server_a, info_a0 = await router.get_next_server(req_a0) + # Do NOT finish req_a0 yet — keep its load active so session B avoids server_a + + req_b0 = CompletionRequest(model="TinyLlama", prompt=[sess_b_tokens]) + server_b, info_b0 = await router.get_next_server(req_b0) + + # Now finish both and populate caches + await router.finish_request(req_a0) + await router.finish_request(req_b0) + router._server_state[server_a].add_blocks(info_a0["block_hashes"][0]) + router._server_state[server_b].add_blocks(info_b0["block_hashes"][0]) + + # Sanity: two sessions should land on different servers + assert server_a != server_b, "Disjoint sessions should land on different servers" + + # Verify block hashes are disjoint between sessions + blocks_a = set(info_a0["block_hashes"][0]) + blocks_b = set(info_b0["block_hashes"][0]) + assert blocks_a.isdisjoint(blocks_b), "Different sessions must not share block hashes" + + # -- Round 2: turn 1 of session A (prefix extends turn 0) ------------ + req_a1 = CompletionRequest(model="TinyLlama", prompt=[sess_a_turn1_tokens]) + server_a1, info_a1 = await router.get_next_server(req_a1) + await router.finish_request(req_a1) + + assert server_a1 == server_a, ( + f"Turn 1 must route to the same server as turn 0 ({server_a}) " + f"due to KV cache prefix hit, but got {server_a1}. " + f"Matches: {info_a1['matches']}" + ) + + # The match count on server_a must equal the prefix overlap + server_a_idx = list(router._server_state.keys()).index(server_a) + expected_prefix_match = len(sess_a_turn0_hids) * tokens_per_block + assert info_a1["matches"][server_a_idx] == expected_prefix_match, ( + f"Expected {expected_prefix_match} matched tokens on server_a, " + f"got {info_a1['matches'][server_a_idx]}" + ) + + # Update server_a cache with new blocks from turn 1 + router._server_state[server_a].add_blocks(info_a1["block_hashes"][0]) + + # -- Round 3: turn 2 of session A (prefix extends turn 1) ------------ + req_a2 = CompletionRequest(model="TinyLlama", prompt=[sess_a_turn2_tokens]) + server_a2, info_a2 = await router.get_next_server(req_a2) + await router.finish_request(req_a2) + + assert server_a2 == server_a, ( + f"Turn 2 must route to the same server as turns 0-1 ({server_a}) " + f"due to KV cache prefix hit, but got {server_a2}. " + f"Matches: {info_a2['matches']}" + ) + + # Turn 2 should match all of turn 0 + turn 1 prefix blocks + expected_full_match = ( + len(sess_a_turn0_hids) + 2 + len(sess_a_turn1_hids) # turn0 + reply + turn1 + ) * tokens_per_block + assert info_a2["matches"][server_a_idx] == expected_full_match, ( + f"Expected {expected_full_match} matched tokens on turn 2, " + f"got {info_a2['matches'][server_a_idx]}" + ) + + # -- Verify session B still routes to its own server ------------------ + req_b1 = CompletionRequest(model="TinyLlama", prompt=[sess_b_tokens]) + server_b1, info_b1 = await router.get_next_server(req_b1) + await router.finish_request(req_b1) + + assert server_b1 == server_b, ( + f"Session B should route to its original server ({server_b}), " + f"but got {server_b1}" + ) + + def test_create_router(servers): default_router = create_router(None, servers) assert isinstance(default_router, RoundRobinRouter) From aed61da7d1d1ae9ec429b44728fc8f6399cb629e Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 19 Mar 2026 19:54:12 -0700 Subject: [PATCH 5/7] [None][test] Parameterize multi-turn router test for completion and chat APIs Add api_type parametrize ("completion", "chat") to test_kv_cache_aware_router_multi_turn_conversation so it exercises both CompletionRequest (prompt with token IDs) and ChatCompletionRequest (prompt_token_ids) code paths through the KvCacheAwareRouter. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tests/unittest/disaggregated/test_router.py | 30 ++++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index c1fd4b18f88f..65251bf40aeb 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -320,7 +320,8 @@ async def test_kv_cache_aware_router(servers): @pytest.mark.asyncio -async def test_kv_cache_aware_router_multi_turn_conversation(): +@pytest.mark.parametrize("api_type", ["completion", "chat"]) +async def test_kv_cache_aware_router_multi_turn_conversation(api_type): """Test that consecutive turns of a multi-turn conversation route to the same server due to KV cache prefix hits. @@ -360,6 +361,22 @@ def build_tokens(hash_ids: list[int]) -> list[int]: tokens.append(0) return tokens + def make_request(token_ids: list[int]): + """Create a CompletionRequest or ChatCompletionRequest with pre-tokenized IDs.""" + if api_type == "completion": + return CompletionRequest(model="TinyLlama", + prompt=[token_ids]) + else: + # Use prompt_token_ids to skip tokenizer (no real model needed) + return ChatCompletionRequest( + model="TinyLlama", + messages=[{ + "role": "user", + "content": "dummy" + }], + prompt_token_ids=token_ids, + ) + # -- dataset-inspired hash_ids per turn (new blocks only) ------------- # Session A (the conversation under test) sess_a_turn0_hids = list(range(10)) # 10 blocks @@ -374,7 +391,6 @@ def build_tokens(hash_ids: list[int]) -> list[int]: sess_a_turn0_tokens = build_tokens(sess_a_turn0_hids) # Turn 1 accumulated: turn 0 tokens + simulated assistant reply + new user tokens - assistant_reply_tokens = [9999] * (tokens_per_block * 2) # 2 blocks of reply sess_a_turn1_tokens = build_tokens( sess_a_turn0_hids + [9990, 9991] + sess_a_turn1_hids ) @@ -391,11 +407,11 @@ def build_tokens(hash_ids: list[int]) -> list[int]: # -- Round 1: initial routing (empty caches) -------------------------- # Route both sessions concurrently so load-balancing spreads them to # different servers (with equal KV cache misses, ties are broken by load). - req_a0 = CompletionRequest(model="TinyLlama", prompt=[sess_a_turn0_tokens]) + req_a0 = make_request(sess_a_turn0_tokens) server_a, info_a0 = await router.get_next_server(req_a0) # Do NOT finish req_a0 yet — keep its load active so session B avoids server_a - req_b0 = CompletionRequest(model="TinyLlama", prompt=[sess_b_tokens]) + req_b0 = make_request(sess_b_tokens) server_b, info_b0 = await router.get_next_server(req_b0) # Now finish both and populate caches @@ -413,7 +429,7 @@ def build_tokens(hash_ids: list[int]) -> list[int]: assert blocks_a.isdisjoint(blocks_b), "Different sessions must not share block hashes" # -- Round 2: turn 1 of session A (prefix extends turn 0) ------------ - req_a1 = CompletionRequest(model="TinyLlama", prompt=[sess_a_turn1_tokens]) + req_a1 = make_request(sess_a_turn1_tokens) server_a1, info_a1 = await router.get_next_server(req_a1) await router.finish_request(req_a1) @@ -435,7 +451,7 @@ def build_tokens(hash_ids: list[int]) -> list[int]: router._server_state[server_a].add_blocks(info_a1["block_hashes"][0]) # -- Round 3: turn 2 of session A (prefix extends turn 1) ------------ - req_a2 = CompletionRequest(model="TinyLlama", prompt=[sess_a_turn2_tokens]) + req_a2 = make_request(sess_a_turn2_tokens) server_a2, info_a2 = await router.get_next_server(req_a2) await router.finish_request(req_a2) @@ -455,7 +471,7 @@ def build_tokens(hash_ids: list[int]) -> list[int]: ) # -- Verify session B still routes to its own server ------------------ - req_b1 = CompletionRequest(model="TinyLlama", prompt=[sess_b_tokens]) + req_b1 = make_request(sess_b_tokens) server_b1, info_b1 = await router.get_next_server(req_b1) await router.finish_request(req_b1) From 24ea90e8a73f462d82452db89e896add7c1963b8 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Thu, 19 Mar 2026 20:03:27 -0700 Subject: [PATCH 6/7] [None][feat] Add TRTLLM_KVCACHE_AWARE_ROUTER_HASH_TOKENS_PER_BLOCK env var Allow overriding the KvCacheAwareRouter tokens_per_block via environment variable. When set, the env var takes precedence over the constructor default (32) and the YAML config value. The effective value is logged at router initialization. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- tensorrt_llm/serve/router.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index f7225d291d2b..03349cae0bfc 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -1,5 +1,6 @@ import asyncio import heapq +import os from abc import ABC, abstractmethod from typing import Awaitable, Callable, Dict, Iterable, List, Optional, Union @@ -642,7 +643,14 @@ def __init__(self, self._tokenizers = {} # TODO: use max_num_tokens? per server? self._max_batch_size = max_batch_size + env_tokens_per_block = os.environ.get( + "TRTLLM_KVCACHE_AWARE_ROUTER_HASH_TOKENS_PER_BLOCK") + if env_tokens_per_block is not None: + tokens_per_block = int(env_tokens_per_block) self._tokens_per_block = tokens_per_block + logger.info( + f"KvCacheAwareRouter: tokens_per_block={self._tokens_per_block}" + ) def _get_tokenizer(self, model: str): if model not in self._tokenizers: From 34b505a4d7fb1756912af4fb32bf7ec7a25a3374 Mon Sep 17 00:00:00 2001 From: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> Date: Tue, 31 Mar 2026 17:53:13 -0700 Subject: [PATCH 7/7] [None][chore] Fix yapf formatting in router and benchmark files Co-Authored-By: Claude Opus 4.6 Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com> --- .../disaggregated/slurm/benchmark/submit.py | 3 +- tensorrt_llm/serve/router.py | 3 +- tests/unittest/disaggregated/test_router.py | 38 ++++++++----------- 3 files changed, 19 insertions(+), 25 deletions(-) diff --git a/examples/disaggregated/slurm/benchmark/submit.py b/examples/disaggregated/slurm/benchmark/submit.py index 68248539f443..cecf30fa9db3 100644 --- a/examples/disaggregated/slurm/benchmark/submit.py +++ b/examples/disaggregated/slurm/benchmark/submit.py @@ -116,7 +116,8 @@ def assign_servers( return allocations -def convert_allocations_to_server_config(allocations, server_port=8333, +def convert_allocations_to_server_config(allocations, + server_port=8333, router_config=None): generation_servers = {} context_servers = {} diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 03349cae0bfc..53685dbc0ad3 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -649,8 +649,7 @@ def __init__(self, tokens_per_block = int(env_tokens_per_block) self._tokens_per_block = tokens_per_block logger.info( - f"KvCacheAwareRouter: tokens_per_block={self._tokens_per_block}" - ) + f"KvCacheAwareRouter: tokens_per_block={self._tokens_per_block}") def _get_tokenizer(self, model: str): if model not in self._tokenizers: diff --git a/tests/unittest/disaggregated/test_router.py b/tests/unittest/disaggregated/test_router.py index 65251bf40aeb..c016be080e83 100644 --- a/tests/unittest/disaggregated/test_router.py +++ b/tests/unittest/disaggregated/test_router.py @@ -364,8 +364,7 @@ def build_tokens(hash_ids: list[int]) -> list[int]: def make_request(token_ids: list[int]): """Create a CompletionRequest or ChatCompletionRequest with pre-tokenized IDs.""" if api_type == "completion": - return CompletionRequest(model="TinyLlama", - prompt=[token_ids]) + return CompletionRequest(model="TinyLlama", prompt=[token_ids]) else: # Use prompt_token_ids to skip tokenizer (no real model needed) return ChatCompletionRequest( @@ -379,7 +378,7 @@ def make_request(token_ids: list[int]): # -- dataset-inspired hash_ids per turn (new blocks only) ------------- # Session A (the conversation under test) - sess_a_turn0_hids = list(range(10)) # 10 blocks + sess_a_turn0_hids = list(range(10)) # 10 blocks sess_a_turn1_hids = list(range(100, 103)) # 3 new blocks sess_a_turn2_hids = list(range(200, 202)) # 2 new blocks @@ -391,16 +390,14 @@ def make_request(token_ids: list[int]): sess_a_turn0_tokens = build_tokens(sess_a_turn0_hids) # Turn 1 accumulated: turn 0 tokens + simulated assistant reply + new user tokens - sess_a_turn1_tokens = build_tokens( - sess_a_turn0_hids + [9990, 9991] + sess_a_turn1_hids - ) + sess_a_turn1_tokens = build_tokens(sess_a_turn0_hids + [9990, 9991] + + sess_a_turn1_hids) # (hash_ids 9990/9991 stand in for the assistant-reply blocks) # Turn 2 accumulated: extends turn 1 further - sess_a_turn2_tokens = build_tokens( - sess_a_turn0_hids + [9990, 9991] + sess_a_turn1_hids + [9992, 9993] - + sess_a_turn2_hids - ) + sess_a_turn2_tokens = build_tokens(sess_a_turn0_hids + [9990, 9991] + + sess_a_turn1_hids + [9992, 9993] + + sess_a_turn2_hids) sess_b_tokens = build_tokens(sess_b_turn0_hids) @@ -426,7 +423,8 @@ def make_request(token_ids: list[int]): # Verify block hashes are disjoint between sessions blocks_a = set(info_a0["block_hashes"][0]) blocks_b = set(info_b0["block_hashes"][0]) - assert blocks_a.isdisjoint(blocks_b), "Different sessions must not share block hashes" + assert blocks_a.isdisjoint( + blocks_b), "Different sessions must not share block hashes" # -- Round 2: turn 1 of session A (prefix extends turn 0) ------------ req_a1 = make_request(sess_a_turn1_tokens) @@ -436,16 +434,14 @@ def make_request(token_ids: list[int]): assert server_a1 == server_a, ( f"Turn 1 must route to the same server as turn 0 ({server_a}) " f"due to KV cache prefix hit, but got {server_a1}. " - f"Matches: {info_a1['matches']}" - ) + f"Matches: {info_a1['matches']}") # The match count on server_a must equal the prefix overlap server_a_idx = list(router._server_state.keys()).index(server_a) expected_prefix_match = len(sess_a_turn0_hids) * tokens_per_block assert info_a1["matches"][server_a_idx] == expected_prefix_match, ( f"Expected {expected_prefix_match} matched tokens on server_a, " - f"got {info_a1['matches'][server_a_idx]}" - ) + f"got {info_a1['matches'][server_a_idx]}") # Update server_a cache with new blocks from turn 1 router._server_state[server_a].add_blocks(info_a1["block_hashes"][0]) @@ -458,17 +454,16 @@ def make_request(token_ids: list[int]): assert server_a2 == server_a, ( f"Turn 2 must route to the same server as turns 0-1 ({server_a}) " f"due to KV cache prefix hit, but got {server_a2}. " - f"Matches: {info_a2['matches']}" - ) + f"Matches: {info_a2['matches']}") # Turn 2 should match all of turn 0 + turn 1 prefix blocks expected_full_match = ( - len(sess_a_turn0_hids) + 2 + len(sess_a_turn1_hids) # turn0 + reply + turn1 + len(sess_a_turn0_hids) + 2 + + len(sess_a_turn1_hids) # turn0 + reply + turn1 ) * tokens_per_block assert info_a2["matches"][server_a_idx] == expected_full_match, ( f"Expected {expected_full_match} matched tokens on turn 2, " - f"got {info_a2['matches'][server_a_idx]}" - ) + f"got {info_a2['matches'][server_a_idx]}") # -- Verify session B still routes to its own server ------------------ req_b1 = make_request(sess_b_tokens) @@ -477,8 +472,7 @@ def make_request(token_ids: list[int]): assert server_b1 == server_b, ( f"Session B should route to its original server ({server_b}), " - f"but got {server_b1}" - ) + f"but got {server_b1}") def test_create_router(servers):