From f739f7d4617ad148f56ce00bffeecc5d2a0dd52d Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 27 May 2026 22:05:19 +0800 Subject: [PATCH] Proxy/runner support for Nixl connector + unified_v3 (offload-decode) policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit scripts/b3_isolated_policy.sh: Recognize unified_v3 as a kv_both-requiring policy; respect explicit KV_CONNECTOR=Nixl override (so unified_v2 / unified_v3 / unified_kv_both can run against either Mooncake or Nixl back-end). When Nixl is selected, skip the bootstrap-ports plumbing — Nixl uses its own UCX side-channel and the proxy forwards kv_transfer_params from the src response body instead of pre-baking engine_id/bootstrap_addr. scripts/cache_aware_proxy.py: - New unified_v3 policy (~250 lines): prefill stays on session-affinity host (preserves intra-session prefix-cache reuse), decode is migrated to a lower-load target when the affinity host is busy with concurrent decodes. KV transfer flows prefill_host → decode_target, opposite of v2. Knobs: v3_min_new_tokens, v3_min_prefill_decode_busy, v3_target_load_ratio, v3_min_load_gap, v3_rotate_affinity, v3_prefer_cache_target. cache_miss_audit found rotation hurts cross- turn locality (9.5% hit with vs ~80% without) so default v3_rotate_affinity=False. - New connector_type setting ("mooncake" | "nixl") gating the PD-sep handshake form: mooncake uses pre-baked kv_transfer_params, nixl forwards them from the response body. Co-Authored-By: Claude Opus 4.7 --- scripts/b3_isolated_policy.sh | 17 +- scripts/cache_aware_proxy.py | 303 ++++++++++++++++++++++++++++++---- 2 files changed, 283 insertions(+), 37 deletions(-) diff --git a/scripts/b3_isolated_policy.sh b/scripts/b3_isolated_policy.sh index 08a888c..9634af0 100755 --- a/scripts/b3_isolated_policy.sh +++ b/scripts/b3_isolated_policy.sh @@ -32,10 +32,14 @@ RUNDIR="${3:?usage: $0 }" # Auto-enable kv_both when the policy requires it. # KV_CONNECTOR (Mooncake|Nixl) selects the underlying connector when KV_BOTH=1. +_KV_CONNECTOR_EXPLICIT="${KV_CONNECTOR:-}" KV_CONNECTOR="${KV_CONNECTOR:-Mooncake}" -if [ "$POLICY" = "unified_v2" ] || [ "$POLICY" = "unified_kv_both" ]; then +if [ "$POLICY" = "unified_v2" ] || [ "$POLICY" = "unified_v3" ] || [ "$POLICY" = "unified_kv_both" ]; then ENABLE_KV_BOTH=1 - KV_CONNECTOR="Mooncake" + # honor explicit KV_CONNECTOR override (e.g. =Nixl); otherwise default Mooncake. + if [ -z "$_KV_CONNECTOR_EXPLICIT" ]; then + KV_CONNECTOR="Mooncake" + fi fi if [ "$POLICY" = "unified_nixl_both" ]; then ENABLE_KV_BOTH=1 @@ -143,7 +147,10 @@ for i in $(seq 0 $((N_INSTANCES - 1))); do combined_args="$combined_args http://127.0.0.1:$((BASE_PORT + i))" done proxy_extra="" -if [ "$ENABLE_KV_BOTH" = "1" ]; then +# Bootstrap ports only needed for Mooncake handshake. Nixl uses its own +# UCX side-channel and the proxy forwards kv_transfer_params from src's +# response body instead of pre-baking engine_id/bootstrap_addr. +if [ "$ENABLE_KV_BOTH" = "1" ] && [ "$KV_CONNECTOR" = "Mooncake" ]; then bp_list="" for i in $(seq 0 $((N_INSTANCES - 1))); do if [ -z "$bp_list" ]; then @@ -154,11 +161,15 @@ if [ "$ENABLE_KV_BOTH" = "1" ]; then done proxy_extra="--bootstrap-ports $bp_list" fi +if [ "$ENABLE_KV_BOTH" = "1" ] && [ "$KV_CONNECTOR" = "Nixl" ]; then + proxy_extra="--connector-type nixl" +fi nohup "$VENV/python" "$ROOT/scripts/cache_aware_proxy.py" \ --port "$PROXY_PORT" \ --combined $combined_args \ --policy "$POLICY" \ $proxy_extra \ + ${EXTRA_PROXY_ARGS:-} \ > "$RUNDIR/proxy.log" 2>&1 & disown tries=0 diff --git a/scripts/cache_aware_proxy.py b/scripts/cache_aware_proxy.py index d70ea9e..ac7d5ae 100644 --- a/scripts/cache_aware_proxy.py +++ b/scripts/cache_aware_proxy.py @@ -82,6 +82,36 @@ class Settings: # Patch 6.6: per-request KV-xfer wall-clock timeout (proxy side). pd_sep_xfer_timeout_s: float = 60.0 + # --- unified_v3 (offload-decode) gating knobs ----------------------- + # v3 differs from v2 in *direction*: prefill stays on the session- + # affinity host (which holds the prefix cache); decode is migrated to + # a less-loaded target. KV transfer flows prefill_host → decode_target. + # The target doesn't need cache — we're shipping the post-prefill KV + # over anyway. After successful migration the session affinity table + # rotates to decode_target so the *next* turn lands where the KV now + # lives. + v3_min_new_tokens: int = 8000 # same as v2: don't migrate tiny prefills + v3_min_prefill_decode_busy: int = 1 # prefill_host must have ≥ this many concurrent decode tokens to justify migrating + v3_target_load_ratio: float = 0.7 # target.num_requests must be < prefill_host.num_requests × this + v3_min_load_gap: int = 1 # target.num_requests must also be ≤ prefill_host - this (absolute slack) + v3_rotate_affinity: bool = True # after migration, set session affinity to decode_target. + # Empirically False is better — see cache_miss_audit (next turn hits 9.5% + # with rotation vs ~80% without), because delay_free_blocks doesn't + # actually preserve cross-turn KV on decode_target. + v3_prefer_cache_target: bool = True # Mechanism B: among low-load candidates, prefer the one + # with the most prefix cache for this prompt — vLLM's connector + # auto-transfers only the missing portion (verified via + # smoke_partial_transfer: cache-rich dst is 77% faster than + # cold dst at 33k tokens, +512 ext). + + # --- KV connector selection (governs PD-sep handshake) ------------- + # "mooncake": pre-baked kv_transfer_params (bootstrap_addr+engine_id+transfer_id). + # Requires --bootstrap-ports and vLLMs launched with MooncakeConnector. + # "nixl" : response-forward handshake. src returns kv_transfer_params via + # response body, proxy forwards to dst. Nixl auto-selects transport + # via UCX (CUDA IPC / NVLink on intra-node, RDMA across nodes). + connector_type: str = "mooncake" + SETTINGS = Settings() @@ -538,6 +568,122 @@ def pick_instance_unified_v2( return chosen, chosen_idx, decision, None +def pick_instance_unified_v3( + instances: list[InstanceState], + token_ids: list[int] | None, + session_id: str | None, + input_length: int, + affinity: dict[str, int], +) -> tuple[InstanceState, int, dict, tuple[InstanceState, int] | None]: + """unified_v3 = unified hybrid + selective DECODE migration. + + Direction-reversed from unified_v2: + - prefill stays on session-affinity host (`prefill_host`) so we keep + the 93%-intra-session prefix-cache reuse intact. + - decode is migrated to a lower-load `decode_target` when the + affinity host is busy with concurrent decodes. + - KV transfer flows prefill_host → decode_target (the opposite of + v2's src → chosen). + - target does NOT need pre-existing cache — we're shipping the + post-prefill KV over anyway. + - On successful migration the *caller* must rotate + `affinity[session_id] = decode_target_idx` so the next turn lands + where the KV now lives (decode_target retains the blocks after + completion, since mooncake defaults to delay_free_blocks=True). + + Decision is purely load-based on the target side: + 1. new_local ≥ v3_min_new_tokens (don't pay RDMA for tiny prefills) + 2. prefill_host.ongoing_decode_tokens ≥ v3_min_prefill_decode_busy + (the host is actually busy decoding; migration buys decode-bw) + 3. ∃ target with both num_requests < prefill_host.num_requests × ratio + and num_requests ≤ prefill_host.num_requests − v3_min_load_gap + + Returns (prefill_host, prefill_idx, decision, migrate). When migrate + is None the request is fully local on prefill_host. When migrate is + (decode_target_inst, decode_target_idx), the handler should run + prefill on prefill_host and ship KV to decode_target for decode. + """ + prefill_host, prefill_idx, decision = pick_instance_unified_hybrid( + instances, token_ids, session_id, input_length, affinity) + + decision["v3_migrate"] = False + decision["v3_decision"] = "local" + decision["v3_reason"] = None + + if not token_ids: + decision["v3_reason"] = "no_token_ids" + return prefill_host, prefill_idx, decision, None + + prefill_cache_hit = prefill_host.estimate_cache_hit(token_ids) + new_local = max(0, input_length - prefill_cache_hit) + decision["v3_prefill_cache_hit"] = prefill_cache_hit + decision["v3_new_local"] = new_local + + # Gate 1: prefill must be large enough to amortise RDMA setup. + if new_local < SETTINGS.v3_min_new_tokens: + decision["v3_reason"] = ( + f"new_local_below_threshold ({new_local} < {SETTINGS.v3_min_new_tokens})" + ) + return prefill_host, prefill_idx, decision, None + + # Gate 2: affinity host must be busy with concurrent decodes — that's + # what migrating decode-traffic-away buys us. If the host is idle + # there's no point. + if prefill_host.ongoing_decode_tokens < SETTINGS.v3_min_prefill_decode_busy: + decision["v3_reason"] = ( + f"prefill_host_not_busy " + f"(ongoing_decode_tokens={prefill_host.ongoing_decode_tokens} < " + f"{SETTINGS.v3_min_prefill_decode_busy})" + ) + return prefill_host, prefill_idx, decision, None + + # Gate 3: pick the lowest-load target that is materially less loaded + # than the prefill_host. Cache content irrelevant — KV ships over. + threshold_loaded = max(1, + int(prefill_host.num_requests * SETTINGS.v3_target_load_ratio)) + candidates = [ + (i, inst) for i, inst in enumerate(instances) + if i != prefill_idx + and inst.num_requests < threshold_loaded + and inst.num_requests <= prefill_host.num_requests - SETTINGS.v3_min_load_gap + ] + if not candidates: + decision["v3_reason"] = ( + f"no_low_load_target " + f"(prefill_host.num_req={prefill_host.num_requests} " + f"threshold={threshold_loaded})" + ) + return prefill_host, prefill_idx, decision, None + + # Mechanism B (v3_prefer_cache_target=True): rank candidates first by + # cache_hit DESC (more cache = less KV to transfer), then by load. vLLM + # auto-skips transferring overlapping prefix when dst's local cache + # matches — verified in smoke_partial_transfer: 77% faster on a 33k + # prompt when dst has the prefix already. + if SETTINGS.v3_prefer_cache_target: + decode_target_idx, decode_target = min( + candidates, + key=lambda x: (-x[1].estimate_cache_hit(token_ids), + x[1].num_requests, x[1].ongoing_tokens)) + else: + decode_target_idx, decode_target = min( + candidates, key=lambda x: (x[1].num_requests, x[1].ongoing_tokens)) + + target_cache_hit = decode_target.estimate_cache_hit(token_ids) + decision["v3_migrate"] = True + decision["v3_decision"] = "migrate_decode" + decision["v3_target_idx"] = decode_target_idx + decision["v3_target_num_req"] = decode_target.num_requests + decision["v3_target_cache_hit"] = target_cache_hit + decision["v3_prefill_num_req"] = prefill_host.num_requests + decision["v3_reason"] = ( + f"prefill_host.num_req={prefill_host.num_requests} busy; " + f"target.num_req={decode_target.num_requests} cache_hit={target_cache_hit}, " + f"transferring KV after prefill" + ) + return prefill_host, prefill_idx, decision, (decode_target, decode_target_idx) + + def _extract_output_token_ids_from_sse( buffer: str, chunk: bytes, @@ -764,10 +910,11 @@ async def lifespan(app: FastAPI): policy = getattr(global_args, 'policy', 'linear') # Mooncake-based modes still need bootstrap discovery; NIXL uses # its own UCX side-channel and doesn't go through our proxy - # bootstrap path (and unified_nixl_both never PD-seps anyway). + # bootstrap path. With --connector-type=nixl, v3 also skips bootstrap. needs_bootstrap = ( global_args.offload - or policy in ("unified_v2", "unified_kv_both") + or (policy in ("unified_v2", "unified_v3", "unified_kv_both") + and getattr(global_args, 'connector_type', 'mooncake') == 'mooncake') ) if needs_bootstrap and bp_list: await init_prefill_bootstrap(combined_instances, app.state.ready) @@ -953,6 +1100,26 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h breakdown.update(decision) if session_id: session_affinity_combined[session_id] = best_idx + elif policy == "unified_v3": + # v3: prefill on affinity (cache reuse), decode migrated to a + # low-load target. KV flows prefill_host → decode_target. + # Reuses _handle_combined_pd_sep_v2 with src=prefill_host, + # dst=decode_target (the handler is direction-agnostic). + chosen, best_idx, decision, pd_sep_v2 = pick_instance_unified_v3( + combined_instances, token_ids, session_id, input_length, + session_affinity_combined) + breakdown.update(decision) + if session_id: + if pd_sep_v2 is not None and SETTINGS.v3_rotate_affinity: + # Migration + rotation: redirect next turn to decode_target, + # assuming KV will live there. (Empirically wrong — see + # cache_miss_audit. Keep behind a flag.) + _decode_target_inst, decode_target_idx = pd_sep_v2 + session_affinity_combined[session_id] = decode_target_idx + else: + # No rotation: keep affinity on prefill_host (where the prefix + # cache lives). This is the empirically correct choice. + session_affinity_combined[session_id] = best_idx else: # linear (default) chosen, best_idx = pick_instance( combined_instances, token_ids, session_id, input_length, @@ -983,13 +1150,35 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h }) if pd_sep_v2 is not None: - src_inst, src_idx = pd_sep_v2 - breakdown["v2_src_url"] = src_inst.url - breakdown["v2_src_idx"] = src_idx - return await _handle_combined_pd_sep_v2( - api, req_data, headers, token_ids, input_length, - src_inst, chosen, breakdown, - request_id=request_id) + # Handler contract: first arg = prefill source (does same-worker + # prefill with do_remote_decode=True, max_tokens=1), second arg = + # decode target (does do_remote_prefill=True, pulls KV via + # Mooncake, decodes). + # + # v2 contract: pd_sep_v2 = (src_inst, src_idx); chosen = decode + # → src does prefill (it has more cache), chosen decodes. + # v3 contract: chosen = prefill_host (affinity, has cache); + # pd_sep_v2 = (decode_target_inst, decode_target_idx) + # → chosen does prefill (cache reuse), decode_target decodes. + if policy == "unified_v3": + decode_target_inst, decode_target_idx = pd_sep_v2 + prefill_inst = chosen + breakdown["v2_src_url"] = prefill_inst.url + breakdown["v2_src_idx"] = best_idx + breakdown["v3_decode_target_url"] = decode_target_inst.url + breakdown["v3_decode_target_idx"] = decode_target_idx + return await _handle_combined_pd_sep_v2( + api, req_data, headers, token_ids, input_length, + prefill_inst, decode_target_inst, breakdown, + request_id=request_id) + else: + src_inst, src_idx = pd_sep_v2 + breakdown["v2_src_url"] = src_inst.url + breakdown["v2_src_idx"] = src_idx + return await _handle_combined_pd_sep_v2( + api, req_data, headers, token_ids, input_length, + src_inst, chosen, breakdown, + request_id=request_id) return await _handle_local_request( api, req_data, headers, token_ids, input_length, @@ -1011,11 +1200,12 @@ async def _handle_combined_pd_sep_v2( of SETTINGS.pd_sep_xfer_timeout_s, so a stuck KV transfer fails the request instead of hanging for 600 s. """ - if src.bootstrap_port is None: + connector = SETTINGS.connector_type + if connector == "mooncake" and src.bootstrap_port is None: raise HTTPException( status_code=500, detail=( - "unified_v2 PD-sep triggered but src instance " + "Mooncake PD-sep triggered but src instance " f"{src.url} has no bootstrap_port; launch with " "kv_role=kv_both and pass --bootstrap-ports" ), @@ -1029,12 +1219,16 @@ async def _handle_combined_pd_sep_v2( src_load_held = True dst_load_held = True + # Build prefill kv_transfer_params per connector. prefill_data = req_data.copy() - prefill_data["kv_transfer_params"] = { - "do_remote_decode": True, - "do_remote_prefill": False, - "transfer_id": f"xfer-{request_id}", - } + if connector == "mooncake": + prefill_data["kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "transfer_id": f"xfer-{request_id}", + } + else: # nixl: src just signals it'll produce KV for remote decode + prefill_data["kv_transfer_params"] = {"do_remote_decode": True} prefill_data["stream"] = False prefill_data["max_tokens"] = 1 prefill_data["min_tokens"] = 1 @@ -1044,11 +1238,23 @@ async def _handle_combined_pd_sep_v2( breakdown["t_prefill_sent"] = _time.monotonic() breakdown["t_prefill_sent_unix"] = _time.time() + forwarded_params: dict | None = None try: resp = await src.client.post(api, json=prefill_data, headers=p_headers) breakdown["t_prefill_done"] = _time.monotonic() breakdown["t_prefill_done_unix"] = _time.time() resp.raise_for_status() + if connector == "nixl": + # Nixl populates kv_transfer_params in the response body with + # remote_block_ids / remote_engine_id / remote_host / remote_port. + # We must read the body BEFORE aclose. + src_resp_json = resp.json() + forwarded_params = src_resp_json.get("kv_transfer_params") + if not forwarded_params or not forwarded_params.get("remote_block_ids"): + raise HTTPException( + status_code=502, + detail=f"Nixl src returned no remote_block_ids: {forwarded_params}", + ) await resp.aclose() src.record_prefix(token_ids) except Exception as e: @@ -1057,11 +1263,16 @@ async def _handle_combined_pd_sep_v2( breakdown["prefill_error"] = True breakdown["error_detail"] = repr(e)[:300] _breakdown_log.append(breakdown) - # Release reservations on failure. - src.ongoing_tokens -= input_length - src.num_requests -= 1 - dst.ongoing_tokens -= input_length - dst.num_requests -= 1 + # Release reservations on failure. Clear load_held flags so the + # finally block below does not double-decrement (CRITICAL audit #1). + if src_load_held: + src.ongoing_tokens -= input_length + src.num_requests -= 1 + src_load_held = False + if dst_load_held: + dst.ongoing_tokens -= input_length + dst.num_requests -= 1 + dst_load_held = False raise HTTPException(status_code=502, detail=f"Prefill failed: {e}") finally: if src_load_held: @@ -1069,17 +1280,19 @@ async def _handle_combined_pd_sep_v2( src.num_requests -= 1 src_load_held = False - parsed = urllib.parse.urlparse(str(src.client.base_url)) - bootstrap_addr = f"http://{parsed.hostname}:{src.bootstrap_port}" - decode_data = req_data.copy() - decode_data["kv_transfer_params"] = { - "do_remote_decode": False, - "do_remote_prefill": True, - "remote_bootstrap_addr": bootstrap_addr, - "remote_engine_id": src.engine_id.get(0, ""), - "transfer_id": f"xfer-{request_id}", - } + if connector == "mooncake": + parsed = urllib.parse.urlparse(str(src.client.base_url)) + bootstrap_addr = f"http://{parsed.hostname}:{src.bootstrap_port}" + decode_data["kv_transfer_params"] = { + "do_remote_decode": False, + "do_remote_prefill": True, + "remote_bootstrap_addr": bootstrap_addr, + "remote_engine_id": src.engine_id.get(0, ""), + "transfer_id": f"xfer-{request_id}", + } + else: # nixl: forward what src returned + decode_data["kv_transfer_params"] = forwarded_params breakdown["t_decode_sent"] = _time.monotonic() breakdown["t_decode_sent_unix"] = _time.time() @@ -1304,7 +1517,8 @@ def parse_args(): p.add_argument("--policy", type=str, default="linear", choices=["linear", "lmetric", "load_only", "sticky", "unified", "unified_kv_both", - "unified_nixl_both", "unified_v2"], + "unified_nixl_both", "unified_v2", + "unified_v3"], help="Routing policy: linear (cache-aware), lmetric (P_tokens × BS), " "load_only (B3 control: pure min-num_requests), " "sticky (B3 control: hard session affinity), " @@ -1317,6 +1531,20 @@ def parse_args(): "or unified_v2 (unified + selective per-request PD-sep " "via Mooncake; requires --bootstrap-ports and " "kv_role=kv_both vLLM launch)") + p.add_argument("--v3-rotate-affinity", type=int, default=1, choices=[0,1], + help="unified_v3 only: 1 = rotate session affinity to decode_target " + "after migration (original behavior, empirically loses prefix cache); " + "0 = keep affinity on prefill_host so next turn hits its cache.") + p.add_argument("--connector-type", type=str, default="mooncake", + choices=["mooncake", "nixl"], + help="PD-sep handshake protocol. 'mooncake' uses pre-baked engine_id" + " + bootstrap_addr (requires --bootstrap-ports). 'nixl' uses" + " response-forward (src returns kv_transfer_params, proxy" + " relays to dst; Nixl/UCX auto-picks NVLink intra-node).") + p.add_argument("--v3-prefer-cache-target", type=int, default=1, choices=[0,1], + help="Mechanism B: unified_v3 picks decode_target with the most" + " prefix cache among low-load candidates (default 1). Set 0" + " to fall back to pure-load tie-break (cache-blind).") p.add_argument("--overload-factor", type=float, default=2.0, help="Break session affinity when instance load > factor * avg") # The four flags below are accepted for bench.sh backward compatibility but @@ -1354,7 +1582,14 @@ if __name__ == "__main__": SETTINGS.max_offload_inflight = global_args.max_offload_inflight SETTINGS.cache_gate_ratio = global_args.cache_gate_ratio SETTINGS.decode_iteration_s = getattr(global_args, 'decode_iteration_s', 0.05) - print("SETTINGS: throughput=%.0f rdma_overhead=%.2f offload=%s" % ( + SETTINGS.v3_rotate_affinity = bool(getattr(global_args, 'v3_rotate_affinity', 1)) + SETTINGS.connector_type = getattr(global_args, 'connector_type', 'mooncake') + SETTINGS.v3_prefer_cache_target = bool(getattr(global_args, 'v3_prefer_cache_target', 1)) + print("SETTINGS: throughput=%.0f rdma_overhead=%.2f offload=%s v3_rotate_affinity=%s " + "connector_type=%s v3_prefer_cache_target=%s" % ( SETTINGS.prefill_throughput, SETTINGS.rdma_overhead_s, - getattr(global_args, 'offload', False))) + getattr(global_args, 'offload', False), + SETTINGS.v3_rotate_affinity, + SETTINGS.connector_type, + SETTINGS.v3_prefer_cache_target)) uvicorn.run(app, host=global_args.host, port=global_args.port)