gahow/agentic-kvc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix multi-turn replay fidelity: track realized output tokens across all components

Browse Source 
The replayer and proxy were building multi-turn prompts from trace tokens,
but the model generates different output tokens. Subsequent turns had wrong
prefix tokens, causing cache misses and invalid experimental measurements.

- replay.py: min_tokens=max_tokens for deterministic length, return_token_ids
  to capture actual output, _apply_realized_prefix for next-turn correction
- proxy: extract output token_ids from SSE, record prompt+output as realized
  prefix in shadow cache, extract _handle_local_request to deduplicate
- bench.sh/launch_elastic_p2p.sh: default elastic mode to unified policy
- mooncake_connector: only send prompt blocks (not stale output blocks),
  track failed_recving_block_ids for error recovery

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-24 14:47:51 +08:00
parent cc4a9c91e7
commit 9cebdb6b9b
5 changed files with 312 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
scripts/bench.sh
View File

@@ -26,7 +26,8 @@ TRACE="${TRACE:-$PROJECT_DIR/traces/w600_r0.0015_st30.jsonl}"
# Defaults
TAG=""
MODE="baseline" # baseline | elastic
POLICY="linear" # linear | lmetric
POLICY="linear" # linear | lmetric | unified
POLICY_SET=false
N_INSTANCES=8
BASE_PORT=8000
PROXY_PORT=9090
@@ -44,7 +45,7 @@ while [[ $# -gt 0 ]]; do
case "$1" in
--tag) TAG="$2"; shift 2 ;;
--mode) MODE="$2"; shift 2 ;;
--policy) POLICY="$2"; shift 2 ;;
--policy) POLICY="$2"; POLICY_SET=true; shift 2 ;;
--instances) N_INSTANCES="$2"; shift 2 ;;
--requests) REQUESTS="$2"; shift 2 ;;
--trace) TRACE="$2"; shift 2 ;;
@@ -60,11 +61,15 @@ while [[ $# -gt 0 ]]; do
done
if [ -z "$TAG" ]; then
echo "Usage: bench.sh --tag NAME --mode {baseline|elastic} [--instances N] [--policy {linear|lmetric}] [--requests N]"
echo "Usage: bench.sh --tag NAME --mode {baseline|elastic} [--instances N] [--policy {linear|lmetric|unified}] [--requests N]"
echo " Trace QPS is controlled by sample_trace.py --sample-ratio, not by bench.sh."
exit 1
fi
if [ "$MODE" = "elastic" ] && [ "$POLICY_SET" = "false" ]; then
POLICY="unified"
fi
OUTDIR="$PROJECT_DIR/outputs/$TAG"
if [ -d "$OUTDIR" ] && [ -f "$OUTDIR/metrics.jsonl" ]; then
echo "[ERROR] Output directory $OUTDIR already exists with data. Use a different --tag."

244
scripts/cache_aware_proxy.py
View File

@@ -14,6 +14,7 @@ Routing policies (--policy):
import argparse
import asyncio
import json
import os
import time as _time
import urllib.parse
@@ -45,6 +46,10 @@ class Settings:
prefill_throughput: float = 7000.0 # tokens/s per GPU (measured on H20)
rdma_overhead_s: float = 0.1 # RDMA PUSH overhead (~10-50ms measured)
cache_capacity_blocks: int = 200000 # per-instance LRU cap on shadow cached_blocks
heavy_threshold: int = 20000
overload_factor: float = 2.0
max_offload_inflight: int = 4
cache_gate_ratio: float = 0.0
SETTINGS = Settings()
@@ -98,7 +103,7 @@ def _p_offload_penalty(inst: InstanceState) -> int:
"""Penalty for PD-sep mode routing (legacy)."""
if inst.active_p_offloads <= 0:
return 0
return inst.active_p_offloads * 20000
return inst.active_p_offloads * SETTINGS.heavy_threshold
def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
@@ -118,7 +123,7 @@ def pick_instance(instances: list[InstanceState], token_ids: list[int] | None,
idx = affinity[session_id]
if idx < len(instances):
inst = instances[idx]
if (inst.ongoing_tokens <= avg_load * 2.0
if (inst.ongoing_tokens <= avg_load * SETTINGS.overload_factor
and inst.active_p_offloads == 0):
return inst, idx
@@ -197,6 +202,54 @@ async def _query_bootstrap_hit(
return None
def _extract_output_token_ids_from_sse(
buffer: str,
chunk: bytes,
) -> tuple[str, list[int]]:
"""Extract vLLM streaming token_ids while preserving the raw stream."""
buffer += chunk.decode("utf-8", errors="ignore")
complete = buffer.endswith("\n") or buffer.endswith("\r")
lines = buffer.splitlines()
if complete:
buffer = ""
elif lines:
buffer = lines.pop()
else:
return buffer, []
output_ids: list[int] = []
for line in lines:
line = line.strip()
if not line.startswith("data:"):
continue
data = line[5:].strip()
if not data or data == "[DONE]":
continue
try:
payload = json.loads(data)
except json.JSONDecodeError:
continue
choices = payload.get("choices", [])
for choice in choices:
token_ids = choice.get("token_ids")
if isinstance(token_ids, list):
output_ids.extend(
int(t) for t in token_ids if isinstance(t, int)
)
return buffer, output_ids
def _realized_tokens(
prompt_token_ids: list[int] | None,
output_token_ids: list[int],
) -> list[int] | None:
if prompt_token_ids is None:
return None
if not output_token_ids:
return prompt_token_ids
return prompt_token_ids + output_token_ids
global_args = None
combined_instances: list[InstanceState] = []
prefill_instances: list[InstanceState] = []
@@ -367,6 +420,54 @@ async def _handle(request: Request, api: str):
input_length, session_id, headers)
async def _handle_local_request(api, req_data, headers, token_ids, input_length,
chosen: InstanceState, estimated_new: int,
breakdown: dict):
breakdown.setdefault("route_class", "LOCAL")
breakdown.setdefault("routed_to", chosen.url)
chosen.ongoing_tokens += input_length
chosen.pending_prefill_tokens += estimated_new
chosen.num_requests += 1
async def generate():
prefill_done = False
sse_buffer = ""
output_token_ids: list[int] = []
try:
for attempt in range(MAX_STREAM_RETRIES):
try:
async with chosen.client.stream("POST", api, json=req_data, headers=headers) as resp:
resp.raise_for_status()
async for chunk in resp.aiter_bytes():
sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
sse_buffer, chunk)
output_token_ids.extend(new_output_ids)
if not prefill_done:
chosen.pending_prefill_tokens -= estimated_new
chosen.ongoing_decode_tokens += input_length
breakdown["t_first_token"] = _time.monotonic()
prefill_done = True
yield chunk
chosen.record_prefix(
_realized_tokens(token_ids, output_token_ids))
break
except (httpx.ConnectError, httpx.RemoteProtocolError):
if prefill_done or attempt >= MAX_STREAM_RETRIES - 1:
raise
await asyncio.sleep(RETRY_DELAY_S)
finally:
if not prefill_done:
chosen.pending_prefill_tokens -= estimated_new
else:
chosen.ongoing_decode_tokens -= input_length
chosen.ongoing_tokens -= input_length
chosen.num_requests -= 1
breakdown["t_done"] = _time.monotonic()
_breakdown_log.append(breakdown)
return StreamingResponse(generate(), media_type="text/event-stream")
async def _handle_combined(api, req_data, token_ids, input_length, session_id, headers):
"""Unified routing: pick the instance with lowest expected latency.
@@ -375,13 +476,57 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
where prefill_time depends on whether the instance has cache (local),
can receive cache via PUSH (remote), or must do cold prefill.
"""
policy = getattr(global_args, 'policy', 'linear')
offload_enabled = getattr(global_args, 'offload', False) and len(combined_instances) >= 2
throughput = SETTINGS.prefill_throughput
if policy in ("linear", "lmetric"):
if policy == "lmetric":
chosen, best_idx = pick_instance_lmetric(
combined_instances, token_ids, session_id, input_length,
session_affinity_combined)
else:
chosen, best_idx = pick_instance(
combined_instances, token_ids, session_id, input_length,
session_affinity_combined)
cache_hit = chosen.estimate_cache_hit(token_ids)
estimated_new = max(0, input_length - cache_hit)
breakdown = {
"request_id": headers.get("X-Request-Id", ""),
"input_length": input_length,
"cache_hit": cache_hit,
"estimated_new_tokens": estimated_new,
"t_proxy_recv": _time.monotonic(),
"policy": policy,
"route_class": "LOCAL",
"routed_to": chosen.url,
}
if session_id and policy == "lmetric":
# LMetric is intentionally per-request; record last target only for
# stats/debugging, not for future decisions.
session_affinity_combined[session_id] = best_idx
return await _handle_local_request(
api, req_data, headers, token_ids, input_length,
chosen, estimated_new, breakdown)
# Compute cache hits for all instances
cache_hits = [inst.estimate_cache_hit(token_ids) for inst in combined_instances]
best_cache_idx = max(range(len(combined_instances)), key=lambda i: cache_hits[i])
best_cache_hit = cache_hits[best_cache_idx]
def _current_offloads() -> int:
return sum(i.active_p_offloads for i in combined_instances)
def _push_allowed(cache_hit: int) -> bool:
if _current_offloads() >= SETTINGS.max_offload_inflight:
return False
push_new = max(0, input_length - cache_hit)
if push_new < SETTINGS.heavy_threshold:
return False
if SETTINGS.cache_gate_ratio > 0:
cache_ratio = cache_hit / max(input_length, 1)
if cache_ratio < SETTINGS.cache_gate_ratio:
return False
return True
def _instance_cost(i: int) -> tuple[float, bool]:
"""Expected latency if this request goes to instance i."""
@@ -391,7 +536,8 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
local_new = max(0, input_length - local_hit)
local_cost = queue + local_new / throughput
if offload_enabled and best_cache_hit > 0 and i != best_cache_idx and local_hit < best_cache_hit:
if (offload_enabled and best_cache_hit > 0 and _push_allowed(best_cache_hit)
and i != best_cache_idx and local_hit < best_cache_hit):
push_new = max(0, input_length - best_cache_hit)
push_cost = queue + push_new / throughput + SETTINGS.rdma_overhead_s
if push_cost < local_cost:
@@ -406,7 +552,7 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
all_costs = [_instance_cost(i) for i in range(len(combined_instances))]
global_best_cost = min(c for c, _ in all_costs)
# Use affinity if it's within 2x of the best option
if affinity_cost <= global_best_cost * 2.0:
if affinity_cost <= global_best_cost * SETTINGS.overload_factor:
best_idx = affinity_idx
best_cost = affinity_cost
best_needs_push = affinity_push
@@ -428,6 +574,7 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
"cache_hit": cache_hit,
"estimated_new_tokens": estimated_new,
"t_proxy_recv": _time.monotonic(),
"policy": policy,
"chosen_cost": round(best_cost, 2),
}
@@ -451,6 +598,23 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
# If real hit > 0, proceed with offload
if push_cache_hit > 0:
push_new = max(0, input_length - push_cache_hit)
cache_ratio = push_cache_hit / max(input_length, 1)
if _current_offloads() >= SETTINGS.max_offload_inflight:
breakdown["push_downgraded"] = "cap_reached"
return await _handle_local_request(
api, req_data, headers, token_ids, input_length,
chosen, estimated_new, breakdown)
if push_new < SETTINGS.heavy_threshold:
breakdown["push_downgraded"] = "below_heavy_threshold"
return await _handle_local_request(
api, req_data, headers, token_ids, input_length,
chosen, estimated_new, breakdown)
if SETTINGS.cache_gate_ratio > 0 and cache_ratio < SETTINGS.cache_gate_ratio:
breakdown["push_downgraded"] = "cache_gate"
return await _handle_local_request(
api, req_data, headers, token_ids, input_length,
chosen, estimated_new, breakdown)
offload_mode = getattr(global_args, 'offload_mode', 'cached_prefill')
breakdown["c_inst"] = c_inst.url
@@ -482,42 +646,9 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
# LOCAL path (also handles downgraded PUSH)
breakdown["route_class"] = "LOCAL"
breakdown["routed_to"] = chosen.url
chosen.ongoing_tokens += input_length
chosen.pending_prefill_tokens += estimated_new
chosen.num_requests += 1
async def generate():
prefill_done = False
try:
for attempt in range(MAX_STREAM_RETRIES):
try:
async with chosen.client.stream("POST", api, json=req_data, headers=headers) as resp:
resp.raise_for_status()
async for chunk in resp.aiter_bytes():
if not prefill_done:
chosen.pending_prefill_tokens -= estimated_new
chosen.ongoing_decode_tokens += input_length
breakdown["t_first_token"] = _time.monotonic()
prefill_done = True
yield chunk
chosen.record_prefix(token_ids)
break
except (httpx.ConnectError, httpx.RemoteProtocolError):
if prefill_done or attempt >= MAX_STREAM_RETRIES - 1:
raise
await asyncio.sleep(RETRY_DELAY_S)
finally:
if not prefill_done:
chosen.pending_prefill_tokens -= estimated_new
else:
chosen.ongoing_decode_tokens -= input_length
chosen.ongoing_tokens -= input_length
chosen.num_requests -= 1
breakdown["t_done"] = _time.monotonic()
_breakdown_log.append(breakdown)
return StreamingResponse(generate(), media_type="text/event-stream")
return await _handle_local_request(
api, req_data, headers, token_ids, input_length,
chosen, estimated_new, breakdown)
PREFILL_TIMEOUT_S = 120 # max seconds to wait for P-instance prefill
@@ -544,6 +675,7 @@ async def _handle_cached_prefill_offload(api, req_data, headers, token_ids,
}
prefill_data["stream"] = False
prefill_data["max_tokens"] = 1
prefill_data["min_tokens"] = 1
prefill_data.pop("max_completion_tokens", None)
prefill_data.pop("stream_options", None)
p_headers = {**headers, "X-data-parallel-rank": "0"}
@@ -591,16 +723,21 @@ async def _handle_cached_prefill_offload(api, req_data, headers, token_ids,
async def generate():
first_token = True
sse_buffer = ""
output_token_ids: list[int] = []
try:
async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
resp.raise_for_status()
async for chunk in resp.aiter_bytes():
sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
sse_buffer, chunk)
output_token_ids.extend(new_output_ids)
if first_token:
d_inst.ongoing_decode_tokens += input_length
breakdown["t_first_token"] = _time.monotonic()
first_token = False
yield chunk
d_inst.record_prefix(token_ids)
d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
finally:
if not first_token:
d_inst.ongoing_decode_tokens -= input_length
@@ -641,17 +778,22 @@ async def _handle_direct_read_offload(api, req_data, headers, token_ids,
async def generate():
first_token = True
sse_buffer = ""
output_token_ids: list[int] = []
try:
async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
resp.raise_for_status()
async for chunk in resp.aiter_bytes():
sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
sse_buffer, chunk)
output_token_ids.extend(new_output_ids)
if first_token:
d_inst.pending_prefill_tokens -= estimated_new
d_inst.ongoing_decode_tokens += input_length
breakdown["t_first_token"] = _time.monotonic()
first_token = False
yield chunk
d_inst.record_prefix(token_ids)
d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
finally:
if first_token:
d_inst.pending_prefill_tokens -= estimated_new
@@ -688,6 +830,7 @@ async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
}
prefill_data["stream"] = False
prefill_data["max_tokens"] = 1
prefill_data["min_tokens"] = 1
prefill_data.pop("max_completion_tokens", None)
prefill_data.pop("stream_options", None)
p_headers = {**headers, "X-data-parallel-rank": "0"}
@@ -726,14 +869,20 @@ async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
async def generate():
first_token = True
sse_buffer = ""
output_token_ids: list[int] = []
try:
async with d_inst.client.stream("POST", api, json=decode_data, headers=headers) as resp:
resp.raise_for_status()
async for chunk in resp.aiter_bytes():
sse_buffer, new_output_ids = _extract_output_token_ids_from_sse(
sse_buffer, chunk)
output_token_ids.extend(new_output_ids)
if first_token:
breakdown["t_first_token"] = _time.monotonic()
first_token = False
yield chunk
d_inst.record_prefix(_realized_tokens(token_ids, output_token_ids))
finally:
breakdown["t_done"] = _time.monotonic()
d_inst.ongoing_tokens -= input_length
@@ -779,8 +928,9 @@ def parse_args():
help="Enable Mooncake KV offload for HEAVY requests (requires kv_both instances)")
p.add_argument("--bootstrap-ports", type=str, default="",
help="Comma-separated bootstrap ports for combined instances (for offload mode)")
p.add_argument("--policy", type=str, default="linear", choices=["linear", "lmetric"],
help="Routing policy: linear (default) or lmetric (P_tokens × BS, OSDI'26)")
p.add_argument("--policy", type=str, default="linear",
choices=["linear", "lmetric", "unified"],
help="Routing policy: linear, lmetric (P_tokens × BS), or unified cost model")
p.add_argument("--overload-factor", type=float, default=2.0,
help="Break session affinity when instance load > factor * avg")
p.add_argument("--max-offload-inflight", type=int, default=4,
@@ -789,7 +939,7 @@ def parse_args():
choices=["direct_read", "cached_prefill"],
help="direct_read: D pulls KV from C (PUSH). "
"cached_prefill: C prefills then D decodes (PD-sep style).")
p.add_argument("--cache-gate-ratio", type=float, default=0.3,
p.add_argument("--cache-gate-ratio", type=float, default=0.0,
help="Min cache_hit/input ratio to allow offload "
"(0.0 disables gate, 1.0 disables offload entirely)")
args = p.parse_args()
@@ -809,6 +959,10 @@ def parse_args():
if __name__ == "__main__":
global_args = parse_args()
SETTINGS.heavy_threshold = global_args.heavy_threshold
SETTINGS.overload_factor = global_args.overload_factor
SETTINGS.max_offload_inflight = global_args.max_offload_inflight
SETTINGS.cache_gate_ratio = global_args.cache_gate_ratio
print("SETTINGS: throughput=%.0f rdma_overhead=%.2f offload=%s" % (
SETTINGS.prefill_throughput, SETTINGS.rdma_overhead_s,
getattr(global_args, 'offload', False)))

1
scripts/launch_elastic_p2p.sh
View File

@@ -90,6 +90,7 @@ $PYTHON "$PROJECT_DIR/scripts/cache_aware_proxy.py" \
--combined $combined_args \
--bootstrap-ports "$bootstrap_ports" \
--offload \
--policy unified \
--heavy-threshold $HEAVY_THRESHOLD \
--port $PROXY_PORT &
sleep 5