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

feat(policy): load-floor bonus for KvAwarePolicy (Q2.B)

Browse Source 
Implements the design proposed and approved in
docs/E1_E2_FIX_DESIGN_ZH.md §Q2.B.

KvAwarePolicy gains a `load_floor_bonus: int = 0` knob. When > 0:

  mean_assigned = sum(assigned[*]) / len(D)
  for each D candidate:
    if not sticky and mean_assigned > 0:
      deficit = max(0, mean_assigned - assigned[D])
      floor_bonus = K * deficit / mean_assigned
    else:
      floor_bonus = 0
    score = (overlap + sticky*α + floor_bonus, sticky, -inflight, -assigned)

Properties (verified by unit-style probe in commit message):
- Default 0 = old behavior preserved
- Sticky-gated: turn-1+ requests of an existing session keep going
  to their original D (cache locality preserved)
- Graduated: bonus magnitude scales with the D's deficit ratio,
  approaches K as deficit/mean → 1, drops to 0 when balanced
- Set above max expected boilerplate overlap (Inferact ~50 → 200)
  so cross-session shared-prefix overlap doesn't pin cold D's idle,
  but real per-session prefix overlap (>K blocks) still wins

Plumbed through ReplayConfig, BenchmarkConfig, and CLI flag
--kvcache-load-floor-bonus on both `replay` and `benchmark-live`.

Empirical verification on synthetic state (same conditions as the
E2 cold-D pathology):
  - OFF (K=0):   route fresh session → decode-0 (boilerplate winner)
  - ON  (K=200): route fresh session → decode-1 (cold D rebalanced)

Validation pass next: scripts/sweep_e3_kvc_v2_loadfloor_rdma.sh
(committed separately).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
tim
2026-05-12 11:45:09 +08:00
parent 905d671135
commit 93fce42747
4 changed files with 86 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/agentic_pd_hybrid/benchmark.py
View File

