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

feat(kvc): session migration with reset-on-success + direct-append threshold tuning

Browse Source 
KVC v2 beats 4DP at ts=1 same-scale on 7/8 metrics:
  TTFT mean -24%, p50 -54%, p90 -64%; lat mean -0.8%, p50 -12.6%, p90 -0.7%.
  Direct-to-D rate jumped 42.8% -> 91.7%. REFACTOR_PLAN_V1 scenario C achieved.

Two-knob fix:
- reset-on-success blacklist decay: clear (sess, D) reject counter on
  successful direct-to-D path. Eliminates v1 thrashing where session 6880
  was stable on decode-1 for 70 turns then collapsed to 75 D-changes after
  cumulative transient pressure tripped the permanent blacklist.
- bump --kvcache-direct-max-uncached-tokens default 2048 -> 8192 via CLI flag.
  41% of v1 fallbacks were 'real-large-append' (>2048 token append); raising
  the threshold lets these go through the direct-to-D fast path.

Code:
- policies.py: RoutingState.session_d_rejects counter + KvAwarePolicy
  migration_reject_threshold; degenerate fallback picks least-rejected D.
- replay.py: record_admission_reject + reset-on-success in _run_request;
  _fallthrough_reason classifies turn-2+ fall-throughs as session-not-resident
  / real-large-append / etc, replacing misleading 'large-append' suffix
  (TEAM_REPORT §2.7).
- cli.py + benchmark.py: --kvcache-migration-reject-threshold flag wiring.

Docs:
- REFACTOR_PLAN_V1_ZH.md: forward-looking plan after ts=1 validation.
- MIGRATION_V1_FINDINGS_ZH.md: v1 thrashing root-cause analysis.
- V2_RESULTS_ZH.md: v2 results, scenario C achievement, attribution.
- TEAM_REPORT_AGENTIC_PD_HYBRID_ZH.md: comprehensive team report.

Scripts:
- sweep_ts1_kvc_n3_plus_dp.sh: ts=1 baseline (KVC 1P3D N=3 + 4DP CA).
- sweep_ts1_migration_v1.sh / v2.sh: validation runs.
- analyze_ts1_validation.py: 4-way comparison analyzer.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
kzlin
2026-05-09 01:18:13 +08:00
parent 1d51704dad
commit 2ec0debef4
12 changed files with 2350 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/agentic_pd_hybrid/benchmark.py
View File

