From 9cebdb6b9b67e58fc50c0b8babe71edc51652124 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sun, 24 May 2026 14:47:51 +0800 Subject: [PATCH] Fix multi-turn replay fidelity: track realized output tokens across all components 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) --- replayer/replay.py | 104 ++++++-- scripts/bench.sh | 11 +- scripts/cache_aware_proxy.py | 244 ++++++++++++++---- scripts/launch_elastic_p2p.sh | 1 + .../v1/mooncake/mooncake_connector.py | 29 ++- 5 files changed, 312 insertions(+), 77 deletions(-) diff --git a/replayer/replay.py b/replayer/replay.py index cc36aca..6b80723 100644 --- a/replayer/replay.py +++ b/replayer/replay.py @@ -84,6 +84,12 @@ class _SessionState: metrics: list[RequestMetrics] = field(default_factory=list) +@dataclass +class _DispatchResult: + metric: RequestMetrics + output_token_ids: list[int] + + _endpoint_counter = 0 @@ -103,14 +109,17 @@ async def _dispatch_request( req: TraceRequest, prompt_token_ids: list[int], sem: asyncio.Semaphore, -) -> RequestMetrics: +) -> _DispatchResult: """Send one request via /v1/completions (streaming) and collect metrics.""" endpoint = _pick_endpoint(config) + target_output_tokens = max(1, req.output_length) payload = { "model": config.model_name, "prompt": prompt_token_ids, - "max_tokens": max(1, req.output_length), + "max_tokens": target_output_tokens, + "min_tokens": target_output_tokens, "temperature": 0, + "return_token_ids": True, "stream": True, "stream_options": {"include_usage": True}, } @@ -122,6 +131,7 @@ async def _dispatch_request( finish_reason = None err = None token_times: list[float] = [] + output_token_ids: list[int] = [] req_headers = {"X-Session-Id": req.session_id} @@ -146,14 +156,24 @@ async def _dispatch_request( except json.JSONDecodeError: continue - now = time.perf_counter() - if ttft_s is None: - ttft_s = now - start - choices = chunk.get("choices", []) if choices: + now = time.perf_counter() delta = choices[0].get("text", "") - if delta: + chunk_token_ids = choices[0].get("token_ids") + if isinstance(chunk_token_ids, list): + clean_ids = [ + int(t) for t in chunk_token_ids + if isinstance(t, int) + ] + if clean_ids: + if ttft_s is None: + ttft_s = now - start + output_token_ids.extend(clean_ids) + token_times.extend([now] * len(clean_ids)) + elif delta: + if ttft_s is None: + ttft_s = now - start token_times.append(now) fr = choices[0].get("finish_reason") if fr: @@ -168,8 +188,15 @@ async def _dispatch_request( end = time.perf_counter() e2e = end - start - if n_output == 0 and token_times: + if output_token_ids: + n_output = len(output_token_ids) + elif n_output == 0 and token_times: n_output = len(token_times) + if err is None and n_output != target_output_tokens: + err = ( + "output_token_mismatch " + f"requested={target_output_tokens} actual={n_output}" + ) tpot = 0.0 if len(token_times) > 1: @@ -177,26 +204,42 @@ async def _dispatch_request( for i in range(len(token_times) - 1)] tpot = sum(inter_token) / len(inter_token) - return RequestMetrics( - request_id=req.request_id, - session_id=req.session_id, - turn_id=req.turn_id, - trace_timestamp_s=req.timestamp_s, - input_length=req.input_length, - output_length=req.output_length, - request_type=req.request_type, - effective_input_length=len(prompt_token_ids), - cached_tokens=cached_tokens, - latency_s=e2e, - ttft_s=ttft_s, - tpot_s=tpot, - actual_output_tokens=n_output, - requested_output_tokens=req.output_length, - finish_reason=finish_reason, - error=err, + return _DispatchResult( + metric=RequestMetrics( + request_id=req.request_id, + session_id=req.session_id, + turn_id=req.turn_id, + trace_timestamp_s=req.timestamp_s, + input_length=req.input_length, + output_length=req.output_length, + request_type=req.request_type, + effective_input_length=len(prompt_token_ids), + cached_tokens=cached_tokens, + latency_s=e2e, + ttft_s=ttft_s, + tpot_s=tpot, + actual_output_tokens=n_output, + requested_output_tokens=req.output_length, + finish_reason=finish_reason, + error=err, + ), + output_token_ids=output_token_ids, ) +def _apply_realized_prefix( + prompt_token_ids: list[int], + realized_context: list[int], +) -> list[int]: + """Replace the reusable session prefix with engine-realized tokens.""" + if not realized_context: + return prompt_token_ids + out = prompt_token_ids.copy() + n = min(len(out), len(realized_context)) + out[:n] = realized_context[:n] + return out + + def _extract_cached_tokens(usage: dict) -> int: ct = 0 details = usage.get("prompt_tokens_details") @@ -220,6 +263,7 @@ async def _run_session( ) -> list[RequestMetrics]: if session_sem is not None: await session_sem.acquire() + realized_context: list[int] = [] try: for req in state.turns: # Wait until this request's trace timestamp @@ -228,13 +272,19 @@ async def _run_session( if elapsed < target_wall: await asyncio.sleep(target_wall - elapsed) - token_ids = _build_prompt_token_ids(req) - metric = await _dispatch_request( + token_ids = _apply_realized_prefix( + _build_prompt_token_ids(req), + realized_context, + ) + result = await _dispatch_request( client=client, config=config, req=req, prompt_token_ids=token_ids, sem=request_sem, ) + metric = result.metric state.metrics.append(metric) await sink.append(metric) + if metric.error is None: + realized_context = token_ids + result.output_token_ids finally: if session_sem is not None: session_sem.release() diff --git a/scripts/bench.sh b/scripts/bench.sh index 771cdf0..77c6013 100755 --- a/scripts/bench.sh +++ b/scripts/bench.sh @@ -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." diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index 7535aeb..64f2a0d 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -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))) diff --git a/scripts/launch_elastic_p2p.sh b/scripts/launch_elastic_p2p.sh index 6291f90..f5e63b3 100755 --- a/scripts/launch_elastic_p2p.sh +++ b/scripts/launch_elastic_p2p.sh @@ -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 diff --git a/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index f30a181..c8d905e 100644 --- a/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/third_party/vllm/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -249,6 +249,10 @@ class MooncakeConnector(KVConnectorBase_V1): assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_block_ids_with_load_errors(self) -> set[int]: + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: assert self.connector_worker is not None assert isinstance(self._connector_metadata, MooncakeConnectorMetadata) @@ -512,10 +516,13 @@ class MooncakeConnectorScheduler: # TODO: check whether block_ids actually ever be 0. If not we could # remove the conditional below - delay_free_blocks = len(block_ids) > 0 + block_size = self.vllm_config.cache_config.block_size + prompt_blocks = (request.num_prompt_tokens + block_size - 1) // block_size + send_block_ids = block_ids[:prompt_blocks] + delay_free_blocks = len(send_block_ids) > 0 if delay_free_blocks: - self._reqs_need_send[request.request_id] = (request, block_ids) + self._reqs_need_send[request.request_id] = (request, send_block_ids) return delay_free_blocks, None @@ -620,6 +627,7 @@ class MooncakeConnectorWorker: self.finished_sending_reqs: set[ReqId] = set() self.finished_recving_reqs: set[ReqId] = set() + self.failed_recving_block_ids: set[int] = set() self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config @@ -1063,6 +1071,11 @@ class MooncakeConnectorWorker: self.finished_recving_reqs = set() return finished_recving_reqs + def get_block_ids_with_load_errors(self) -> set[int]: + failed = self.failed_recving_block_ids + self.failed_recving_block_ids = set() + return failed + async def fetch_finished_sending_reqs(self) -> set[ReqId]: finished_sending_reqs = self.finished_sending_reqs self.finished_sending_reqs = set() @@ -1176,6 +1189,10 @@ class MooncakeConnectorWorker: logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.") except Exception as e: logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e) + for req_id in req_ids: + pull_meta = pull_metas[req_id] + self.failed_recving_block_ids.update(pull_meta.local_block_ids) + self.finished_recving_reqs.add(pull_meta.d_req_id) return def process_pulling_result( @@ -1201,6 +1218,12 @@ class MooncakeConnectorWorker: response.err_reqs, response.err_msg, ) + for req_id in response.err_reqs: + pull_meta = pull_metas.get(req_id) + if pull_meta is None: + continue + self.failed_recving_block_ids.update(pull_meta.local_block_ids) + self.finished_recving_reqs.add(pull_meta.d_req_id) async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str): url = remote_bootstrap_addr + "/query" @@ -1322,11 +1345,13 @@ class MooncakeConnectorWorker: logger.info("direct_push %s: %d blocks pushed from C", req_id, matched) else: logger.debug("direct_push %s: %d matched, pushed=%s", req_id, matched, pushed) + self.failed_recving_block_ids.update(local_block_ids) self.finished_recving_reqs.add(req_id) except Exception as e: logger.error("direct_push %s failed: %s", req_id, e) + self.failed_recving_block_ids.update(pm.local_block_ids) self.finished_recving_reqs.add(req_id) async def _start_load_kv(