@@ -48,6 +48,7 @@ class BenchmarkConfig:
enable_backpressure: bool = False
backpressure_max_pause_s: float = 2.0
kvcache_migration_reject_threshold: int = 3
kvcache_load_floor_bonus: int = 0
sample_profile: str = "default"
min_initial_input_tokens: int | None = None
max_initial_input_tokens: int | None = None
@@ -200,6 +201,7 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
enable_backpressure=config.enable_backpressure,
backpressure_max_pause_s=config.backpressure_max_pause_s,
kvcache_migration_reject_threshold=config.kvcache_migration_reject_threshold,
kvcache_load_floor_bonus=config.kvcache_load_floor_bonus,
)
if config.request_timeout_s is not None:
replay_config = replace(
@@ -261,6 +263,7 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
"enable_backpressure": config.enable_backpressure,
"backpressure_max_pause_s": config.backpressure_max_pause_s,
"kvcache_migration_reject_threshold": config.kvcache_migration_reject_threshold,
"kvcache_load_floor_bonus": config.kvcache_load_floor_bonus,
"sample_profile": config.sample_profile,
"min_initial_input_tokens": config.min_initial_input_tokens,
"max_initial_input_tokens": config.max_initial_input_tokens,

28
src/agentic_pd_hybrid/cli.py
View File

@@ -270,6 +270,19 @@ def main() -> None:
"See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1."
),
)
replay.add_argument(
"--kvcache-load-floor-bonus",
type=int,
default=0,
help=(
"Graduated bonus added to lex-score position 0 for under-loaded D "
"workers (gated on not-sticky so turn-1+ requests still stick). "
"Magnitude scales as K * max(0, mean - assigned[D]) / mean. "
"Set above max expected cross-session boilerplate overlap "
"(Inferact ~50 → use 200). 0 disables. "
"See docs/E1_E2_FIX_DESIGN_ZH.md §Q2."
),
)
sample = subparsers.add_parser(
"sample-sessions",
@@ -521,6 +534,19 @@ def main() -> None:
"See REFACTOR_PLAN_V1 §6.2 / TEAM_REPORT §2.1."
),
)
benchmark.add_argument(
"--kvcache-load-floor-bonus",
type=int,
default=0,
help=(
"Graduated bonus added to lex-score position 0 for under-loaded D "
"workers (gated on not-sticky so turn-1+ requests still stick). "
"Magnitude scales as K * max(0, mean - assigned[D]) / mean. "
"Set above max expected cross-session boilerplate overlap "
"(Inferact ~50 → use 200). 0 disables. "
"See docs/E1_E2_FIX_DESIGN_ZH.md §Q2."
),
)
benchmark.add_argument(
"--sample-profile",
choices=["default", "small-append"],
@@ -607,6 +633,7 @@ def main() -> None:
enable_backpressure=args.enable_backpressure,
backpressure_max_pause_s=args.backpressure_max_pause_s,
kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
kvcache_load_floor_bonus=args.kvcache_load_floor_bonus,
)
results = asyncio.run(replay_trace(config))
print(
@@ -754,6 +781,7 @@ def main() -> None:
enable_backpressure=args.enable_backpressure,
backpressure_max_pause_s=args.backpressure_max_pause_s,
kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
kvcache_load_floor_bonus=args.kvcache_load_floor_bonus,
sample_profile=args.sample_profile,
min_initial_input_tokens=args.min_initial_input_tokens,
max_initial_input_tokens=args.max_initial_input_tokens,

53
src/agentic_pd_hybrid/policies.py
View File

@@ -161,6 +161,28 @@ class KvAwarePolicy:
# 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
# Load-floor bonus: graduated boost added to lex-score position 0 for
# under-loaded D workers, gated on `not sticky` so turn-1+ requests of an
# existing session continue to stick to their original D. The boost
# magnitude scales linearly with the D's deficit relative to the running
# mean of `decode_assignment_counts`:
# floor_bonus = K * max(0, mean - assigned[D]) / max(1, mean)
# When mean=0 (warmup), bonus is 0 for all workers (lex tiebreak by
# iteration order). Once any D has been assigned, under-loaded D's get a
# bonus that approaches K as their deficit-to-mean ratio approaches 1.
# The bonus naturally decays as load equalises, leaving the original
# overlap+sticky scoring in charge of steady-state selection.
#
# Set this above the maximum cross-session boilerplate overlap you expect
# so that fresh sessions are routed to under-loaded D's even when those
# D's currently have 0 overlap, but below the magnitude of "real" prefix
# overlap (e.g., a session with 800-block per-session prefix on an
# already-warm D should still go there).
#
# 0 disables. See docs/E1_E2_FIX_DESIGN_ZH.md §Q2 for the full design and
# docs/E1_E2_RESULTS_ZH.md §5d for why this is needed on Inferact-shaped
# workloads where boilerplate overlap pins D2 cold forever.
load_floor_bonus: int = 0
def select(
self,
@@ -172,6 +194,12 @@ class KvAwarePolicy:
prefill_worker_id = state.next_prefill_worker_id(topology)
session = state.session_state.get(request.session_id)
# Pre-compute the running mean of decode assignments. Used by the
# load-floor bonus inside the candidate loop.
n_route_workers = max(1, len(topology.route_workers))
total_assigned = sum(state.decode_assignment_counts.values())
mean_assigned = total_assigned / n_route_workers
best_decode_worker_id: str | None = None
best_score: tuple[int, int, int, int] | None = None
candidates_considered = 0
@@ -189,9 +217,18 @@ class KvAwarePolicy:
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)
assignment_penalty = -state.decode_assignment_counts.get(worker.worker_id, 0)
worker_assigned = state.decode_assignment_counts.get(worker.worker_id, 0)
assignment_penalty = -worker_assigned
# Load-floor bonus: only for fresh placements (not sticky), and
# only when the knob is enabled. See docstring above.
floor_bonus = 0
if self.load_floor_bonus > 0 and not sticky and mean_assigned > 0:
deficit = max(0.0, mean_assigned - worker_assigned)
floor_bonus = int(self.load_floor_bonus * deficit / mean_assigned)
score = (
overlap + sticky * self.sticky_bonus,
overlap + sticky * self.sticky_bonus + floor_bonus,
sticky,
inflight_penalty,
assignment_penalty,
@@ -223,14 +260,22 @@ class KvAwarePolicy:
)
def create_policy(name: str, *, migration_reject_threshold: int = 3) -> RoutingPolicy:
def create_policy(
name: str,
*,
migration_reject_threshold: int = 3,
load_floor_bonus: int = 0,
) -> 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(migration_reject_threshold=migration_reject_threshold)
return KvAwarePolicy(
migration_reject_threshold=migration_reject_threshold,
load_floor_bonus=load_floor_bonus,
)
raise ValueError(f"Unsupported policy: {name}")

6
src/agentic_pd_hybrid/replay.py
View File

@@ -111,6 +111,11 @@ class ReplayConfig:
# 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
# Load-floor bonus magnitude for KvAwarePolicy: graduated boost added to
# under-loaded D workers to break overlap-pinning imbalance on workloads
# with shared cross-session prefix. 0 disables. See
# docs/E1_E2_FIX_DESIGN_ZH.md §Q2.
kvcache_load_floor_bonus: int = 0
structural_log_dir: Path | None = None
@@ -198,6 +203,7 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]:
policy = create_policy(
config.policy_name,
migration_reject_threshold=config.kvcache_migration_reject_threshold,
load_floor_bonus=config.kvcache_load_floor_bonus,
)
state = RoutingState.create(config.topology)
state_lock = asyncio.Lock()