diff --git a/src/aituner/harness.py b/src/aituner/harness.py index f55e51b..112f18e 100644 --- a/src/aituner/harness.py +++ b/src/aituner/harness.py @@ -24,6 +24,13 @@ _RUNTIME_KEYS = { _STRONG_INCUMBENT_MIN_GAIN = 1.8 _MIN_POST_INCUMBENT_VALIDATION_TRIALS = 2 _VALIDATION_TRIALS_WITHOUT_FAMILY_COVERAGE = 3 +# Decode-bound throughput is frequently KV-cache limited, so more gpu-memory-utilization +# yields more KV blocks and more concurrent decode. Hill-climb in small steps toward a +# safe ceiling and let measurement find the real peak: a too-high target regresses or +# fails to launch and is rejected by the incumbent guard, and its tested signature then +# blocks re-proposal so the climb terminates. +_GMU_STEP = 0.02 +_GMU_SAFE_CEILING = 0.97 def build_harness_context( @@ -383,14 +390,17 @@ def _knob_harnesses( "knob_family": "gpu-memory-utilization", "use_when": [ "The engine launches cleanly but memory headroom limits batching.", + "A decode-bound incumbent (decode_tpot) is KV-cache limited and could sustain more concurrent decode with more KV blocks.", ], "procedure": [ "Make small adjustments only after topology and batching knobs are stable.", + "Raise gpu-memory-utilization one small step at a time and keep the step only if request_rate_per_gpu improves and the engine still launches.", ], "guards": [ "Treat launch OOM as hard negative evidence and back off immediately.", + "Do not exceed a safe utilization ceiling; stop climbing once a higher target regresses or fails to launch.", ], - "active_now": False, + "active_now": active_bottleneck in {"decode_tpot", "admission_or_queueing"}, } ) return harnesses @@ -1184,6 +1194,15 @@ def _runtime_candidate_actions( topology_patch = _preserve_topology_patch(study, anchor_flags) actions: list[dict[str, Any]] = [] + base_tp = _parse_int_like(study.engine.base_flags.get("tensor-parallel-size"), default=1) + base_dp = _parse_int_like(study.engine.base_flags.get("data-parallel-size"), default=1) + cur_tp = _parse_int_like(anchor_flags.get("tensor-parallel-size"), default=base_tp) + cur_dp = _parse_int_like(anchor_flags.get("data-parallel-size"), default=base_dp) + # Topology-before-runtime: gpu-mem-util / raising max-num-seqs are micro-tuning that is + # only justified once topology has moved off the baseline. At the baseline a latency + # bottleneck must still be answered with a topology change, not a runtime tweak. + topology_settled = cur_tp > base_tp or cur_dp > base_dp + if "max-num-batched-tokens" in tunable: current_mbt = _parse_int_like(anchor_flags.get("max-num-batched-tokens"), default=0) mbt_targets: list[tuple[str, int]] = [] @@ -1226,8 +1245,17 @@ def _runtime_candidate_actions( if top_bottleneck == "admission_or_queueing": target = max(8, int(current_mns * 1.5)) if current_mns > 0 else 64 mns_targets.append(("raise_max_num_seqs", _round_up_to_multiple(target, 8))) - elif top_bottleneck == "decode_tpot" and current_mns > 8: - mns_targets.append(("lower_max_num_seqs", max(8, current_mns // 2))) + elif top_bottleneck == "decode_tpot": + if current_mns > 8: + mns_targets.append(("lower_max_num_seqs", max(8, current_mns // 2))) + # Decode concurrency can also be too low: once topology is settled, raising + # max-num-seqs exploits decode parallelism when the incumbent has SLO headroom. + # The incumbent guard keeps it only if per-GPU rate improves. + if topology_settled: + raise_target = _round_up_to_multiple( + max(16, int(current_mns * 1.5)) if current_mns > 0 else 48, 8 + ) + mns_targets.append(("raise_max_num_seqs", raise_target)) for action_id, target in mns_targets: patch = {**topology_patch, "max-num-seqs": target} signature = _config_signature({"env_patch": {}, "flag_patch": patch}) @@ -1273,6 +1301,37 @@ def _runtime_candidate_actions( ], ) ) + + if ( + "gpu-memory-utilization" in tunable + and topology_settled + and top_bottleneck in {"decode_tpot", "admission_or_queueing"} + ): + current_gmu = _parse_float_like( + anchor_flags.get("gpu-memory-utilization"), default=0.9 + ) + if 0.0 < current_gmu < _GMU_SAFE_CEILING: + target = round(min(_GMU_SAFE_CEILING, current_gmu + _GMU_STEP), 4) + if target > current_gmu: + patch = {**topology_patch, "gpu-memory-utilization": target} + signature = _config_signature({"env_patch": {}, "flag_patch": patch}) + if signature not in tested_signatures: + actions.append( + _runtime_action( + action_id="raise_gpu_memory_utilization", + knob_family="gpu-memory-utilization", + score=0.4 + _information_gain(bottleneck_hypotheses, "runtime"), + patch=patch, + hypothesis=( + "Raise gpu-memory-utilization to add KV-cache headroom so the " + "decode-bound incumbent can sustain more concurrent decode." + ), + expected_effects=[ + "add KV-cache blocks for higher decode concurrency on the incumbent topology", + "reject if the higher memory target regresses request_rate_per_gpu or fails to launch", + ], + ) + ) return actions @@ -2252,6 +2311,19 @@ def _parse_int_like(value: Any, *, default: int) -> int: return default +def _parse_float_like(value: Any, *, default: float) -> float: + if value is None or isinstance(value, bool): + return default + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str) and value.strip(): + try: + return float(value.strip()) + except ValueError: + return default + return default + + def _config_signature(config_patch: Any) -> str: if not isinstance(config_patch, dict): config_patch = {} diff --git a/tests/test_core_flow.py b/tests/test_core_flow.py index c90b55e..eee56fe 100644 --- a/tests/test_core_flow.py +++ b/tests/test_core_flow.py @@ -1318,6 +1318,113 @@ class CoreFlowTests(unittest.TestCase): }, ) + def test_harness_raises_gpu_mem_util_on_settled_decode_bound_incumbent(self) -> None: + """Regression for the coverage gap that let the naive baseline beat the harness: + a settled TP incumbent that is decode_tpot-bound must get a gpu-memory-utilization + raise (KV-cache headroom) before the harness is allowed to stop.""" + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + study_path = _write_study_assets( + tmp_path, + slo_overrides={ + "ttft_rule": {"kind": "fixed_ms", "threshold_ms": 4000}, + "tpot_rule": {"kind": "fixed_ms", "threshold_ms": 50}, + }, + engine_overrides={ + "tunable_flags": [ + "tensor-parallel-size", + "gpu-memory-utilization", + ], + "topology_constraints": { + "allowed_tensor_parallel_sizes": [1, 2, 4], + "allowed_data_parallel_sizes": [1], + "allowed_tp_dp_products": [1, 2, 4], + }, + }, + ) + study = load_study_spec(study_path) + result_path = tmp_path / "trial-0002.json" + result_path.write_text( + json.dumps( + { + "status": "completed", + "best_sampling_u": 0.074, + "best_request_rate": 2.6, + "best_pass_rate": 0.97, + "probes": [ + { + "threshold": 0.074, + "feasible": True, + "payload": { + "request_count": 300, + "pass_rate": 0.97, + "request_rate": 2.6, + "latency_summary": {"failed_reason_counts": {}}, + }, + }, + { + "threshold": 0.09, + "feasible": False, + "payload": { + "request_count": 300, + "pass_rate": 0.6, + "request_rate": 3.2, + "early_stop_reason": "slo_pass_rate_unrecoverable", + "latency_summary": { + "failed_reason_counts": {"tpot_ms>50.0": 90} + }, + }, + }, + ], + } + ), + encoding="utf-8", + ) + state = StudyState( + study_id=study.study_id, + best_trial_id="trial-0002", + best_request_rate=2.6, + best_request_rate_per_gpu=0.65, + trials=[ + TrialSummary( + trial_id="trial-0001", + status="completed", + best_request_rate=1.1, + best_request_rate_per_gpu=0.275, + config_patch={"env_patch": {}, "flag_patch": {"tensor-parallel-size": 2}}, + ), + TrialSummary( + trial_id="trial-0002", + status="completed", + best_request_rate=2.6, + best_request_rate_per_gpu=0.65, + result_path=str(result_path), + config_patch={ + "env_patch": {}, + "flag_patch": { + "tensor-parallel-size": 4, + "gpu-memory-utilization": 0.9, + }, + }, + ), + ], + ) + context = build_harness_context( + study=study, window_summary={"prompt_tokens_p95": 1500}, state=state + ) + proposal = build_harness_guided_proposal(context) + self.assertIsNotNone(proposal) + self.assertFalse(proposal.should_stop) + # TP4 preserved; gpu-memory-utilization hill-climbed one step (0.9 -> 0.92). + self.assertEqual( + proposal.config_patch.flag_patch.get("tensor-parallel-size"), 4 + ) + self.assertEqual( + proposal.config_patch.flag_patch.get("gpu-memory-utilization"), 0.92 + ) + # And the harness must NOT authorize a stop while that knob is untried. + self.assertIsNone(build_harness_stop_proposal(context)) + def test_harness_validates_unmeasured_tp_frontier_before_runtime_refinement(self) -> None: with tempfile.TemporaryDirectory() as tmp: tmp_path = Path(tmp)