@@ -47,6 +47,7 @@ class BenchmarkConfig:
pool_poll_include_sessions: bool = True
enable_backpressure: bool = False
backpressure_max_pause_s: float = 2.0
kvcache_migration_reject_threshold: int = 3
sample_profile: str = "default"
min_initial_input_tokens: int | None = None
max_initial_input_tokens: int | None = None
@@ -198,6 +199,7 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
pool_poll_include_sessions=config.pool_poll_include_sessions,
enable_backpressure=config.enable_backpressure,
backpressure_max_pause_s=config.backpressure_max_pause_s,
kvcache_migration_reject_threshold=config.kvcache_migration_reject_threshold,
)
if config.request_timeout_s is not None:
replay_config = replace(
@@ -258,6 +260,7 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
"pool_poll_include_sessions": config.pool_poll_include_sessions,
"enable_backpressure": config.enable_backpressure,
"backpressure_max_pause_s": config.backpressure_max_pause_s,
"kvcache_migration_reject_threshold": config.kvcache_migration_reject_threshold,
"sample_profile": config.sample_profile,
"min_initial_input_tokens": config.min_initial_input_tokens,
"max_initial_input_tokens": config.max_initial_input_tokens,

22
src/agentic_pd_hybrid/cli.py
View File

@@ -260,6 +260,16 @@ def main() -> None:
default=2.0,
help="Cap on per-request backpressure sleep, regardless of D hint.",
)
replay.add_argument(
"--kvcache-migration-reject-threshold",
type=int,
default=3,
help=(
"Per-(session, D) admission-reject count after which KvAwarePolicy "
"skips that D for the session (forces migration). 0 disables. "
"See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1."
),
)
sample = subparsers.add_parser(
"sample-sessions",
@@ -501,6 +511,16 @@ def main() -> None:
default=2.0,
help="Cap on per-request backpressure sleep, regardless of D hint.",
)
benchmark.add_argument(
"--kvcache-migration-reject-threshold",
type=int,
default=3,
help=(
"Per-(session, D) admission-reject count after which KvAwarePolicy "
"skips that D for the session (forces migration). 0 disables. "
"See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1."
),
)
benchmark.add_argument(
"--sample-profile",
choices=["default", "small-append"],
@@ -586,6 +606,7 @@ def main() -> None:
pool_poll_include_sessions=not args.pool_poll_no_sessions,
enable_backpressure=args.enable_backpressure,
backpressure_max_pause_s=args.backpressure_max_pause_s,
kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
)
results = asyncio.run(replay_trace(config))
print(
@@ -732,6 +753,7 @@ def main() -> None:
pool_poll_include_sessions=not args.pool_poll_no_sessions,
enable_backpressure=args.enable_backpressure,
backpressure_max_pause_s=args.backpressure_max_pause_s,
kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
sample_profile=args.sample_profile,
min_initial_input_tokens=args.min_initial_input_tokens,
max_initial_input_tokens=args.max_initial_input_tokens,

42
src/agentic_pd_hybrid/policies.py
View File

@@ -44,6 +44,10 @@ class RoutingState:
inflight_decode: Counter[str] = field(default_factory=Counter)
decode_assignment_counts: Counter[str] = field(default_factory=Counter)
decode_resident_blocks: dict[str, set[int]] = field(default_factory=dict)
# Migration support: per-(session_id, decode_worker_id) admission reject counter.
# KvAwarePolicy uses this to skip D's that have repeatedly rejected this session
# (avoids the structural starvation observed in TEAM_REPORT §2.1).
session_d_rejects: Counter[tuple[str, str]] = field(default_factory=Counter)
@classmethod
def create(cls, topology: SingleNodeTopology) -> "RoutingState":
@@ -66,6 +70,12 @@ class RoutingState:
self.decode_cursor += 1
return worker.worker_id
def record_admission_reject(self, session_id: str, decode_worker_id: str) -> int:
"""Increment per-(session, D) rejection counter. Returns new count."""
key = (session_id, decode_worker_id)
self.session_d_rejects[key] += 1
return self.session_d_rejects[key]
def finish(self, request: TraceRequest, decision: RoutingDecision) -> None:
session = self.session_state.setdefault(request.session_id, SessionRouteState())
session.last_decode_worker = decision.decode_worker_id
@@ -146,6 +156,11 @@ class StickyDecodePolicy:
class KvAwarePolicy:
name: str = "kv-aware"
sticky_bonus: int = 1
# Session migration: when (session, D) has been rejected this many times,
# skip D entirely for this session (force migration to another D).
# 0 disables the mechanism. Default 3 picked empirically to allow brief
# transient saturation without panicking, but to reroute persistent starvation.
migration_reject_threshold: int = 3
def select(
self,
@@ -158,8 +173,19 @@ class KvAwarePolicy:
session = state.session_state.get(request.session_id)
best_decode_worker_id: str | None = None
best_score: tuple[int, int, int] | None = None
best_score: tuple[int, int, int, int] | None = None
candidates_considered = 0
for worker in topology.route_workers:
# Migration: skip workers that have rejected this session too many times.
# If all candidates get filtered (degenerate case), fall through to
# un-filtered selection below.
if self.migration_reject_threshold > 0:
rejects = state.session_d_rejects.get(
(request.session_id, worker.worker_id), 0
)
if rejects >= self.migration_reject_threshold:
continue
candidates_considered += 1
overlap = _overlap_blocks(request, state, worker.worker_id)
sticky = int(session is not None and session.last_decode_worker == worker.worker_id)
inflight_penalty = -state.inflight_decode.get(worker.worker_id, 0)
@@ -174,6 +200,16 @@ class KvAwarePolicy:
best_score = score
best_decode_worker_id = worker.worker_id
# Degenerate fallback: every D was filtered. Pick the least-rejected D.
if best_decode_worker_id is None:
best_decode_worker_id = min(
(w.worker_id for w in topology.route_workers),
key=lambda wid: state.session_d_rejects.get(
(request.session_id, wid), 0
),
)
best_score = (0, 0, 0, 0)
assert best_decode_worker_id is not None
reuse_expected = bool(best_score and best_score[0] > 0)
return _build_decision(
@@ -187,14 +223,14 @@ class KvAwarePolicy:
)
def create_policy(name: str) -> RoutingPolicy:
def create_policy(name: str, *, migration_reject_threshold: int = 3) -> RoutingPolicy:
normalized = name.strip().lower()
if normalized == "default":
return DefaultPolicy()
if normalized == "sticky":
return StickyDecodePolicy()
if normalized in {"kv-aware", "kv_aware", "kv"}:
return KvAwarePolicy()
return KvAwarePolicy(migration_reject_threshold=migration_reject_threshold)
raise ValueError(f"Unsupported policy: {name}")

100
src/agentic_pd_hybrid/replay.py
View File

@@ -106,6 +106,11 @@ class ReplayConfig:
pool_poll_include_sessions: bool = True
enable_backpressure: bool = False
backpressure_max_pause_s: float = 2.0
# Session migration via per-(sess, D) admission reject memory.
# When a session has been admission-rejected this many times on a given D,
# KvAwarePolicy skips that D for the session (forcing migration). Default 3.
# Set 0 to disable. See REFACTOR_PLAN_V1 §6.2.
kvcache_migration_reject_threshold: int = 3
structural_log_dir: Path | None = None
@@ -190,7 +195,10 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]:
if turn_count > 1
),
)
policy = create_policy(config.policy_name)
policy = create_policy(
config.policy_name,
migration_reject_threshold=config.kvcache_migration_reject_threshold,
)
state = RoutingState.create(config.topology)
state_lock = asyncio.Lock()
semaphore = asyncio.Semaphore(config.concurrency_limit)
@@ -350,6 +358,22 @@ async def _run_request(
async with state_lock:
state.finish(request, decision)
# Migration feedback: if this request was forced into a fallback path
# because the chosen D rejected admission, record the (session, D)
# rejection so KvAwarePolicy can migrate this session next turn.
if _is_admission_rejection_mode(execution.execution_mode):
state.record_admission_reject(
request.session_id,
decision.decode_worker_id,
)
# Reset-on-success: a successful direct-to-D path proves D-X can
# currently serve this session — clear the cumulative reject counter
# so that brief past saturation doesn't permanently blacklist the D.
# (MIGRATION_V1_FINDINGS §4.1: blacklist-permanence bug fix.)
elif execution.execution_mode == "kvcache-direct-to-d-session":
state.session_d_rejects[
(request.session_id, decision.decode_worker_id)
] = 0
return RequestMetrics.from_decision(
request,
@@ -1349,6 +1373,49 @@ def _is_stale_decode_session_error(exc: Exception) -> bool:
)
# execution_mode substrings that signal D-side admission rejected this request.
# Used by _run_request to update state.session_d_rejects so KvAwarePolicy can
# migrate persistently-starved sessions to a different D next turn.
_ADMISSION_REJECTION_SUBSTRINGS = (
"session-cap",
"no-d-capacity",
"d-backpressure",
)
def _is_admission_rejection_mode(execution_mode: str) -> bool:
return any(token in execution_mode for token in _ADMISSION_REJECTION_SUBSTRINGS)
def _fallthrough_reason(
*,
request: TraceRequest,
config: ReplayConfig,
decision,
direct_append_length: int | None,
direct_session_reused: bool,
direct_session_reset: bool,
) -> str:
"""Classify why a turn-2+ KVC request fell through to the seed/large-append branch.
Returns a short label suffix used in execution_mode strings to replace the
misleading 'large-append' label (TEAM_REPORT §2.7). In particular,
'session-not-resident' is the §1 starvation signature — direct_session_reused
is False because the session was never opened on the policy-chosen D.
"""
if not direct_session_reused:
return "session-not-resident"
if direct_session_reset:
return "session-was-evicted"
if direct_append_length is None:
return "no-direct-info"
if direct_append_length > config.kvcache_direct_max_uncached_tokens:
return "real-large-append"
if not _should_bypass_prefill(request=request, config=config, decision=decision):
return "policy-no-bypass"
return "other-large-append"
def _dynamic_decode_headroom_tokens(
*,
residency: DecodeResidencyState,
@@ -2510,6 +2577,17 @@ async def _execute_request(
decode_residency=decode_residency,
)
# TEAM_REPORT §2.7: 'large-append' is misleading — most fallthroughs are
# actually 'session-not-resident-on-pinned-D' (§1 starvation). Classify
# the real reason and embed it in the execution_mode label.
fallthrough = _fallthrough_reason(
request=request,
config=config,
decision=decision,
direct_append_length=direct_append_length,
direct_session_reused=direct_session_reused,
direct_session_reset=direct_session_reset,
)
seed_filter_reason = _seed_filter_reason(
request=request,
config=config,
@@ -2521,7 +2599,7 @@ async def _execute_request(
client=client,
config=config,
decision=decision,
execution_mode=f"pd-router-fallback-large-append-{seed_filter_reason}",
execution_mode=f"pd-router-fallback-{fallthrough}-{seed_filter_reason}",
decode_residency=decode_residency,
)
async with direct_session_lock:
@@ -2566,7 +2644,7 @@ async def _execute_request(
client=client,
config=config,
decision=decision,
execution_mode="pd-router-fallback-large-append-session-cap",
execution_mode=f"pd-router-fallback-{fallthrough}-session-cap",
decode_residency=decode_residency,
)
if can_seed:
@@ -2582,23 +2660,27 @@ async def _execute_request(
decode_residency=decode_residency,
reserved_tokens=reserved_tokens,
execution_mode=(
"pd-router-large-append-reseed"
f"pd-router-{fallthrough}-reseed"
+ _eviction_suffix(
evicted_sessions,
prefill_backed_evictions,
)
),
)
# Preserve seed_reason in the label so migration feedback fires for
# 'd-no-space' / 'd-*-backpressure' (matched via _is_admission_rejection_mode).
if _is_decode_backpressure_reason(seed_reason):
mode_label = f"pd-router-fallback-{fallthrough}-d-backpressure"
elif seed_reason == "d-no-space":
mode_label = f"pd-router-fallback-{fallthrough}-no-d-capacity"
else:
mode_label = f"pd-router-fallback-{fallthrough}"
return await _invoke_plain_router(
request=request,
client=client,
config=config,
decision=decision,
execution_mode=(
"pd-router-fallback-d-backpressure"
if _is_decode_backpressure_reason(seed_reason)
else "pd-router-fallback-large-append"
),
execution_mode=mode_label,
decode_residency=decode_residency,
)