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

Ablation 2: fire-and-forget vs await-prefill scheduling

Browse Source 
Added --fire-and-forget flag to cache_aware_proxy.py for async prefill dispatch.

Results on 6P+2D config:
  Await:  TTFT=1.48s  TPOT=0.066s  E2E=5.95s  94% success
  FnF:    TTFT=5.32s  TPOT=0.037s  E2E=11.9s  85% success

Fire-and-forget improves TPOT by 44% (pipeline overlap) but degrades
TTFT by 260% (decode internally waits for KV, less efficiently than
proxy-level await) and increases errors from KV race conditions.

Full 4-way ablation summary in analyze_ablations.py.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-21 23:02:42 +08:00
parent 9dee25907b
commit c7afdc5074
2 changed files with 149 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
scripts/cache_aware_proxy.py
View File

@@ -196,36 +196,55 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
return StreamingResponse(generate(), media_type="text/event-stream")
async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
session_id, headers):
"""PD-Sep mode: await prefill, then stream decode."""
p_inst, _ = pick_instance(prefill_instances, token_ids, session_id,
input_length, session_affinity)
d_inst = min(decode_instances, key=lambda x: x.ongoing_tokens)
# Await prefill
p_inst.ongoing_tokens += input_length
async def _send_prefill_async(p_inst, api, prefill_data, p_headers, token_ids, input_length):
"""Fire-and-forget prefill: send and don't block caller."""
try:
prefill_data = req_data.copy()
prefill_data["kv_transfer_params"] = {
"do_remote_decode": True, "do_remote_prefill": False,
"transfer_id": f"xfer-{request_id}",
}
prefill_data["stream"] = False
prefill_data["max_tokens"] = 1
prefill_data.pop("max_completion_tokens", None)
prefill_data.pop("stream_options", None)
p_headers = {**headers, "X-data-parallel-rank": "0"}
resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers)
resp.raise_for_status()
await resp.aclose()
p_inst.record_prefix(token_ids)
except Exception as e:
raise HTTPException(status_code=502, detail=f"Prefill failed: {e}")
except Exception:
pass
finally:
p_inst.ongoing_tokens -= input_length
async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
session_id, headers):
"""PD-Sep mode. --fire-and-forget controls prefill waiting behavior."""
p_inst, _ = pick_instance(prefill_instances, token_ids, session_id,
input_length, session_affinity)
d_inst = min(decode_instances, key=lambda x: x.ongoing_tokens)
prefill_data = req_data.copy()
prefill_data["kv_transfer_params"] = {
"do_remote_decode": True, "do_remote_prefill": False,
"transfer_id": f"xfer-{request_id}",
}
prefill_data["stream"] = False
prefill_data["max_tokens"] = 1
prefill_data.pop("max_completion_tokens", None)
prefill_data.pop("stream_options", None)
p_headers = {**headers, "X-data-parallel-rank": "0"}
p_inst.ongoing_tokens += input_length
if global_args.fire_and_forget:
# Fire-and-forget: send prefill async, immediately proceed to decode
asyncio.create_task(_send_prefill_async(
p_inst, api, prefill_data, p_headers, token_ids, input_length))
else:
# Await: block until prefill completes, then send decode
try:
resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers)
resp.raise_for_status()
await resp.aclose()
p_inst.record_prefix(token_ids)
except Exception as e:
raise HTTPException(status_code=502, detail=f"Prefill failed: {e}")
finally:
p_inst.ongoing_tokens -= input_length
# Stream decode
d_inst.ongoing_tokens += input_length
parsed = urllib.parse.urlparse(str(p_inst.client.base_url))
@@ -260,6 +279,8 @@ def parse_args():
help="PD-Sep prefill: URL [bootstrap_port]")
p.add_argument("--decode", nargs=1, action="append", dest="decode_raw",
help="PD-Sep decode: URL")
p.add_argument("--fire-and-forget", action="store_true",
help="Send prefill async, don't await before decode")
args = p.parse_args()
args.prefill = []