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

Proxy/runner support for Nixl connector + unified_v3 (offload-decode) policy

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-27 22:05:19 +08:00
parent da39ab6804
commit f739f7d461
2 changed files with 283 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

303
scripts/cache_aware_proxy.py
View File

@@ -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)