diff --git a/src/agentic_pd_hybrid/benchmark.py b/src/agentic_pd_hybrid/benchmark.py index 055247c..0811331 100644 --- a/src/agentic_pd_hybrid/benchmark.py +++ b/src/agentic_pd_hybrid/benchmark.py @@ -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, diff --git a/src/agentic_pd_hybrid/cli.py b/src/agentic_pd_hybrid/cli.py index ca2777c..2c37d73 100644 --- a/src/agentic_pd_hybrid/cli.py +++ b/src/agentic_pd_hybrid/cli.py @@ -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, diff --git a/src/agentic_pd_hybrid/policies.py b/src/agentic_pd_hybrid/policies.py index 162bf84..61e8c13 100644 --- a/src/agentic_pd_hybrid/policies.py +++ b/src/agentic_pd_hybrid/policies.py @@ -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}") diff --git a/src/agentic_pd_hybrid/replay.py b/src/agentic_pd_hybrid/replay.py index 9e065ef..4c0be51 100644 --- a/src/agentic_pd_hybrid/replay.py +++ b/src/agentic_pd_hybrid/replay.py @@ -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()