gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Harness: explore gpu-memory-utilization (and raise max-num-seqs) before Stop-B

Browse Source 
The harness defined a gpu-memory-utilization family but hard-coded active_now=False
and never generated a candidate for it, and only ever *lowered* max-num-seqs for
decode_tpot. So on the decode-bound 27B incumbent it stopped at TP4=0.648 while the
naive (use_harness=false) baseline freely found gpu-memory-utilization=0.94 -> 0.873
(+35%) and max-num-seqs=48. That made the harness look worse than naive -- a real
coverage gap, not bad luck.

Fix in _runtime_candidate_actions (topology-before-runtime gated: only once topology
has moved off the baseline, so a baseline latency bottleneck still gets a TP change):
- Add a gpu-memory-utilization hill-climb candidate (+0.02/step toward a 0.97 safe
  ceiling) for decode_tpot/admission incumbents, scored high enough (>=0.35) to block
  a premature Stop-B until it is tried; the incumbent guard keeps the step only if
  per-GPU rate improves and the engine launches, and the tested signature terminates
  the climb (so 0.96 OOM/regression backs off to 0.94 automatically).
- Let max-num-seqs *rise* for decode_tpot (not only fall) to exploit decode parallelism.
- Activate the gpu-memory-utilization harness family for decode_tpot/admission.

Verified: new unit test asserts a settled TP4 decode-bound incumbent gets a
gpu-memory-utilization raise (0.9->0.92) and no stop while untried. 115 tests pass.
Empirical reliability (harness recovers ~0.87 and stops) to be confirmed by re-run.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-19 10:25:47 +08:00
parent 95c02d7dd9
commit a3523f5601
2 changed files with 182 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
src/aituner/harness.py
View File

@@ -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 = {}

107
tests/test_core_flow.py
View File

@@ -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)