diff --git a/benchmarks/di-dsv4/README.md b/benchmarks/di-dsv4/README.md new file mode 100644 index 000000000000..ce972f28008b --- /dev/null +++ b/benchmarks/di-dsv4/README.md @@ -0,0 +1,30 @@ +# DSv4-Flash benchmark harness + +Scripts used to evaluate dynamo disagg / KV migration / engine configs against +the production fleet. + +| script | purpose | +| --- | --- | +| `migtest.sh` | aiperf sweep vs a second fleet; warms both endpoints first | +| `mirror_h2h.sh` | mirror the same prod shards to two fleets, compare N windows | +| `compare.sh` | one fleet vs prod, N windows | +| `h2h.sh` | synthetic 7v7 aiperf | +| `kvxfer_force.py` | force a KV migration and check output determinism at temp=0 | +| `kvxfer_cross.py` | migration-vs-cold-prefill crossover sweep | +| `summarize.py` | median-per-window summary of compare/mirror logs | + +## Hard-won rules + +1. **Warm before the first measured point.** The first aiperf concurrency level + after a fleet boot reads ~5.8x slow (1.53 vs 8.90 req/s on identical config). + This produced three false "engine regression" verdicts. +2. **Never trust pod logs for rates.** The container log holds only ~900 lines, + so `--since=1h` and `--since=1m` return nearly the same count. Use + `dynamo_frontend_requests_total{status,error_type}` instead. +3. **Preflight both endpoints.** Port-forwards die when a pod rolls, and the + harness silently records `req/s=0.00` rather than failing. +4. **Give each synthetic prompt a distinct first 256 tokens**, or they collide + on one session id and the seed request 400s. +5. **Mirror the same shards to both candidates.** Synthetic traffic that is too + uniform/cacheable flatters KV-aware routing; a shard-sliced mirror is a + load-based sample, so both sides must get the identical slice. diff --git a/benchmarks/di-dsv4/compare.sh b/benchmarks/di-dsv4/compare.sh new file mode 100644 index 000000000000..2b367f73165e --- /dev/null +++ b/benchmarks/di-dsv4/compare.sh @@ -0,0 +1,27 @@ +set -u +LABEL="$1"; WINDOWS="${2:-7}"; GAP="${3:-300}" +cd /data/home/pernekhan/backend +export PATH=$HOME/miniconda3/bin:$PATH +PY=$HOME/miniconda3/envs/di-main/bin/python +D='deepseek-ai/DeepSeek-V4-Flash-0731-roce-disagg' +P='deepseek-ai/DeepSeek-V4-Flash-0731' + +q() { $PY -m scripts.cli vm-query --instant --label='' --query "$1" 2>/dev/null | tail -1 | awk -F, '{print $NF+0}'; } +tq() { q "histogram_quantile($1, sum(rate(vllm:time_to_first_token_seconds_bucket{model_name=\"$2\"}[5m])) by (le))"; } +hit() { q "sum(rate(vllm:prefix_cache_hits_total{model_name=\"$1\"$2}[5m]))/sum(rate(vllm:prefix_cache_queries_total{model_name=\"$1\"$2}[5m]))"; } +gen() { q "sum(rate(vllm:generation_tokens_total{model_name=\"$1\"}[5m]))"; } +req() { q "sum(rate(vllm:request_success_total{model_name=\"$1\"}[5m]))"; } + +echo "CONFIG=$LABEL windows=$WINDOWS gap=${GAP}s" +for i in $(seq 1 "$WINDOWS"); do + echo "W$i t=$(date -u +%H:%M)" \ + "hitD=$(hit "$D" ',dynamo_component="backend"')" \ + "hitP=$(hit "$P" '')" \ + "t50D=$(tq 0.50 "$D") t90D=$(tq 0.90 "$D") t99D=$(tq 0.99 "$D")" \ + "t50P=$(tq 0.50 "$P") t90P=$(tq 0.90 "$P") t99P=$(tq 0.99 "$P")" \ + "genD=$(gen "$D") genP=$(gen "$P")" \ + "reqD=$(req "$D") reqP=$(req "$P")" \ + "eng=$(q "count(vllm:num_requests_running{model_name=\"$P\"})")" + [ "$i" -lt "$WINDOWS" ] && sleep "$GAP" +done +echo "DONE_$LABEL" diff --git a/benchmarks/di-dsv4/h2h.sh b/benchmarks/di-dsv4/h2h.sh new file mode 100644 index 000000000000..dc7456552b36 --- /dev/null +++ b/benchmarks/di-dsv4/h2h.sh @@ -0,0 +1,48 @@ +set -u +T=/home/pernekhan/.claude/jobs/3aed5ed4/tmp +export PATH=$HOME/miniconda3/bin:$PATH +TOKDIR=$T/ds4tok/tokenizers/di--deepseek-ai--DeepSeek-V4-Flash-0731--D9jmKKxS +KEY=$(cat $T/.tok) +OUT=$T/h2h; mkdir -p $OUT + +# prod-shaped: long shared prefixes (cacheable) + unique tail, ~300 output tokens +PREFIX_N=64 +PREFIX_LEN=6000 +IN_MEAN=3000 +IN_STD=2000 +OSL=300 + +run() { + NAME=$1; MODEL=$2; URL=$3; C=$4; N=$5; shift 5 + D=$OUT/${NAME}_c${C} + rm -rf $D + timeout 2400 conda run -n di-main --no-capture-output aiperf profile \ + -m "$MODEL" --url "$URL" --endpoint-type chat --streaming \ + --tokenizer "$TOKDIR" --tokenizer-trust-remote-code \ + --num-prefix-prompts $PREFIX_N --prefix-prompt-length $PREFIX_LEN \ + --synthetic-input-tokens-mean $IN_MEAN --synthetic-input-tokens-stddev $IN_STD \ + --output-tokens-mean $OSL --output-tokens-stddev 0 \ + --concurrency "$C" --request-count "$N" --num-warmup-requests 8 \ + --random-seed 8800 --output-artifact-dir "$D" "$@" >/dev/null 2>&1 + python3 - "$D/profile_export_aiperf.json" "$NAME" "$C" <<'PY' +import json,sys +try: d=json.load(open(sys.argv[1])) +except Exception: + print(f" {sys.argv[2]:<8} c={sys.argv[3]:<4} FAILED"); raise SystemExit +g=lambda k,f='avg': (d.get(k) or {}).get(f) or 0 +print(f" {sys.argv[2]:<8} c={sys.argv[3]:<4} req/s={g('request_throughput'):6.2f} " + f"TTFT p50={g('time_to_first_token','p50'):7.0f} p90={g('time_to_first_token','p90'):7.0f} p99={g('time_to_first_token','p99'):8.0f}ms " + f"ITL p50={g('inter_token_latency','p50'):5.1f}ms out_tok/s={g('output_token_throughput'):7.0f}") +PY +} + +echo "=== 7 GPU vs 7 GPU, identical synthetic traffic (seed 8800) ===" +echo " prefix pool=$PREFIX_N x ${PREFIX_LEN}tok, unique tail mean=$IN_MEAN, out=$OSL" +for C in 8 24 48 96; do + N=$(( C * 12 )); [ $N -lt 200 ] && N=200 + run vllm7 "Pernekhan/DeepSeek-V4-Flash-0731-test" "http://localhost:18002" $C $N \ + --custom-endpoint /v1/openai/chat/completions --api-key "$KEY" + run dynamo7 "deepseek-ai/DeepSeek-V4-Flash-0731-roce-disagg" "http://localhost:18001" $C $N + echo "" +done +echo H2H_COMPLETE diff --git a/benchmarks/di-dsv4/kvxfer_cross.py b/benchmarks/di-dsv4/kvxfer_cross.py new file mode 100644 index 000000000000..3758541f4bac --- /dev/null +++ b/benchmarks/di-dsv4/kvxfer_cross.py @@ -0,0 +1,33 @@ +import json, urllib.request, time, sys +URL="http://localhost:80/v1/chat/completions"; MODEL="deepseek-ai/DeepSeek-V4-Flash-0731-roce-disagg" +A,B,REPS = sys.argv[1], sys.argv[2], int(sys.argv[3]) +UNIT="Timing probe segment for decode to decode key value transfer measurement. " +def body_of(prompt): + return {"model":MODEL,"temperature":0.0,"max_tokens":16, + "messages":[{"role":"user","content":prompt}]} +def post(prompt, worker): + r=urllib.request.Request(URL,data=json.dumps(body_of(prompt)).encode(), + headers={"Content-Type":"application/json"}) + if worker: r.add_header("x-dynamo-worker-instance-id",worker) + t=time.time() + try: + with urllib.request.urlopen(r,timeout=600) as resp: d=json.load(resp) + except Exception as e: return time.time()-t,{"error":str(e)} + return time.time()-t,d +def mk(tag,mult): return (f"Probe {tag} segment for decode to decode transfer measurement. "*mult) + "\n\nQuestion: say ok." +for mult,label in ((1200,"21k"),(3200,"67k")): + cold=[]; mig=[] + for i in range(REPS): + stamp=f"{label}-{i}-{int(time.time())}" + el,d=post(mk("cold-"+stamp,mult),B) + pt=(d.get("usage") or {}).get("prompt_tokens") + cold.append(el); print(f"COLD {label} rep{i} sec={el:.3f} tokens={pt} err={d.get('error')}"); sys.stdout.flush() + p=mk("mig-"+stamp,mult) + el2,d2=post(p,A); print(f"SEED {label} rep{i} sec={el2:.3f} tokens={(d2.get('usage') or {}).get('prompt_tokens')} err={d2.get('error')}"); sys.stdout.flush() + time.sleep(75) + el3,d3=post(p,B) + mig.append(el3); print(f"MIGR {label} rep{i} sec={el3:.3f} tokens={(d3.get('usage') or {}).get('prompt_tokens')} err={d3.get('error')}"); sys.stdout.flush() + if cold and mig: + import statistics as st + c=st.median(cold); m=st.median(mig) + print(f"RESULT {label} cold med={c:.3f} min={min(cold):.3f} | migrated med={m:.3f} min={min(mig):.3f} | speedup_med={c/m:.2f}x speedup_min={min(cold)/min(mig):.2f}x") diff --git a/benchmarks/di-dsv4/kvxfer_force.py b/benchmarks/di-dsv4/kvxfer_force.py new file mode 100644 index 000000000000..49b8ff2c1fe7 --- /dev/null +++ b/benchmarks/di-dsv4/kvxfer_force.py @@ -0,0 +1,39 @@ +import json, urllib.request, time, sys + +URL = "http://localhost:80/v1/chat/completions" +MODEL = "deepseek-ai/DeepSeek-V4-Flash-0731-roce-disagg" +A, B, WAIT = sys.argv[1], sys.argv[2], int(sys.argv[3]) +MULT = int(sys.argv[4]) if len(sys.argv) > 4 else 260 +FILLER = ("The migration harness verifies decode-to-decode key value transfer. " * MULT) +PROMPT = FILLER + "\n\nQuestion: Reply with exactly the five words: alpha bravo charlie delta echo." + +def post(prompt, worker, max_tokens=64): + body = {"model": MODEL, "temperature": 0.0, "max_tokens": max_tokens, + "messages": [{"role": "user", "content": prompt}]} + req = urllib.request.Request(URL, data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + if worker: + req.add_header("x-dynamo-worker-instance-id", worker) + try: + with urllib.request.urlopen(req, timeout=180) as r: + return json.load(r) + except Exception as e: + return {"error": str(e)} + +def summarize(tag, d): + ch = d.get("choices", [{}])[0].get("message", {}).get("content") + u = d.get("usage", {}) or {} + print(tag, json.dumps({ + "out": ch, + "prompt_tokens": u.get("prompt_tokens"), + "cached": (u.get("prompt_tokens_details") or {}).get("cached_tokens"), + "completion_tokens": u.get("completion_tokens"), + "err": d.get("error"), + })) + return ch + +r1 = post(PROMPT, A); o1 = summarize("SEED_A", r1) +print("WAIT", WAIT); sys.stdout.flush() +time.sleep(WAIT) +r2 = post(PROMPT, B); o2 = summarize("FORCED_B", r2) +print("IDENTICAL", json.dumps(o1 is not None and o1 == o2)) diff --git a/benchmarks/di-dsv4/migtest.sh b/benchmarks/di-dsv4/migtest.sh new file mode 100644 index 000000000000..f20a900afdaa --- /dev/null +++ b/benchmarks/di-dsv4/migtest.sh @@ -0,0 +1,90 @@ +set -u +T=/home/pernekhan/.claude/jobs/3aed5ed4/tmp +export PATH=$HOME/miniconda3/bin:$PATH +PY=$HOME/miniconda3/envs/di-main/bin/python +TOKDIR=$T/ds4tok/tokenizers/di--deepseek-ai--DeepSeek-V4-Flash-0731--D9jmKKxS +KEY=$(cat $T/.tok) +OUT=$T/migtest; mkdir -p $OUT +NS=deepinfra +D='deepseek-ai/DeepSeek-V4-Flash-0731-roce-disagg' +V='Pernekhan/DeepSeek-V4-Flash-0731-test' + +# Prompts must clear the migration gate: overlap_blocks*256 >= 16384 tokens. +# 12 shared prefixes x 24k tokens, ~2k unique tail => ~26k prompts, ~24k cacheable. +PREFIX_N=12 +PREFIX_LEN=24000 +IN_MEAN=2000 +OSL=300 + +fe() { kubectl -n $NS get pods --no-headers | grep 'roce-disagg-fron' | grep -v Terminating | head -1 | awk '{print $1}'; } + +run() { + NAME=$1; MODEL=$2; URL=$3; C=$4; N=$5; shift 5 + DIR=$OUT/${NAME}_c${C}; rm -rf $DIR + timeout 3600 conda run -n di-main --no-capture-output aiperf profile \ + -m "$MODEL" --url "$URL" --endpoint-type chat --streaming \ + --tokenizer "$TOKDIR" --tokenizer-trust-remote-code \ + --num-prefix-prompts $PREFIX_N --prefix-prompt-length $PREFIX_LEN \ + --synthetic-input-tokens-mean $IN_MEAN --synthetic-input-tokens-stddev 500 \ + --output-tokens-mean $OSL --output-tokens-stddev 0 \ + --concurrency "$C" --request-count "$N" --num-warmup-requests 6 \ + --random-seed 4242 --output-artifact-dir "$DIR" "$@" >/dev/null 2>&1 + $PY - "$DIR/profile_export_aiperf.json" "$NAME" "$C" <<'PY' +import json,sys +try: d=json.load(open(sys.argv[1])) +except Exception: + print(f" {sys.argv[2]:<9} c={sys.argv[3]:<4} FAILED"); raise SystemExit +g=lambda k,f='avg': (d.get(k) or {}).get(f) or 0 +print(f" {sys.argv[2]:<9} c={sys.argv[3]:<4} req/s={g('request_throughput'):6.2f} " + f"TTFT p50={g('time_to_first_token','p50'):7.0f} p90={g('time_to_first_token','p90'):7.0f} p99={g('time_to_first_token','p99'):8.0f}ms " + f"ITL={g('inter_token_latency','p50'):5.1f}ms tok/s={g('output_token_throughput'):7.0f}") +PY +} + +# Warm both fleets before the first measured point. Without this the first +# concurrency level absorbs all CUDA-graph/JIT warmup and reads ~5x slow -- +# it produced three bogus "engine regression" verdicts before it was caught. +warmup() { + timeout 900 conda run -n di-main --no-capture-output aiperf profile \ + -m "$2" --url "$3" --endpoint-type chat --streaming \ + --tokenizer "$TOKDIR" --tokenizer-trust-remote-code \ + --num-prefix-prompts $PREFIX_N --prefix-prompt-length $PREFIX_LEN \ + --synthetic-input-tokens-mean $IN_MEAN --output-tokens-mean $OSL \ + --concurrency 16 --request-count 64 --num-warmup-requests 4 \ + --random-seed 4242 --output-artifact-dir "$OUT/warm_$1" "${@:4}" >/dev/null 2>&1 + echo " warmed $1" +} +warmup dynamo7 "$D" "http://localhost:18041" +warmup vllm7 "$V" "http://localhost:18042" --custom-endpoint /v1/openai/chat/completions --api-key "$KEY" +echo "WARMUP_DONE" + +echo "=== 7 dynamo GPU vs 7 standalone vLLM GPU, no mirror, seed 4242 ===" +echo " prompts: ${PREFIX_N} shared prefixes x ${PREFIX_LEN} tok + ~${IN_MEAN} tail (clears the 16384 gate)" +echo "" + +for C in 24 48 128; do + N=$(( C * 8 )); [ $N -lt 150 ] && N=150 + + # capture the dynamo frontend log for the whole dynamo run so migration + # counts are complete (container log holds only ~900 lines) + FP=$(fe) + LOG=$OUT/fe_c${C}.log + ( kubectl -n $NS logs -f "$FP" --since=5s > "$LOG" 2>/dev/null ) & + TAILPID=$! + sleep 3 + run dynamo7 "$D" "http://localhost:18041" $C $N + sleep 5 + kill $TAILPID 2>/dev/null; wait $TAILPID 2>/dev/null + + SP=$(grep -c 'KVMIGRATE_SPILL' "$LOG" 2>/dev/null || echo 0) + HS=$(grep -c 'handshake obtained' "$LOG" 2>/dev/null || echo 0) + MG=$(grep -c 'KVMIGRATE: decoding on target' "$LOG" 2>/dev/null || echo 0) + PF=$(grep -c 'producer dispatch failed' "$LOG" 2>/dev/null || echo 0) + echo " migration@c${C}: spills=$SP handshakes=$HS migrated=$MG producer_failed=$PF (of $N requests)" + grep -oE 'migratable_tokens=[0-9]+' "$LOG" 2>/dev/null | head -3 | sed 's/^/ sample /' + + run vllm7 "$V" "http://localhost:18042" $C $N \ + --custom-endpoint /v1/openai/chat/completions --api-key "$KEY" + echo "" +done +echo MIGTEST_COMPLETE diff --git a/benchmarks/di-dsv4/mirror_h2h.sh b/benchmarks/di-dsv4/mirror_h2h.sh new file mode 100644 index 000000000000..9dca05718c58 --- /dev/null +++ b/benchmarks/di-dsv4/mirror_h2h.sh @@ -0,0 +1,29 @@ +set -u +LABEL="${1:-run}"; WINDOWS="${2:-7}"; GAP="${3:-300}" +cd /data/home/pernekhan/backend +export PATH=$HOME/miniconda3/bin:$PATH +PY=$HOME/miniconda3/envs/di-main/bin/python +D='deepseek-ai/DeepSeek-V4-Flash-0731-roce-disagg' +V='Pernekhan/DeepSeek-V4-Flash-0731-test' + +q() { $PY -m scripts.cli vm-query --instant --label='' --query "$1" 2>/dev/null | tail -1 | awk -F, '{print $NF+0}'; } +tq() { q "histogram_quantile($1, sum(rate(vllm:time_to_first_token_seconds_bucket{model_name=\"$2\"}[5m])) by (le))"; } +hit() { q "sum(rate(vllm:prefix_cache_hits_total{model_name=\"$1\"}[5m]))/sum(rate(vllm:prefix_cache_queries_total{model_name=\"$1\"}[5m]))"; } +gen() { q "sum(rate(vllm:generation_tokens_total{model_name=\"$1\"}[5m]))"; } +req() { q "sum(rate(vllm:request_success_total{model_name=\"$1\"}[5m]))"; } +eng() { q "count(vllm:num_requests_running{model_name=\"$1\"})"; } +kv() { q "avg(vllm:kv_cache_usage_perc{model_name=\"$1\"})"; } + +echo "CONFIG=$LABEL windows=$WINDOWS (same 7 prod shards mirrored to both)" +for i in $(seq 1 "$WINDOWS"); do + echo "W$i t=$(date -u +%H:%M)" \ + "hitD=$(hit "$D") hitV=$(hit "$V")" \ + "t50D=$(tq 0.50 "$D") t90D=$(tq 0.90 "$D") t99D=$(tq 0.99 "$D")" \ + "t50V=$(tq 0.50 "$V") t90V=$(tq 0.90 "$V") t99V=$(tq 0.99 "$V")" \ + "genD=$(gen "$D") genV=$(gen "$V")" \ + "reqD=$(req "$D") reqV=$(req "$V")" \ + "engD=$(eng "$D") engV=$(eng "$V")" \ + "kvD=$(kv "$D") kvV=$(kv "$V")" + [ "$i" -lt "$WINDOWS" ] && sleep "$GAP" +done +echo "DONE_$LABEL" diff --git a/benchmarks/di-dsv4/summarize.py b/benchmarks/di-dsv4/summarize.py new file mode 100644 index 000000000000..869d4fe7f6a3 --- /dev/null +++ b/benchmarks/di-dsv4/summarize.py @@ -0,0 +1,61 @@ +import sys, statistics as st + +GPUS_DISAGG = 13 + +def load(path): + rows = [] + for line in open(path): + if not line.startswith('W'): + continue + d = {} + for kv in line.split(): + if '=' in kv: + k, v = kv.split('=', 1) + d[k] = v + try: + rows.append({k: float(v) for k, v in d.items() if k != 't'}) + except ValueError: + pass + return rows + +def med(rows, key): + vals = [r[key] for r in rows if key in r and r[key] > 0] + return st.median(vals) if vals else float('nan') + +def report(label, path): + rows = load(path) + if not rows: + print(f" {label}: no windows"); return None + eng = med(rows, 'eng') + out = { + 'n': len(rows), + 'hitD': med(rows, 'hitD'), 'hitP': med(rows, 'hitP'), + 't50D': med(rows, 't50D') * 1000, 't90D': med(rows, 't90D') * 1000, 't99D': med(rows, 't99D') * 1000, + 't50P': med(rows, 't50P') * 1000, 't90P': med(rows, 't90P') * 1000, 't99P': med(rows, 't99P') * 1000, + 'genD_gpu': med(rows, 'genD') / GPUS_DISAGG, 'genP_gpu': med(rows, 'genP') / eng, + 'reqD_gpu': med(rows, 'reqD') / GPUS_DISAGG, 'reqP_gpu': med(rows, 'reqP') / eng, + 'eng': eng, + } + print(f"\n === {label} (n={out['n']} windows, prod engines={eng:.0f}) ===") + print(f" {'metric':<22} {'disagg':>10} {'prod':>10} {'ratio':>9}") + def row(name, d, p, better_low=False, fmt="{:.0f}"): + r = (p / d) if better_low else (d / p) + mark = " <-- win" if r > 1.0 else "" + print(f" {name:<22} {fmt.format(d):>10} {fmt.format(p):>10} {r:>8.2f}x{mark}") + row('cache hit', out['hitD'], out['hitP'], fmt="{:.3f}") + row('TTFT p50 (ms)', out['t50D'], out['t50P'], better_low=True) + row('TTFT p90 (ms)', out['t90D'], out['t90P'], better_low=True) + row('TTFT p99 (ms)', out['t99D'], out['t99P'], better_low=True) + row('gen tok/s per GPU', out['genD_gpu'], out['genP_gpu']) + row('req/s per GPU', out['reqD_gpu'], out['reqP_gpu']) + return out + +a = report('A: queue gate ON', sys.argv[1]) +b = report('B: queue gate OFF', sys.argv[2]) if len(sys.argv) > 2 else None + +if a and b: + print("\n === A vs B (disagg only, median of windows) ===") + for k, name, low in (('t50D','TTFT p50',True), ('t90D','TTFT p90',True), ('t99D','TTFT p99',True), + ('hitD','cache hit',False), ('genD_gpu','gen tok/s/GPU',False), ('reqD_gpu','req/s/GPU',False)): + delta = (a[k] / b[k]) if low else (b[k] / a[k]) + print(f" {name:<16} A={a[k]:>10.3f} B={b[k]:>10.3f} B is {delta:.2f}x {'better' if delta>1 else 'worse'}") diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index b3e3e60eb5d6..f15a87d715ae 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/components/src/dynamo/vllm/handlers.py @@ -2740,6 +2740,8 @@ async def generate_tokens( request_id, data_parallel_rank=None, lora_request=None, + emit_kv_transfer_params: bool = False, + embedding_sequence_length=None, trace_headers=None, priority=0, reasoning_ended=None, @@ -2848,6 +2850,14 @@ async def generate_tokens( if finish_reason: out["finish_reason"] = normalize_finish_reason(finish_reason) + # KV migration source: publish the producer handshake so + # the target decode worker can pull these blocks. + if emit_kv_transfer_params: + kv_tp = getattr(res, "kv_transfer_params", None) + if kv_tp is not None: + out["disaggregated_params"] = { + "kv_transfer_params": kv_tp + } out[ "completion_usage" ] = BaseWorkerHandler._build_completion_usage( @@ -3150,6 +3160,31 @@ async def _generate_token_mode(self, request, context, request_id): logger.debug( f"Using disaggregated params from prefill for request {request_id}" ) + + # KV migration source (frontend sets kv_migration_source): this worker + # already holds the conversation's prefix, so when the router picks it + # as a migration SOURCE we behave like a prefill worker for one step -- + # emit producer kv_transfer_params and generate a single token -- so the + # TARGET decode worker can pull the blocks over NIXL instead of + # recomputing them. Requires kv_role=kv_both, which our decode workers + # already run. + if request.get("kv_migration_source"): + if sampling_params.extra_args is None: + sampling_params.extra_args = {} + sampling_params.extra_args["kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "remote_engine_id": None, + "remote_block_ids": None, + "remote_host": None, + "remote_port": None, + } + sampling_params.max_tokens = 1 + sampling_params.min_tokens = 1 + logger.info( + f"KV migration source: serving {request_id} as producer " + f"(1 token, blocks offered over NIXL)" + ) prefill_prompt_tokens_details = ( prefill_result.get("prompt_tokens_details") if prefill_result else None ) @@ -3215,6 +3250,7 @@ async def _generate_token_mode(self, request, context, request_id): request_id, data_parallel_rank=dp_rank, lora_request=lora_request, + emit_kv_transfer_params=bool(request.get("kv_migration_source")), trace_headers=trace_headers, priority=priority, reasoning_ended=reasoning_ended, diff --git a/dist/ai_dynamo_runtime-1.4.1-cp310-abi3-manylinux_2_35_x86_64.whl b/dist/ai_dynamo_runtime-1.4.1-cp310-abi3-manylinux_2_35_x86_64.whl new file mode 100644 index 000000000000..24011206ccf0 Binary files /dev/null and b/dist/ai_dynamo_runtime-1.4.1-cp310-abi3-manylinux_2_35_x86_64.whl differ diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 7ca5847f8b21..effefbecc7da 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -1456,9 +1456,15 @@ impl< None => (None, None), }; + let (best_overlap_worker, best_overlap_blocks) = match request.max_overlap_worker() { + Some((worker, blocks)) => (Some(worker), blocks), + None => (None, 0.0), + }; let response = SchedulingResponse { best_worker: selection.worker, effective_overlap_blocks: selection.effective_overlap_blocks, + best_overlap_worker, + best_overlap_blocks, cached_tokens: selection.cached_tokens, selected_worker_tiers, request_progress, diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index c33955635627..e36bfcbae5cb 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -74,6 +74,10 @@ impl KvSchedulerError { pub struct SchedulingResponse { pub best_worker: WorkerWithDpRank, pub effective_overlap_blocks: f64, + /// Worker holding the most of this request's blocks, selected or not. + /// Lets a caller find the blocks for any request, with no session concept. + pub best_overlap_worker: Option, + pub best_overlap_blocks: f64, pub cached_tokens: usize, pub selected_worker_tiers: SelectedWorkerTierSnapshot, pub request_progress: Option, @@ -273,6 +277,16 @@ impl SchedulingRequest { .unwrap_or(0) } + /// Argmax over the per-worker overlap map: (worker, blocks). + pub(crate) fn max_overlap_worker(&self) -> Option<(WorkerWithDpRank, f64)> { + self.overlap + .effective_overlap_blocks + .iter() + .filter(|(_, blocks)| **blocks > 0.0) + .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(worker, blocks)| (*worker, *blocks)) + } + pub(crate) fn effective_overlap_blocks_for(&self, worker: WorkerWithDpRank) -> f64 { self.overlap .effective_overlap_blocks diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index b6a187fa56b6..0903dee63cbe 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -136,6 +136,9 @@ pub enum FindBestMatchOutcome { effective_overlap_blocks: f64, cached_tokens: usize, routing_hashes: Option, + /// Worker holding the most of this request's blocks, selected or not. + best_overlap_worker: Option, + best_overlap_blocks: f64, }, QueueRejected { rejection: scheduling::QueueRejection, @@ -911,6 +914,8 @@ where effective_overlap_blocks: response.effective_overlap_blocks, cached_tokens: response.cached_tokens, routing_hashes, + best_overlap_worker: response.best_overlap_worker, + best_overlap_blocks: response.best_overlap_blocks, }, lifecycle, )) diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index afbba1a0b4e4..98c433524793 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -194,7 +194,7 @@ impl next: ServerStreamingEngine>, ) -> Result>> { // Extract request data while preserving context - let (mut req, context) = request.into_parts(); + let (mut req, mut context) = request.into_parts(); let request_id = context.id().to_string(); let metadata = context.metadata().clone(); let engine_ctx = context.context(); @@ -209,6 +209,131 @@ impl return next.generate(context.map(|_| req)).await; } + // Conditional disaggregation: prompts below DYN_MIN_REMOTE_PREFILL_TOKENS + // are served entirely on the decode worker (aggregate-style local + // prefill). Remote prefill only pays for itself when the prefill + // compute exceeds the handoff cost (route + prefill queue + KV + // transfer + decode admission); below that the handoff dominates + // TTFT. 0 (default) disables the gate. + static MIN_REMOTE_PREFILL_TOKENS: std::sync::OnceLock = + std::sync::OnceLock::new(); + let min_remote = *MIN_REMOTE_PREFILL_TOKENS.get_or_init(|| { + std::env::var("DYN_MIN_REMOTE_PREFILL_TOKENS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0) + }); + + // Derive a session-affinity key from the prompt's leading tokens when + // the caller did not send one. A conversation's rendered history grows + // by appending, so its first N tokens are identical across turns; the + // affinity map then pins every turn to the same decode worker, which + // is what makes its prefix cache reusable. Overlap-credit scoring + // cannot do this at high load: the decode-load term (thousands of + // blocks) dwarfs any per-request overlap credit. 0 (default) disables. + static SESSION_FROM_PREFIX_TOKENS: std::sync::OnceLock = + std::sync::OnceLock::new(); + let prefix_key_tokens = *SESSION_FROM_PREFIX_TOKENS.get_or_init(|| { + std::env::var("DYN_SESSION_FROM_PREFIX_TOKENS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0) + }); + // Hash a fixed-length prefix window, never `min(len, N)`: a shorter + // window would change as the conversation grows, so every turn would + // look like a brand new session and land on a worker that has none of + // its blocks. Requests below the window get no session id and fall back + // to plain KV-overlap routing. + if prefix_key_tokens > 0 + && req.token_ids.len() >= prefix_key_tokens + && context + .get_optional::(SESSION_AFFINITY_CONTEXT_KEY) + .ok() + .flatten() + .is_none() + { + use std::hash::{Hash, Hasher}; + let mut hasher = std::hash::DefaultHasher::new(); + req.token_ids[..prefix_key_tokens].hash(&mut hasher); + context.insert( + SESSION_AFFINITY_CONTEXT_KEY, + SessionAffinityId::new(format!("prefix-{:016x}", hasher.finish())), + ); + } + + // Uncached-length gating (DYN_LOCAL_CONTINUATIONS=1): a conversation + // continuation's history is already cached on its pinned decode + // worker, so only its NEW tail needs prefill -- serve it locally no + // matter how long the raw prompt is; only cold large prefills gain + // from a dedicated prefill worker. Approximated with a bounded + // recently-seen set keyed by the session prefix hash. Loss on restart + // is benign: a continuation goes remote once, then re-registers. + static LOCAL_CONTINUATIONS: std::sync::OnceLock = std::sync::OnceLock::new(); + static SEEN_SESSIONS: std::sync::OnceLock< + std::sync::Mutex>, + > = std::sync::OnceLock::new(); + let local_continuations = *LOCAL_CONTINUATIONS.get_or_init(|| { + std::env::var("DYN_LOCAL_CONTINUATIONS") + .map(|v| v == "1" || v.eq_ignore_ascii_case("true")) + .unwrap_or(false) + }); + let session_seen = if local_continuations && prefix_key_tokens > 0 { + match context.get_optional::(SESSION_AFFINITY_CONTEXT_KEY) { + Ok(Some(sid)) => { + use std::hash::{Hash, Hasher}; + let mut hasher = std::hash::DefaultHasher::new(); + sid.as_str().hash(&mut hasher); + let key = hasher.finish(); + let now = std::time::Instant::now(); + let mut seen = SEEN_SESSIONS + .get_or_init(Default::default) + .lock() + .expect("seen-session map poisoned"); + if seen.len() > 4_000_000 { + seen.clear(); + } + seen.insert(key, now) + .is_some_and(|t| now.duration_since(t).as_secs() < 3600) + } + _ => false, + } + } else { + false + }; + // KV migration (DYN_KV_MIGRATION_SAMPLE > 0): for a sampled fraction + // of continuation requests, run the prefill leg on the session's OLD + // pinned decode worker (whose prefix cache makes it nearly free) and + // decode on a different worker, moving the conversation's KV via the + // existing bidirectional NIXL path. Mechanism validation for + // load-relief-that-preserves-cache; production distress wiring + // replaces the random sample once the mechanism is proven. + static MIGRATION_SAMPLE: std::sync::OnceLock = std::sync::OnceLock::new(); + let migration_sample = *MIGRATION_SAMPLE.get_or_init(|| { + std::env::var("DYN_KV_MIGRATION_SAMPLE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0.0) + }); + let migrate = migration_sample > 0.0 + && session_seen + && self.lifecycle_state() == PrefillLifecycleState::Active + && rand::random::() < migration_sample; + if migrate { + // Telemetry-only for now: the D->D dispatch leg lands in the next + // increment (needs backend-instance targeting for the bootstrap + // handshake); this quantifies migration rate and candidate shape + // in production traffic without touching the data path. + tracing::info!( + request_id = %request_id, + tokens = req.token_ids.len(), + "KVMIGRATE_CANDIDATE: continuation would migrate (prefill on pinned worker, decode elsewhere)" + ); + } + + if min_remote > 0 && (req.token_ids.len() < min_remote || session_seen) { + return next.generate(context.map(|_| req)).await; + } + let session_affinity = context .get_optional::(SESSION_AFFINITY_CONTEXT_KEY) .map_err(|message| anyhow::anyhow!("invalid session affinity context: {message}"))?; @@ -306,10 +431,23 @@ impl error = %error, "request rejected by prefill worker (at capacity)" ); - } else { - tracing::error!(error = %error, "Remote prefill failed, failing request"); + return Err(error); + } + // Transient dispatch failures (e.g. the selected instance + // missing from the transport map during a discovery re-list) + // must not fail the request: the decode worker can always run + // the prefill locally, exactly like the passthrough taken when + // the prefill router is inactive. + tracing::warn!( + error = %error, + "Remote prefill dispatch failed, falling back to local prefill" + ); + if let Some(ref tracker) = req.tracker { + let _decode_permit = tracker.set_phase(RequestPhase::Decode).await; } - return Err(error); + let mut local_req = req; + local_req.stop_conditions.max_tokens = original_max_tokens; + return next.generate(context.map(|_| local_req)).await; } }; diff --git a/lib/llm/src/kv_router/push_router.rs b/lib/llm/src/kv_router/push_router.rs index 0311f46f1f77..87c994ecba68 100644 --- a/lib/llm/src/kv_router/push_router.rs +++ b/lib/llm/src/kv_router/push_router.rs @@ -108,6 +108,10 @@ pub struct KvPushRouter { pub chooser: Arc, request_metrics: Arc, affinity: Option, + /// session_id -> (worker that last served it, when). Observational only: + /// it never influences selection, it only names a migration source when the + /// router's KV index has not caught up with a recent placement. + kv_home: Arc>, } impl KvPushRouter { @@ -139,6 +143,7 @@ impl KvPushRouter { chooser, request_metrics, affinity, + kv_home: Arc::new(dashmap::DashMap::new()), } } @@ -325,6 +330,56 @@ impl KvPushRouter { Ok(guard) } + /// Record which worker just served this session, so a later turn that lands + /// elsewhere knows where the blocks are. Bounded: swept back to entries seen + /// in the last 15 minutes whenever it grows past the cap. + fn remember_kv_home(&self, session_id: &str, worker_id: u64) { + const KV_HOME_TTL: std::time::Duration = std::time::Duration::from_secs(900); + const KV_HOME_CAP: usize = 200_000; + let now = std::time::Instant::now(); + self.kv_home.insert(session_id.to_string(), (worker_id, now)); + if self.kv_home.len() > KV_HOME_CAP { + self.kv_home + .retain(|_, (_, seen)| now.duration_since(*seen) < KV_HOME_TTL); + } + } + + /// Run one producer step on `source` so it publishes a NIXL handshake for + /// this conversation's blocks. The caller hands the result to the worker + /// that will actually decode, which then pulls the blocks instead of + /// recomputing the prefix. Returns None if the source cannot produce one. + async fn kv_migration_handshake( + &self, + req: &PreprocessedRequest, + ctx_id: &str, + metadata: std::collections::BTreeMap, + source: u64, + ) -> Option { + let mut src_req = req.clone(); + src_req.kv_migration_source = Some(true); + src_req.stop_conditions.max_tokens = Some(1); + let src_ctx = dynamo_runtime::pipeline::Context::with_id_and_metadata( + src_req, + format!("{ctx_id}-kvmig"), + metadata, + ); + let mut stream = match self.inner.direct(src_ctx, source).await { + Ok(stream) => stream, + Err(error) => { + tracing::warn!(%error, source, "KVMIGRATE: producer dispatch failed"); + return None; + } + }; + let first = stream.next().await?; + let data = first.data.as_ref()?; + let disaggregated_params = data.disaggregated_params.clone()?; + tracing::info!(source, "KVMIGRATE: producer handshake obtained"); + Some(crate::protocols::common::preprocessor::PrefillResult { + disaggregated_params, + prompt_tokens_details: None, + }) + } + async fn dispatch_selection( &self, request: SingleIn, @@ -572,6 +627,83 @@ impl AsyncEngine, ManyOut = std::sync::OnceLock::new(); + let migration_mode = *KV_MIGRATION.get_or_init(|| { + match std::env::var("DYN_KV_MIGRATION").unwrap_or_default().as_str() { + "1" | "true" | "TRUE" => 2, + "dry" | "DRY" => 1, + _ => 0, + } + }); + static KV_MIGRATION_MIN_TOKENS: std::sync::OnceLock = std::sync::OnceLock::new(); + let migration_min_tokens = *KV_MIGRATION_MIN_TOKENS.get_or_init(|| { + std::env::var("DYN_KV_MIGRATION_MIN_TOKENS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(16384) + }); + let migratable_tokens = + (selection.best_overlap_blocks * self.chooser.block_size() as f64) as usize; + let request = if migration_mode > 0 + && phase != RequestPhase::Prefill + && migratable_tokens >= migration_min_tokens + { + let session_key = affinity_id(&request).ok().flatten(); + let cached_worker = selection + .best_overlap_worker + .or_else(|| { + session_key + .as_ref() + .and_then(|sid| self.kv_home.get(sid.as_str())) + .map(|e| e.value().0) + }) + .filter(|w| *w != selection.instance_id); + if let Some(sid) = session_key.as_ref() { + self.remember_kv_home(sid.as_str(), selection.instance_id); + } + match cached_worker { + Some(source) => { + let ctx_id = request.context().id().to_string(); + tracing::info!( + request_id = %ctx_id, + source, + target = selection.instance_id, + tokens = request.token_ids.len(), + migratable_tokens, + "KVMIGRATE_SPILL: request routed off the worker holding its KV" + ); + if migration_mode < 2 { + request + } else { + let metadata = request.metadata().clone(); + let (mut req, ctx) = request.into_parts(); + if req.prefill_result.is_none() { + if let Some(result) = self + .kv_migration_handshake(&req, &ctx_id, metadata, source) + .await + { + tracing::info!( + request_id = %ctx_id, source, target = selection.instance_id, + "KVMIGRATE: decoding on target with migrated KV" + ); + req.prefill_result = Some(result); + } + } + ctx.map(|_| req) + } + } + None => request, + } + } else { + request + }; + let stream = match self .dispatch_selection(request, selection, guard, operation.is_some()) .await diff --git a/lib/llm/src/kv_router/push_router/selection.rs b/lib/llm/src/kv_router/push_router/selection.rs index 72884b19d20f..b1d6efc43369 100644 --- a/lib/llm/src/kv_router/push_router/selection.rs +++ b/lib/llm/src/kv_router/push_router/selection.rs @@ -28,6 +28,10 @@ pub(super) struct WorkerSelection { pub(super) cached_tokens: usize, pub(super) routing_hashes: Option, pub(super) lifecycle: Option<(RequestProgressUpdater, RequestLifecycleLease)>, + /// Worker holding the most of this request's blocks, selected or not, and + /// how many. Drives KV migration with no session concept involved. + pub(super) best_overlap_worker: Option, + pub(super) best_overlap_blocks: f64, } #[derive(Clone, Copy)] @@ -101,7 +105,11 @@ impl KvPushRouter { effective_overlap_blocks, cached_tokens, routing_hashes, + best_overlap_worker, + best_overlap_blocks, } => Ok(WorkerSelection { + best_overlap_worker: best_overlap_worker.map(|w| w.worker_id), + best_overlap_blocks, instance_id: worker.worker_id, dp_rank: worker.dp_rank, overlap_amount: overlap_blocks, diff --git a/lib/llm/src/protocols/common/preprocessor.rs b/lib/llm/src/protocols/common/preprocessor.rs index 5e05e5751e16..d7493940a4db 100644 --- a/lib/llm/src/protocols/common/preprocessor.rs +++ b/lib/llm/src/protocols/common/preprocessor.rs @@ -244,6 +244,13 @@ pub struct PreprocessedRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub prefill_result: Option, + /// Ask this worker to act as a KV migration SOURCE: serve one token as a + /// producer and publish the NIXL handshake, so another decode worker can + /// pull the conversation's blocks instead of recomputing them. + #[builder(default)] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub kv_migration_source: Option, + /// Multimodal encoder handoff payload, set by the frontend when /// forwarding a request from an Encode worker to a downstream /// Prefill/Aggregated peer. Engine-opaque JSON object;