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

66
src/agentic_pd_hybrid/benchmark.py
View File

@@ -27,10 +27,20 @@ class BenchmarkConfig:
time_scale: float = 1.0
concurrency_limit: int = 32
timeout_s: float = 1200.0
request_timeout_s: float | None = None
stream: bool = True
stream_idle_timeout_s: float | None = 900.0
kvcache_direct_max_uncached_tokens: int = 2048
kvcache_admission_mode: str = "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_prefill_backup_policy: str = "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
sample_profile: str = "default"
min_initial_input_tokens: int | None = None
max_initial_input_tokens: int | None = None
@@ -59,10 +69,17 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
topology = config.topology
if config.mechanism_name == "kvcache-centric":
prefill_extra_server_args = topology.prefill_extra_server_args + (
"--enable-streaming-session",
)
if config.kvcache_prefill_priority_eviction:
prefill_extra_server_args = prefill_extra_server_args + (
"--radix-eviction-policy",
"priority",
)
topology = replace(
topology,
prefill_extra_server_args=topology.prefill_extra_server_args
+ ("--enable-streaming-session",),
prefill_extra_server_args=prefill_extra_server_args,
decode_extra_server_args=topology.decode_extra_server_args
+ (
"--enable-streaming-session",
@@ -107,6 +124,11 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
prefill_policy="round_robin",
decode_policy=_decode_policy_for(config.policy_name),
timeout_s=config.timeout_s,
router_request_timeout_s=(
config.request_timeout_s
if config.request_timeout_s is not None
else config.timeout_s
),
include_router=(
config.mechanism_name in {"pd-disaggregation", "kvcache-centric"}
),
@@ -142,7 +164,27 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
stream_idle_timeout_s=config.stream_idle_timeout_s,
kvcache_direct_max_uncached_tokens=config.kvcache_direct_max_uncached_tokens,
kvcache_admission_mode=config.kvcache_admission_mode, # type: ignore[arg-type]
kvcache_seed_max_resident_tokens=config.kvcache_seed_max_resident_tokens,
kvcache_seed_max_output_tokens=config.kvcache_seed_max_output_tokens,
kvcache_seed_min_turn_id=config.kvcache_seed_min_turn_id,
kvcache_seed_only_multiturn_sessions=(
config.kvcache_seed_only_multiturn_sessions
),
kvcache_prefill_backup_policy=config.kvcache_prefill_backup_policy, # type: ignore[arg-type]
kvcache_seed_max_inflight_decode=(
config.kvcache_seed_max_inflight_decode
),
kvcache_prefill_priority_eviction=(
config.kvcache_prefill_priority_eviction
),
kvcache_prefill_direct_priority=config.kvcache_prefill_direct_priority,
kvcache_prefill_normal_priority=config.kvcache_prefill_normal_priority,
)
if config.request_timeout_s is not None:
replay_config = replace(
replay_config,
timeout_s=config.request_timeout_s,
)
asyncio.run(replay_trace(replay_config))
finally:
signal.signal(signal.SIGINT, previous_sigint)
@@ -163,10 +205,30 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
"time_scale": config.time_scale,
"concurrency_limit": config.concurrency_limit,
"timeout_s": config.timeout_s,
"request_timeout_s": config.request_timeout_s,
"stream": config.stream,
"stream_idle_timeout_s": config.stream_idle_timeout_s,
"kvcache_direct_max_uncached_tokens": config.kvcache_direct_max_uncached_tokens,
"kvcache_admission_mode": config.kvcache_admission_mode,
"kvcache_seed_max_resident_tokens": config.kvcache_seed_max_resident_tokens,
"kvcache_seed_max_output_tokens": config.kvcache_seed_max_output_tokens,
"kvcache_seed_min_turn_id": config.kvcache_seed_min_turn_id,
"kvcache_seed_only_multiturn_sessions": (
config.kvcache_seed_only_multiturn_sessions
),
"kvcache_prefill_backup_policy": config.kvcache_prefill_backup_policy,
"kvcache_seed_max_inflight_decode": (
config.kvcache_seed_max_inflight_decode
),
"kvcache_prefill_priority_eviction": (
config.kvcache_prefill_priority_eviction
),
"kvcache_prefill_direct_priority": (
config.kvcache_prefill_direct_priority
),
"kvcache_prefill_normal_priority": (
config.kvcache_prefill_normal_priority
),
"sample_profile": config.sample_profile,
"min_initial_input_tokens": config.min_initial_input_tokens,
"max_initial_input_tokens": config.max_initial_input_tokens,