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

Add normalized prefill scheduler harness

Browse Source
This commit is contained in:
Gahow Wang
2026-06-29 01:12:19 +08:00
parent 7ad439730e
commit 36c301c128
3 changed files with 650 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
src/aituner/harness.py
View File

@@ -39,6 +39,9 @@ _VALIDATION_TRIALS_WITHOUT_FAMILY_COVERAGE = 3
_GMU_STEP = 0.02
_GMU_NOMINAL_FLOOR = 0.9
_GMU_SAFE_CEILING = 0.97
_PREFILL_QUANTUM_HEAD_OF_LINE_RATIO = 1.0
_PREFILL_QUANTUM_FRAGMENTATION_RATIO = 0.5
_ADMISSION_PRESSURE_STEP_UP = 1.5
def build_harness_context(
@@ -355,19 +358,27 @@ def _knob_harnesses(
}
)
if "enable-chunked-prefill" in tunable:
prefill_scheduler_active = (
active_bottleneck in {"ttft_prefill", "admission_or_queueing"}
and _prefill_scheduler_workload_applies(study, window_summary)
)
harnesses.append(
{
"knob_family": "enable-chunked-prefill",
"use_when": [
"The L profile has a long tail and long prefills block shorter requests.",
"Treat chunked prefill as part of a scheduler interaction with prefill quantum and admission pressure, not as a standalone magic flag.",
],
"procedure": [
"Keep chunked prefill enabled for heavy-tail chat windows unless history shows chunking overhead dominates.",
"Enable chunked prefill when the measured bottleneck indicates prefill head-of-line blocking.",
"Move max-num-batched-tokens by relative trust-region steps in normalized prefill_quantum_ratio space.",
"Move max-num-seqs only as a relative admission-pressure correction, and preserve topology while testing this scheduler hypothesis.",
],
"guards": [
"Do not disable chunked prefill on a heavy-tail workload without direct evidence from a nearby trial.",
"Do not use fixed absolute MBT/MNS tables; derive the next concrete flag values from the incumbent and workload scale.",
],
"active_now": False,
"active_now": prefill_scheduler_active,
}
)
if "expert-parallel-size" in tunable or "enable-expert-parallel" in tunable:
@@ -1280,6 +1291,20 @@ def _runtime_candidate_actions(
)
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,
)
)
if (
"max-num-batched-tokens" in tunable
and _anchor_has_topology_patch(anchor)
@@ -1636,11 +1661,216 @@ def _next_gpu_memory_utilization_target(
return target
def _prefill_scheduler_candidate_actions(
study: StudySpec,
window_summary: dict[str, Any],
anchor_flags: dict[str, Any],
runtime_base_patch: dict[str, Any],
top_bottleneck: str,
bottleneck_hypotheses: list[dict[str, Any]],
*,
topology_settled: bool,
seen_signatures: set[str],
blocked_candidates: list[dict[str, Any]],
) -> list[dict[str, Any]]:
tunable = set(study.engine.tunable_flags)
if not topology_settled:
return []
if top_bottleneck not in {"ttft_prefill", "admission_or_queueing"}:
return []
if "enable-chunked-prefill" not in tunable or "max-num-batched-tokens" not in tunable:
return []
if not _prefill_scheduler_workload_applies(study, window_summary):
return []
prompt_scale = _prefill_prompt_scale(window_summary)
if prompt_scale <= 0:
return []
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)
current_chunked = bool(anchor_flags.get("enable-chunked-prefill", False))
quantum_step = _next_prefill_quantum_step(
current_mbt,
prompt_scale,
top_bottleneck=top_bottleneck,
)
admission_step = (
_next_admission_pressure_step(
study,
current_mns,
top_bottleneck=top_bottleneck,
quantum_direction=quantum_step["direction"],
)
if "max-num-seqs" in tunable
else None
)
if current_chunked and quantum_step["target"] is None and admission_step 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
signature = _effective_config_signature(study, {"env_patch": {}, "flag_patch": patch})
action_id = _prefill_scheduler_action_id(quantum_step["direction"], admission_step)
if signature in seen_signatures:
blocked_candidates.append(
_blocked_candidate(
action_id=action_id,
knob_family="prefill-scheduler-interaction",
config_patch={"env_patch": {}, "flag_patch": patch},
blocked_reason="blocked_noop_or_repeat_effective_full_config",
effective_config_signature=signature,
)
)
return []
current_ratio = current_mbt / prompt_scale if current_mbt > 0 else None
target_mbt = quantum_step["target"] if quantum_step["target"] is not None else current_mbt
target_ratio = target_mbt / prompt_scale if target_mbt > 0 else None
confidence = _hypothesis_confidence(bottleneck_hypotheses, top_bottleneck)
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
score = relief * max(confidence, 0.35) + _information_gain(bottleneck_hypotheses, "runtime") + 0.08
factors = {
"expected_bottleneck_relief": round(relief, 4),
"bottleneck_confidence": round(confidence, 4),
"information_gain": round(_information_gain(bottleneck_hypotheses, "runtime"), 4),
"launch_safety": 0.08,
"regression_risk": 0.06 if current_chunked else 0.1,
"prefill_quantum_ratio_current": (
round(current_ratio, 4) if current_ratio is not None else None
),
"prefill_quantum_ratio_target": (
round(target_ratio, 4) if target_ratio is not None else None
),
"admission_pressure_current": current_mns or None,
"admission_pressure_target": admission_step,
}
actions = [
_runtime_action(
action_id=action_id,
knob_family="prefill-scheduler-interaction",
score=score,
score_factors=factors,
patch=patch,
hypothesis=(
"Test the prefill scheduler hypothesis in normalized control space: "
"chunked prefill changes the scheduler mode, max-num-batched-tokens "
"controls prefill_quantum_ratio, and max-num-seqs controls admission pressure."
),
expected_effects=[
"preserve the settled topology while perturbing scheduler controls",
"reduce long-prefill head-of-line blocking when the prefill quantum is too large",
"reject this scheduler hypothesis if request_rate_per_gpu does not improve under the configured SLO",
],
)
]
seen_signatures.add(signature)
return actions
def _prefill_scheduler_workload_applies(
study: StudySpec,
window_summary: dict[str, Any],
) -> bool:
if study.trace.request_mode == "decode_only":
return False
prompt_p95 = _as_float(window_summary.get("prompt_tokens_p95"))
prompt_p99 = _as_float(window_summary.get("prompt_tokens_p99"))
tail_ratio = _as_float(window_summary.get("prompt_tail_ratio_p95_p50"))
if prompt_p95 <= 0 and prompt_p99 > 0:
prompt_p95 = prompt_p99
if _length_regime(prompt_p95, tail_ratio) == "short_or_moderate":
return False
prefix_cache = window_summary.get("prefix_cache")
cache_ratio = 0.0
if isinstance(prefix_cache, dict):
cache_ratio = _as_float(prefix_cache.get("repeated_token_ratio_estimate"))
return _cache_regime(cache_ratio) != "high_prefix_reuse"
def _prefill_prompt_scale(window_summary: dict[str, Any]) -> float:
prompt_p95 = _as_float(window_summary.get("prompt_tokens_p95"))
prompt_p99 = _as_float(window_summary.get("prompt_tokens_p99"))
if prompt_p95 > 0:
return prompt_p95
return prompt_p99
def _next_prefill_quantum_step(
current_mbt: int,
prompt_scale: float,
*,
top_bottleneck: str,
) -> dict[str, Any]:
if current_mbt <= 0:
return {
"direction": "seed",
"target": _round_up_to_multiple(int(prompt_scale), 1024),
}
ratio = current_mbt / prompt_scale if prompt_scale > 0 else 0.0
if top_bottleneck == "ttft_prefill" and ratio > _PREFILL_QUANTUM_HEAD_OF_LINE_RATIO:
target = int((current_mbt * prompt_scale) ** 0.5)
target = _round_up_to_multiple(target, 1024)
if target < current_mbt:
return {"direction": "lower", "target": target}
if ratio < _PREFILL_QUANTUM_FRAGMENTATION_RATIO:
target = int((current_mbt * prompt_scale) ** 0.5)
target = _round_up_to_multiple(target, 1024)
if target > current_mbt:
return {"direction": "raise", "target": target}
return {"direction": "hold", "target": None}
def _next_admission_pressure_step(
study: StudySpec,
current_mns: int,
*,
top_bottleneck: str,
quantum_direction: str,
) -> int | None:
if current_mns <= 0:
return None
target_concurrency = max(int(study.trace.max_concurrency), 1)
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)
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)
return None
def _prefill_scheduler_action_id(quantum_direction: str, admission_target: int | 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"
return "enable_chunked_prefill_scheduler_mode"
def _runtime_action(
*,
action_id: str,
knob_family: str,
score: float,
score_factors: dict[str, Any] | None = None,
patch: dict[str, Any],
hypothesis: str,
expected_effects: list[str],
@@ -1649,7 +1879,8 @@ def _runtime_action(
"action_id": action_id,
"knob_family": knob_family,
"score": round(score, 4),
"score_factors": {
"score_factors": score_factors
or {
"expected_bottleneck_relief": round(max(score - 0.1, 0.0), 4),
"information_gain": 0.1,
"launch_safety": 0.05,