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

Harden prefill scheduler harness

Browse Source
This commit is contained in:
Gahow Wang
2026-06-29 01:54:02 +08:00
parent bfd85793f3
commit ee101a7c24
3 changed files with 624 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
src/aituner/harness.py
View File

@@ -1280,30 +1280,31 @@ def _runtime_candidate_actions(
# only justified once no untested TP increase remains. At an intermediate TP (e.g. TP2
# while TP4 is still reachable and untried) a latency bottleneck must still be answered
# by climbing TP, not a runtime tweak -- otherwise runtime tuning preempts the frontier.
_next_tp = _next_allowed_tp(study, current_tp=cur_tp, current_dp=cur_dp)
higher_tp_patch = _higher_tp_frontier_patch(
study,
current_tp=cur_tp,
current_dp=cur_dp,
)
tp_frontier_open = (
_next_tp is not None
and _effective_config_signature(
study,
{"env_patch": {}, "flag_patch": {"tensor-parallel-size": _next_tp}}
)
higher_tp_patch is not None
and _effective_config_signature(study, {"env_patch": {}, "flag_patch": higher_tp_patch})
not in tested_signatures
)
topology_settled = not tp_frontier_open
actions.extend(
_prefill_scheduler_candidate_actions(
study,
window_summary,
anchor_flags,
runtime_base_patch,
top_bottleneck,
bottleneck_hypotheses,
topology_settled=topology_settled,
seen_signatures=seen_signatures,
blocked_candidates=blocked_candidates,
)
prefill_scheduler_actions = _prefill_scheduler_candidate_actions(
study,
window_summary,
anchor_flags,
runtime_base_patch,
top_bottleneck,
bottleneck_hypotheses,
topology_settled=topology_settled,
seen_signatures=seen_signatures,
blocked_candidates=blocked_candidates,
)
actions.extend(prefill_scheduler_actions)
prefill_scheduler_candidate_available = bool(prefill_scheduler_actions)
if (
"max-num-batched-tokens" in tunable
@@ -1312,6 +1313,7 @@ def _runtime_candidate_actions(
and recent_diagnostics[-1].get("trial_id") == anchor.get("trial_id")
and cur_tp > 1
and not bottleneck_hypotheses
and not prefill_scheduler_candidate_available
):
current_mbt = _parse_int_like(anchor_flags.get("max-num-batched-tokens"), default=0)
target_mbt = (
@@ -1361,7 +1363,7 @@ def _runtime_candidate_actions(
)
seen_signatures.add(signature)
if "max-num-batched-tokens" in tunable:
if "max-num-batched-tokens" in tunable and not prefill_scheduler_candidate_available:
current_mbt = _parse_int_like(anchor_flags.get("max-num-batched-tokens"), default=0)
mbt_targets: list[tuple[str, int]] = []
if top_bottleneck == "ttft_prefill":
@@ -1484,6 +1486,7 @@ def _runtime_candidate_actions(
and "max-num-batched-tokens" in tunable
and "max-num-seqs" in tunable
and max_num_seqs_tested
and not prefill_scheduler_candidate_available
):
current_mbt = _parse_int_like(anchor_flags.get("max-num-batched-tokens"), default=0)
current_mns = _parse_int_like(anchor_flags.get("max-num-seqs"), default=0)
@@ -1540,7 +1543,11 @@ def _runtime_candidate_actions(
)
)
if "enable-chunked-prefill" in tunable and top_bottleneck == "ttft_prefill":
if (
"enable-chunked-prefill" in tunable
and top_bottleneck == "ttft_prefill"
and not prefill_scheduler_candidate_available
):
current = bool(anchor_flags.get("enable-chunked-prefill", False))
if not current:
patch = {**runtime_base_patch, "enable-chunked-prefill": True}
@@ -1706,14 +1713,15 @@ def _prefill_scheduler_candidate_actions(
else None
)
if current_chunked and quantum_step["target"] is None and admission_step is None:
admission_target = admission_step["target"] if admission_step is not None else None
if current_chunked and quantum_step["target"] is None and admission_target is None:
return []
patch = {**runtime_base_patch, "enable-chunked-prefill": True}
if quantum_step["target"] is not None:
patch["max-num-batched-tokens"] = quantum_step["target"]
if admission_step is not None:
patch["max-num-seqs"] = admission_step
if admission_target is not None:
patch["max-num-seqs"] = admission_target
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
action_id = _prefill_scheduler_action_id(quantum_step["direction"], admission_step)
@@ -1736,12 +1744,12 @@ def _prefill_scheduler_candidate_actions(
relief = 0.56 if quantum_step["direction"] == "lower" else 0.42
if quantum_step["direction"] == "seed":
relief = 0.38
if admission_step is not None:
relief += 0.06
if admission_target is not None:
relief += 0.08 if admission_step and admission_step["direction"] == "lower" else 0.06
coverage_bonus = 0.0
if quantum_step["direction"] == "seed" or not current_chunked:
coverage_bonus = 0.28
elif quantum_step["target"] is not None or admission_step is not None:
elif quantum_step["target"] is not None or admission_target is not None:
coverage_bonus = 0.14
information_gain = _information_gain(bottleneck_hypotheses, "runtime")
score = relief * max(confidence, 0.35) + information_gain + 0.08 + coverage_bonus
@@ -1759,7 +1767,14 @@ def _prefill_scheduler_candidate_actions(
round(target_ratio, 4) if target_ratio is not None else None
),
"admission_pressure_current": current_mns or None,
"admission_pressure_target": admission_step,
"admission_pressure_target": admission_target,
"admission_pressure_direction": admission_step["direction"] if admission_step else "hold",
"admission_pressure_ratio_current": (
round(admission_step["current_ratio"], 4) if admission_step else None
),
"admission_pressure_ratio_target": (
round(admission_step["target_ratio"], 4) if admission_step else None
),
}
actions = [
_runtime_action(
@@ -1843,32 +1858,59 @@ def _next_admission_pressure_step(
*,
top_bottleneck: str,
quantum_direction: str,
) -> int | None:
) -> dict[str, Any] | None:
if current_mns <= 0:
return None
target_concurrency = max(int(study.trace.max_concurrency), 1)
current_ratio = current_mns / target_concurrency
if top_bottleneck == "admission_or_queueing" and current_mns < target_concurrency:
target = min(target_concurrency, int(current_mns * _ADMISSION_PRESSURE_STEP_UP))
return _round_up_to_multiple(target, 8)
target = _round_up_to_multiple(target, 8)
return {
"direction": "raise",
"target": target,
"current_ratio": current_ratio,
"target_ratio": target / target_concurrency,
}
if (
top_bottleneck == "ttft_prefill"
and quantum_direction in {"hold", "raise"}
and current_mns < target_concurrency
):
target = min(target_concurrency, int(current_mns * _ADMISSION_PRESSURE_STEP_UP))
return _round_up_to_multiple(target, 8)
target = _round_up_to_multiple(target, 8)
return {
"direction": "raise",
"target": target,
"current_ratio": current_ratio,
"target_ratio": target / target_concurrency,
}
if top_bottleneck == "ttft_prefill" and current_mns > target_concurrency:
target = min(current_mns - 8, _round_up_to_multiple(target_concurrency, 8))
if target > 0 and target < current_mns:
return {
"direction": "lower",
"target": target,
"current_ratio": current_ratio,
"target_ratio": target / target_concurrency,
}
return None
def _prefill_scheduler_action_id(quantum_direction: str, admission_target: int | None) -> str:
def _prefill_scheduler_action_id(
quantum_direction: str,
admission_step: dict[str, Any] | None,
) -> str:
if quantum_direction == "lower":
return "lower_prefill_quantum_with_chunked_prefill"
if quantum_direction == "raise":
return "raise_prefill_quantum_with_chunked_prefill"
if quantum_direction == "seed":
return "seed_chunked_prefill_quantum"
if admission_target is not None:
return "adjust_admission_pressure_with_chunked_prefill"
if admission_step is not None and admission_step["direction"] == "lower":
return "lower_admission_pressure_with_chunked_prefill"
if admission_step is not None and admission_step["direction"] == "raise":
return "raise_admission_pressure_with_chunked_prefill"
return "enable_chunked_prefill_scheduler_mode"
@@ -2393,8 +2435,12 @@ def _topology_frontier_status(
flags = _effective_flags_for_item(study, best)
current_tp = _parse_int_like(flags.get("tensor-parallel-size"), default=1)
current_dp = _parse_int_like(flags.get("data-parallel-size"), default=1)
next_tp = _next_allowed_tp(study, current_tp=current_tp, current_dp=current_dp)
if next_tp is None:
flag_patch = _higher_tp_frontier_patch(
study,
current_tp=current_tp,
current_dp=current_dp,
)
if flag_patch is None:
return {
**default,
"reason": "no_legal_higher_tp_frontier",
@@ -2402,11 +2448,8 @@ def _topology_frontier_status(
"current_tp": current_tp,
"current_dp": current_dp,
}
next_tp = _parse_int_like(flag_patch.get("tensor-parallel-size"), default=current_tp)
flag_patch: dict[str, Any] = {"tensor-parallel-size": next_tp}
base_dp = _parse_int_like(study.engine.base_flags.get("data-parallel-size"), default=1)
if current_dp != base_dp:
flag_patch["data-parallel-size"] = current_dp
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": flag_patch})
if signature in _state_tested_signatures(study, state):
return {
@@ -2430,6 +2473,22 @@ def _topology_frontier_status(
}
def _higher_tp_frontier_patch(
study: StudySpec,
*,
current_tp: int,
current_dp: int,
) -> dict[str, Any] | None:
next_tp = _next_allowed_tp(study, current_tp=current_tp, current_dp=current_dp)
if next_tp is None:
return None
flag_patch: dict[str, Any] = {"tensor-parallel-size": next_tp}
base_dp = _parse_int_like(study.engine.base_flags.get("data-parallel-size"), default=1)
if current_dp != base_dp:
flag_patch["data-parallel-size"] = current_dp
return flag_patch
def _effective_flags_for_item(study: StudySpec, item: dict[str, Any]) -> dict[str, Any]:
flags = dict(study.engine.base_flags)
patch = item.get("config_patch")