Add per-request breakdown profiling, identify KV cache memory bottleneck
Breakdown profiling at proxy level captures: t_proxy_recv → t_prefill_sent → t_prefill_done → t_decode_sent → t_first_token Key finding: 87.7% of TTFT is spent in kv+decode phase, NOT prefill. Root cause: decode instance KV cache memory saturation (97.1% usage). With 6P+2D config, 2 decode GPUs have only ~56GB total KV cache. Large agentic requests (avg 33.6k tokens) fill this quickly. Small requests (49 tokens, prefill=0.044s) wait 114s for KV cache to be freed by large requests completing decode. vLLM log confirms: Running=0, Waiting=6, KV cache=97.1% GPU is idle but requests queue for KV cache memory, not compute. This is the fundamental bottleneck of single-machine PD separation for long-context agentic workloads: concentrating decode onto fewer GPUs creates a KV cache memory wall. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -196,25 +196,42 @@ async def _handle_combined(api, req_data, token_ids, input_length, session_id, h
|
||||
return StreamingResponse(generate(), media_type="text/event-stream")
|
||||
|
||||
|
||||
async def _send_prefill_async(p_inst, api, prefill_data, p_headers, token_ids, input_length):
|
||||
import time as _time
|
||||
|
||||
# Per-request breakdown log (append-only)
|
||||
_breakdown_log: list[dict] = []
|
||||
|
||||
|
||||
async def _send_prefill_async(p_inst, api, prefill_data, p_headers, token_ids,
|
||||
input_length, breakdown):
|
||||
"""Fire-and-forget prefill: send and don't block caller."""
|
||||
try:
|
||||
resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers)
|
||||
breakdown["t_prefill_done"] = _time.monotonic()
|
||||
resp.raise_for_status()
|
||||
await resp.aclose()
|
||||
p_inst.record_prefix(token_ids)
|
||||
except Exception:
|
||||
pass
|
||||
breakdown["t_prefill_done"] = _time.monotonic()
|
||||
breakdown["prefill_error"] = True
|
||||
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."""
|
||||
"""PD-Sep mode with per-stage breakdown profiling."""
|
||||
breakdown = {
|
||||
"request_id": request_id,
|
||||
"input_length": input_length,
|
||||
"t_proxy_recv": _time.monotonic(),
|
||||
}
|
||||
|
||||
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)
|
||||
breakdown["p_inst"] = p_inst.url
|
||||
breakdown["d_inst"] = d_inst.url
|
||||
|
||||
prefill_data = req_data.copy()
|
||||
prefill_data["kv_transfer_params"] = {
|
||||
@@ -228,24 +245,27 @@ async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
|
||||
p_headers = {**headers, "X-data-parallel-rank": "0"}
|
||||
|
||||
p_inst.ongoing_tokens += input_length
|
||||
breakdown["t_prefill_sent"] = _time.monotonic()
|
||||
|
||||
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))
|
||||
p_inst, api, prefill_data, p_headers, token_ids, input_length, breakdown))
|
||||
else:
|
||||
# Await: block until prefill completes, then send decode
|
||||
try:
|
||||
resp = await p_inst.client.post(api, json=prefill_data, headers=p_headers)
|
||||
breakdown["t_prefill_done"] = _time.monotonic()
|
||||
resp.raise_for_status()
|
||||
await resp.aclose()
|
||||
p_inst.record_prefix(token_ids)
|
||||
except Exception as e:
|
||||
breakdown["t_prefill_done"] = _time.monotonic()
|
||||
breakdown["prefill_error"] = True
|
||||
_breakdown_log.append(breakdown)
|
||||
raise HTTPException(status_code=502, detail=f"Prefill failed: {e}")
|
||||
finally:
|
||||
p_inst.ongoing_tokens -= input_length
|
||||
|
||||
# Stream decode
|
||||
# Send decode
|
||||
d_inst.ongoing_tokens += input_length
|
||||
parsed = urllib.parse.urlparse(str(p_inst.client.base_url))
|
||||
bootstrap_addr = f"http://{parsed.hostname}:{p_inst.bootstrap_port}"
|
||||
@@ -258,18 +278,32 @@ async def _handle_pd_sep(api, req_data, request_id, token_ids, input_length,
|
||||
"transfer_id": f"xfer-{request_id}",
|
||||
}
|
||||
|
||||
breakdown["t_decode_sent"] = _time.monotonic()
|
||||
|
||||
async def generate():
|
||||
first_token = True
|
||||
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():
|
||||
if first_token:
|
||||
breakdown["t_first_token"] = _time.monotonic()
|
||||
first_token = False
|
||||
yield chunk
|
||||
finally:
|
||||
breakdown["t_done"] = _time.monotonic()
|
||||
d_inst.ongoing_tokens -= input_length
|
||||
_breakdown_log.append(breakdown)
|
||||
|
||||
return StreamingResponse(generate(), media_type="application/json")
|
||||
|
||||
|
||||
@app.get("/breakdown")
|
||||
async def get_breakdown():
|
||||
"""Return per-request breakdown data for analysis."""
|
||||
return _breakdown_log
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="Unified cache-aware global scheduler")
|
||||
p.add_argument("--port", type=int, default=8000)
|
||||
|
||||
Reference in New Issue
Block a user