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

Add kvcache-centric profiling and admission controls

Browse Source
This commit is contained in:
Gahow Wang
2026-04-25 16:00:52 +00:00
parent 08b13d22bc
commit 13bb31a446
9 changed files with 1044 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

207
src/agentic_pd_hybrid/replay.py
View File

@@ -3,7 +3,8 @@ from __future__ import annotations
import asyncio
import json
import time
from dataclasses import dataclass, field
from collections import Counter
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, Literal
@@ -26,6 +27,8 @@ from agentic_pd_hybrid.trace import (
HeaderMode = Literal["none", "routing-key", "target-worker", "auto"]
KvCacheAdmissionMode = Literal["router", "worker"]
KvCachePrefillBackupPolicy = Literal["release-after-transfer", "capacity-backup"]
_ADMISSION_PROBE_TIMEOUT_S = 2.0
@dataclass(frozen=True)
@@ -47,6 +50,18 @@ class ReplayConfig:
stream_idle_timeout_s: float | None = 900.0
kvcache_direct_max_uncached_tokens: int = 2048
kvcache_admission_mode: KvCacheAdmissionMode = "router"
kvcache_seed_max_resident_tokens: int | None = None
kvcache_seed_max_output_tokens: int | None = None
kvcache_seed_min_turn_id: int = 1
kvcache_seed_only_multiturn_sessions: bool = False
kvcache_seed_allowed_session_ids: frozenset[str] | None = None
kvcache_prefill_backup_policy: KvCachePrefillBackupPolicy = (
"release-after-transfer"
)
kvcache_seed_max_inflight_decode: int | None = 3
kvcache_prefill_priority_eviction: bool = False
kvcache_prefill_direct_priority: int = -100
kvcache_prefill_normal_priority: int = 100
@dataclass
@@ -104,11 +119,23 @@ class ExecutionResult:
latency_s: float | None
ttft_s: float | None
tpot_s: float | None
prefill_request_priority: int | None = None
decode_request_priority: int | None = None
error: str | None = None
async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]:
requests = load_trace(config.trace_path, request_limit=config.request_limit)
if config.kvcache_seed_only_multiturn_sessions:
session_turns = Counter(request.session_id for request in requests)
config = replace(
config,
kvcache_seed_allowed_session_ids=frozenset(
session_id
for session_id, turn_count in session_turns.items()
if turn_count > 1
),
)
policy = create_policy(config.policy_name)
state = RoutingState.create(config.topology)
state_lock = asyncio.Lock()
@@ -242,6 +269,8 @@ async def _run_request(
latency_s=execution.latency_s,
ttft_s=execution.ttft_s,
tpot_s=execution.tpot_s,
prefill_request_priority=execution.prefill_request_priority,
decode_request_priority=execution.decode_request_priority,
error=execution.error,
)
@@ -253,6 +282,8 @@ async def _invoke_router(
config: ReplayConfig,
decode_worker_index: int,
session_id: str | None = None,
prefill_request_priority: int | None = None,
decode_request_priority: int | None = None,
) -> tuple[float, float | None, float | None, int]:
headers = _build_headers(
request=request,
@@ -274,6 +305,10 @@ async def _invoke_router(
}
if session_id is not None:
payload["session_params"] = {"id": session_id}
if prefill_request_priority is not None:
payload["smg_prefill_priority"] = prefill_request_priority
if decode_request_priority is not None:
payload["smg_decode_priority"] = decode_request_priority
return await _invoke_generate(
client=client,
@@ -462,6 +497,7 @@ async def _open_streaming_session(
"session_id": session_id,
"streaming": True,
},
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
response.raise_for_status()
opened_session_id = response.json()
@@ -481,6 +517,7 @@ async def _close_streaming_session(
response = await client.post(
f"{server_url.rstrip('/')}/close_session",
json={"session_id": session_id},
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
if response.is_success:
return
@@ -538,7 +575,10 @@ async def _fetch_decode_server_state(
server_url: str,
) -> tuple[dict[str, Any], int, int]:
try:
response = await client.get(f"{server_url.rstrip('/')}/server_info")
response = await client.get(
f"{server_url.rstrip('/')}/server_info",
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
response.raise_for_status()
payload = response.json()
except Exception:
@@ -567,6 +607,7 @@ async def _query_decode_direct_admission(
"uncached_input_tokens": max(0, uncached_input_tokens),
"output_tokens": max(0, output_tokens),
},
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
response.raise_for_status()
payload = response.json()
@@ -643,6 +684,51 @@ def _estimate_session_resident_tokens(request: TraceRequest) -> int:
return request.input_length + request.output_length
def _seed_filter_reason(
*,
request: TraceRequest,
config: ReplayConfig,
inflight_decode_load: int | None = None,
) -> str | None:
if request.turn_id < config.kvcache_seed_min_turn_id:
return "seed-filter-early-turn"
if (
config.kvcache_seed_max_inflight_decode is not None
and inflight_decode_load is not None
and inflight_decode_load > config.kvcache_seed_max_inflight_decode
):
return "seed-filter-inflight-decode-load"
if (
config.kvcache_seed_allowed_session_ids is not None
and request.session_id not in config.kvcache_seed_allowed_session_ids
):
return "seed-filter-single-turn-session"
resident_tokens = _estimate_session_resident_tokens(request)
if (
config.kvcache_seed_max_resident_tokens is not None
and resident_tokens > config.kvcache_seed_max_resident_tokens
):
return "seed-filter-resident-tokens"
if (
config.kvcache_seed_max_output_tokens is not None
and request.output_length > config.kvcache_seed_max_output_tokens
):
return "seed-filter-output-tokens"
return None
def _prefill_priority_for_router_request(
*,
config: ReplayConfig,
direct_to_d_predicted: bool,
) -> int | None:
if not config.kvcache_prefill_priority_eviction:
return None
if direct_to_d_predicted:
return config.kvcache_prefill_direct_priority
return config.kvcache_prefill_normal_priority
def _inspect_direct_request(
*,
request: TraceRequest,
@@ -802,6 +888,7 @@ async def _fetch_decode_load_snapshot(
response = await client.get(
f"{server_url.rstrip('/')}/v1/loads",
params={"include": "core,disagg"},
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
response.raise_for_status()
payload = response.json()
@@ -865,6 +952,13 @@ def _is_decode_backpressure_reason(reason: str | None) -> bool:
}
def _is_stale_decode_session_error(exc: Exception) -> bool:
return (
isinstance(exc, httpx.HTTPStatusError)
and exc.response.status_code == 400
)
def _dynamic_decode_headroom_tokens(
*,
residency: DecodeResidencyState,
@@ -1469,17 +1563,23 @@ async def _invoke_plain_router(
decision,
execution_mode: str,
) -> ExecutionResult:
prefill_priority = _prefill_priority_for_router_request(
config=config,
direct_to_d_predicted=False,
)
latency_s, ttft_s, tpot_s, cached_tokens = await _invoke_router(
client=client,
request=request,
config=config,
decode_worker_index=decision.decode_worker_index,
prefill_request_priority=prefill_priority,
)
return ExecutionResult(
execution_mode=execution_mode,
actual_kv_transfer_blocks=decision.kv_transfer_blocks,
effective_input_length=request.input_length,
cached_tokens=cached_tokens,
prefill_request_priority=prefill_priority,
session_reused=False,
session_reset=False,
latency_s=latency_s,
@@ -1502,17 +1602,20 @@ async def _invoke_kvcache_seeded_router(
reserved_tokens: int,
execution_mode: str,
) -> ExecutionResult:
keep_prefill_backup = False
prefill_reserved_tokens = 0
async with direct_session_lock:
keep_prefill_backup, prefill_reserved_tokens, _prefill_evicted = (
await _reserve_prefill_backup_capacity(
client=client,
request=request,
prefill_url=prefill_url,
session=decode_session,
direct_sessions=direct_sessions,
residency=decode_residency,
if config.kvcache_prefill_backup_policy == "capacity-backup":
keep_prefill_backup, prefill_reserved_tokens, _prefill_evicted = (
await _reserve_prefill_backup_capacity(
client=client,
request=request,
prefill_url=prefill_url,
session=decode_session,
direct_sessions=direct_sessions,
residency=decode_residency,
)
)
)
if (
decode_session.prefill_opened
and decode_session.prefill_server_url != prefill_url
@@ -1538,6 +1641,10 @@ async def _invoke_kvcache_seeded_router(
decode_session_newly_opened = False
try:
prefill_priority = _prefill_priority_for_router_request(
config=config,
direct_to_d_predicted=True,
)
async with direct_session_lock:
if not decode_session.opened:
await _open_streaming_session(
@@ -1555,6 +1662,7 @@ async def _invoke_kvcache_seeded_router(
config=config,
decode_worker_index=decision.decode_worker_index,
session_id=request.session_id,
prefill_request_priority=prefill_priority,
)
except Exception:
async with direct_session_lock:
@@ -1615,6 +1723,7 @@ async def _invoke_kvcache_seeded_router(
actual_kv_transfer_blocks=decision.kv_transfer_blocks,
effective_input_length=request.input_length,
cached_tokens=cached_tokens,
prefill_request_priority=prefill_priority,
session_reused=False,
session_reset=False,
latency_s=latency_s,
@@ -1697,6 +1806,19 @@ async def _execute_request(
)
if request.turn_id == 1:
seed_filter_reason = _seed_filter_reason(
request=request,
config=config,
inflight_decode_load=decision.inflight_decode_load,
)
if seed_filter_reason is not None:
return await _invoke_plain_router(
client=client,
request=request,
config=config,
decision=decision,
execution_mode=f"pd-router-turn1-{seed_filter_reason}",
)
async with direct_session_lock:
admit_new_decode_session = _should_admit_new_decode_session(
residency=decode_residency,
@@ -1800,16 +1922,33 @@ async def _execute_request(
)
)
if can_direct:
return await _invoke_decode_session_direct(
client=client,
request=request,
config=config,
decision=decision,
direct_sessions=direct_sessions,
direct_session_lock=direct_session_lock,
decode_residency=decode_residency,
reserved_tokens=direct_reserved_tokens,
)
try:
return await _invoke_decode_session_direct(
client=client,
request=request,
config=config,
decision=decision,
direct_sessions=direct_sessions,
direct_session_lock=direct_session_lock,
decode_residency=decode_residency,
reserved_tokens=direct_reserved_tokens,
)
except Exception as exc:
if not _is_stale_decode_session_error(exc):
raise
async with direct_session_lock:
await _close_decode_session(
client=client,
session=decode_session,
residency=decode_residency,
)
return await _invoke_plain_router(
client=client,
request=request,
config=config,
decision=decision,
execution_mode="pd-router-fallback-stale-d-session",
)
if _is_decode_backpressure_reason(direct_reason):
return await _invoke_plain_router(
client=client,
@@ -1819,6 +1958,19 @@ async def _execute_request(
execution_mode="pd-router-fallback-d-backpressure",
)
seed_filter_reason = _seed_filter_reason(
request=request,
config=config,
inflight_decode_load=decision.inflight_decode_load,
)
if seed_filter_reason is not None:
return await _invoke_plain_router(
client=client,
request=request,
config=config,
decision=decision,
execution_mode=f"pd-router-fallback-{seed_filter_reason}",
)
async with direct_session_lock:
admit_new_decode_session = _should_admit_new_decode_session(
residency=decode_residency,
@@ -1894,6 +2046,19 @@ async def _execute_request(
),
)
seed_filter_reason = _seed_filter_reason(
request=request,
config=config,
inflight_decode_load=decision.inflight_decode_load,
)
if seed_filter_reason is not None:
return await _invoke_plain_router(
request=request,
client=client,
config=config,
decision=decision,
execution_mode=f"pd-router-fallback-large-append-{seed_filter_reason}",
)
async with direct_session_lock:
admit_new_decode_session = _should_admit_new_decode_session(
residency=decode_residency,