From cd25c41c5341b4cc5f251a816b6885ac692e7a3a Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Tue, 7 Jul 2026 11:28:25 -0700 Subject: [PATCH 01/24] autoresearch: generate non-catalog levers --- docs/autoresearch.md | 230 +++++++++ gitm/agents/autoresearch.py | 888 ++++++++++++++++++++++++++++++++ gitm/scheduler/loop.py | 89 ++++ tests/test_autoresearch.py | 681 ++++++++++++++++++++++++ tests/test_run_loop_workload.py | 32 ++ 5 files changed, 1920 insertions(+) create mode 100644 docs/autoresearch.md create mode 100644 gitm/agents/autoresearch.py create mode 100644 tests/test_autoresearch.py diff --git a/docs/autoresearch.md b/docs/autoresearch.md new file mode 100644 index 0000000..eb5b845 --- /dev/null +++ b/docs/autoresearch.md @@ -0,0 +1,230 @@ +# Autoresearch + +Autoresearch is the agentic half of the optimizer's search. The curated +intervention library (`gitm/kernels/library.yaml`) is finite and reviewed; +autoresearch proposes **non-catalog** levers — real runtime knobs that are not in +the library — constrained to the bottleneck class attribution already identified, +then routes each proposal through the *same* gate and rollback path as a catalog +lever. + +It is a **candidate source, not a new trust path**. Nothing it proposes can +bypass the selection gate or be kept without a measured win. + +Code: `gitm/agents/autoresearch.py`. + +## Pipeline + +``` +trace ──▶ classify_bottleneck ──▶ propose ──▶ select_interventions ──▶ apply_intervention + (idle/memory/compute) (candidate (safety gate + (snapshot → apply → + specs) replay ranking) measure → keep|rollback) +``` + +Every stage is one it shares with the catalog: + +1. **`select_interventions`** (`gitm/agents/policy.py`) — pre-filters on the + safety tier and qualification commit, then ranks survivors by counterfactual + replay (`predict_delta`). A rejected proposal is recorded and dropped. +2. **`apply_intervention`** (`gitm/optimizer/apply.py`) — snapshots, applies, + measures, and keeps the change only if the measured delta clears + `min_keep_delta`; otherwise it restores. A proposal that doesn't measurably + help is rolled back. + +## Bottleneck classification + +`classify_bottleneck(trace) -> str` maps a captured trace to one of three classes +using two signals read straight from the trace: + +| Signal | Meaning | Threshold | +|--------|---------|-----------| +| serialized-concurrency fraction | kernels ran back-to-back on one stream instead of overlapping → scheduling gaps / launch-bound idle | `_SC_THRESHOLD = 0.5` | +| memcpy fraction | data movement dominates the op mix → memory bound | `_MEMCPY_THRESHOLD = 0.25` | + +Each signal is scored against its threshold; the stronger signal wins. If neither +crosses its threshold the workload is **`compute_bound`**. An empty trace has no +stall/movement signal and defaults to `compute_bound`. + +``` +sc_score = serialized_fraction / 0.5 +mem_score = memcpy_fraction / 0.25 +class = compute_bound if max(sc_score, mem_score) < 1 + idle_stall if sc_score >= mem_score + memory_bound otherwise +``` + +This is a deliberately simple v0 heuristic, not a tuned model — the thresholds +only have to route the search into the right candidate table. The rollback gate, +not the classifier, is what protects a wrong route: a bad proposal is measured +and reverted. + +## Repoint at the largest residual + +The class picks *which* levers to try; the residual target picks *where* to aim +them. `largest_residual(res)` takes the residuals the attribution phase already +computed (`gitm/optimizer/monitor.residuals` → the gap-vs-ceiling +`r_kt = (t_obs − t_pred)/t_pred`), aggregates them per op, and returns the op +whose kernels run furthest **over** the predicted ceiling — the biggest gap. + +Two things this deliberately is **not**: +- not `measure_trace`'s residual, which is within-kernel-type jitter (a uniformly + slow, dominant op has ~0 jitter); +- not the top Granger hypothesis, which is ranked by p-value (significance), not + gap magnitude. + +When a target op is found *and it matches a real kernel name* in the trace, each +proposal is scoped to it via `applies_to_kernels`, so the existing ranking gate +(`predict_delta`) weights the lever by that op's share of trace time. If the op +isn't present (the residual label comes from the predicted graph, which may not +match captured kernel names), the target is still **recorded** as the +justification but proposals are left unscoped — tagging an absent op would zero +the coverage and make a relevant lever look worthless. + +**Honest ceiling:** repointing changes what the search aims at, how proposals are +prioritized, and the recorded "why" (`target_op` on every result). It does **not** +by itself produce genuinely different knobs per op — the lever set still comes +from the class table. Real per-op lever affinity would need hand-authored +op→knob mappings, which v0 does not claim. + +## Proposal sources (the Proposer seam) + +Candidates come from a `Proposer` behind a common seam; each returns +`list[InterventionSpec]` and feeds the *same* selection + rollback gate, so a +proposer is a candidate **source**, not a new trust path. + +- **`GenerativeProposer`** (the loop's active source) — *generates* candidates by + searching a workload's knob surface instead of reading a frozen list. It pulls + knobs from a **`KnobSource`**, drops any that duplicate `library.yaml`, keeps + the ones affine to the bottleneck class, and searches a small **value grid** per + knob. This is what makes autoresearch a *search* rather than a lookup: each knob + is tried at several values (`½×`/`2×`/`4×` of the default, a bool flip, the + other members of an enum, or an explicit grid like `2048`/`4096`), emitted as + `autoresearch::=` — and only a measured win survives the gate. +- **`StochasticProposer`** — entropy-guided sampling of the same knob surface. + The class affinity *weights the dice* (affine knobs get most of the mass) but a + nonzero floor (`epsilon`) keeps every knob reachable, so the search can wander + off-class and surface a lever the keyword heuristic would never pick. A seeded + RNG draws the `(knob, value)` samples — reproducible for a given seed, varied by + changing it — and the rollback gate makes unbounded entropy safe. `epsilon=0` + collapses to pure heuristic; higher values explore more widely. (Available + behind the seam; the loop still runs the generative proposer by default.) +- **`TableProposer`** (the fallback) — the static per-class `_RULES` table below. + `FallbackProposer(EngineArgsProposer(), TableProposer())` uses it only when the + active proposer has nothing for a class (e.g. an unknown class), so the reviewed + catalog's levers still apply. + +Neither source can emit a hallucinated knob: the generated names come from the +workload's real config surface, and both paths exclude anything already in +`library.yaml`. + +### Fitting different workloads + +Versatility comes from the `KnobSource`, **not** a `{workload: knobs}` table: + +- **`VLLMKnobSource`** (used by `EngineArgsProposer`, the vLLM binding of + `GenerativeProposer`) introspects `dataclasses.fields(EngineArgs)` when vLLM is + importable — exposing `Literal` fields as searchable enums — and yields a frozen + catalog otherwise, so the path runs offline and deterministically. +- Another workload plugs in by yielding its own `Knob` list and a `workload` + label: `GenerativeProposer(MyKnobSource(), workload="triton-serve")`. Candidates + carry that label in their applicability; nothing enumerates workloads centrally. +- **Class affinity** is per-knob, not per-workload: a `KnobSource` can tag each + `Knob` with the bottleneck classes it targets (`Knob.classes`), and only knobs + with no tag fall back to the vLLM-flavoured keyword heuristic on the name. + +## Fallback table + +The static `_RULES` table — one small, fixed `(knob, value, rationale)` per +class, every knob a **real, current vLLM argument** (verified against +docs.vllm.ai) that is **not** in `library.yaml`: + +| Class | Knob | Value | Rationale | +|-------|------|-------|-----------| +| `idle_stall` | `max_num_partial_prefills` | 4 | raise partial-prefill concurrency so prefill overlaps decode | +| `idle_stall` | `long_prefill_token_threshold` | 2048 | chunk big prompts so they interleave, closing decode gaps | +| `memory_bound` | `cpu_offload_gb` | 4 | offload cold weights to host RAM, freeing HBM for KV cache | +| `memory_bound` | `preemption_mode` | `swap` | swap preempted KV blocks instead of recomputing under pressure | +| `compute_bound` | `compilation_config` | 3 | torch.compile level 3: kernel fusion + piecewise CUDA graphs | + +The **rationales are plausibility arguments, not measured claims**. Each proposal +carries `expected_delta_mean = 0.05` (`lo = 0.0`, `hi = 0.15`) — a modest, +honestly-unproven range — and a `source` that says so. Only the measured A/B +turns a proposal into a real number. + +## Safety posture + +- Proposals are always `moderate` tier — never `high_risk`. Topology and weight + changes stay in the reviewed catalog. +- The rollback gate is the whole safety story for a proposal: applied behind a + snapshot, kept only on a measured win, reverted otherwise. + +## Loop integration + +The 24-hour loop (`gitm/scheduler/loop.py`) runs autoresearch as **Phase 4b** on +the `vllm-decode` path, after the catalog pass: + +- Guarded by remaining budget — if the catalog pass spent the budget, the phase + is skipped (the class it *would* have searched is still reported). +- With no live engine attached the loop uses a `DryRunApplicator`, so autoresearch + claims land **unverified** (`measured_delta = None`) — exactly like unverified + catalog claims, never counted as a proven gain. +- It reuses the attribution phase's already-computed residuals (`res`) to repoint + at the largest-residual op — no recomputation. +- Candidates are **generated** by `FallbackProposer(EngineArgsProposer(), + TableProposer())`: the EngineArgs value-grid search is the active source, with + the static table as the fallback. +- Results are written to `/autoresearch.json` (including the `target` op + and per-result `target_op`); the summary carries `bottleneck_class` and + `n_autoresearch`; generated candidates appear in the report as + `autoresearch::=` claims and rejected ones in the rejected + list. + +## API + +```python +from gitm.agents.autoresearch import ( + EngineArgsProposer, FallbackProposer, TableProposer, + autoresearch, classify_bottleneck, propose, +) + +# End-to-end: classify, find the largest-residual op, generate, gate, apply/rollback. +# Without a proposer it uses the static table; the loop passes the generative one. +proposer = FallbackProposer(EngineArgsProposer(), TableProposer()) +run = autoresearch(trace, applicator=applicator, policy=policy, residuals=res, + proposer=proposer) +run.bottleneck_class # "idle_stall" | "memory_bound" | "compute_bound" +run.target # ResidualTarget(op, residual, n_kernels) | None +run.results # list[AutoresearchResult] + +# Lower-level pieces. +cls = classify_bottleneck(trace) +target = largest_residual(res) # ResidualTarget | None +specs = EngineArgsProposer().propose(cls, target_op=target.op if target else None) +specs = propose(cls, target_op=target.op if target else None) # the static table +``` + +`AutoresearchResult` carries `spec`, `bottleneck_class`, `predicted_delta`, +`applicable`, `rejected_reason`, `measured_delta`, `rolled_back`, and `target_op`. + +## v0 limits and roadmap + +The lever set is now *generated* from the `EngineArgs` surface with a value grid +per knob, not a frozen table — but the generation is still crude and honest about +it: + +- **Class→knob affinity is a keyword heuristic** (`prefill` / `cache` / `compil` + …) matched against the knob name, not a learned or hand-authored mapping. +- **Value grids are blind** — a generic `½×`/`2×`/`4×` ladder off the default (or + a hand-seeded grid where that would be nonsensical, e.g. token thresholds), not + tuned per knob. Per-class candidate counts are bounded by `max_candidates`. +- **Live `EngineArgs` typing is best-effort** — introspected fields are typed + coarsely from their annotation; enums the introspector can't detect fall back + to "no grid" and are simply skipped. +- **The stochastic sampler isn't wired into the loop yet.** `StochasticProposer` + exists behind the seam (seeded, reproducible), but Phase 4b still runs the + generative proposer; flipping the loop onto it (or a `FallbackProposer` chain) + is the remaining wiring. +- **No cross-run feedback yet.** The remaining next steps: an effect model that + reweights the search from realized measured deltas — still blocked because the + loop runs dry until a live applicator is attached — plus sourcing the + bottleneck-class vocabulary from the shared attribution layer rather than the + local `BOTTLENECK_CLASSES` heuristic here. diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py new file mode 100644 index 0000000..6d7be92 --- /dev/null +++ b/gitm/agents/autoresearch.py @@ -0,0 +1,888 @@ +"""Autoresearch — propose non-catalog levers within the attributed bottleneck class. + +The curated library (``library.yaml``) is finite. Autoresearch is the agentic +half of the README's "select from a library of known optimizations *and* run +agentic search for novel ones": it proposes real vLLM config knobs *outside* +that catalog, constrained to the bottleneck class the attribution layer +identified (idle / memory / compute). + +Every proposal is then routed through the exact same path as a catalog lever: + +1. the selection gate — :func:`gitm.agents.policy.select_interventions` — which + pre-filters on the safety tier and qualification commit, then ranks the + survivors by counterfactual replay (:func:`gitm.optimizer.replay.predict_delta`); +2. the rollback-gated live apply — :func:`gitm.optimizer.apply.apply_intervention` — + which snapshots, applies, measures, and keeps only on a measured win. + +A proposal the gate rejects is recorded and dropped; one that applies but does +not measurably help is rolled back. Autoresearch is a *candidate source*, not a +new trust path — nothing it proposes can bypass the gate or be kept without a +measured win. + +The proposed knobs are real, current vLLM arguments (verified against +docs.vllm.ai); their expected deltas, however, are unproven estimates. The +``source`` field says so, and only the measured A/B keeps or discards them. + +v0 classifies the bottleneck from coarse trace telemetry (:func:`classify_bottleneck`) +and repoints the search at the largest-residual op. Candidates come from one of +two sources behind the :class:`Proposer` seam: the static per-class table +(:func:`propose`) or :class:`GenerativeProposer`, which searches a workload's knob +surface (supplied by a :class:`KnobSource` — vLLM's ``EngineArgs`` by default) at +a small value grid per knob. The seam is workload-agnostic: a new workload plugs +in a KnobSource rather than a per-workload table. The loop runs the generative +proposer with the table as a fallback. Later versions learn an effect model from +realized deltas and sample the knob space stochastically. +""" + +from __future__ import annotations + +import random +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Protocol + +from gitm.agents.policy import Policy, select_interventions +from gitm.kernels.library import load_library +from gitm.kernels.spec import Applicability, InterventionSpec, SafetyGate +from gitm.optimizer.apply import Applicator, apply_intervention +from gitm.optimizer.monitor import Residuals, _serialized_fraction +from gitm.tracer.schema import Trace + +if TYPE_CHECKING: + from gitm.optimizer.preconditions import GateContext + +#: The module's public surface: the bottleneck vocabulary, the ``KnobSource`` and +#: ``Proposer`` seams (so any workload can plug in), and the entry points. Helpers +#: prefixed with ``_`` (the fallback table, value grids, EngineArgs introspection) +#: are internal. +__all__ = [ + "BOTTLENECK_CLASSES", + "IDLE_STALL", + "MEMORY_BOUND", + "COMPUTE_BOUND", + "classify_bottleneck", + "ResidualTarget", + "largest_residual", + "Knob", + "KnobSource", + "VLLMKnobSource", + "Proposer", + "TableProposer", + "GenerativeProposer", + "EngineArgsProposer", + "StochasticProposer", + "FallbackProposer", + "propose", + "AutoresearchResult", + "AutoresearchRun", + "autoresearch", + "autoresearch_v0", +] + +# --- bottleneck classification ---------------------------------------------- +# +# The attribution vocabulary autoresearch searches within. Nothing upstream +# emits these labels yet, so v0 derives them from coarse trace telemetry. These +# are deliberately simple heuristics, not a tuned model — the thresholds only +# have to route the search into the right candidate table; the rollback gate is +# what actually protects a wrong route (a bad proposal is measured and reverted). + +#: The bottleneck classes autoresearch searches within — the single source of +#: truth shared by classify_bottleneck (the producer), the keyword-affinity map, +#: and the fallback table, so the three can't drift (guarded by a test). These are +#: workload-agnostic GPU-execution categories, not a per-workload vocabulary. +IDLE_STALL = "idle_stall" +MEMORY_BOUND = "memory_bound" +COMPUTE_BOUND = "compute_bound" +BOTTLENECK_CLASSES = (IDLE_STALL, MEMORY_BOUND, COMPUTE_BOUND) + +#: Serialized-concurrency fraction above this ⇒ kernels ran back-to-back on one +#: stream instead of overlapping: scheduling gaps / launch-bound idle time. +_SC_THRESHOLD = 0.5 +#: memcpy share of GPU operations above this ⇒ data movement dominates. +_MEMCPY_THRESHOLD = 0.25 + + +def classify_bottleneck(trace: Trace) -> str: + """Map a captured trace to one of ``idle_stall`` / ``memory_bound`` / ``compute_bound``. + + v0 heuristic on two signals read straight from the trace: the + serialized-concurrency fraction (poor kernel overlap ⇒ idle/scheduling gaps) + and the memcpy fraction (data movement dominating ⇒ memory bound). Each is + scored against its threshold; the stronger signal wins, and if neither + crosses its threshold the workload is treated as compute bound. An empty + trace has no stall/movement signal, so it defaults to ``compute_bound``. + """ + kernels = trace.kernels() + if not kernels: + return COMPUTE_BOUND + + memcpys = [e for e in trace.events if e.kind == "memcpy"] + sc = _serialized_fraction(kernels) + kernel_ns = sum(max(0, k.end_ns - k.start_ns) for k in kernels) + memcpy_ns = sum(max(0, e.end_ns - e.start_ns) for e in memcpys) + gpu_op_ns = kernel_ns + memcpy_ns + memcpy_frac = memcpy_ns / gpu_op_ns if gpu_op_ns else 0.0 + + sc_score = sc / _SC_THRESHOLD + mem_score = memcpy_frac / _MEMCPY_THRESHOLD + if max(sc_score, mem_score) < 1.0: + return COMPUTE_BOUND + return IDLE_STALL if sc_score >= mem_score else MEMORY_BOUND + + +# --- residual targeting ----------------------------------------------------- +# +# "Repoint at the largest residual": instead of aiming the search at the whole +# trace, aim it at the single op whose kernels run furthest *over* the predicted +# ceiling — the biggest gap the attribution layer found. This is the gap residual +# ``r_kt = (t_obs - t_pred)/t_pred`` from gitm.optimizer.monitor.residuals(), NOT +# the within-kernel-type jitter that measure_trace() computes, and NOT the Granger +# p-value rank (significance, not magnitude). + + +@dataclass +class ResidualTarget: + """The op with the largest kernel-time gap vs its predicted ceiling.""" + + op: str + residual: float # mean r_kt over the op's kernels (fraction over the ceiling) + n_kernels: int + + +def largest_residual(res: Residuals) -> ResidualTarget | None: + """The op whose kernels run furthest over the predicted ceiling. + + Aggregates the per-kernel gap residual by op (mean ``r_kt``) and returns the + op with the largest *positive* mean — the biggest bottleneck, not the + jitteriest op. Returns ``None`` when there is no residual data or nothing runs + over its ceiling (all means ≤ 0). + """ + if not res.per_kernel: + return None + by_op: dict[str, list[float]] = {} + for kr in res.per_kernel: + by_op.setdefault(kr.op, []).append(kr.r_kt) + op, values = max(by_op.items(), key=lambda kv: sum(kv[1]) / len(kv[1])) + mean = sum(values) / len(values) + if mean <= 0: + return None + return ResidualTarget(op=op, residual=mean, n_kernels=len(values)) + + +def _op_present(trace: Trace, op: str) -> bool: + """True if ``op`` names (a substring of) a real kernel in the trace. + + Guards the ``applies_to_kernels`` tagging: the residual op label comes from + the predicted graph, so only tag a proposal with it when it actually matches + a captured kernel name — otherwise ``predict_delta`` coverage would be 0 and + the proposal would rank as worthless rather than untargeted. + """ + return any(op in k.name for k in trace.kernels()) + + +# --- candidate table -------------------------------------------------------- +# +# Per-bottleneck candidate perturbations: (knob, value, one-line rationale). +# Every knob is a real, current vLLM argument (docs.vllm.ai) that is NOT in the +# curated library.yaml — autoresearch proposes *outside* the catalog. The +# rationales are plausibility arguments, not measured claims. +_RULES: dict[str, list[tuple[str, object, str]]] = { + "idle_stall": [ + ("max_num_partial_prefills", 4, + "raise partial-prefill concurrency so prefill overlaps decode instead of stalling it"), + ("long_prefill_token_threshold", 2048, + "lower the long-prefill threshold so big prompts chunk and interleave, closing decode gaps"), + ], + "memory_bound": [ + ("cpu_offload_gb", 4, + "offload cold weights to host RAM to free HBM for a larger KV cache"), + ("preemption_mode", "swap", + "swap preempted KV blocks to host instead of recomputing them under memory pressure"), + ], + "compute_bound": [ + ("compilation_config", 3, + "raise torch.compile to level 3 for kernel fusion + piecewise CUDA graphs"), + ], +} + + +@dataclass +class AutoresearchResult: + spec: InterventionSpec + bottleneck_class: str + predicted_delta: float + applicable: bool + rejected_reason: str | None + measured_delta: float | None + rolled_back: bool + target_op: str | None = None # the largest-residual op this proposal aimed at + + +@dataclass +class AutoresearchRun: + """One end-to-end autoresearch pass: the classified bottleneck + its results.""" + + bottleneck_class: str + results: list[AutoresearchResult] = field(default_factory=list) + target: ResidualTarget | None = None # the largest-residual op the search aimed at + + +#: The honest, unproven delta band every candidate carries until the measured A/B +#: replaces it. One place to tune what "proposed, not measured" means. +_DELTA_MEAN, _DELTA_LO, _DELTA_HI = 0.05, 0.0, 0.15 + + +def _candidate_spec( + *, + name: str, + summary: str, + knob: str, + value: object, + applies_to_kernels: list[str], + bottleneck_class: str, + workload: str, + source: str, +) -> InterventionSpec: + """Build a candidate spec with the fields every proposer forces. + + The single place for the honest-but-unproven delta band, the workload + applicability, and the moderate/rollback-gated safety posture — so the table + and generative proposers can't drift apart on what a *candidate* is. Only the + caller-varying parts (name, summary, knob/value, source) are parameters. + """ + return InterventionSpec( + name=name, + summary=summary, + knob=knob, + value=value, # int | float | str | bool + applies_to_kernels=applies_to_kernels, + # Proposed, not measured: an honest, modest range. The measured A/B is + # what turns this into a real number. + expected_delta_mean=_DELTA_MEAN, + expected_delta_lo=_DELTA_LO, + expected_delta_hi=_DELTA_HI, + source=source, + applicability=Applicability(workloads=[workload], other=f"targets {bottleneck_class}"), + # Unproven ⇒ never high-risk (topology/weights changes stay in the + # reviewed catalog). Moderate + the rollback gate is the whole safety + # story for a candidate. + safety=SafetyGate( + tier="moderate", + notes="autoresearch candidate — kept only on a measured, rollback-gated win.", + ), + ) + + +def propose(bottleneck_class: str, *, target_op: str | None = None) -> list[InterventionSpec]: + """Emit candidate specs for a bottleneck class (empty if the class is unknown). + + The static ``_RULES`` table — the offline fallback. When ``target_op`` is + given, each proposal is scoped to that op via ``applies_to_kernels`` so the + ranking gate (``predict_delta``) weights it by that op's share of trace time — + this is how "repoint at the largest residual" reaches the selection. The op is + the caller's job to validate against the trace (see :func:`_op_present`); an + off-trace op would zero the coverage. + """ + applies = [target_op] if target_op else [] + aim = f" (targeting {target_op})" if target_op else "" + return [ + _candidate_spec( + name=f"autoresearch:{bottleneck_class}:{knob}", + summary=why + aim, + knob=knob, + value=value, + applies_to_kernels=applies, + bottleneck_class=bottleneck_class, + workload="vllm-decode", + source="autoresearch-v0 (proposed knob, not catalog; verified real vLLM arg)", + ) + for knob, value, why in _RULES.get(bottleneck_class, []) + ] + + +# --- proposal sources (the "Proposer" seam) --------------------------------- +# +# ``propose`` above is the static per-class table. ``GenerativeProposer`` is the +# generative counterpart: instead of a frozen list it searches a workload's knob +# surface — supplied by a ``KnobSource`` — keeping the knobs affine to the +# attributed bottleneck class and trying a small value grid per knob. +# ``VLLMKnobSource`` (introspect ``EngineArgs``) is one source; another workload +# plugs in by yielding its own knobs, so the mechanism is workload-agnostic with +# no ``{workload: knobs}`` table. Both proposers return ``list[InterventionSpec]`` +# and feed the exact same selection + rollback gate — a Proposer is a candidate +# *source*, not a new trust path. Nothing here can propose a knob outside the +# workload's real surface, or one that duplicates the curated library. + + +class Proposer(Protocol): + """A source of candidate specs for a bottleneck class. The gate is source-agnostic.""" + + def propose( + self, bottleneck_class: str, *, target_op: str | None = None + ) -> list[InterventionSpec]: ... + + +class TableProposer: + """The static ``_RULES`` table (see :func:`propose`) as a Proposer. + + The offline default and the fallback under :class:`FallbackProposer` when the + generative proposer has nothing to offer for a class. + """ + + def propose( + self, bottleneck_class: str, *, target_op: str | None = None + ) -> list[InterventionSpec]: + return propose(bottleneck_class, target_op=target_op) + + +@dataclass(frozen=True) +class Knob: + """A workload config knob and how to search its value. + + ``grid`` is an explicit set of search points, used where a derived grid would + be nonsensical (e.g. token thresholds); when empty, the grid is derived from + ``kind``/``default``. ``classes`` optionally tags which bottleneck classes the + knob is affine to — a :class:`KnobSource` can declare this for any class + vocabulary, so affinity need not rely on the vLLM-flavoured keyword heuristic. + """ + + name: str + kind: str # "int" | "float" | "bool" | "enum" | "str" + default: object = None + choices: tuple = () + grid: tuple = () + classes: tuple = () # bottleneck classes this knob is affine to (optional) + + +#: Frozen fallback catalog: real, current vLLM EngineArgs (docs.vllm.ai) that are +#: NOT in library.yaml. Used when vLLM can't be imported (air-gapped operator, no +#: GPU stack) so the generative path still runs offline and deterministically. +_FALLBACK_KNOBS: tuple[Knob, ...] = ( + Knob("max_num_partial_prefills", "int", default=1), + Knob("long_prefill_token_threshold", "int", default=0, grid=(2048, 4096)), + Knob("max_long_partial_prefills", "int", default=1), + Knob("cpu_offload_gb", "int", default=0), + Knob("preemption_mode", "enum", default="recompute", choices=("recompute", "swap")), + Knob("compilation_config", "int", default=0, grid=(2, 3)), +) + +#: Which knobs to search for each bottleneck class, matched against the knob NAME +#: (substring). Deliberately a keyword heuristic, and honest about being one: it +#: biases the search toward the attributed class; the gate does the proving. +_CLASS_KEYWORDS: dict[str, tuple[str, ...]] = { + "idle_stall": ("prefill", "partial", "schedul", "chunk", "overlap"), + "memory_bound": ("cache", "swap", "offload", "block", "gpu_memory", "kv", "preempt", "cpu"), + "compute_bound": ("compil", "cudagraph", "cuda_graph", "graph", "quant", "fus", "eager"), +} + + +def _affine(knob: Knob, bottleneck_class: str, keywords: tuple[str, ...]) -> bool: + """Is ``knob`` worth searching for this class? + + An explicit ``Knob.classes`` tag wins (a source can declare affinity for any + class vocabulary); otherwise fall back to a keyword-substring match on the + name — the honest, vLLM-flavoured heuristic for sources that don't self-tag. + """ + if knob.classes: + return bottleneck_class in knob.classes + lname = knob.name.lower() + return any(k in lname for k in keywords) + + +#: Multipliers applied to a numeric knob's default to derive its search points. +_GRID_MULTIPLIERS = (0.5, 2.0, 4.0) + + +def _value_grid(knob: Knob) -> list[object]: + """A small set of candidate values to search for ``knob`` (excludes the default). + + An explicit ``grid`` wins; otherwise the grid is derived from the kind: flip a + bool, the other members of an enum, or a ½×/2×/4× ladder for a number. This is + what turns the search from a single-value lookup into an actual value search. + """ + if knob.grid: + return [v for v in knob.grid if v != knob.default] + if knob.kind == "bool": + return [not bool(knob.default)] + if knob.kind == "enum": + return [c for c in knob.choices if c != knob.default] + if knob.kind in ("int", "float"): + d = knob.default + base = d if isinstance(d, int | float) and d not in (0, False) else 1 + raw = [base * m for m in _GRID_MULTIPLIERS] + if knob.kind == "int": + vals = sorted({int(round(x)) for x in raw if round(x) >= 1}) + else: + vals = sorted({round(x, 3) for x in raw if x > 0}) + return [v for v in vals if v != d] + return [] # unknown / free-form (str): nothing safe to search + + +def _annotation_kind(annotation: object) -> str: + """Coarsely map a dataclass-field annotation to a value-grid kind (best-effort).""" + text = str(annotation).lower() + if "bool" in text: + return "bool" + if "int" in text: + return "int" + if "float" in text: + return "float" + return "str" + + +def _field_kind_and_choices(annotation: object) -> tuple[str, tuple]: + """(kind, choices) for a field annotation. + + A ``Literal[...]`` becomes an enum with its members as choices, so those knobs + become searchable (a value grid = the other members). Anything else falls back + to the coarse string match with no choices. Best-effort: vLLM's stringised + annotations (``from __future__ import annotations``) won't resolve here, so + Literal extraction only fires when the annotation is a real typing object. + """ + try: + import typing + + if typing.get_origin(annotation) is typing.Literal: + return "enum", tuple(typing.get_args(annotation)) + except Exception: + pass + return _annotation_kind(annotation), () + + +#: EngineArgs field-name fragments that are never runtime *performance* knobs — +#: model identity, I/O paths, logging, RNG. Mutating them wouldn't close a +#: bottleneck (and could be nonsensical), so the introspected surface excludes +#: them even though they're typed int/bool. vLLM-specific, so applied only here; +#: the curated fallback catalog is authoritative and left untouched. +_NON_TUNABLE_HINTS = ( + "model", + "tokenizer", + "seed", + "log", + "name", + "path", + "dir", + "revision", + "trust_remote_code", + "config_format", + "download", +) + + +def _is_tunable(field_name: str) -> bool: + """False for EngineArgs fields that aren't runtime performance knobs.""" + lname = field_name.lower() + return not any(h in lname for h in _NON_TUNABLE_HINTS) + + +def _engine_arg_knobs() -> list[Knob]: + """Enumerate real vLLM EngineArgs, or fall back to the frozen catalog. + + Best-effort introspection: when vLLM is importable, each tunable dataclass + field is a candidate knob (typed as best we can from its annotation, with + ``Literal`` fields exposed as searchable enums). Non-performance fields (model + identity, paths, logging, RNG — see :data:`_NON_TUNABLE_HINTS`) are skipped. + When vLLM isn't importable (no GPU stack / air-gapped), the frozen catalog + keeps the generative path working offline. + """ + try: + import dataclasses + + from vllm import EngineArgs # type: ignore + except Exception: + return list(_FALLBACK_KNOBS) + + knobs: list[Knob] = [] + for f in dataclasses.fields(EngineArgs): + if not _is_tunable(f.name): + continue + default = None if f.default is dataclasses.MISSING else f.default + kind, choices = _field_kind_and_choices(f.type) + knobs.append(Knob(name=f.name, kind=kind, default=default, choices=choices)) + return knobs + + +class KnobSource(Protocol): + """Yields the knob surface to search — a workload's config namespace. + + vLLM's is :class:`VLLMKnobSource` (introspect ``EngineArgs``). Another workload + plugs in by yielding its own ``Knob`` list; there is deliberately no + ``{workload: knobs}`` table — versatility comes from the source, not a map. + """ + + def knobs(self) -> list[Knob]: ... + + +class VLLMKnobSource: + """The real vLLM ``EngineArgs`` surface (frozen fallback when vLLM is absent).""" + + def knobs(self) -> list[Knob]: + return _engine_arg_knobs() + + +@dataclass(frozen=True) +class _ListKnobSource: + """A fixed knob list as a KnobSource (the ``knobs=`` convenience and tests).""" + + _knobs: tuple[Knob, ...] + + def knobs(self) -> list[Knob]: + return list(self._knobs) + + +class _ProposerBase: + """Shared setup + eligibility for knob-surface proposers. + + Holds the source, workload label, catalog exclusion, and affinity map, and + exposes the two pieces every knob-surface proposer needs: the searchable knob + set (outside the catalog, with something to search) and the per-candidate spec + builder. Subclasses decide *how* to pick from the searchable knobs — exhaustive + value grid (:class:`GenerativeProposer`) or weighted sampling + (:class:`StochasticProposer`). + """ + + def __init__( + self, + knob_source: KnobSource, + *, + workload: str = "vllm-decode", + catalog_knobs: set[str] | None = None, + affinity_keywords: dict[str, tuple[str, ...]] | None = None, + ) -> None: + self._source = knob_source + self._workload = workload + # Keyword affinity is the fallback for knobs that don't self-tag via + # ``Knob.classes``. Default is the vLLM-flavoured vocabulary; a workload + # with its own knob naming can supply its own (still no per-workload table). + self._affinity = _CLASS_KEYWORDS if affinity_keywords is None else affinity_keywords + self._catalog = ( + set(catalog_knobs) + if catalog_knobs is not None + else {s.knob for s in load_library()} + ) + + def _searchable(self) -> list[Knob]: + """Knobs worth searching: outside the catalog and with a non-empty grid.""" + return [ + k for k in self._source.knobs() if k.name not in self._catalog and _value_grid(k) + ] + + def _spec( + self, + bottleneck_class: str, + knob: Knob, + value: object, + *, + target_op: str | None, + verb: str, + source: str, + ) -> InterventionSpec: + aim = f" (targeting {target_op})" if target_op else "" + return _candidate_spec( + name=f"autoresearch:{bottleneck_class}:{knob.name}={value}", + summary=f"{verb} {knob.name}={value} for {bottleneck_class}{aim}", + knob=knob.name, + value=value, + applies_to_kernels=[target_op] if target_op else [], + bottleneck_class=bottleneck_class, + workload=self._workload, + source=source, + ) + + +class GenerativeProposer(_ProposerBase): + """Search a workload's knob surface exhaustively, gated like any spec. + + Pulls knobs from ``knob_source`` (any workload's config namespace), drops any + that duplicate the curated library, keeps the ones affine to the bottleneck + class (an explicit ``Knob.classes`` tag, else a keyword heuristic on the name), + and emits one candidate per value-grid point. Forced to ``moderate`` tier with + an honest, unproven delta band — it can only *widen* the candidate set; the + selection + rollback gate is what keeps anything. ``workload`` labels the + candidates' applicability, so one mechanism serves any workload without a table. + """ + + def __init__( + self, + knob_source: KnobSource, + *, + workload: str = "vllm-decode", + catalog_knobs: set[str] | None = None, + max_candidates: int | None = None, + affinity_keywords: dict[str, tuple[str, ...]] | None = None, + ) -> None: + super().__init__( + knob_source, + workload=workload, + catalog_knobs=catalog_knobs, + affinity_keywords=affinity_keywords, + ) + self._max = max_candidates + + def propose( + self, bottleneck_class: str, *, target_op: str | None = None + ) -> list[InterventionSpec]: + keywords = self._affinity.get(bottleneck_class, ()) + out = [ + self._spec( + bottleneck_class, + knob, + value, + target_op=target_op, + verb="search", + source="autoresearch-v0 (generated candidate; real workload knob, unproven delta)", + ) + for knob in self._searchable() + if _affine(knob, bottleneck_class, keywords) + for value in _value_grid(knob) + ] + # Bound the per-class candidate count so a large config surface can't + # flood the gate; the rollback gate still ranks and proves what survives. + return out if self._max is None else out[: self._max] + + +class EngineArgsProposer(GenerativeProposer): + """vLLM binding of :class:`GenerativeProposer`: the EngineArgs surface + vllm-decode. + + The loop's convenience entry point. ``knobs=`` overrides the surface (used by + tests); otherwise it introspects EngineArgs, falling back to the frozen catalog + offline. Defaults to a candidate cap since the real ``EngineArgs`` surface is + large — the offline fallback catalog is well under it, so counts are unchanged. + """ + + def __init__( + self, + *, + knobs: list[Knob] | None = None, + catalog_knobs: set[str] | None = None, + max_candidates: int | None = 24, + ) -> None: + source: KnobSource = ( + _ListKnobSource(tuple(knobs)) if knobs is not None else VLLMKnobSource() + ) + super().__init__( + source, + workload="vllm-decode", + catalog_knobs=catalog_knobs, + max_candidates=max_candidates, + ) + + +class FallbackProposer: + """Try ``primary``; use ``secondary`` only when primary yields nothing. + + Wires the generative proposer as the active source with the static table as + the genuine fallback — a class the EngineArgs surface can't populate (or an + unknown class) still gets the reviewed catalog's levers. + """ + + def __init__(self, primary: Proposer, secondary: Proposer) -> None: + self._primary = primary + self._secondary = secondary + + def propose( + self, bottleneck_class: str, *, target_op: str | None = None + ) -> list[InterventionSpec]: + specs = self._primary.propose(bottleneck_class, target_op=target_op) + return specs or self._secondary.propose(bottleneck_class, target_op=target_op) + + +class StochasticProposer(_ProposerBase): + """Entropy-guided sampling of a workload's knob surface (reproducible by seed). + + The heuristic weights the dice: knobs affine to the bottleneck class carry most + of the mass, but every eligible knob keeps a nonzero floor (``epsilon``), so the + search can wander off-class and surface a lever the keyword heuristic would + never pick. A seeded RNG draws the actual (knob, value) candidates — + reproducible for a given seed, varied by changing it — and the rollback gate + makes unbounded entropy safe. Same seam as the others: a candidate *source*, not + a trust path. ``epsilon=0`` collapses to pure heuristic (affine knobs only); + higher ``epsilon`` explores more widely. + """ + + def __init__( + self, + knob_source: KnobSource, + *, + workload: str = "vllm-decode", + catalog_knobs: set[str] | None = None, + n_samples: int = 6, + seed: int = 0, + epsilon: float = 0.15, + affinity_keywords: dict[str, tuple[str, ...]] | None = None, + ) -> None: + super().__init__( + knob_source, + workload=workload, + catalog_knobs=catalog_knobs, + affinity_keywords=affinity_keywords, + ) + self._n = n_samples + self._seed = seed + self._epsilon = epsilon + + def propose( + self, bottleneck_class: str, *, target_op: str | None = None + ) -> list[InterventionSpec]: + keywords = self._affinity.get(bottleneck_class, ()) + eligible = [(k, _value_grid(k)) for k in self._searchable()] + eligible = [(k, grid) for k, grid in eligible if grid] + # Bias the dice toward the attributed class; the floor keeps off-class knobs + # reachable — that's the entropy. All-zero weight (epsilon=0, nothing affine) + # means no heuristic signal *and* no entropy, so nothing to sample. + weights = [ + 1.0 if _affine(k, bottleneck_class, keywords) else self._epsilon + for k, _grid in eligible + ] + if not any(weights): + return [] + + rng = random.Random(self._seed) # noqa: S311 - reproducible search, not security + seen: set[tuple[str, object]] = set() + out: list[InterventionSpec] = [] + max_attempts = max(self._n * 4, len(eligible) * 4) + for _ in range(max_attempts): + if len(out) >= self._n: + break + knob, grid = rng.choices(eligible, weights=weights, k=1)[0] + value = rng.choice(grid) + if (knob.name, value) in seen: # don't gate the same candidate twice + continue + seen.add((knob.name, value)) + out.append( + self._spec( + bottleneck_class, + knob, + value, + target_op=target_op, + verb="sample", + source="autoresearch-v0 (stochastic sample; real workload knob, unproven delta)", + ) + ) + return out + + +def autoresearch_v0( + trace: Trace, + bottleneck_class: str, + *, + applicator: Applicator, + policy: Policy | None = None, + min_keep_delta: float = 0.0, + target: ResidualTarget | None = None, + proposer: Proposer | None = None, + ctx: GateContext | None = None, + reject: Callable[[InterventionSpec], str | None] | None = None, +) -> list[AutoresearchResult]: + """Propose → gate → (apply + measure + rollback) for one bottleneck class. + + Proposals are ranked and pre-filtered by :func:`select_interventions` (the + same gate the catalog goes through), then each survivor is applied behind the + rollback gate so a proposal that doesn't clear ``min_keep_delta`` is reverted. + + ``proposer`` chooses the candidate source. Default (``None``) is the static + :func:`propose` table; pass an :class:`EngineArgsProposer` (or a + :class:`FallbackProposer`) to generate candidates from the real EngineArgs + surface. Either way the ranking + rollback gate below is identical. + + ``ctx`` is the precondition gate context, forwarded to + :func:`select_interventions` so candidates face the *same* applicability gate + as the catalog. ``reject`` is an optional per-candidate veto applied after the + gate but before apply (the loop uses it for the live structural-knob-needs- + restart guard) — it keeps engine-specific policy out of this workload-agnostic + core. + + A ``target`` (the largest-residual op) repoints the search: when that op is + present in the trace, proposals are scoped to it so the gate prioritizes + levers hitting the biggest gap, and every result records the op it aimed at. + """ + # Only tag with the op when it matches a real kernel — otherwise coverage is 0. + target_op = target.op if (target is not None and _op_present(trace, target.op)) else None + if proposer is None: + proposals = propose(bottleneck_class, target_op=target_op) + else: + proposals = proposer.propose(bottleneck_class, target_op=target_op) + if not proposals: + return [] + + ranked = select_interventions( + trace, proposals, policy or Policy(), top_n=len(proposals), ctx=ctx + ) + aimed_at = target.op if target is not None else None + + results: list[AutoresearchResult] = [] + for c in ranked: + # Gate rejection wins; else the caller's veto (e.g. a live structural knob + # with no restart hook) can reject before we touch the engine. Rejected + # candidates are recorded but never applied; survivors go through the + # rollback-gated apply. Both land in one result shape. + reason = c.rejected_reason + if reason is None and reject is not None: + reason = reject(c.spec) + applied = ( + apply_intervention(c.spec, applicator, min_keep_delta=min_keep_delta) + if reason is None + else None + ) + results.append( + AutoresearchResult( + spec=c.spec, + bottleneck_class=bottleneck_class, + predicted_delta=c.predicted_delta, + applicable=applied is not None, + rejected_reason=reason, + measured_delta=applied.measured_delta if applied else None, + rolled_back=applied.rolled_back if applied else False, + target_op=aimed_at, + ) + ) + return results + + +def autoresearch( + trace: Trace, + *, + applicator: Applicator, + policy: Policy | None = None, + min_keep_delta: float = 0.0, + residuals: Residuals | None = None, + proposer: Proposer | None = None, + ctx: GateContext | None = None, + reject: Callable[[InterventionSpec], str | None] | None = None, +) -> AutoresearchRun: + """Classify the trace's bottleneck, then run the full propose→gate→apply pass. + + This is the end-to-end entry point: hand it a captured trace and a live + applicator and it decides which class to search, proposes non-catalog levers + for that class, and routes each through the selection + rollback gates. + + ``proposer`` selects the candidate source (default: the static table); the + loop passes an :class:`EngineArgsProposer` so the search generates candidates + from the real EngineArgs surface rather than a frozen list. ``ctx`` forwards + the precondition gate context (same applicability gate as the catalog); + ``reject`` is an optional per-candidate veto (the loop's structural-knob + guard). See :func:`autoresearch_v0`. + + When ``residuals`` (from :func:`gitm.optimizer.monitor.residuals`) are passed, + the search is repointed at the largest-residual op — the biggest gap vs the + predicted ceiling — rather than the whole trace. The loop already computes + these in its attribution phase, so it passes them straight through. + """ + bottleneck_class = classify_bottleneck(trace) + target = largest_residual(residuals) if residuals is not None else None + return AutoresearchRun( + bottleneck_class=bottleneck_class, + target=target, + results=autoresearch_v0( + trace, + bottleneck_class, + applicator=applicator, + policy=policy, + min_keep_delta=min_keep_delta, + target=target, + proposer=proposer, + ctx=ctx, + reject=reject, + ), + ) diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index a3fed2e..9634e8f 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -18,6 +18,14 @@ from typing import Any from gitm._paths import runs_dir, traces_dir +from gitm.agents.autoresearch import ( + AutoresearchRun, + EngineArgsProposer, + FallbackProposer, + TableProposer, + autoresearch, + classify_bottleneck, +) from gitm.agents.policy import Policy, select_interventions from gitm.kernels.library import load_library from gitm.optimizer.apply import ( @@ -536,6 +544,85 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: if time.time_ns() - started_ns >= int(budget_s * 1e9): break + + # Phase 4b - agentic autoresearch through the catalog gate/rollback path. + if time.time_ns() - started_ns < int(budget_s * 1e9): + proposer = FallbackProposer(EngineArgsProposer(), TableProposer()) + + def _unenactable(spec: Any) -> str | None: + if ( + cfg.engine is not None + and live_restart_fn is None + and knob_kind(spec.knob) == "structural" + ): + return "structural knob: needs engine restart, no restart_fn" + return None + + ar_run = autoresearch( + trace, + applicator=applicator, + policy=policy, + residuals=res, + proposer=proposer, + ctx=pctx.gate, + reject=_unenactable, + ) + else: + ar_run = AutoresearchRun(bottleneck_class=classify_bottleneck(trace), results=[]) + + ar_causal_evidence = ", ".join( + f"{h.cause_op}→{h.effect_op} (p={h.p_value:.2g})" for h in hypotheses.top(2) + ) or "no strong causal signal" + for r in ar_run.results: + if not r.applicable: + rejected.append(f"{r.spec.name} ({r.rejected_reason})") + continue + if r.rolled_back: + rolled_back.append(r.spec.name) + claims.append( + Claim( + summary=r.spec.summary, + residual_invariant="kernel_time", + residual_value=0.0, + causal_evidence=ar_causal_evidence, + intervention_name=r.spec.name, + predicted_delta=r.predicted_delta, + measured_delta=r.measured_delta, + rolled_back=r.rolled_back, + ) + ) + (run_dir / "autoresearch.json").write_text( + json.dumps( + { + "bottleneck_class": ar_run.bottleneck_class, + "target": ( + { + "op": ar_run.target.op, + "residual": ar_run.target.residual, + "n_kernels": ar_run.target.n_kernels, + } + if ar_run.target is not None + else None + ), + "results": [ + { + "name": r.spec.name, + "knob": r.spec.knob, + "value": r.spec.value, + "applicable": r.applicable, + "rejected_reason": r.rejected_reason, + "predicted_delta": r.predicted_delta, + "measured_delta": r.measured_delta, + "rolled_back": r.rolled_back, + "target_op": r.target_op, + } + for r in ar_run.results + ], + }, + indent=2, + ) + ) + # Phase 5 — stabilize + write report provenance = build_provenance( workload_id=workload, @@ -572,6 +659,8 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: "n_claims": len(claims), "n_rolled_back": len(rolled_back), "n_rejected": len(rejected), + "bottleneck_class": ar_run.bottleneck_class, + "n_autoresearch": len(ar_run.results), "scheduler_stats": asdict(sched_summary) if sched_stats.samples else None, "report_path": str(run_dir / "report.md"), } diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py new file mode 100644 index 0000000..f682dbc --- /dev/null +++ b/tests/test_autoresearch.py @@ -0,0 +1,681 @@ +"""Tests for the autoresearch agent — classify → propose → gate → apply/rollback.""" + +from __future__ import annotations + +import gitm.agents.autoresearch as ar +from gitm.agents.autoresearch import ( + AutoresearchRun, + autoresearch, + autoresearch_v0, + classify_bottleneck, + largest_residual, + propose, +) +from gitm.agents.policy import Policy +from gitm.kernels.library import load_library +from gitm.kernels.spec import InterventionSpec, SafetyGate +from gitm.optimizer.apply import DictApplicator +from gitm.optimizer.monitor import KernelResidual, Residuals +from gitm.tracer.schema import Trace + +from .conftest import make_kernel, make_memcpy, make_sync, make_trace + + +def _residuals(pairs: list[tuple[str, float]]) -> Residuals: + """Build a Residuals from (op, r_kt) pairs.""" + res = Residuals() + res.per_kernel = [KernelResidual(op=op, layer=None, r_kt=r, r_mt=None) for op, r in pairs] + return res + + +def test_module_imports() -> None: + """Regression guard: the shipped stub imported a nonexistent module.""" + assert ar.propose is propose + + +def test_public_api_names_all_resolve() -> None: + """Every name in __all__ is actually defined (no stale/typo'd exports).""" + assert ar.__all__, "the module should declare a public surface" + missing = [name for name in ar.__all__ if not hasattr(ar, name)] + assert not missing, f"__all__ names undefined symbols: {missing}" + + +# --- propose ----------------------------------------------------------------- + + +def test_propose_known_class_returns_specs() -> None: + specs = propose("idle_stall") + assert specs, "known class should yield proposals" + assert all(isinstance(s, InterventionSpec) for s in specs) + # Never high-risk: unproven proposals stay at moderate and lean on the gate. + assert all(s.safety.tier != "high_risk" for s in specs) + + +def test_propose_unknown_class_is_empty() -> None: + assert propose("no_such_bottleneck") == [] + + +def test_proposed_knobs_are_disjoint_from_catalog() -> None: + """The whole point of autoresearch is proposing *outside* the library.""" + catalog_knobs = {spec.knob for spec in load_library()} + assert catalog_knobs, "catalog should be non-empty" + for cls in ("idle_stall", "memory_bound", "compute_bound"): + for spec in propose(cls): + assert spec.knob not in catalog_knobs, f"{spec.knob} duplicates the catalog" + + +# --- classify_bottleneck ----------------------------------------------------- + + +def test_classify_idle_stall_from_serialized_kernels() -> None: + # Back-to-back kernels on one stream, no overlap ⇒ serialized concurrency = 1. + events = [ + make_kernel("k", start_ns=i * 100, end_ns=i * 100 + 90, stream_id=0) + for i in range(6) + ] + assert classify_bottleneck(make_trace(events=events)) == "idle_stall" + + +def test_classify_memory_bound_from_memcpy_heavy_trace() -> None: + # Overlapping kernels (no stall) but memcpys dominate GPU-op time. + kernels = [ + make_kernel("a", start_ns=0, end_ns=100, stream_id=0), + make_kernel("b", start_ns=0, end_ns=100, stream_id=1), + ] + memcpys = [make_memcpy(start_ns=i * 100, end_ns=i * 100 + 100) for i in range(4)] + assert classify_bottleneck(make_trace(events=kernels + memcpys)) == "memory_bound" + + +def test_classify_compute_bound_when_transfers_do_not_dominate_gpu_time() -> None: + # Counting events would overreact to tiny transfers; the classifier uses GPU-op time. + kernels = [ + make_kernel("a", start_ns=0, end_ns=1000, stream_id=0), + make_kernel("b", start_ns=0, end_ns=1000, stream_id=1), + ] + memcpys = [make_memcpy(start_ns=i * 10, end_ns=i * 10 + 5) for i in range(4)] + syncs = [make_sync(start_ns=i * 20, end_ns=i * 20 + 1) for i in range(8)] + assert classify_bottleneck(make_trace(events=kernels + memcpys + syncs)) == "compute_bound" + + +def test_classify_compute_bound_when_overlapped_and_no_memcpy() -> None: + # Two kernels that temporally overlap ⇒ not serialized, no memcpy ⇒ compute. + events = [ + make_kernel("a", start_ns=0, end_ns=100, stream_id=0), + make_kernel("b", start_ns=50, end_ns=150, stream_id=1), + ] + assert classify_bottleneck(make_trace(events=events)) == "compute_bound" + + +def test_classify_empty_trace_defaults_to_compute() -> None: + assert classify_bottleneck(make_trace(events=[])) == "compute_bound" + + +# --- autoresearch_v0: apply / rollback --------------------------------------- + + +def _trace() -> Trace: + return make_trace(events=[make_kernel("paged_attention", start_ns=0, end_ns=500)]) + + +def test_applies_and_keeps_on_measured_win() -> None: + config: dict = {} + applicator = DictApplicator(config, measure_fn=lambda spec: 0.10) # +10% ⇒ keep + + results = autoresearch_v0(_trace(), "idle_stall", applicator=applicator) + + assert results, "idle_stall has proposals" + kept = [r for r in results if r.applicable and not r.rolled_back] + assert kept, "a positive delta must be kept" + for r in kept: + assert r.measured_delta == 0.10 + assert config.get(r.spec.knob) == r.spec.value # the knob was actually set + + +def test_rolls_back_on_regression() -> None: + config: dict = {} + applicator = DictApplicator(config, measure_fn=lambda spec: -0.10) # slower ⇒ revert + + results = autoresearch_v0(_trace(), "memory_bound", applicator=applicator) + + assert results + assert all(r.applicable and r.rolled_back for r in results) + assert config == {}, "a regressing proposal must leave the config untouched" + + +def test_gate_rejection_is_recorded_not_applied(monkeypatch) -> None: + """A proposal the gate rejects is reported, never applied.""" + high_risk = InterventionSpec( + name="autoresearch:test:danger", + summary="hypothetical high-risk proposal", + knob="some_risky_knob", + value=1, + expected_delta_mean=0.05, + expected_delta_lo=0.0, + expected_delta_hi=0.15, + source="test", + safety=SafetyGate(tier="high_risk"), + ) + monkeypatch.setattr(ar, "propose", lambda _cls, target_op=None: [high_risk]) + + config: dict = {} + applicator = DictApplicator(config, measure_fn=lambda spec: 0.10) + results = autoresearch_v0( + _trace(), "idle_stall", applicator=applicator, policy=Policy(skip_high_risk=True) + ) + + assert len(results) == 1 + r = results[0] + assert not r.applicable + assert r.rejected_reason == "policy.skip_high_risk" + assert config == {}, "a rejected proposal must never touch the config" + + +def test_unknown_class_yields_no_results() -> None: + applicator = DictApplicator({}, measure_fn=lambda spec: 0.10) + assert autoresearch_v0(_trace(), "no_such_class", applicator=applicator) == [] + + +# --- end-to-end entry point -------------------------------------------------- + + +def test_autoresearch_end_to_end_classifies_and_runs() -> None: + events = [ + make_kernel("k", start_ns=i * 100, end_ns=i * 100 + 90, stream_id=0) + for i in range(6) + ] + config: dict = {} + applicator = DictApplicator(config, measure_fn=lambda spec: 0.08) + + run = autoresearch(make_trace(events=events), applicator=applicator) + + assert isinstance(run, AutoresearchRun) + assert run.bottleneck_class == "idle_stall" + assert run.results + assert all(r.bottleneck_class == "idle_stall" for r in run.results) + + +# --- repoint at the largest residual ----------------------------------------- + + +def test_largest_residual_picks_biggest_gap_not_jitter() -> None: + """The target is the op furthest over its ceiling (largest mean r_kt), not the + jitteriest op (high variance, ~0 mean) or the smaller-gap op.""" + res = _residuals([ + ("A", 0.80), ("A", 0.60), # mean +0.70 — the biggest gap + ("B", 0.10), ("B", 0.20), # mean +0.15 — smaller gap + ("C", 1.50), ("C", -1.50), # mean 0.00 — pure jitter, must NOT win + ]) + target = largest_residual(res) + assert target is not None + assert target.op == "A" + assert target.n_kernels == 2 + + +def test_largest_residual_none_when_nothing_over_ceiling() -> None: + # Everything runs at or under the predicted ceiling ⇒ no bottleneck to chase. + assert largest_residual(_residuals([("A", -0.2), ("B", -0.1)])) is None + assert largest_residual(Residuals()) is None + + +def test_autoresearch_targets_largest_residual_op() -> None: + # Kernels named so the residual op "paged_attention" matches a real kernel. + events = [ + make_kernel("paged_attention", start_ns=i * 100, end_ns=i * 100 + 90, stream_id=0) + for i in range(6) + ] + res = _residuals([("paged_attention", 0.9), ("linear", 0.2)]) + applicator = DictApplicator({}, measure_fn=lambda spec: 0.05) + + run = autoresearch(make_trace(events=events), applicator=applicator, residuals=res) + + assert run.target is not None and run.target.op == "paged_attention" + assert run.results + assert all(r.target_op == "paged_attention" for r in run.results) + # The op is present in the trace, so proposals are scoped to it. + assert all(r.spec.applies_to_kernels == ["paged_attention"] for r in run.results) + + +def test_autoresearch_without_residuals_has_no_target() -> None: + events = [make_kernel("k", start_ns=i * 100, end_ns=i * 100 + 90) for i in range(4)] + run = autoresearch(make_trace(events=events), applicator=DictApplicator({})) + assert run.target is None + assert all(r.target_op is None for r in run.results) + # Untargeted proposals stay unscoped (coverage 1.0 in predict_delta). + assert all(r.spec.applies_to_kernels == [] for r in run.results) + + +def test_off_trace_residual_op_is_not_tagged() -> None: + """If the largest-residual op doesn't match any kernel name, don't scope to it + (that would zero predict_delta coverage) — but still record it as the target.""" + events = [make_kernel("gemm", start_ns=i * 100, end_ns=i * 100 + 90) for i in range(4)] + res = _residuals([("paged_attention", 0.9)]) # op not present in the trace + run = autoresearch(make_trace(events=events), applicator=DictApplicator({}), residuals=res) + + assert run.target is not None and run.target.op == "paged_attention" + assert all(r.target_op == "paged_attention" for r in run.results) # recorded + assert all(r.spec.applies_to_kernels == [] for r in run.results) # but not scoped + + +# --- EngineArgs-driven generative proposer ----------------------------------- + +from gitm.agents.autoresearch import ( # noqa: E402 + _CLASS_KEYWORDS, + _RULES, + BOTTLENECK_CLASSES, + EngineArgsProposer, + FallbackProposer, + GenerativeProposer, + Knob, + StochasticProposer, + TableProposer, + VLLMKnobSource, + _candidate_spec, + _field_kind_and_choices, + _is_tunable, + _value_grid, +) + + +class _ListSource: + """A minimal KnobSource for tests: yields a fixed knob list.""" + + def __init__(self, knobs: list[Knob]) -> None: + self._knobs = knobs + + def knobs(self) -> list[Knob]: + return list(self._knobs) + + +def test_value_grid_flips_bool() -> None: + assert _value_grid(Knob("enforce_eager", "bool", default=False)) == [True] + + +def test_value_grid_enum_returns_other_members() -> None: + grid = _value_grid( + Knob("preemption_mode", "enum", default="recompute", choices=("recompute", "swap")) + ) + assert grid == ["swap"] + + +def test_value_grid_int_searches_multiple_distinct_positive_values() -> None: + grid = _value_grid(Knob("max_num_partial_prefills", "int", default=1)) + assert len(grid) >= 2, "a value grid must actually search several values" + assert len(set(grid)) == len(grid) # distinct + assert all(isinstance(v, int) and v >= 1 and v != 1 for v in grid) + + +def test_value_grid_explicit_grid_wins() -> None: + grid = _value_grid( + Knob("long_prefill_token_threshold", "int", default=0, grid=(2048, 4096)) + ) + assert grid == [2048, 4096] + + +def _stub_knobs() -> list[Knob]: + return [ + Knob("max_num_partial_prefills", "int", default=1), # idle_stall-affine + Knob("cpu_offload_gb", "int", default=0), # memory_bound-affine + Knob("block_size", "int", default=16), # memory-affine but IN the catalog + ] + + +def test_engineargs_proposer_searches_multiple_values_per_knob() -> None: + p = EngineArgsProposer( + knobs=[Knob("max_num_partial_prefills", "int", default=1)], catalog_knobs=set() + ) + specs = p.propose("idle_stall") + assert len(specs) >= 2, "value-grid search should try several values" + assert len({s.value for s in specs}) == len(specs) # one spec per distinct value + assert all(s.knob == "max_num_partial_prefills" for s in specs) + assert all(isinstance(s, InterventionSpec) for s in specs) + # Generated candidates are never high-risk; they lean on the rollback gate. + assert all(s.safety.tier == "moderate" for s in specs) + # The value-grid naming (knob=value) distinguishes them from table proposals. + assert all("=" in s.name for s in specs) + + +def test_engineargs_proposer_excludes_catalog_knobs() -> None: + """block_size is a real EngineArgs knob, but it's in the curated library — + autoresearch proposes *outside* the catalog, so it must be dropped.""" + p = EngineArgsProposer(knobs=_stub_knobs(), catalog_knobs={"block_size"}) + specs = p.propose("memory_bound") + assert specs, "cpu_offload_gb is affine to memory_bound" + assert all(s.knob != "block_size" for s in specs) + assert {s.knob for s in specs} == {"cpu_offload_gb"} + + +def test_engineargs_proposer_scopes_candidates_to_bottleneck_class() -> None: + p = EngineArgsProposer(knobs=_stub_knobs(), catalog_knobs=set()) + idle = {s.knob for s in p.propose("idle_stall")} + mem = {s.knob for s in p.propose("memory_bound")} + assert idle == {"max_num_partial_prefills"} + assert mem == {"cpu_offload_gb", "block_size"} # both memory-affine by name + + +def test_engineargs_proposer_unknown_class_is_empty() -> None: + assert EngineArgsProposer(knobs=_stub_knobs()).propose("no_such_class") == [] + + +def test_engineargs_proposer_scopes_specs_to_target_op() -> None: + p = EngineArgsProposer( + knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + ) + specs = p.propose("memory_bound", target_op="paged_attention") + assert specs + assert all(s.applies_to_kernels == ["paged_attention"] for s in specs) + + +def test_engineargs_offline_fallback_runs_without_vllm() -> None: + """vLLM isn't importable in CI; the frozen catalog must still yield candidates + for every class, and every candidate stays outside the curated library.""" + catalog = {s.knob for s in load_library()} + p = EngineArgsProposer() # default knobs → _engine_arg_knobs() → frozen fallback + for cls in ("idle_stall", "memory_bound", "compute_bound"): + specs = p.propose(cls) + assert specs, f"{cls} should yield fallback candidates offline" + assert all(s.safety.tier == "moderate" for s in specs) + assert all(s.knob not in catalog for s in specs) # disjoint from the library + + +def test_engineargs_candidates_route_through_gate_and_rollback() -> None: + """A generated candidate is kept on a measured win and reverted on a regression — + the same gate the catalog goes through, nothing special for generated specs.""" + proposer = EngineArgsProposer( + knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + ) + grid = _value_grid(Knob("cpu_offload_gb", "int", default=0)) + + kept_cfg: dict = {} + kept = autoresearch_v0( + _trace(), + "memory_bound", + applicator=DictApplicator(kept_cfg, measure_fn=lambda s: 0.10), + proposer=proposer, + ) + assert kept and all(r.applicable and not r.rolled_back for r in kept) + assert kept_cfg.get("cpu_offload_gb") in grid # the generated knob was set + + revert_cfg: dict = {} + reverted = autoresearch_v0( + _trace(), + "memory_bound", + applicator=DictApplicator(revert_cfg, measure_fn=lambda s: -0.10), + proposer=proposer, + ) + assert reverted and all(r.rolled_back for r in reverted) + assert revert_cfg == {}, "a regressing candidate must leave the config untouched" + + +def test_fallback_proposer_uses_table_only_when_primary_is_empty() -> None: + table = TableProposer() + # Primary has no compute_bound knob here → falls back to the table's compute lever. + fb = FallbackProposer( + EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set()), + table, + ) + # memory_bound: primary yields cpu_offload_gb candidates → table NOT consulted. + mem = fb.propose("memory_bound") + assert {s.knob for s in mem} == {"cpu_offload_gb"} + assert all("=" in s.name for s in mem) # generated, not table + # compute_bound: primary empty → table lever surfaces. + compute = fb.propose("compute_bound") + assert compute == table.propose("compute_bound") + assert compute, "the table has a compute_bound lever" + + +def test_autoresearch_end_to_end_with_engineargs_proposer() -> None: + events = [ + make_kernel("k", start_ns=i * 100, end_ns=i * 100 + 90, stream_id=0) for i in range(6) + ] # serialized → idle_stall + proposer = EngineArgsProposer( + knobs=[Knob("max_num_partial_prefills", "int", default=1)], catalog_knobs=set() + ) + run = autoresearch( + make_trace(events=events), + applicator=DictApplicator({}, measure_fn=lambda s: 0.05), + proposer=proposer, + ) + assert run.bottleneck_class == "idle_stall" + assert run.results and all(r.spec.knob == "max_num_partial_prefills" for r in run.results) + assert all("=" in r.spec.name for r in run.results) # generated value-grid candidates + + +def test_table_proposer_matches_module_propose() -> None: + assert [s.knob for s in TableProposer().propose("idle_stall")] == [ + s.knob for s in propose("idle_stall") + ] + + +# --- workload-agnostic generative proposer (KnobSource seam) ------------------ + + +def test_generative_proposer_serves_a_non_vllm_workload() -> None: + """A non-vLLM workload plugs in via a KnobSource + workload label — no table. + Knobs declare their class affinity explicitly, so the vLLM keyword table is + irrelevant here.""" + source = _ListSource([Knob("cache_prefetch_depth", "int", default=2, classes=("memory_bound",))]) + p = GenerativeProposer(source, workload="triton-serve", catalog_knobs=set()) + + specs = p.propose("memory_bound") + assert specs and all(s.knob == "cache_prefetch_depth" for s in specs) + # The candidate is labelled for the caller's workload, not hardcoded vLLM. + assert all(s.applicability.workloads == ["triton-serve"] for s in specs) + assert p.propose("idle_stall") == [] # not tagged for idle_stall + + +def test_knob_explicit_classes_override_keyword_affinity() -> None: + # A knob whose NAME wouldn't keyword-match idle_stall, but is tagged for it. + source = _ListSource([Knob("weird_lever", "int", default=1, classes=("idle_stall",))]) + p = GenerativeProposer(source, catalog_knobs=set()) + assert {s.knob for s in p.propose("idle_stall")} == {"weird_lever"} + assert p.propose("memory_bound") == [] + + +def test_engineargs_proposer_is_a_vllm_bound_generative_proposer() -> None: + """EngineArgsProposer is just GenerativeProposer bound to the vLLM surface.""" + p = EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set()) + assert isinstance(p, GenerativeProposer) + specs = p.propose("memory_bound") + assert specs and all(s.applicability.workloads == ["vllm-decode"] for s in specs) + + +def test_vllm_knob_source_yields_offline_fallback_without_vllm() -> None: + # vLLM isn't importable in CI → the source yields the frozen fallback catalog. + knobs = VLLMKnobSource().knobs() + assert knobs and all(isinstance(k, Knob) for k in knobs) + assert "max_num_partial_prefills" in {k.name for k in knobs} + + +def test_is_tunable_excludes_non_perf_engine_args() -> None: + # Identity / IO / logging / RNG fields are not performance knobs. + for name in ("model", "served_model_name", "tokenizer", "seed", + "disable_log_stats", "download_dir", "revision"): + assert not _is_tunable(name), f"{name} should be excluded" + # Real optimization knobs pass — including tricky substrings. + for name in ("max_num_seqs", "gpu_memory_utilization", "compilation_config", + "long_prefill_token_threshold", "cpu_offload_gb"): + assert _is_tunable(name), f"{name} should be kept" + + +def test_field_kind_and_choices_extracts_literal_enum() -> None: + import typing + + kind, choices = _field_kind_and_choices(typing.Literal["recompute", "swap"]) + assert kind == "enum" + assert choices == ("recompute", "swap") + # Plain scalars fall back to kind-only (no choices) and stay robust. + assert _field_kind_and_choices(bool) == ("bool", ()) + assert _field_kind_and_choices(int) == ("int", ()) + + +def test_candidate_specs_share_the_forced_fields() -> None: + """DRY guard: table and generated candidates agree on delta band, safety tier, + and applicability — the fields _candidate_spec centralizes.""" + table = propose("idle_stall")[0] + generated = EngineArgsProposer( + knobs=[Knob("max_num_partial_prefills", "int", default=1)], catalog_knobs=set() + ).propose("idle_stall")[0] + for s in (table, generated): + assert s.safety.tier == "moderate" + assert (s.expected_delta_mean, s.expected_delta_lo, s.expected_delta_hi) == (0.05, 0.0, 0.15) + assert s.applicability.workloads == ["vllm-decode"] + + +def test_bottleneck_vocabulary_is_single_sourced() -> None: + """classify emits only known classes; the fallback table and keyword-affinity + map key off the same authoritative vocabulary, so the three can't drift.""" + assert set(_RULES) == set(BOTTLENECK_CLASSES) + assert set(_CLASS_KEYWORDS) == set(BOTTLENECK_CLASSES) + + +def test_classify_always_returns_a_known_class() -> None: + traces = [ + make_trace(events=[]), # empty → compute default + make_trace(events=[ + make_kernel("k", start_ns=i * 100, end_ns=i * 100 + 90) for i in range(6) + ]), # serialized → idle + make_trace(events=[make_kernel("k", start_ns=0, end_ns=1000)] + + [make_memcpy(start_ns=i * 10, end_ns=i * 10 + 5) for i in range(4)]), + ] + assert all(classify_bottleneck(t) in BOTTLENECK_CLASSES for t in traces) + + +def test_generative_proposer_accepts_custom_affinity_keywords() -> None: + """A workload whose knobs use their own naming supplies its own affinity + keywords — no reliance on the vLLM-flavoured defaults, and still no table.""" + source = _ListSource([Knob("shard_rebalance_interval", "int", default=4)]) + p = GenerativeProposer( + source, + workload="mesh-serve", + catalog_knobs=set(), + affinity_keywords={"memory_bound": ("shard", "rebalance")}, + ) + assert {s.knob for s in p.propose("memory_bound")} == {"shard_rebalance_interval"} + # The default vLLM keyword vocabulary wouldn't have matched this name. + assert GenerativeProposer(source, catalog_knobs=set()).propose("memory_bound") == [] + + +def test_generative_proposer_is_uncapped_by_default() -> None: + knobs = [Knob(f"prefill_{i}", "bool", default=False) for i in range(30)] + p = GenerativeProposer(_ListSource(knobs), catalog_knobs=set()) + assert len(p.propose("idle_stall")) == 30 # all idle-affine, one value each + + +def test_generative_proposer_caps_candidate_count() -> None: + knobs = [Knob(f"prefill_{i}", "bool", default=False) for i in range(50)] + p = GenerativeProposer(_ListSource(knobs), catalog_knobs=set(), max_candidates=10) + assert len(p.propose("idle_stall")) == 10 # bounded search over a large surface + + +def test_engineargs_proposer_bounds_a_large_surface_by_default() -> None: + # The vLLM binding caps by default so the real ~100-field surface can't flood. + knobs = [Knob(f"prefill_{i}", "bool", default=False) for i in range(100)] + specs = EngineArgsProposer(knobs=knobs, catalog_knobs=set()).propose("idle_stall") + assert 0 < len(specs) <= 24 + + +def test_candidate_spec_helper_forces_safety_and_delta() -> None: + spec = _candidate_spec( + name="autoresearch:test:x=1", + summary="s", + knob="x", + value=1, + applies_to_kernels=[], + bottleneck_class="idle_stall", + workload="some-workload", + source="test", + ) + assert spec.safety.tier == "moderate" + assert spec.expected_delta_hi == 0.15 + assert spec.applicability.workloads == ["some-workload"] + + +# --- stochastic (entropy-guided) proposer ------------------------------------ + + +def _mixed_source() -> _ListSource: + # One idle-affine knob (name contains "prefill") and one off-class knob. + return _ListSource( + [Knob("prefill_slots", "int", default=1), Knob("zzz_widget", "int", default=1)] + ) + + +def test_stochastic_is_reproducible_for_a_given_seed() -> None: + src = _mixed_source() + p = StochasticProposer(src, catalog_knobs=set(), n_samples=6, seed=42, epsilon=0.4) + first = [s.name for s in p.propose("idle_stall")] + second = [s.name for s in p.propose("idle_stall")] + assert first and first == second # same seed → identical draws + + +def test_stochastic_seed_actually_varies_exploration() -> None: + src = _mixed_source() + variants = { + tuple(s.name for s in StochasticProposer( + src, catalog_knobs=set(), n_samples=6, seed=k, epsilon=0.5 + ).propose("idle_stall")) + for k in range(6) + } + assert len(variants) > 1 # the seed is a real exploration control + + +def test_stochastic_epsilon_zero_is_pure_heuristic() -> None: + # No entropy floor ⇒ only class-affine knobs are ever sampled, any seed. + src = _mixed_source() + p = StochasticProposer(src, catalog_knobs=set(), n_samples=20, seed=7, epsilon=0.0) + specs = p.propose("idle_stall") + assert specs and all(s.knob == "prefill_slots" for s in specs) + + +def test_stochastic_epsilon_lets_the_search_wander_off_class() -> None: + # A nonzero floor ⇒ the off-class knob is reachable (the entropy source). + src = _mixed_source() + p = StochasticProposer(src, catalog_knobs=set(), n_samples=20, seed=0, epsilon=1.0) + knobs = {s.knob for s in p.propose("idle_stall")} + assert "prefill_slots" in knobs and "zzz_widget" in knobs + + +def test_stochastic_empty_when_no_signal_and_no_entropy() -> None: + # Off-class knob only, epsilon=0 ⇒ no heuristic signal and no floor ⇒ nothing. + src = _ListSource([Knob("zzz_widget", "int", default=1)]) + p = StochasticProposer(src, catalog_knobs=set(), n_samples=8, seed=0, epsilon=0.0) + assert p.propose("idle_stall") == [] + + +def test_stochastic_retries_duplicates_but_returns_unique_candidates() -> None: + src = _ListSource([Knob("prefill_enabled", "bool", default=False)]) + p = StochasticProposer(src, catalog_knobs=set(), n_samples=6, seed=0, epsilon=0.0) + + specs = p.propose("idle_stall") + + assert len(specs) == 1 + assert specs[0].knob == "prefill_enabled" + assert specs[0].value is True + + +def test_stochastic_skips_knobs_without_value_grid() -> None: + src = _ListSource([ + Knob("prefill_freeform", "str", default="auto"), + Knob("prefill_enabled", "bool", default=False), + ]) + p = StochasticProposer(src, catalog_knobs=set(), n_samples=3, seed=0, epsilon=0.0) + + specs = p.propose("idle_stall") + + assert specs + assert {s.knob for s in specs} == {"prefill_enabled"} + + +def test_stochastic_candidates_route_through_gate_and_rollback() -> None: + proposer = StochasticProposer( + _mixed_source(), catalog_knobs=set(), n_samples=6, seed=1, epsilon=0.3 + ) + cfg: dict = {} + results = autoresearch_v0( + _trace(), + "idle_stall", + applicator=DictApplicator(cfg, measure_fn=lambda s: 0.10), + proposer=proposer, + ) + assert results and all(r.applicable and not r.rolled_back for r in results) + assert all(s.safety.tier == "moderate" for s in (r.spec for r in results)) diff --git a/tests/test_run_loop_workload.py b/tests/test_run_loop_workload.py index ddc7aff..9c390d1 100644 --- a/tests/test_run_loop_workload.py +++ b/tests/test_run_loop_workload.py @@ -129,6 +129,38 @@ def test_vllm_workload_still_uses_intervention_path(tmp_path: Path, monkeypatch) assert result["summary"]["mode"] == "intervention" +def test_vllm_loop_runs_autoresearch(tmp_path: Path, monkeypatch): + """The vllm path runs agentic autoresearch: it classifies the bottleneck, then + *generates* non-catalog levers from the real EngineArgs surface and searches a + value grid per knob (EngineArgsProposer), surfacing them in the summary + + report. The serialized same-stream kernels above classify as idle_stall; the + frozen fallback catalog yields three idle knobs across six value-grid points.""" + import gitm.scheduler.loop as loop + + monkeypatch.setattr(loop, "capture", _fake_capture_with_kernels("paged_attention")) + monkeypatch.setattr(loop, "sync_device", lambda: None) + + from gitm import optimize + + # budget="30s" (not 1s): the catalog pass must not exhaust the budget and skip + # the autoresearch phase on a slow box. + result = optimize( + workload="vllm-decode", budget="30s", scratch=str(tmp_path), workload_runner=lambda: {} + ) + s = result["summary"] + assert s["bottleneck_class"] == "idle_stall" + # Generative, not the 2-row table: three idle knobs × their value grids = 6. + assert s["n_autoresearch"] == 6 + assert "autoresearch:" in result["report_md"] # candidates reach the report + + ar_json = (Path(result["run_dir"]) / "autoresearch.json").read_text(encoding="utf-8") + # The value-grid naming (knob=value) proves generation ran, not the flat table. + assert "max_num_partial_prefills=" in ar_json + assert all(name in ar_json for name in ("=2048", "=4096")) # explicit-grid search points + # Dry-run (no live engine) mutates nothing, so no safety trail is written. + assert not (Path(result["run_dir"]) / "audit.jsonl").exists() + + def test_cli_run_returns_nonzero_on_no_data(tmp_path: Path, capsys): """Automation must see a failure exit when a run measures nothing.""" from gitm.cli import main From 889d5d9e5f162ab7d06082d4c3788b222990a45b Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Tue, 7 Jul 2026 15:15:05 -0700 Subject: [PATCH 02/24] autoresearch: source valid knob domains from EngineArgs CLI args MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The value grid was blind: a ½×/2×/4× ladder that invents values for enum knobs and list-valued args, which then fail to apply and waste rollback-gated trials. Read each field's valid domain from its vLLM CLI argument instead: - argparse `choices=` → the real enum domain, so the grid proposes only values that can apply (not an invented ladder point). - argparse `type=` → sharpens int/float when the annotation is opaque. - `nargs` → skip list-valued args (not scalar knobs). Split the extraction into `_knobs_from_engine_args` / `_argparse_domains` / `_knob_domain` so it's tested against a fake EngineArgs-like class without vLLM present. Honest limit: argparse carries no numeric min/max, so unconstrained ints/floats still use the ladder. Offline (no vLLM) path unchanged — frozen fallback catalog. Docs reconciled. Co-Authored-By: Claude Opus 4.8 --- docs/autoresearch.md | 23 ++++--- gitm/agents/autoresearch.py | 110 +++++++++++++++++++++++++++---- gitm/workloads.py | 5 ++ tests/test_autoresearch.py | 72 ++++++++++++++++++++ tests/test_workload_bootstrap.py | 58 ++++++++++++++++ 5 files changed, 246 insertions(+), 22 deletions(-) diff --git a/docs/autoresearch.md b/docs/autoresearch.md index eb5b845..4cb52a4 100644 --- a/docs/autoresearch.md +++ b/docs/autoresearch.md @@ -122,8 +122,11 @@ Versatility comes from the `KnobSource`, **not** a `{workload: knobs}` table: - **`VLLMKnobSource`** (used by `EngineArgsProposer`, the vLLM binding of `GenerativeProposer`) introspects `dataclasses.fields(EngineArgs)` when vLLM is - importable — exposing `Literal` fields as searchable enums — and yields a frozen - catalog otherwise, so the path runs offline and deterministically. + importable and reads each field's **valid domain from its CLI argument** — + argparse `choices=` (the real enum domain) and `type=` — via + `EngineArgs.add_cli_args`, so the search only proposes values that can apply; + list-valued args (`nargs`) and non-performance fields are skipped. Yields a + frozen catalog otherwise, so the path runs offline and deterministically. - Another workload plugs in by yielding its own `Knob` list and a `workload` label: `GenerativeProposer(MyKnobSource(), workload="triton-serve")`. Candidates carry that label in their applicability; nothing enumerates workloads centrally. @@ -213,12 +216,16 @@ it: - **Class→knob affinity is a keyword heuristic** (`prefill` / `cache` / `compil` …) matched against the knob name, not a learned or hand-authored mapping. -- **Value grids are blind** — a generic `½×`/`2×`/`4×` ladder off the default (or - a hand-seeded grid where that would be nonsensical, e.g. token thresholds), not - tuned per knob. Per-class candidate counts are bounded by `max_candidates`. -- **Live `EngineArgs` typing is best-effort** — introspected fields are typed - coarsely from their annotation; enums the introspector can't detect fall back - to "no grid" and are simply skipped. +- **Enum domains come from vLLM; numeric ranges don't.** When vLLM is importable, + each knob's valid **choices** and **type** are read from its `EngineArgs` CLI + argument (argparse `choices=`/`type=`), and list-valued args are skipped — so + the search stops proposing enum values or shapes that can't apply. But argparse + carries no numeric min/max, so unconstrained ints/floats still use the generic + `½×`/`2×`/`4×` ladder (or a hand-seeded grid). Per-class candidate counts are + bounded by `max_candidates`. +- **Live typing is otherwise best-effort** — fields with no CLI `choices`/`type` + are classified coarsely from their annotation; ones the introspector can't type + fall back to "no grid" and are simply skipped. - **The stochastic sampler isn't wired into the loop yet.** `StochasticProposer` exists behind the seam (seeded, reproducible), but Phase 4b still runs the generative proposer; flipping the loop onto it (or a `FallbackProposer` chain) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 6d7be92..9d83c5b 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -476,33 +476,115 @@ def _is_tunable(field_name: str) -> bool: return not any(h in lname for h in _NON_TUNABLE_HINTS) -def _engine_arg_knobs() -> list[Knob]: - """Enumerate real vLLM EngineArgs, or fall back to the frozen catalog. +@dataclass(frozen=True) +class _ArgDomain: + """Valid-value metadata for one EngineArgs field, read from its CLI argument. + + ``choices`` is the argparse ``choices=`` set (the *real* enum domain); ``type`` + is the argparse ``type=`` callable; ``is_list`` flags list-valued args + (``nargs`` +/*/N) which aren't scalar knobs. argparse carries no numeric + min/max, so ranges still fall to the value-grid ladder. + """ + + type: object | None = None + choices: tuple | None = None + default: object = None + is_list: bool = False + + +def _argparse_domains(engine_args_cls: object) -> dict[str, _ArgDomain]: + """Read each EngineArgs field's valid domain from ``add_cli_args`` (best-effort). - Best-effort introspection: when vLLM is importable, each tunable dataclass - field is a candidate knob (typed as best we can from its annotation, with - ``Literal`` fields exposed as searchable enums). Non-performance fields (model - identity, paths, logging, RNG — see :data:`_NON_TUNABLE_HINTS`) are skipped. - When vLLM isn't importable (no GPU stack / air-gapped), the frozen catalog - keeps the generative path working offline. + vLLM builds its CLI from the same dataclass, so the argparse actions are the + authoritative source of choices/types — far better than string-matching the + annotation. Returns ``{}`` if the class has no ``add_cli_args`` or it raises. """ try: - import dataclasses + import argparse - from vllm import EngineArgs # type: ignore + parser = engine_args_cls.add_cli_args(argparse.ArgumentParser()) # type: ignore[attr-defined] except Exception: - return list(_FALLBACK_KNOBS) + return {} + out: dict[str, _ArgDomain] = {} + for action in getattr(parser, "_actions", []): + dest = getattr(action, "dest", "") + if not dest or dest == "help": + continue + nargs = getattr(action, "nargs", None) + is_list = nargs in ("+", "*") or (isinstance(nargs, int) and nargs > 1) + choices = getattr(action, "choices", None) + out[dest] = _ArgDomain( + type=getattr(action, "type", None), + choices=tuple(choices) if choices else None, + default=getattr(action, "default", None), + is_list=is_list, + ) + return out + + +def _knob_domain(annotation: object, domain: _ArgDomain | None) -> tuple[str, tuple]: + """(kind, choices) for a field, preferring argparse metadata over the annotation. + + argparse ``choices`` gives the exact enum domain (so the search only proposes + values that can apply); its ``type`` sharpens int/float when the annotation is + opaque. Everything else falls back to the annotation-based classification. + """ + if domain is not None and domain.choices: + return "enum", domain.choices + kind, choices = _field_kind_and_choices(annotation) + if kind == "str" and domain is not None: + kind = {int: "int", float: "float", bool: "bool"}.get(domain.type, kind) # type: ignore[arg-type] + return kind, choices + +def _knobs_from_engine_args(engine_args_cls: object) -> list[Knob]: + """Build the searchable knob list from an EngineArgs-like dataclass. + + Split out from :func:`_engine_arg_knobs` (which handles the vLLM import + + fallback) so the extraction is testable without vLLM present. Skips + non-performance fields (:data:`_NON_TUNABLE_HINTS`) and list-valued args (not + scalar knobs), and sources each field's valid domain from its CLI argument. + """ + import dataclasses + + domains = _argparse_domains(engine_args_cls) knobs: list[Knob] = [] - for f in dataclasses.fields(EngineArgs): + for f in dataclasses.fields(engine_args_cls): # type: ignore[arg-type] if not _is_tunable(f.name): continue - default = None if f.default is dataclasses.MISSING else f.default - kind, choices = _field_kind_and_choices(f.type) + domain = domains.get(f.name) + if domain is not None and domain.is_list: + continue # list-valued arg: not a scalar knob, can't grid-search + kind, choices = _knob_domain(f.type, domain) + if f.default is not dataclasses.MISSING: + default = f.default + else: + default = domain.default if domain is not None else None knobs.append(Knob(name=f.name, kind=kind, default=default, choices=choices)) return knobs +def _engine_arg_knobs() -> list[Knob]: + """Enumerate real vLLM EngineArgs, or fall back to the frozen catalog. + + Best-effort introspection: when vLLM is importable, each tunable, scalar field + is a candidate knob, with its valid domain (enum ``choices``, type) read from + the field's CLI argument (:func:`_argparse_domains`) so the search only + proposes values that can actually apply. Non-performance fields + (:data:`_NON_TUNABLE_HINTS`) and list-valued args are skipped. When vLLM isn't + importable (no GPU stack / air-gapped), the frozen catalog keeps the generative + path working offline. + """ + try: + from vllm import EngineArgs # type: ignore + except Exception: + return list(_FALLBACK_KNOBS) + try: + return _knobs_from_engine_args(EngineArgs) or list(_FALLBACK_KNOBS) + except Exception: + return list(_FALLBACK_KNOBS) + + class KnobSource(Protocol): """Yields the knob surface to search — a workload's config namespace. diff --git a/gitm/workloads.py b/gitm/workloads.py index fc9c6d1..8127d46 100644 --- a/gitm/workloads.py +++ b/gitm/workloads.py @@ -639,6 +639,11 @@ def _restart(_old_engine: Any, knob: str, value: Any) -> Any: f"two-engine V1 distributed clash: {exc}" ) from exc + run.engine = llm + llm.gitm_throughput_fn = _throughput + llm.gitm_restart_fn = _restart + return run + def _vllm_synthetic_runner(n_prompts: int, max_tokens: int) -> WorkloadRunner: """A CPU-only stand-in for the decode loop (no vLLM, no GPU). diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index f682dbc..da82ca7 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -269,9 +269,11 @@ def test_off_trace_residual_op_is_not_tagged() -> None: StochasticProposer, TableProposer, VLLMKnobSource, + _argparse_domains, _candidate_spec, _field_kind_and_choices, _is_tunable, + _knobs_from_engine_args, _value_grid, ) @@ -508,6 +510,76 @@ def test_field_kind_and_choices_extracts_literal_enum() -> None: assert _field_kind_and_choices(int) == ("int", ()) +# --- valid domains sourced from EngineArgs' CLI args ------------------------- + +import argparse # noqa: E402 +import dataclasses # noqa: E402 + + +@dataclasses.dataclass +class _FakeEngineArgs: + """An EngineArgs-like dataclass with a CLI builder — exercises the argparse + domain extraction without importing vLLM (absent in CI).""" + + kv_cache_dtype: str = "auto" + max_num_seqs: int = 256 + gpu_frac: float = 0.9 + enforce_eager: bool = False + served_model_name: str = "m" # non-perf (model/name) → excluded + middleware: tuple = () # list-valued (nargs=+) → skipped + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + parser.add_argument( + "--kv-cache-dtype", dest="kv_cache_dtype", + choices=["auto", "fp8", "fp8_e5m2"], default="auto", + ) + parser.add_argument("--max-num-seqs", dest="max_num_seqs", type=int, default=256) + parser.add_argument("--gpu-frac", dest="gpu_frac", type=float, default=0.9) + parser.add_argument("--enforce-eager", dest="enforce_eager", action="store_true") + parser.add_argument("--served-model-name", dest="served_model_name", type=str, default="m") + parser.add_argument("--middleware", dest="middleware", nargs="+", type=str, default=[]) + return parser + + +def test_knobs_from_engine_args_sources_valid_enum_domain() -> None: + """The enum domain comes from argparse ``choices`` — so the value grid only + contains values that can actually apply, never an invented ladder point.""" + knobs = {k.name: k for k in _knobs_from_engine_args(_FakeEngineArgs)} + kv = knobs["kv_cache_dtype"] + assert kv.kind == "enum" + assert set(kv.choices) == {"auto", "fp8", "fp8_e5m2"} + assert set(_value_grid(kv)) == {"fp8", "fp8_e5m2"} # valid non-default choices only + + +def test_knobs_from_engine_args_types_from_argparse() -> None: + knobs = {k.name: k for k in _knobs_from_engine_args(_FakeEngineArgs)} + assert knobs["max_num_seqs"].kind == "int" + assert knobs["gpu_frac"].kind == "float" + assert knobs["enforce_eager"].kind == "bool" + + +def test_knobs_from_engine_args_skips_nonperf_and_list_args() -> None: + names = {k.name for k in _knobs_from_engine_args(_FakeEngineArgs)} + assert "served_model_name" not in names # non-performance field + assert "middleware" not in names # list-valued arg (nargs=+): not a scalar knob + + +def test_argparse_domains_reads_choices_type_and_nargs() -> None: + d = _argparse_domains(_FakeEngineArgs) + assert d["kv_cache_dtype"].choices == ("auto", "fp8", "fp8_e5m2") + assert d["max_num_seqs"].type is int + assert d["middleware"].is_list is True + assert d["kv_cache_dtype"].is_list is False + + +def test_argparse_domains_empty_when_no_cli_builder() -> None: + class _Bare: + pass + + assert _argparse_domains(_Bare) == {} + + def test_candidate_specs_share_the_forced_fields() -> None: """DRY guard: table and generated candidates agree on delta band, safety tier, and applicability — the fields _candidate_spec centralizes.""" diff --git a/tests/test_workload_bootstrap.py b/tests/test_workload_bootstrap.py index d197983..cdadd61 100644 --- a/tests/test_workload_bootstrap.py +++ b/tests/test_workload_bootstrap.py @@ -96,3 +96,61 @@ def test_generate_importable_from_package(): from gitm.benchmarks.hft import generate assert generate.generate and generate.GenConfig + + +def test_vllm_factory_returns_runner_with_live_engine_hooks(monkeypatch): + """The vLLM factory must return the decode runner, not just build the engine.""" + import sys + import types + + import gitm.workloads as workloads + + class FakeSamplingParams: + def __init__(self, *, max_tokens: int, temperature: float): + self.max_tokens = max_tokens + self.temperature = temperature + + class FakeLLM: + instances = [] + + def __init__(self, *, model: str, **kwargs): + self.model = model + self.kwargs = kwargs + self.calls = 0 + FakeLLM.instances.append(self) + + def generate(self, prompts, params): + self.calls += 1 + return [ + types.SimpleNamespace( + outputs=[types.SimpleNamespace(token_ids=list(range(params.max_tokens)))] + ) + for _ in prompts + ] + + monkeypatch.setitem( + sys.modules, + "vllm", + types.SimpleNamespace(LLM=FakeLLM, SamplingParams=FakeSamplingParams), + ) + monkeypatch.setattr(workloads, "sync_device", lambda: None) + monkeypatch.setenv("GITM_VLLM_MODEL", "fake/model") + monkeypatch.setenv("GITM_VLLM_PROMPTS", "2") + monkeypatch.setenv("GITM_VLLM_MAX_TOKENS", "4") + monkeypatch.setenv("GITM_VLLM_GPU_MEM", "0.45") + + runner = workloads.get_factory("vllm-decode")(None) + + assert callable(runner) + assert runner.engine is FakeLLM.instances[0] + assert callable(runner.engine.gitm_throughput_fn) + assert callable(runner.engine.gitm_restart_fn) + + summary = runner() + assert summary["generated_tokens"] == 8 + assert FakeLLM.instances[0].calls == 1 + assert runner.engine.gitm_throughput_fn(runner.engine) > 0 + + restarted = runner.engine.gitm_restart_fn(runner.engine, "swap_space", 2) + assert restarted.kwargs["gpu_memory_utilization"] == 0.45 + assert restarted.kwargs["swap_space"] == 2 From 9b1ce9403d56d9f3941579335b34552dd96cd995 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Tue, 7 Jul 2026 20:13:04 -0700 Subject: [PATCH 03/24] vllm-decode: add GITM_VLLM_ENFORCE_EAGER to disable CUDA graphs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit On platforms whose CUPTI doesn't attribute CUDA-graph-replayed kernels, the tracer records an empty trace and the loop honestly reports no-data. Setting GITM_VLLM_ENFORCE_EAGER=1 builds the engine (and every restart-A/B candidate, via the shared _base_kwargs) with enforce_eager=True — decode kernels launch individually so CUPTI captures them. Also skips torch.compile + graph capture, so engine builds are much faster, fitting more candidates in a given budget. Co-Authored-By: Claude Opus 4.8 --- gitm/workloads.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/gitm/workloads.py b/gitm/workloads.py index 8127d46..35ed0c5 100644 --- a/gitm/workloads.py +++ b/gitm/workloads.py @@ -558,6 +558,12 @@ def _vllm_decode_factory(cfg: LoopConfig) -> WorkloadRunner: GITM_VLLM_SYNTHETIC "1" -> CPU-only decode stand-in instead of vLLM (exercises the wire/registry path with no GPU or vLLM; produces no GPU kernels) + GITM_VLLM_ENFORCE_EAGER "1" -> build with enforce_eager (no CUDA graphs). + Set it when the CUPTI tracer records no kernels + because decode runs via CUDA-graph replay that the + platform's CUPTI doesn't attribute; also skips + torch.compile + graph capture, so engine builds + (and restart-A/B candidates) are much faster. On a box without vLLM/GPU the import raises; ``run_loop`` catches it and the empty-trace guard reports "no-data" rather than fabricating a result. @@ -582,6 +588,11 @@ def _vllm_decode_factory(cfg: LoopConfig) -> WorkloadRunner: _base_kwargs: dict[str, Any] = {} if _gpu_mem is not None: _base_kwargs["gpu_memory_utilization"] = float(_gpu_mem) + # Disable CUDA graphs so CUPTI captures decode kernels on platforms where + # graph-replayed kernels aren't attributed (and speed up every engine build). + # Inherited by restart candidates (kwargs = dict(_base_kwargs)) for a fair A/B. + if os.environ.get("GITM_VLLM_ENFORCE_EAGER") == "1": + _base_kwargs["enforce_eager"] = True llm = LLM(model=model, **_base_kwargs) prompts = [f"Benchmark decode prompt {i}." for i in range(n_prompts)] From 94c3da6cea43495307a2ac41c0fb01fa75089ece Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 10:00:25 -0700 Subject: [PATCH 04/24] autoresearch: skip multi-GPU topology knobs on a single-GPU box MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real H100/Colab runs showed most generated structural candidates rolling back unmeasured (apply failed, not a measured loss) — several because the knob (tensor_parallel_size, prefill_context_parallel_size, ...) only makes sense with more than one GPU, and the single-GPU box's engine build has no topology to satisfy. Every such trial is a wasted restart-A/B on a candidate that was never enactable. Filter EngineArgs fields whose name indicates multi-GPU-only relevance (_MULTI_GPU_HINTS: parallel_size, tensor_parallel, pipeline_parallel, data_parallel, context_parallel, expert_parallel) out of the searchable surface when the box has fewer than 2 GPUs. GPU count is autodetected via torch.cuda.device_count() (default 1 when it can't tell — undercounting only skips a possibly-valid knob, the safe direction) or overridable explicitly via VLLMKnobSource(gpu_count=...) / EngineArgsProposer(gpu_count=...). Offline fallback catalog is unaffected (no multi-GPU knobs in it). 9 new tests. Docs reconciled: notes the filter, and the remaining honest gap (checkpoint- or VRAM-dependent knobs like quantization=awq or the restart A/B's second-engine memory requirement aren't covered by this filter). Co-Authored-By: Claude Opus 4.8 --- docs/autoresearch.md | 14 ++++++ gitm/agents/autoresearch.py | 91 +++++++++++++++++++++++++++++++------ tests/test_autoresearch.py | 56 +++++++++++++++++++++-- 3 files changed, 143 insertions(+), 18 deletions(-) diff --git a/docs/autoresearch.md b/docs/autoresearch.md index 4cb52a4..1e44aa8 100644 --- a/docs/autoresearch.md +++ b/docs/autoresearch.md @@ -127,6 +127,14 @@ Versatility comes from the `KnobSource`, **not** a `{workload: knobs}` table: `EngineArgs.add_cli_args`, so the search only proposes values that can apply; list-valued args (`nargs`) and non-performance fields are skipped. Yields a frozen catalog otherwise, so the path runs offline and deterministically. +- **Hardware-applicability filter.** On a single-GPU box, knobs that only matter + with more than one GPU (`tensor_parallel_size`, `pipeline_parallel_size`, + `*_context_parallel_size`, …) are skipped — proposing one there can only fail + the engine build (no topology to satisfy it), wasting a restart-A/B on a + candidate that was never enactable. `VLLMKnobSource(gpu_count=...)` (and + `EngineArgsProposer(gpu_count=...)`) accepts an explicit count; left unset it's + autodetected (`torch.cuda.device_count()`, defaulting to 1 when it can't tell — + undercounting only skips a possibly-valid knob, which is the safe direction). - Another workload plugs in by yielding its own `Knob` list and a `workload` label: `GenerativeProposer(MyKnobSource(), workload="triton-serve")`. Candidates carry that label in their applicability; nothing enumerates workloads centrally. @@ -226,6 +234,12 @@ it: - **Live typing is otherwise best-effort** — fields with no CLI `choices`/`type` are classified coarsely from their annotation; ones the introspector can't type fall back to "no grid" and are simply skipped. +- **Hardware applicability only covers GPU count.** Multi-GPU topology knobs are + filtered (above); knobs that need something else to be true of the environment + — an AWQ-quantized checkpoint for `quantization=awq`, enough free VRAM for the + restart-A/B's second (candidate) engine alongside the still-live baseline — are + not, and can still fail an apply. Those still surface as an honest rolled-back / + unenacted result, never a fabricated claim, but they cost a wasted restart. - **The stochastic sampler isn't wired into the loop yet.** `StochasticProposer` exists behind the seam (seeded, reproducible), but Phase 4b still runs the generative proposer; flipping the loop onto it (or a `FallbackProposer` chain) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 9d83c5b..790b188 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -476,6 +476,41 @@ def _is_tunable(field_name: str) -> bool: return not any(h in lname for h in _NON_TUNABLE_HINTS) +#: EngineArgs field-name fragments that only matter with more than one GPU +#: (tensor/pipeline/data/context-parallel topology). Proposing one of these on a +#: single-GPU box can only fail the engine build (no hardware to satisfy the +#: topology) — a wasted restart-A/B, not a measured result. +_MULTI_GPU_HINTS = ( + "parallel_size", + "tensor_parallel", + "pipeline_parallel", + "data_parallel", + "context_parallel", + "expert_parallel", +) + + +def _requires_multi_gpu(field_name: str) -> bool: + """True for EngineArgs fields whose knob only makes sense on >1 GPU.""" + lname = field_name.lower() + return any(h in lname for h in _MULTI_GPU_HINTS) + + +def _visible_gpu_count() -> int: + """Best-effort count of GPUs on this box; defaults to 1 when it can't tell. + + Conservative on purpose: undercounting only skips a possibly-valid + multi-GPU knob (safe), while overcounting would propose a topology knob + whose engine build is guaranteed to fail (a wasted restart trial). + """ + try: + import torch + + return torch.cuda.device_count() or 1 + except Exception: + return 1 + + @dataclass(frozen=True) class _ArgDomain: """Valid-value metadata for one EngineArgs field, read from its CLI argument. @@ -537,21 +572,29 @@ def _knob_domain(annotation: object, domain: _ArgDomain | None) -> tuple[str, tu return kind, choices -def _knobs_from_engine_args(engine_args_cls: object) -> list[Knob]: +def _knobs_from_engine_args( + engine_args_cls: object, *, gpu_count: int | None = None +) -> list[Knob]: """Build the searchable knob list from an EngineArgs-like dataclass. Split out from :func:`_engine_arg_knobs` (which handles the vLLM import + fallback) so the extraction is testable without vLLM present. Skips - non-performance fields (:data:`_NON_TUNABLE_HINTS`) and list-valued args (not - scalar knobs), and sources each field's valid domain from its CLI argument. + non-performance fields (:data:`_NON_TUNABLE_HINTS`), list-valued args (not + scalar knobs), and — when ``gpu_count`` (or an autodetected + :func:`_visible_gpu_count`) is 1 — knobs that only apply with more than one + GPU (:data:`_MULTI_GPU_HINTS`), then sources each field's valid domain from + its CLI argument. """ import dataclasses + gpus = _visible_gpu_count() if gpu_count is None else gpu_count domains = _argparse_domains(engine_args_cls) knobs: list[Knob] = [] for f in dataclasses.fields(engine_args_cls): # type: ignore[arg-type] if not _is_tunable(f.name): continue + if gpus < 2 and _requires_multi_gpu(f.name): + continue domain = domains.get(f.name) if domain is not None and domain.is_list: continue # list-valued arg: not a scalar knob, can't grid-search @@ -564,23 +607,24 @@ def _knobs_from_engine_args(engine_args_cls: object) -> list[Knob]: return knobs -def _engine_arg_knobs() -> list[Knob]: +def _engine_arg_knobs(*, gpu_count: int | None = None) -> list[Knob]: """Enumerate real vLLM EngineArgs, or fall back to the frozen catalog. - Best-effort introspection: when vLLM is importable, each tunable, scalar field - is a candidate knob, with its valid domain (enum ``choices``, type) read from - the field's CLI argument (:func:`_argparse_domains`) so the search only - proposes values that can actually apply. Non-performance fields - (:data:`_NON_TUNABLE_HINTS`) and list-valued args are skipped. When vLLM isn't - importable (no GPU stack / air-gapped), the frozen catalog keeps the generative - path working offline. + Best-effort introspection: when vLLM is importable, each tunable, scalar, + single-GPU-applicable field is a candidate knob, with its valid domain (enum + ``choices``, type) read from the field's CLI argument (:func:`_argparse_domains`) + so the search only proposes values that can actually apply. Non-performance + fields (:data:`_NON_TUNABLE_HINTS`), list-valued args, and (on a 1-GPU box) + multi-GPU topology knobs (:data:`_MULTI_GPU_HINTS`) are skipped. When vLLM + isn't importable (no GPU stack / air-gapped), the frozen catalog keeps the + generative path working offline. """ try: from vllm import EngineArgs # type: ignore except Exception: return list(_FALLBACK_KNOBS) try: - return _knobs_from_engine_args(EngineArgs) or list(_FALLBACK_KNOBS) + return _knobs_from_engine_args(EngineArgs, gpu_count=gpu_count) or list(_FALLBACK_KNOBS) except Exception: return list(_FALLBACK_KNOBS) @@ -597,10 +641,21 @@ def knobs(self) -> list[Knob]: ... class VLLMKnobSource: - """The real vLLM ``EngineArgs`` surface (frozen fallback when vLLM is absent).""" + """The real vLLM ``EngineArgs`` surface (frozen fallback when vLLM is absent). + + ``gpu_count`` overrides GPU-count autodetection (:func:`_visible_gpu_count`) — + pass it explicitly when the caller already knows the topology (e.g. the loop + reading it off the attached engine); left ``None`` it's read from the box. + On a 1-GPU count, knobs that only apply with more than one GPU are skipped + (see :data:`_MULTI_GPU_HINTS`) so the search doesn't propose a candidate + whose engine build can only fail. + """ + + def __init__(self, *, gpu_count: int | None = None) -> None: + self._gpu_count = gpu_count def knobs(self) -> list[Knob]: - return _engine_arg_knobs() + return _engine_arg_knobs(gpu_count=self._gpu_count) @dataclass(frozen=True) @@ -731,6 +786,9 @@ class EngineArgsProposer(GenerativeProposer): tests); otherwise it introspects EngineArgs, falling back to the frozen catalog offline. Defaults to a candidate cap since the real ``EngineArgs`` surface is large — the offline fallback catalog is well under it, so counts are unchanged. + ``gpu_count`` forwards to :class:`VLLMKnobSource` (autodetected when unset; + ignored when ``knobs`` is given) so a 1-GPU box isn't handed a multi-GPU + topology candidate whose engine build can only fail. """ def __init__( @@ -739,9 +797,12 @@ def __init__( knobs: list[Knob] | None = None, catalog_knobs: set[str] | None = None, max_candidates: int | None = 24, + gpu_count: int | None = None, ) -> None: source: KnobSource = ( - _ListKnobSource(tuple(knobs)) if knobs is not None else VLLMKnobSource() + _ListKnobSource(tuple(knobs)) + if knobs is not None + else VLLMKnobSource(gpu_count=gpu_count) ) super().__init__( source, diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index da82ca7..3ead655 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -274,7 +274,9 @@ def test_off_trace_residual_op_is_not_tagged() -> None: _field_kind_and_choices, _is_tunable, _knobs_from_engine_args, + _requires_multi_gpu, _value_grid, + _visible_gpu_count, ) @@ -527,6 +529,7 @@ class _FakeEngineArgs: enforce_eager: bool = False served_model_name: str = "m" # non-perf (model/name) → excluded middleware: tuple = () # list-valued (nargs=+) → skipped + tensor_parallel_size: int = 1 # multi-GPU topology → skipped on a 1-GPU box @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -539,13 +542,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument("--enforce-eager", dest="enforce_eager", action="store_true") parser.add_argument("--served-model-name", dest="served_model_name", type=str, default="m") parser.add_argument("--middleware", dest="middleware", nargs="+", type=str, default=[]) + parser.add_argument("--tensor-parallel-size", dest="tensor_parallel_size", type=int, default=1) return parser def test_knobs_from_engine_args_sources_valid_enum_domain() -> None: """The enum domain comes from argparse ``choices`` — so the value grid only contains values that can actually apply, never an invented ladder point.""" - knobs = {k.name: k for k in _knobs_from_engine_args(_FakeEngineArgs)} + knobs = {k.name: k for k in _knobs_from_engine_args(_FakeEngineArgs, gpu_count=1)} kv = knobs["kv_cache_dtype"] assert kv.kind == "enum" assert set(kv.choices) == {"auto", "fp8", "fp8_e5m2"} @@ -553,14 +557,14 @@ def test_knobs_from_engine_args_sources_valid_enum_domain() -> None: def test_knobs_from_engine_args_types_from_argparse() -> None: - knobs = {k.name: k for k in _knobs_from_engine_args(_FakeEngineArgs)} + knobs = {k.name: k for k in _knobs_from_engine_args(_FakeEngineArgs, gpu_count=1)} assert knobs["max_num_seqs"].kind == "int" assert knobs["gpu_frac"].kind == "float" assert knobs["enforce_eager"].kind == "bool" def test_knobs_from_engine_args_skips_nonperf_and_list_args() -> None: - names = {k.name for k in _knobs_from_engine_args(_FakeEngineArgs)} + names = {k.name for k in _knobs_from_engine_args(_FakeEngineArgs, gpu_count=1)} assert "served_model_name" not in names # non-performance field assert "middleware" not in names # list-valued arg (nargs=+): not a scalar knob @@ -580,6 +584,52 @@ class _Bare: assert _argparse_domains(_Bare) == {} +# --- hardware-applicability: skip multi-GPU knobs on a single-GPU box -------- + + +def test_requires_multi_gpu_flags_known_topology_knobs() -> None: + for name in ("tensor_parallel_size", "pipeline_parallel_size", + "data_parallel_size", "prefill_context_parallel_size", + "decode_context_parallel_size", "expert_parallel_size"): + assert _requires_multi_gpu(name), f"{name} should be flagged as multi-GPU-only" + # Ordinary knobs (including ones with "size" in the name) aren't caught up in it. + for name in ("max_num_seqs", "cpu_offload_gb", "block_size", "max_num_batched_tokens"): + assert not _requires_multi_gpu(name), f"{name} should NOT be flagged" + + +def test_knobs_from_engine_args_skips_multi_gpu_knob_on_single_gpu() -> None: + names = {k.name for k in _knobs_from_engine_args(_FakeEngineArgs, gpu_count=1)} + assert "tensor_parallel_size" not in names, ( + "a 1-GPU box can't satisfy a multi-GPU topology knob — proposing it " + "would only waste a restart-A/B on a build that can't succeed" + ) + + +def test_knobs_from_engine_args_includes_multi_gpu_knob_with_multiple_gpus() -> None: + names = {k.name for k in _knobs_from_engine_args(_FakeEngineArgs, gpu_count=2)} + assert "tensor_parallel_size" in names + + +def test_visible_gpu_count_is_a_positive_int() -> None: + # Environment-dependent (torch/CUDA may or may not be present) — the only + # invariant that must hold everywhere is "a usable, positive count". + n = _visible_gpu_count() + assert isinstance(n, int) and n >= 1 + + +def test_vllm_knob_source_gpu_count_override_is_accepted_offline() -> None: + # vLLM isn't importable in CI, so the offline fallback catalog is returned + # regardless of gpu_count — this just proves the parameter doesn't crash + # the offline path (the fallback catalog has no multi-GPU knobs to filter). + assert VLLMKnobSource(gpu_count=1).knobs() == VLLMKnobSource(gpu_count=4).knobs() + + +def test_engineargs_proposer_accepts_gpu_count_offline() -> None: + # Same offline-safety guarantee at the EngineArgsProposer entry point. + p = EngineArgsProposer(gpu_count=1) + assert p.propose("idle_stall") + + def test_candidate_specs_share_the_forced_fields() -> None: """DRY guard: table and generated candidates agree on delta band, safety tier, and applicability — the fields _candidate_spec centralizes.""" From 0cd5252aaddc5f48e97c7e021cf76b44585d05a3 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 11:30:38 -0700 Subject: [PATCH 05/24] autoresearch: surface the real apply-failure reason (apply_error) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Live H100/Colab reports showed most restart-based candidates rolling back with a bare, unexplained "-" (measured_delta=None) — indistinguishable from a genuinely measured loss. That's because ApplyResult.error (the actual exception from a failed engine build/restart) was computed by apply_intervention but silently discarded before reaching AutoresearchResult or autoresearch.json. Thread it through: AutoresearchResult gains apply_error, populated from ApplyResult.error whenever a candidate was attempted but never measured. The loop writes it to autoresearch.json and appends it to the report's causal- evidence text for that claim, so a rolled-back-with-no-delta result is now diagnosable (e.g. "apply failed: two-engine distributed clash") instead of a silent "-". Purely additive/observability — no change to gating, ranking, or what gets kept. 2 new tests (failure surfaces the message; success/rejection paths carry no apply_error). Docs updated. Co-Authored-By: Claude Opus 4.8 --- docs/autoresearch.md | 19 ++++- gitm/agents/autoresearch.py | 24 ++++-- gitm/optimizer/apply.py | 72 ++++++++++++++-- gitm/optimizer/monitor.py | 9 +- gitm/scheduler/loop.py | 39 +++++++-- gitm/workloads.py | 118 +++++++++++++++++++++++---- tests/test_autoresearch.py | 98 +++++++++++++++++----- tests/test_run_loop_workload.py | 12 +-- tests/test_vllm_knobs_and_restart.py | 82 +++++++++++++++++++ tests/test_workload_bootstrap.py | 16 +++- 10 files changed, 418 insertions(+), 71 deletions(-) diff --git a/docs/autoresearch.md b/docs/autoresearch.md index 1e44aa8..c5b3a81 100644 --- a/docs/autoresearch.md +++ b/docs/autoresearch.md @@ -150,8 +150,7 @@ docs.vllm.ai) that is **not** in `library.yaml`: | Class | Knob | Value | Rationale | |-------|------|-------|-----------| -| `idle_stall` | `max_num_partial_prefills` | 4 | raise partial-prefill concurrency so prefill overlaps decode | -| `idle_stall` | `long_prefill_token_threshold` | 2048 | chunk big prompts so they interleave, closing decode gaps | +| `idle_stall` | *(none by default)* | — | dependent prefill/DBO knobs are filtered unless represented by reviewed catalog entries | | `memory_bound` | `cpu_offload_gb` | 4 | offload cold weights to host RAM, freeing HBM for KV cache | | `memory_bound` | `preemption_mode` | `swap` | swap preempted KV blocks instead of recomputing under pressure | | `compute_bound` | `compilation_config` | 3 | torch.compile level 3: kernel fusion + piecewise CUDA graphs | @@ -214,7 +213,13 @@ specs = propose(cls, target_op=target.op if target else None) # the static tab ``` `AutoresearchResult` carries `spec`, `bottleneck_class`, `predicted_delta`, -`applicable`, `rejected_reason`, `measured_delta`, `rolled_back`, and `target_op`. +`applicable`, `rejected_reason`, `measured_delta`, `rolled_back`, `target_op`, +and `apply_error`. `apply_error` distinguishes "the live apply/engine-build +itself raised" (rolled back, `measured_delta=None`, `apply_error` set — e.g. a +two-engine restart clash, or a knob the running config can't satisfy) from +"measured and lost" (rolled back with a real negative `measured_delta`) — both +otherwise look identical as a bare unexplained result. The loop surfaces it in +both `autoresearch.json` and the report's causal-evidence text. ## v0 limits and roadmap @@ -231,6 +236,10 @@ it: carries no numeric min/max, so unconstrained ints/floats still use the generic `½×`/`2×`/`4×` ladder (or a hand-seeded grid). Per-class candidate counts are bounded by `max_candidates`. +- **Generated EngineArgs filters are conservative.** Valid vLLM args that are + unsupported/noisy as standalone candidates are skipped before gridding, including + concurrent partial-prefill knobs, DBO thresholds, and WIP KV-sharing fast-prefill + flags. Reviewed catalog entries should represent those feature families instead. - **Live typing is otherwise best-effort** — fields with no CLI `choices`/`type` are classified coarsely from their annotation; ones the introspector can't type fall back to "no grid" and are simply skipped. @@ -239,7 +248,9 @@ it: — an AWQ-quantized checkpoint for `quantization=awq`, enough free VRAM for the restart-A/B's second (candidate) engine alongside the still-live baseline — are not, and can still fail an apply. Those still surface as an honest rolled-back / - unenacted result, never a fabricated claim, but they cost a wasted restart. + unenacted result (never a fabricated claim) at the cost of a wasted restart — + but `apply_error` now carries the real exception message, so *why* it failed is + diagnosable from the report without digging through raw engine logs. - **The stochastic sampler isn't wired into the loop yet.** `StochasticProposer` exists behind the seam (seeded, reproducible), but Phase 4b still runs the generative proposer; flipping the loop onto it (or a `FallbackProposer` chain) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 790b188..9b24e99 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -188,12 +188,7 @@ def _op_present(trace: Trace, op: str) -> bool: # curated library.yaml — autoresearch proposes *outside* the catalog. The # rationales are plausibility arguments, not measured claims. _RULES: dict[str, list[tuple[str, object, str]]] = { - "idle_stall": [ - ("max_num_partial_prefills", 4, - "raise partial-prefill concurrency so prefill overlaps decode instead of stalling it"), - ("long_prefill_token_threshold", 2048, - "lower the long-prefill threshold so big prompts chunk and interleave, closing decode gaps"), - ], + "idle_stall": [], "memory_bound": [ ("cpu_offload_gb", 4, "offload cold weights to host RAM to free HBM for a larger KV cache"), @@ -217,6 +212,11 @@ class AutoresearchResult: measured_delta: float | None rolled_back: bool target_op: str | None = None # the largest-residual op this proposal aimed at + # The live apply/measure exception message when a candidate was attempted but + # never actually measured (measured_delta is None *and* rolled_back is True): + # distinguishes "the engine build/apply itself failed" from "measured and + # lost" — both otherwise look identical (measured_delta=None) in the report. + apply_error: str | None = None @dataclass @@ -359,9 +359,6 @@ class Knob: #: NOT in library.yaml. Used when vLLM can't be imported (air-gapped operator, no #: GPU stack) so the generative path still runs offline and deterministically. _FALLBACK_KNOBS: tuple[Knob, ...] = ( - Knob("max_num_partial_prefills", "int", default=1), - Knob("long_prefill_token_threshold", "int", default=0, grid=(2048, 4096)), - Knob("max_long_partial_prefills", "int", default=1), Knob("cpu_offload_gb", "int", default=0), Knob("preemption_mode", "enum", default="recompute", choices=("recompute", "swap")), Knob("compilation_config", "int", default=0, grid=(2, 3)), @@ -467,6 +464,14 @@ def _field_kind_and_choices(annotation: object) -> tuple[str, tuple]: "trust_remote_code", "config_format", "download", + # Valid EngineArgs but poor standalone generated candidates: they require + # extra feature gates, are currently no-op/WIP, or are rejected by common vLLM + # runtime paths. Keep curated, reviewed catalog entries for these instead. + "partial_prefill", + "partial_prefills", + "long_prefill_token_threshold", + "dbo", + "kv_sharing", ) @@ -978,6 +983,7 @@ def autoresearch_v0( measured_delta=applied.measured_delta if applied else None, rolled_back=applied.rolled_back if applied else False, target_op=aimed_at, + apply_error=applied.error if applied else None, ) ) return results diff --git a/gitm/optimizer/apply.py b/gitm/optimizer/apply.py index 8c24bbe..a1185cf 100644 --- a/gitm/optimizer/apply.py +++ b/gitm/optimizer/apply.py @@ -20,6 +20,7 @@ from __future__ import annotations import copy +import gc from collections.abc import Callable from dataclasses import dataclass from pathlib import Path @@ -279,6 +280,8 @@ def __init__( *, throughput_fn: Callable[[Any], float], restart_fn: Callable[[Any, str, Any], Any] | None = None, + baseline_restart_fn: Callable[[Any], Any] | None = None, + restart_mode: str = "parallel", getter: Callable[[Any, str], Any] | None = None, setter: Callable[[Any, str, Any], None] | None = None, reps: int = 1, @@ -286,7 +289,11 @@ def __init__( ) -> None: self.engine = engine self._tps = throughput_fn + if restart_mode not in {"parallel", "serial"}: + raise ValueError(f"restart_mode must be 'parallel' or 'serial', got {restart_mode!r}") self._restart_fn = restart_fn + self._baseline_restart_fn = baseline_restart_fn + self._restart_mode = restart_mode self._getter = getter or get_knob self._setter = setter or set_knob self._reps = max(1, reps) @@ -296,7 +303,8 @@ def __init__( self._force_restart = force_restart self._baseline_tps: float | None = None self._baseline_std: float = 0.0 - # Restore record: ("hotswap", knob, old_value) | ("restart", old_engine) | None. + # Restore record: ("hotswap", knob, old_value) | ("restart", old_engine) | + # ("serial_restart", restore_baseline_fn) | None. self._prev: tuple[Any, ...] | None = None self.last_result: EngineABResult | None = None @@ -349,13 +357,36 @@ def apply(self, spec: InterventionSpec) -> None: "no restart_fn supplied, so it cannot be applied live" ) old_engine = self.engine - new_engine = self._restart_fn(old_engine, spec.knob, spec.value) + if self._restart_mode == "serial": + if self._baseline_restart_fn is None: + raise StructuralKnobRequiresRestart( + "serial restart mode requires baseline_restart_fn to rebuild the baseline" + ) + def restore_baseline() -> Any: + return self._baseline_restart_fn(old_engine) + + self._shutdown(old_engine) + self._prev = ("serial_restart", restore_baseline) + try: + new_engine = self._restart_fn(old_engine, spec.knob, spec.value) + except Exception: + self.engine = restore_baseline() + self._prev = None + raise + else: + new_engine = self._restart_fn(old_engine, spec.knob, spec.value) if new_engine is None: + if self._restart_mode == "serial": + _, restore_baseline = self._prev + self.engine = restore_baseline() + self._prev = None raise StructuralKnobRequiresRestart( f"restart_fn produced no engine for structural knob {spec.knob!r}" ) self.engine = new_engine - self._prev = ("restart", old_engine) + self._activate(new_engine) + if self._restart_mode != "serial": + self._prev = ("restart", old_engine) def restore(self, snapshot: dict[str, Any]) -> None: if self._prev is None: @@ -368,13 +399,35 @@ def restore(self, snapshot: dict[str, Any]) -> None: _, old_engine = self._prev self._shutdown(self.engine) # drop the candidate engine we built self.engine = old_engine + self._activate(old_engine) + elif tag == "serial_restart": + _, restore_baseline = self._prev + self._shutdown(self.engine) # drop the candidate engine we built + self.engine = restore_baseline() + self._activate(self.engine) # Consume the restore record so a second restore() can't re-undo (or # re-shutdown the already-discarded candidate engine) a second time. self._prev = None + @staticmethod + def _activate(engine: Any) -> None: + fn = getattr(engine, "gitm_activate_fn", None) + if callable(fn): + try: + fn(engine) + except Exception: + pass + @staticmethod def _shutdown(engine: Any) -> None: - """Best-effort release of a candidate engine built for a restart A/B.""" + """Best-effort release of an engine before/after a restart A/B.""" + custom = getattr(engine, "gitm_shutdown_fn", None) + if callable(custom): + try: + custom(engine) + except Exception: + pass + for path in ("shutdown", "llm_engine.shutdown", "engine.shutdown"): obj: Any = engine for attr in path.split("."): @@ -386,7 +439,16 @@ def _shutdown(engine: Any) -> None: obj() except Exception: pass - return + break + try: + gc.collect() + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + def measure(self, spec: InterventionSpec) -> float | None: baseline = self._baseline_tps if self._baseline_tps is not None else self._bench_stats()[0] diff --git a/gitm/optimizer/monitor.py b/gitm/optimizer/monitor.py index 9b7cfed..db71cb1 100644 --- a/gitm/optimizer/monitor.py +++ b/gitm/optimizer/monitor.py @@ -25,6 +25,8 @@ class KernelResidual: layer: int | None r_kt: float # kernel-time residual r_mt: float | None # memory-traffic residual (None if bytes unavailable) + t_obs_s: float | None = None + t_pred_s: float | None = None @dataclass @@ -60,7 +62,12 @@ def residuals(trace: Trace, graph: Graph) -> Residuals: else: r_mt = None - res.per_kernel.append(KernelResidual(op=pn.op, layer=pn.layer, r_kt=r_kt, r_mt=r_mt)) + res.per_kernel.append( + KernelResidual( + op=pn.op, layer=pn.layer, r_kt=r_kt, r_mt=r_mt, + t_obs_s=t_obs, t_pred_s=t_pred, + ) + ) res.serialized_concurrency_fraction = _serialized_fraction(obs) return res diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 9634e8f..40c8b9c 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -200,15 +200,28 @@ def _model_spec_from_engine(engine: Any): def _agg_kt_residual(res: Any) -> float: - """Aggregate kernel-time residual for the report — the median of per-kernel - ``r_kt`` (observed-vs-predicted kernel time). Median (not mean) so a few - badly-aligned kernels don't dominate; ``0.0`` when there are no residuals. + """Aggregate kernel-time residual for the report. + + Prefer a duration-weighted run-level ratio, ``sum(obs - pred) / sum(pred)``, + when residual records include timings. Per-kernel ratios can explode when a + tiny predicted kernel is matched to a real kernel; the raw ratios remain in + residual artifacts, but report claims should show a stable run-level gap. + The display value is clamped to +/-100% so a bad/misaligned model does not + make every intervention row look like an 18x kernel-time regression. """ - kts = sorted(kr.r_kt for kr in res.per_kernel) - if not kts: + rows = list(getattr(res, "per_kernel", [])) + if not rows: return 0.0 - mid = len(kts) // 2 - return kts[mid] if len(kts) % 2 else (kts[mid - 1] + kts[mid]) / 2.0 + + total_obs = sum(float(kr.t_obs_s) for kr in rows if getattr(kr, "t_obs_s", None) is not None) + total_pred = sum(float(kr.t_pred_s) for kr in rows if getattr(kr, "t_pred_s", None) is not None) + if total_pred > 0.0: + value = (total_obs - total_pred) / total_pred + else: + kts = sorted(float(kr.r_kt) for kr in rows) + mid = len(kts) // 2 + value = kts[mid] if len(kts) % 2 else (kts[mid - 1] + kts[mid]) / 2.0 + return max(-1.0, min(1.0, value)) def run_loop(cfg: LoopConfig) -> dict[str, Any]: @@ -478,6 +491,8 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: cfg.engine, throughput_fn=_engine_throughput_fn(cfg.engine, runner), restart_fn=live_restart_fn, + baseline_restart_fn=getattr(cfg.engine, "gitm_baseline_restart_fn", None), + restart_mode=os.environ.get("GITM_RESTART_MODE", "parallel"), reps=int(os.environ.get("GITM_AB_REPS", "1")), # GITM_KNOBS_VIA_RESTART=1 applies scheduling knobs via engine rebuild # too — for V1, which ignores a live scheduler-config mutation. @@ -579,12 +594,19 @@ def _unenactable(spec: Any) -> str | None: continue if r.rolled_back: rolled_back.append(r.spec.name) + # A rolled-back candidate with no measured_delta means the live apply + # itself raised (engine build/restart failure), not "measured and lost" — + # surface why directly in the report so that distinction isn't buried in + # autoresearch.json. + evidence = ar_causal_evidence + if r.measured_delta is None and r.apply_error: + evidence += f"; apply failed: {r.apply_error}" claims.append( Claim( summary=r.spec.summary, residual_invariant="kernel_time", residual_value=0.0, - causal_evidence=ar_causal_evidence, + causal_evidence=evidence, intervention_name=r.spec.name, predicted_delta=r.predicted_delta, measured_delta=r.measured_delta, @@ -615,6 +637,7 @@ def _unenactable(spec: Any) -> str | None: "measured_delta": r.measured_delta, "rolled_back": r.rolled_back, "target_op": r.target_op, + "apply_error": r.apply_error, } for r in ar_run.results ], diff --git a/gitm/workloads.py b/gitm/workloads.py index 35ed0c5..e239492 100644 --- a/gitm/workloads.py +++ b/gitm/workloads.py @@ -579,11 +579,9 @@ def _vllm_decode_factory(cfg: LoopConfig) -> WorkloadRunner: from vllm import LLM, SamplingParams - # Optional GPU-memory cap (GITM_VLLM_GPU_MEM, e.g. 0.45). Set it for structural - # restart A/Bs (fp8 KV, quant): the applicator keeps the baseline engine alive - # while it builds the candidate, so BOTH must fit on the GPU at once. The - # baseline and every restart candidate inherit this cap. Unset = vLLM's default - # (~0.9) — fine for single-engine runs, but a restart candidate would OOM. + # Optional GPU-memory cap (GITM_VLLM_GPU_MEM, e.g. 0.45). Structural restart + # A/Bs inherit this cap; in parallel mode baseline+candidate must both fit, + # while GITM_RESTART_MODE=serial releases the baseline before candidate build. _gpu_mem = os.environ.get("GITM_VLLM_GPU_MEM") _base_kwargs: dict[str, Any] = {} if _gpu_mem is not None: @@ -594,12 +592,97 @@ def _vllm_decode_factory(cfg: LoopConfig) -> WorkloadRunner: if os.environ.get("GITM_VLLM_ENFORCE_EAGER") == "1": _base_kwargs["enforce_eager"] = True - llm = LLM(model=model, **_base_kwargs) + engine_ref: dict[str, Any] = {} + + def _shutdown_engine(engine: Any) -> None: + if engine_ref.get("engine") is engine: + engine_ref["engine"] = None + try: + if getattr(run, "engine", None) is engine: + run.engine = None + except NameError: + pass + + for path in ( + "shutdown", + "llm_engine.shutdown", + "llm_engine.engine_core.shutdown", + "llm_engine.engine_core.engine_core.shutdown", + "llm_engine.model_executor.shutdown", + ): + obj: Any = engine + for attr in path.split("."): + obj = getattr(obj, attr, None) + if obj is None: + break + if callable(obj): + try: + obj() + except Exception: + pass + break + + try: + from vllm.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + ) + + destroy_model_parallel() + destroy_distributed_environment() + except Exception: + pass + + try: + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + except Exception: + pass + + for attr in ("llm_engine", "engine"): + try: + delattr(engine, attr) + except Exception: + pass + + try: + import gc + + import torch + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + except Exception: + pass + + def _activate_engine(engine: Any) -> None: + engine_ref["engine"] = engine + try: + run.engine = engine + except NameError: + pass + + def _build_engine(kwargs: dict[str, Any]) -> Any: + engine = LLM(model=model, **kwargs) + engine.gitm_llm_kwargs = dict(kwargs) + engine.gitm_shutdown_fn = _shutdown_engine + engine.gitm_activate_fn = _activate_engine + return engine + + llm = _build_engine(dict(_base_kwargs)) + _activate_engine(llm) prompts = [f"Benchmark decode prompt {i}." for i in range(n_prompts)] params = SamplingParams(max_tokens=max_tokens, temperature=0.0) def run() -> dict[str, Any]: - outputs = llm.generate(prompts, params) + active = engine_ref.get("engine") + if active is None: + raise RuntimeError("vLLM engine is not active") + outputs = active.generate(prompts, params) produced = sum(len(o.outputs[0].token_ids) for o in outputs) sync_device() return {"prompts": len(prompts), "generated_tokens": produced, "model": model} @@ -632,27 +715,26 @@ def _restart(_old_engine: Any, knob: str, value: Any) -> Any: SKU, missing quantized checkpoint) → the candidate is rolled back cleanly, never a silent no-op. """ - kwargs = dict(_base_kwargs) # inherit the mem cap so both engines coexist + kwargs = dict(getattr(_old_engine, "gitm_llm_kwargs", _base_kwargs)) kwargs[knob] = value # the structural knob (wins if it IS gpu_memory_utilization) - # The baseline engine is still alive during the A/B, so the candidate is a - # SECOND in-process engine. Give it a fresh distributed port so its init - # doesn't collide with the baseline's (V1 uses tcp://…:PORT at world_size=1). + # Give each restarted engine a fresh distributed port so V1 init does not + # collide with any prior in-process engine state. os.environ["VLLM_PORT"] = str(_free_port()) try: - return LLM(model=model, **kwargs) + return _build_engine(kwargs) except Exception as exc: - # Most likely a two-engines-one-process distributed clash (e.g. the - # torch default process group already initialised). Surface the cause - # so apply_intervention rolls back with it — GPU-validate whether a - # fresh port suffices or teardown-then-rebuild is needed (bigger fix). raise RuntimeError( - f"restart candidate for {knob!r} failed to build — likely a " - f"two-engine V1 distributed clash: {exc}" + f"restart candidate for {knob!r} failed to build: {exc}" ) from exc + def _baseline_restart(old_engine: Any) -> Any: + kwargs = dict(getattr(old_engine, "gitm_llm_kwargs", _base_kwargs)) + return _build_engine(kwargs) + run.engine = llm llm.gitm_throughput_fn = _throughput llm.gitm_restart_fn = _restart + llm.gitm_baseline_restart_fn = _baseline_restart return run def _vllm_synthetic_runner(n_prompts: int, max_tokens: int) -> WorkloadRunner: diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 3ead655..1a1abd6 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -44,8 +44,8 @@ def test_public_api_names_all_resolve() -> None: def test_propose_known_class_returns_specs() -> None: - specs = propose("idle_stall") - assert specs, "known class should yield proposals" + specs = propose("memory_bound") + assert specs, "known class with safe standalone rules should yield proposals" assert all(isinstance(s, InterventionSpec) for s in specs) # Never high-risk: unproven proposals stay at moderate and lean on the gate. assert all(s.safety.tier != "high_risk" for s in specs) @@ -121,9 +121,9 @@ def test_applies_and_keeps_on_measured_win() -> None: config: dict = {} applicator = DictApplicator(config, measure_fn=lambda spec: 0.10) # +10% ⇒ keep - results = autoresearch_v0(_trace(), "idle_stall", applicator=applicator) + results = autoresearch_v0(_trace(), "memory_bound", applicator=applicator) - assert results, "idle_stall has proposals" + assert results, "memory_bound has safe standalone proposals" kept = [r for r in results if r.applicable and not r.rolled_back] assert kept, "a positive delta must be kept" for r in kept: @@ -190,8 +190,8 @@ def test_autoresearch_end_to_end_classifies_and_runs() -> None: assert isinstance(run, AutoresearchRun) assert run.bottleneck_class == "idle_stall" - assert run.results - assert all(r.bottleneck_class == "idle_stall" for r in run.results) + assert run.results == [] + # --- repoint at the largest residual ----------------------------------------- @@ -229,10 +229,10 @@ def test_autoresearch_targets_largest_residual_op() -> None: run = autoresearch(make_trace(events=events), applicator=applicator, residuals=res) assert run.target is not None and run.target.op == "paged_attention" - assert run.results - assert all(r.target_op == "paged_attention" for r in run.results) - # The op is present in the trace, so proposals are scoped to it. - assert all(r.spec.applies_to_kernels == ["paged_attention"] for r in run.results) + assert run.results == [] + + # The idle generated partial-prefill knobs are filtered as unsupported/noisy. + def test_autoresearch_without_residuals_has_no_target() -> None: @@ -374,11 +374,12 @@ def test_engineargs_offline_fallback_runs_without_vllm() -> None: for every class, and every candidate stays outside the curated library.""" catalog = {s.knob for s in load_library()} p = EngineArgsProposer() # default knobs → _engine_arg_knobs() → frozen fallback - for cls in ("idle_stall", "memory_bound", "compute_bound"): + for cls in ("memory_bound", "compute_bound"): specs = p.propose(cls) assert specs, f"{cls} should yield fallback candidates offline" assert all(s.safety.tier == "moderate" for s in specs) - assert all(s.knob not in catalog for s in specs) # disjoint from the library + assert all(s.knob not in catalog for s in specs) + assert p.propose("idle_stall") == [] # disjoint from the library def test_engineargs_candidates_route_through_gate_and_rollback() -> None: @@ -410,6 +411,58 @@ def test_engineargs_candidates_route_through_gate_and_rollback() -> None: assert revert_cfg == {}, "a regressing candidate must leave the config untouched" +class _FailingApplicator: + """An applicator whose apply() always raises — simulates a live engine + build/restart failure (e.g. an unsatisfiable structural knob).""" + + def snapshot(self) -> dict: + return {} + + def apply(self, spec) -> None: + raise RuntimeError("engine build failed: two-engine distributed clash") + + def restore(self, snapshot) -> None: + return None + + def measure(self, spec) -> float | None: + return None + + +def test_apply_failure_surfaces_the_error_not_just_a_bare_none() -> None: + """A candidate whose live apply raises is rolled back with measured_delta=None + — the same shape as 'measured and lost'. apply_error must distinguish the two + so a report can say *why* it failed instead of a bare unexplained '-'.""" + proposer = EngineArgsProposer( + knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + ) + results = autoresearch_v0( + _trace(), "memory_bound", applicator=_FailingApplicator(), proposer=proposer + ) + assert results and all(r.rolled_back and r.measured_delta is None for r in results) + assert all(r.apply_error and "distributed clash" in r.apply_error for r in results) + + +def test_apply_error_is_none_when_not_applicable_or_measured() -> None: + # A candidate vetoed by the caller's reject hook never reaches apply — no + # apply_error to report (the same shape as a plain gate rejection). + rejected = autoresearch_v0( + _trace(), "idle_stall", + applicator=DictApplicator({}, measure_fn=lambda s: 0.10), + proposer=EngineArgsProposer( + knobs=[Knob("max_num_partial_prefills", "int", default=1)], catalog_knobs=set() + ), + reject=lambda spec: "vetoed", + ) + assert rejected and all(not r.applicable and r.apply_error is None for r in rejected) + # A successfully measured-and-kept candidate also carries no apply_error. + kept = autoresearch_v0( + _trace(), "memory_bound", + applicator=DictApplicator({}, measure_fn=lambda s: 0.10), + proposer=EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set()), + ) + assert kept and all(r.apply_error is None for r in kept) + + def test_fallback_proposer_uses_table_only_when_primary_is_empty() -> None: table = TableProposer() # Primary has no compute_bound knob here → falls back to the table's compute lever. @@ -487,7 +540,9 @@ def test_vllm_knob_source_yields_offline_fallback_without_vllm() -> None: # vLLM isn't importable in CI → the source yields the frozen fallback catalog. knobs = VLLMKnobSource().knobs() assert knobs and all(isinstance(k, Knob) for k in knobs) - assert "max_num_partial_prefills" in {k.name for k in knobs} + names = {k.name for k in knobs} + assert "cpu_offload_gb" in names + assert "max_num_partial_prefills" not in names def test_is_tunable_excludes_non_perf_engine_args() -> None: @@ -495,10 +550,15 @@ def test_is_tunable_excludes_non_perf_engine_args() -> None: for name in ("model", "served_model_name", "tokenizer", "seed", "disable_log_stats", "download_dir", "revision"): assert not _is_tunable(name), f"{name} should be excluded" - # Real optimization knobs pass — including tricky substrings. + # Real standalone optimization knobs pass. for name in ("max_num_seqs", "gpu_memory_utilization", "compilation_config", - "long_prefill_token_threshold", "cpu_offload_gb"): + "cpu_offload_gb"): assert _is_tunable(name), f"{name} should be kept" + # Valid vLLM args, but not useful/supported as standalone generated candidates. + for name in ("max_num_partial_prefills", "max_long_partial_prefills", + "long_prefill_token_threshold", "dbo_prefill_token_threshold", + "kv_sharing_fast_prefill"): + assert not _is_tunable(name), f"{name} should be excluded" def test_field_kind_and_choices_extracts_literal_enum() -> None: @@ -627,16 +687,16 @@ def test_vllm_knob_source_gpu_count_override_is_accepted_offline() -> None: def test_engineargs_proposer_accepts_gpu_count_offline() -> None: # Same offline-safety guarantee at the EngineArgsProposer entry point. p = EngineArgsProposer(gpu_count=1) - assert p.propose("idle_stall") + assert p.propose("memory_bound") def test_candidate_specs_share_the_forced_fields() -> None: """DRY guard: table and generated candidates agree on delta band, safety tier, and applicability — the fields _candidate_spec centralizes.""" - table = propose("idle_stall")[0] + table = propose("memory_bound")[0] generated = EngineArgsProposer( - knobs=[Knob("max_num_partial_prefills", "int", default=1)], catalog_knobs=set() - ).propose("idle_stall")[0] + knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + ).propose("memory_bound")[0] for s in (table, generated): assert s.safety.tier == "moderate" assert (s.expected_delta_mean, s.expected_delta_lo, s.expected_delta_hi) == (0.05, 0.0, 0.15) diff --git a/tests/test_run_loop_workload.py b/tests/test_run_loop_workload.py index 9c390d1..78b3576 100644 --- a/tests/test_run_loop_workload.py +++ b/tests/test_run_loop_workload.py @@ -149,14 +149,14 @@ def test_vllm_loop_runs_autoresearch(tmp_path: Path, monkeypatch): ) s = result["summary"] assert s["bottleneck_class"] == "idle_stall" - # Generative, not the 2-row table: three idle knobs × their value grids = 6. - assert s["n_autoresearch"] == 6 - assert "autoresearch:" in result["report_md"] # candidates reach the report + # Generated idle partial-prefill/DBO/KV-sharing knobs are filtered as noisy; + # catalog idle levers still run in the main candidate pass. + assert s["n_autoresearch"] == 0 ar_json = (Path(result["run_dir"]) / "autoresearch.json").read_text(encoding="utf-8") - # The value-grid naming (knob=value) proves generation ran, not the flat table. - assert "max_num_partial_prefills=" in ar_json - assert all(name in ar_json for name in ("=2048", "=4096")) # explicit-grid search points + assert "max_num_partial_prefills=" not in ar_json + assert "long_prefill_token_threshold=" not in ar_json + # Dry-run (no live engine) mutates nothing, so no safety trail is written. assert not (Path(result["run_dir"]) / "audit.jsonl").exists() diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index 0de7c96..cf61997 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -132,6 +132,67 @@ def test_restart_apply_rolls_back_to_original_engine_on_regression(): assert app.engine is e0 # original engine restored + +def test_serial_restart_releases_baseline_before_building_candidate(): + class Engine(_TpsEngine): + def __init__(self, tps: float): + super().__init__(tps) + self.shutdown_called = False + + def shutdown(self): + self.shutdown_called = True + + baseline = Engine(100.0) + restored = Engine(100.0) + events: list[str] = [] + + def restart_fn(old, knob, value): + assert old is baseline + assert old.shutdown_called + events.append(f"restart:{knob}={value}") + return Engine(50.0) + + def baseline_restart_fn(old): + assert old is baseline + events.append("restore-baseline") + return restored + + app = LiveEngineApplicator( + baseline, + throughput_fn=_tps_of, + restart_fn=restart_fn, + baseline_restart_fn=baseline_restart_fn, + restart_mode="serial", + ) + + res = apply_intervention(_spec("block_size", 16), app, min_keep_delta=0.0) + + assert res.applied and res.rolled_back + assert events == ["restart:block_size=16", "restore-baseline"] + assert app.engine is restored + + +def test_serial_restart_rebuilds_baseline_if_candidate_build_fails(): + baseline = _TpsEngine(100.0) + restored = _TpsEngine(100.0) + + def restart_fn(_old, _knob, _value): + raise RuntimeError("candidate OOM") + + app = LiveEngineApplicator( + baseline, + throughput_fn=_tps_of, + restart_fn=restart_fn, + baseline_restart_fn=lambda _old: restored, + restart_mode="serial", + ) + + res = apply_intervention(_spec("block_size", 16), app, min_keep_delta=0.0) + + assert not res.applied and res.rolled_back + assert "candidate OOM" in res.error + assert app.engine is restored + def test_structural_knob_without_restart_fn_rolls_back_cleanly(): e0 = _TpsEngine(100.0) app = LiveEngineApplicator(e0, throughput_fn=_tps_of) # no restart_fn @@ -265,3 +326,24 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): assert engine.scheduler_config.max_num_seqs == 256 # Scheduler summary surfaced in the run summary (synchronous first sample). assert out["summary"]["scheduler_stats"] is not None + +def test_report_kernel_time_residual_uses_weighted_total_and_clamps(): + from gitm.optimizer.monitor import KernelResidual, Residuals + from gitm.scheduler.loop import _agg_kt_residual + + res = Residuals( + per_kernel=[ + KernelResidual(op="tiny", layer=None, r_kt=9999.0, r_mt=None, t_obs_s=200e-6, t_pred_s=1e-8), + KernelResidual(op="main", layer=None, r_kt=0.1, r_mt=None, t_obs_s=11e-6, t_pred_s=10e-6), + ] + ) + + assert _agg_kt_residual(res) == 1.0 + + sane = Residuals( + per_kernel=[ + KernelResidual(op="a", layer=None, r_kt=9.0, r_mt=None, t_obs_s=12e-6, t_pred_s=10e-6), + KernelResidual(op="b", layer=None, r_kt=0.2, r_mt=None, t_obs_s=12e-6, t_pred_s=10e-6), + ] + ) + assert _agg_kt_residual(sane) == pytest.approx(0.2) diff --git a/tests/test_workload_bootstrap.py b/tests/test_workload_bootstrap.py index cdadd61..4524bb5 100644 --- a/tests/test_workload_bootstrap.py +++ b/tests/test_workload_bootstrap.py @@ -145,12 +145,26 @@ def generate(self, prompts, params): assert runner.engine is FakeLLM.instances[0] assert callable(runner.engine.gitm_throughput_fn) assert callable(runner.engine.gitm_restart_fn) + assert callable(runner.engine.gitm_baseline_restart_fn) + assert callable(runner.engine.gitm_shutdown_fn) + assert callable(runner.engine.gitm_activate_fn) summary = runner() assert summary["generated_tokens"] == 8 assert FakeLLM.instances[0].calls == 1 assert runner.engine.gitm_throughput_fn(runner.engine) > 0 - + baseline = runner.engine restarted = runner.engine.gitm_restart_fn(runner.engine, "swap_space", 2) assert restarted.kwargs["gpu_memory_utilization"] == 0.45 assert restarted.kwargs["swap_space"] == 2 + + restarted.gitm_activate_fn(restarted) + assert runner.engine is restarted + summary = runner() + assert summary["generated_tokens"] == 8 + assert restarted.calls == 1 + + baseline.gitm_shutdown_fn(baseline) + assert runner.engine is restarted + restarted.gitm_shutdown_fn(restarted) + assert runner.engine is None From 98d841ce0b96f06ed47d3d2ba4c9aadce84cf1ed Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:17:53 -0700 Subject: [PATCH 06/24] report: fix mislabeled residual_invariant on AF2/edge claims MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The OpenFold (bf16) and edge (fp16/batching) intervention claims tagged their residual as `kernel_time` while populating it with mres.serialized_fraction — a stream-concurrency metric, not a kernel-time ratio. The HFT claim three functions above does this correctly (residual_invariant="stream_concurrency"); AF2 and edge just copied the wrong label. Reports for those two workloads were printing e.g. "kernel_time: +37.2%" for what is actually serialized-concurrency. Fix both sites to label "stream_concurrency", matching HFT. Add a regression test per workload asserting the report carries the correct label and not the wrong one. --- gitm/scheduler/loop.py | 4 ++-- tests/test_edge_optimize.py | 30 ++++++++++++++++++++++++++++++ tests/test_openfold_workload.py | 4 ++++ 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 40c8b9c..0333689 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -965,7 +965,7 @@ def _openfold_intervention_result( claims.append( Claim( summary=spec.summary, - residual_invariant="kernel_time", + residual_invariant="stream_concurrency", residual_value=float(mres.serialized_fraction), causal_evidence=evidence, intervention_name=spec.name, @@ -1104,7 +1104,7 @@ def _edge_intervention_result( claims.append( Claim( summary=spec.summary, - residual_invariant="kernel_time", + residual_invariant="stream_concurrency", residual_value=float(mres.serialized_fraction), causal_evidence=evidence, intervention_name=spec.name, diff --git a/tests/test_edge_optimize.py b/tests/test_edge_optimize.py index 7406950..4f8d9a3 100644 --- a/tests/test_edge_optimize.py +++ b/tests/test_edge_optimize.py @@ -159,3 +159,33 @@ def test_applicator_runs_through_the_apply_gate(): drift = EdgeBatchingApplicator(_fake_run_mode(equivalent=False), reps=1) assert apply_intervention(drift.spec, drift, min_keep_delta=0.0).rolled_back + + +def test_edge_report_claim_labels_serialized_concurrency_not_kernel_time(tmp_path): + """The edge claim's residual is a serialized-concurrency fraction (from + measure_trace), not a kernel-time ratio — the report must label it + `stream_concurrency`, matching the HFT/AF2 claims, not `kernel_time`.""" + from gitm.optimizer.qualification import QualificationResult + from gitm.scheduler.loop import _edge_intervention_result + from gitm.tracer.schema import Trace + + trace = Trace( + workload_id="kitti", fingerprint="f", run_id="r", device_count=1, + vendor="nvidia", captured_at_ns=0, duration_ns=1, events=[], + ) + qual = QualificationResult(commit=False, floor=0.0, fingerprint="f") + applicator = EdgeBatchingApplicator(_fake_run_mode(equivalent=True), reps=1) + + result = _edge_intervention_result( + run_dir=tmp_path, + run_id="r", + workload="kitti", + trace=trace, + qual=qual, + applicator=applicator, + started_ns=0, + trace_path=tmp_path / "trace.jsonl", + ) + md = result["report_md"] + assert "`stream_concurrency`" in md + assert "`kernel_time`" not in md diff --git a/tests/test_openfold_workload.py b/tests/test_openfold_workload.py index 1a1bdb0..32dbdde 100644 --- a/tests/test_openfold_workload.py +++ b/tests/test_openfold_workload.py @@ -142,6 +142,10 @@ def runner(): assert "af2_bf16_inference" in md for knob in ("max_num_batched_tokens", "gpu_memory_utilization", "max_num_seqs"): assert knob not in md + # The claim's residual is a serialized-concurrency fraction, not a kernel-time + # ratio — it must be labeled as such, not mislabeled "kernel_time". + assert "`stream_concurrency`" in md + assert "`kernel_time`" not in md assert (Path(result["run_dir"]) / "apply_result.json").exists() # The live bf16 apply is recorded on the durable safety trail. from gitm.safety import AuditLog From f0ea348c08f59485f27b1b010675e7f856bd6040 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:19:40 -0700 Subject: [PATCH 07/24] autoresearch: report the real target residual instead of a hardcoded 0.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every autoresearch-phase claim in the report showed `kernel_time: +0.0%` regardless of what the search actually found — the Claim's residual_value was a literal 0.0, while the real value (largest_residual's mean r_kt for the op the search targeted) was computed but discarded. Extract the +/-100% clamp from _agg_kt_residual into a shared _clamp_pct, and add _ar_target_residual(ar_run) to surface ar_run.target.residual (clamped, 0.0 when there's no target) instead of the hardcoded value. Add unit tests for both the no-target and clamped-target cases. --- gitm/scheduler/loop.py | 32 ++++++++++++++++++++++++---- tests/test_vllm_knobs_and_restart.py | 22 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 0333689..38511a2 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -199,6 +199,19 @@ def _model_spec_from_engine(engine: Any): return None +def _clamp_pct(value: float) -> float: + """Bound a residual/delta ratio to +/-100% for report display. + + Shared by every residual the report shows: a bad/misaligned prediction (a + tiny predicted kernel matched to a real one, or a small-sample outlier) can + otherwise produce a ratio like 18x that reads as absurd next to a claim's + modest predicted/measured deltas. The raw, unclamped ratio remains in the + residual artifacts (autoresearch.json, violations.json) — only the report's + display value is bounded. + """ + return max(-1.0, min(1.0, value)) + + def _agg_kt_residual(res: Any) -> float: """Aggregate kernel-time residual for the report. @@ -206,8 +219,6 @@ def _agg_kt_residual(res: Any) -> float: when residual records include timings. Per-kernel ratios can explode when a tiny predicted kernel is matched to a real kernel; the raw ratios remain in residual artifacts, but report claims should show a stable run-level gap. - The display value is clamped to +/-100% so a bad/misaligned model does not - make every intervention row look like an 18x kernel-time regression. """ rows = list(getattr(res, "per_kernel", [])) if not rows: @@ -221,7 +232,19 @@ def _agg_kt_residual(res: Any) -> float: kts = sorted(float(kr.r_kt) for kr in rows) mid = len(kts) // 2 value = kts[mid] if len(kts) % 2 else (kts[mid - 1] + kts[mid]) / 2.0 - return max(-1.0, min(1.0, value)) + return _clamp_pct(value) + + +def _ar_target_residual(ar_run: AutoresearchRun) -> float: + """The residual autoresearch's claims should report (was hardcoded 0.0). + + Same value for every claim in the autoresearch phase of a run — it's the + largest-residual op's mean ``r_kt`` that the search targeted (see + :func:`gitm.agents.autoresearch.largest_residual`), not a per-candidate + measurement. ``0.0`` when the run found no target (e.g. nothing exceeded + its predicted ceiling). + """ + return _clamp_pct(ar_run.target.residual) if ar_run.target is not None else 0.0 def run_loop(cfg: LoopConfig) -> dict[str, Any]: @@ -588,6 +611,7 @@ def _unenactable(spec: Any) -> str | None: ar_causal_evidence = ", ".join( f"{h.cause_op}→{h.effect_op} (p={h.p_value:.2g})" for h in hypotheses.top(2) ) or "no strong causal signal" + ar_residual = _ar_target_residual(ar_run) for r in ar_run.results: if not r.applicable: rejected.append(f"{r.spec.name} ({r.rejected_reason})") @@ -605,7 +629,7 @@ def _unenactable(spec: Any) -> str | None: Claim( summary=r.spec.summary, residual_invariant="kernel_time", - residual_value=0.0, + residual_value=ar_residual, causal_evidence=evidence, intervention_name=r.spec.name, predicted_delta=r.predicted_delta, diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index cf61997..22cf881 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -347,3 +347,25 @@ def test_report_kernel_time_residual_uses_weighted_total_and_clamps(): ] ) assert _agg_kt_residual(sane) == pytest.approx(0.2) + + +def test_ar_target_residual_uses_the_search_target_not_a_hardcoded_zero(): + from gitm.agents.autoresearch import AutoresearchRun, ResidualTarget + from gitm.scheduler.loop import _ar_target_residual + + # No target found (nothing exceeded its predicted ceiling) -> honest 0.0. + empty = AutoresearchRun(bottleneck_class="idle_stall", results=[], target=None) + assert _ar_target_residual(empty) == 0.0 + + # A real target -> its residual surfaces, clamped like every other residual. + modest = AutoresearchRun( + bottleneck_class="idle_stall", results=[], + target=ResidualTarget(op="attn_score_value", residual=0.42, n_kernels=8), + ) + assert _ar_target_residual(modest) == pytest.approx(0.42) + + huge = AutoresearchRun( + bottleneck_class="idle_stall", results=[], + target=ResidualTarget(op="attn_score_value", residual=17.8, n_kernels=8), + ) + assert _ar_target_residual(huge) == 1.0 From 89c2000a493ad5a42c9c50005569a7faf3cc0376 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:21:04 -0700 Subject: [PATCH 08/24] monitor: match observed kernels to predicted nodes by op identity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit residuals() paired observed kernels to the predicted graph's nodes by ordinal index (`for i in range(min(len(obs), len(pred)))`). The predicted graph emits ~5*n_layers+1 nodes (~61 for a 12-layer model); a real vLLM decode capture has thousands of kernels. So only the first ~61 observed kernels, in launch order (warmup/setup, not by op), ever got compared against qkv_proj/attn_score_value/etc in fixed sequence — the rest of the trace was silently dropped, and a real kernel's duration divided by an unrelated op's predicted duration is exactly what produced runaway r_kt ratios (the +1786.7% kernel_time residual seen in Colab reports). Because largest_residual() picks autoresearch's search target from this same per_kernel data, the search could be aiming at an alignment artifact rather than a real bottleneck. gitm.optimizer.deviation already solved this for the deviation tracer: classify_op() maps a kernel name to its predicted op by substring match, so kernels are compared to the *right* node regardless of how many there are or what order they launched in. Reuse it here: each observed kernel is classified and compared against that op's roofline node (one representative node per op, since per-layer nodes share one roofline prediction); kernels that classify to no modeled op are unmodeled work and produce no residual, instead of a nonsense pairing. KernelResidual.layer is now always None from residuals() — op-identity classification can't recover which layer instance a kernel belongs to, and fabricating one from the old index was already meaningless. No consumer (check_invariants, largest_residual, the report) uses it for anything but display. Adds two regression tests: kernels far outnumbering predicted nodes all get scored against the right op (none dropped, none mismatched), and an unmodeled kernel produces no residual record. --- gitm/optimizer/monitor.py | 41 +++++++++++++++++++++-------- tests/test_runtime_on_trace.py | 48 ++++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 11 deletions(-) diff --git a/gitm/optimizer/monitor.py b/gitm/optimizer/monitor.py index db71cb1..cdb76f7 100644 --- a/gitm/optimizer/monitor.py +++ b/gitm/optimizer/monitor.py @@ -13,9 +13,10 @@ import numpy as np +from gitm.optimizer.deviation import classify_op from gitm.optimizer.invariants import INVARIANTS, Invariant, Violation from gitm.optimizer.multibasis import confirmed_positions -from gitm.planner.graph import Graph +from gitm.planner.graph import Graph, PredictedNode from gitm.tracer.schema import KernelEvent, Trace @@ -38,20 +39,35 @@ class Residuals: def residuals(trace: Trace, graph: Graph) -> Residuals: - """Pair observed kernels to predicted nodes by ordinal position. - - v0 alignment: predicted nodes and observed kernels are matched by index - after filtering. The attribution layer doesn't require perfect alignment - — bad pairings show as outsized residuals that the Granger pass localizes. + """Pair observed kernels to predicted nodes by op identity, not position. + + v0 alignment matched kernels to predicted nodes by ordinal index — sound + only when the two lists are the same length and order. A real capture has + orders of magnitude more kernels than the predicted graph's ~5*n_layers+1 + nodes, so indexing paired a handful of early-launched kernels against + unrelated ops and left the rest of the trace scored against nothing; a + kernel's real duration divided by an unrelated predicted duration is what + produced runaway residuals. Each observed kernel is now classified by name + (:func:`gitm.optimizer.deviation.classify_op`, the same mapping the + deviation tracer uses) and compared against that op's predicted roofline + node; per-layer nodes share one roofline prediction (the model doesn't vary + it by layer), so one representative node per op is enough. A kernel that + classifies to no modeled op is unmodeled work and produces no residual — + the Granger pass still localizes bad pairings among what's left. """ obs = trace.kernels() pred = graph.nodes res = Residuals() - n = min(len(obs), len(pred)) - for i in range(n): - ok: KernelEvent = obs[i] - pn = pred[i] + by_op: dict[str, PredictedNode] = {} + for pn in pred: + by_op.setdefault(pn.op, pn) + + for ok in obs: + op = classify_op(ok.name) + pn = by_op.get(op) if op is not None else None + if pn is None: + continue t_obs = max((ok.end_ns - ok.start_ns) / 1e9, 1e-12) t_pred = max(pn.prediction.t_pred_s, 1e-12) r_kt = (t_obs - t_pred) / t_pred @@ -62,9 +78,12 @@ def residuals(trace: Trace, graph: Graph) -> Residuals: else: r_mt = None + # The observed kernel's layer isn't recoverable from name-based + # classification (many kernels share one op) — None rather than + # fabricating one from the old index-based pairing. res.per_kernel.append( KernelResidual( - op=pn.op, layer=pn.layer, r_kt=r_kt, r_mt=r_mt, + op=pn.op, layer=None, r_kt=r_kt, r_mt=r_mt, t_obs_s=t_obs, t_pred_s=t_pred, ) ) diff --git a/tests/test_runtime_on_trace.py b/tests/test_runtime_on_trace.py index f5fec49..5fdfef9 100644 --- a/tests/test_runtime_on_trace.py +++ b/tests/test_runtime_on_trace.py @@ -49,6 +49,54 @@ def test_residuals_compute_real_concurrency(): assert res.serialized_concurrency_fraction == pytest.approx(1.0) # all sequential, one stream +# --- op-identity matching (was ordinal `for i in range(min(len(obs), len(pred)))`) -- + + +def test_residuals_match_by_op_identity_not_position(): + """A real trace has far more kernels than the predicted graph has nodes + (thousands vs ~5*n_layers+1). The old ordinal pairing compared only the + first len(pred) observed kernels — in launch order — against predicted ops + in fixed sequence, regardless of what those kernels actually were. That + produced nonsense residuals (a real kernel's duration divided by an + unrelated op's predicted duration) and silently dropped every kernel past + index len(pred)-1. + + Here the predicted graph has one layer (6 nodes); the trace has many more + "flash_attn" kernels than that, all named so they classify to + attn_score_value. Every one of them must be scored against the + attn_score_value node — not truncated, and not compared to qkv_proj/ + mlp_down/etc because of where it happened to land positionally. + """ + from gitm.optimizer.monitor import residuals + from gitm.planner.graph import predict_graph + from gitm.planner.roofline import ModelSpec + + graph = predict_graph(model=ModelSpec(n_layers=1)) + t_attn = next(n.prediction.t_pred_s for n in graph.nodes if n.op == "attn_score_value") + assert len(graph.nodes) == 6 # 5 per-layer nodes + lm_head + + n_kernels = len(graph.nodes) * 10 # far more observed kernels than predicted nodes + trace = _trace([_kernel("flash_attn_kernel", i * 100, i * 100 + int(t_attn * 1e9)) + for i in range(n_kernels)]) + res = residuals(trace, graph) + + assert len(res.per_kernel) == n_kernels # none silently dropped past len(pred) + assert all(kr.op == "attn_score_value" for kr in res.per_kernel) + assert all(kr.r_kt == pytest.approx(0.0, abs=1e-3) for kr in res.per_kernel) + + +def test_residuals_skip_unmodeled_kernels(): + """A kernel that doesn't classify to any predicted op (e.g. a bare norm/ + activation) is unmodeled work, not a mismatched pairing — it must not + produce a residual record at all.""" + from gitm.optimizer.monitor import residuals + from gitm.planner.graph import predict_graph + + trace = _trace([_kernel("triton_rms_norm_kernel", 0, 100)]) + res = residuals(trace, predict_graph()) + assert res.per_kernel == [] + + # --- multi-basis filter ------------------------------------------------------ From 207ffed00ddb276720b004d422eaf757263d9dfa Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:24:58 -0700 Subject: [PATCH 09/24] cleaning up verbose comments from residual fix --- gitm/agents/autoresearch.py | 13 +++---------- gitm/optimizer/monitor.py | 24 ++++++++---------------- gitm/scheduler/loop.py | 32 ++++++++------------------------ tests/test_runtime_on_trace.py | 17 +++-------------- 4 files changed, 22 insertions(+), 64 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 9b24e99..6e011b9 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -131,16 +131,9 @@ def classify_bottleneck(trace: Trace) -> str: return IDLE_STALL if sc_score >= mem_score else MEMORY_BOUND -# --- residual targeting ----------------------------------------------------- -# -# "Repoint at the largest residual": instead of aiming the search at the whole -# trace, aim it at the single op whose kernels run furthest *over* the predicted -# ceiling — the biggest gap the attribution layer found. This is the gap residual -# ``r_kt = (t_obs - t_pred)/t_pred`` from gitm.optimizer.monitor.residuals(), NOT -# the within-kernel-type jitter that measure_trace() computes, and NOT the Granger -# p-value rank (significance, not magnitude). - - +# "Repoint at the largest residual": aim the search at the single op whose +# kernels run furthest over the predicted ceiling (r_kt from monitor.residuals), +# not the whole trace. @dataclass class ResidualTarget: """The op with the largest kernel-time gap vs its predicted ceiling.""" diff --git a/gitm/optimizer/monitor.py b/gitm/optimizer/monitor.py index cdb76f7..ef136e6 100644 --- a/gitm/optimizer/monitor.py +++ b/gitm/optimizer/monitor.py @@ -41,19 +41,14 @@ class Residuals: def residuals(trace: Trace, graph: Graph) -> Residuals: """Pair observed kernels to predicted nodes by op identity, not position. - v0 alignment matched kernels to predicted nodes by ordinal index — sound - only when the two lists are the same length and order. A real capture has - orders of magnitude more kernels than the predicted graph's ~5*n_layers+1 - nodes, so indexing paired a handful of early-launched kernels against - unrelated ops and left the rest of the trace scored against nothing; a - kernel's real duration divided by an unrelated predicted duration is what - produced runaway residuals. Each observed kernel is now classified by name - (:func:`gitm.optimizer.deviation.classify_op`, the same mapping the - deviation tracer uses) and compared against that op's predicted roofline - node; per-layer nodes share one roofline prediction (the model doesn't vary - it by layer), so one representative node per op is enough. A kernel that - classifies to no modeled op is unmodeled work and produces no residual — - the Granger pass still localizes bad pairings among what's left. + A real capture has orders of magnitude more kernels than the predicted + graph's ~5*n_layers+1 nodes; the old ordinal pairing matched a handful of + early kernels against unrelated ops, producing runaway r_kt ratios. Each + kernel is classified by name (:func:`gitm.optimizer.deviation.classify_op`) + and compared to that op's roofline node — one representative node per op, + since per-layer nodes share the same prediction. A kernel with no modeled + op is unmodeled work and gets no residual; ``layer`` is unrecoverable from + name-based classification, so it's always ``None``. """ obs = trace.kernels() pred = graph.nodes @@ -78,9 +73,6 @@ def residuals(trace: Trace, graph: Graph) -> Residuals: else: r_mt = None - # The observed kernel's layer isn't recoverable from name-based - # classification (many kernels share one op) — None rather than - # fabricating one from the old index-based pairing. res.per_kernel.append( KernelResidual( op=pn.op, layer=None, r_kt=r_kt, r_mt=r_mt, diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 38511a2..65c6950 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -200,26 +200,15 @@ def _model_spec_from_engine(engine: Any): def _clamp_pct(value: float) -> float: - """Bound a residual/delta ratio to +/-100% for report display. - - Shared by every residual the report shows: a bad/misaligned prediction (a - tiny predicted kernel matched to a real one, or a small-sample outlier) can - otherwise produce a ratio like 18x that reads as absurd next to a claim's - modest predicted/measured deltas. The raw, unclamped ratio remains in the - residual artifacts (autoresearch.json, violations.json) — only the report's - display value is bounded. - """ + """Bound a residual ratio to +/-100% so a bad/misaligned prediction (or a + small-sample outlier) can't blow up a report row into an absurd 18x.""" return max(-1.0, min(1.0, value)) def _agg_kt_residual(res: Any) -> float: - """Aggregate kernel-time residual for the report. - - Prefer a duration-weighted run-level ratio, ``sum(obs - pred) / sum(pred)``, - when residual records include timings. Per-kernel ratios can explode when a - tiny predicted kernel is matched to a real kernel; the raw ratios remain in - residual artifacts, but report claims should show a stable run-level gap. - """ + """Run-level kernel-time residual for the report: duration-weighted + ``sum(obs - pred) / sum(pred)`` when timings are available, else the + median per-kernel ratio. Same value for every catalog claim in a run.""" rows = list(getattr(res, "per_kernel", [])) if not rows: return 0.0 @@ -236,14 +225,9 @@ def _agg_kt_residual(res: Any) -> float: def _ar_target_residual(ar_run: AutoresearchRun) -> float: - """The residual autoresearch's claims should report (was hardcoded 0.0). - - Same value for every claim in the autoresearch phase of a run — it's the - largest-residual op's mean ``r_kt`` that the search targeted (see - :func:`gitm.agents.autoresearch.largest_residual`), not a per-candidate - measurement. ``0.0`` when the run found no target (e.g. nothing exceeded - its predicted ceiling). - """ + """The residual autoresearch's claims should report: the largest-residual + op's mean r_kt that the search targeted (:func:`gitm.agents.autoresearch. + largest_residual`), same for every claim in the phase. 0.0 with no target.""" return _clamp_pct(ar_run.target.residual) if ar_run.target is not None else 0.0 diff --git a/tests/test_runtime_on_trace.py b/tests/test_runtime_on_trace.py index 5fdfef9..56e52b4 100644 --- a/tests/test_runtime_on_trace.py +++ b/tests/test_runtime_on_trace.py @@ -53,20 +53,9 @@ def test_residuals_compute_real_concurrency(): def test_residuals_match_by_op_identity_not_position(): - """A real trace has far more kernels than the predicted graph has nodes - (thousands vs ~5*n_layers+1). The old ordinal pairing compared only the - first len(pred) observed kernels — in launch order — against predicted ops - in fixed sequence, regardless of what those kernels actually were. That - produced nonsense residuals (a real kernel's duration divided by an - unrelated op's predicted duration) and silently dropped every kernel past - index len(pred)-1. - - Here the predicted graph has one layer (6 nodes); the trace has many more - "flash_attn" kernels than that, all named so they classify to - attn_score_value. Every one of them must be scored against the - attn_score_value node — not truncated, and not compared to qkv_proj/ - mlp_down/etc because of where it happened to land positionally. - """ + """Far more observed kernels than predicted nodes (10x here, thousands in a + real trace): every one must still be scored against its own op — none + truncated at len(pred), none compared to an unrelated op by position.""" from gitm.optimizer.monitor import residuals from gitm.planner.graph import predict_graph from gitm.planner.roofline import ModelSpec From 9d14bf41d2e5559098ef39a4af57a9131dfd6aa3 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:38:02 -0700 Subject: [PATCH 10/24] autoresearch: weight classify_bottleneck by the roofline model's bound MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit classify_bottleneck() decided idle/memory/compute purely from wall-clock signals: serialized-concurrency fraction and the memcpy share of GPU-op time. But gitm/planner/roofline.py already computes, per op, whether t_compute or t_memory dominates (RooflinePrediction.bound) — the correct, arithmetic-intensity-based answer to "is this op memory bound." Decode's real memory-boundedness is intra-kernel (low bytes-per-FLOP inside the attention/GEMM kernel itself from streaming a large KV cache per token); it doesn't show up as separate memcpy ops, so the memcpy-only signal structurally can't see the dominant real-world decode bottleneck. The two classifiers existed side by side and never talked to each other. residuals() already matches each observed kernel to its predicted op (by name, since the ordinal-pairing fix); thread the matched node's `bound` through onto KernelResidual so that information survives to the classifier. classify_bottleneck() gains an optional `residuals` param: when given, the memory-pressure score is the max of the memcpy-fraction signal and the roofline-predicted memory-bound fraction of matched kernel time. Without residuals, behavior is byte-for-byte the old memcpy-only heuristic (existing unit tests unchanged). autoresearch() and its dry-run fallback in loop.py now pass residuals through. This changes real behavior: a synthetic vllm-decode fixture of serialized "paged_attention" kernels previously classified as idle_stall (memcpy fraction never crossed threshold) now classifies as memory_bound, because attn_score_value at batch=1 is genuinely bandwidth-bound per the roofline model — a stronger signal than serialization alone. Updated test_vllm_loop_runs_autoresearch to assert the corrected classification and the resulting cpu_offload_gb/ preemption_mode candidates instead of the old idle_stall/zero-candidate expectation. Adds unit tests for _roofline_memory_fraction (None with no data, correct time-weighted fraction) and for classify_bottleneck catching intrinsic memory-boundedness that the memcpy-only path misses. --- gitm/agents/autoresearch.py | 60 ++++++++++++++++++++++++++------- gitm/optimizer/monitor.py | 8 +++-- gitm/scheduler/loop.py | 2 +- tests/test_autoresearch.py | 42 +++++++++++++++++++++++ tests/test_run_loop_workload.py | 17 ++++++---- 5 files changed, 107 insertions(+), 22 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 6e011b9..e27cf97 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -23,8 +23,10 @@ docs.vllm.ai); their expected deltas, however, are unproven estimates. The ``source`` field says so, and only the measured A/B keeps or discards them. -v0 classifies the bottleneck from coarse trace telemetry (:func:`classify_bottleneck`) -and repoints the search at the largest-residual op. Candidates come from one of +v0 classifies the bottleneck from trace telemetry (:func:`classify_bottleneck`) — +weighted by the roofline model's per-op compute/memory prediction when residuals +are available, not just wall-clock heuristics — and repoints the search at the +largest-residual op. Candidates come from one of two sources behind the :class:`Proposer` seam: the static per-class table (:func:`propose`) or :class:`GenerativeProposer`, which searches a workload's knob surface (supplied by a :class:`KnobSource` — vLLM's ``EngineArgs`` by default) at @@ -103,15 +105,41 @@ _MEMCPY_THRESHOLD = 0.25 -def classify_bottleneck(trace: Trace) -> str: +def _roofline_memory_fraction(residuals: Residuals | None) -> float | None: + """Fraction of matched kernel time the roofline model predicts is memory-bound. + + ``None`` when there's nothing to compute from (no residuals passed, or no + observed kernel matched a predicted op) — ``classify_bottleneck`` then falls + back to the memcpy-only signal. Unlike memcpy share, this catches the actual + decode bottleneck: low arithmetic intensity *inside* an attention/GEMM kernel + (roofline's ``t_memory > t_compute``), not just standalone copy ops — decode + is rarely dominated by explicit memcpys, so the memcpy signal alone misses it. + """ + if residuals is None or not residuals.per_kernel: + return None + total = mem = 0.0 + for kr in residuals.per_kernel: + t = kr.t_obs_s or 0.0 + total += t + if kr.bound == "memory": + mem += t + return mem / total if total > 0 else None + + +def classify_bottleneck(trace: Trace, residuals: Residuals | None = None) -> str: """Map a captured trace to one of ``idle_stall`` / ``memory_bound`` / ``compute_bound``. - v0 heuristic on two signals read straight from the trace: the - serialized-concurrency fraction (poor kernel overlap ⇒ idle/scheduling gaps) - and the memcpy fraction (data movement dominating ⇒ memory bound). Each is - scored against its threshold; the stronger signal wins, and if neither - crosses its threshold the workload is treated as compute bound. An empty - trace has no stall/movement signal, so it defaults to ``compute_bound``. + Two signals, scored against a threshold each; the stronger wins, and neither + crossing its threshold defaults to compute bound (an empty trace too, having + no stall/movement signal at all): + + - serialized-concurrency fraction: poor kernel overlap ⇒ idle/scheduling gaps. + - memory pressure: the memcpy share of GPU-op time (data movement dominating), + widened by the roofline-predicted memory-bound fraction of matched kernel + time when ``residuals`` (from :func:`gitm.optimizer.monitor.residuals`) is + passed — real arithmetic intensity per op, not just standalone copies. + Whichever of the two reads higher memory pressure wins; without + ``residuals`` this is exactly the old memcpy-only heuristic. """ kernels = trace.kernels() if not kernels: @@ -126,6 +154,9 @@ def classify_bottleneck(trace: Trace) -> str: sc_score = sc / _SC_THRESHOLD mem_score = memcpy_frac / _MEMCPY_THRESHOLD + roofline_frac = _roofline_memory_fraction(residuals) + if roofline_frac is not None: + mem_score = max(mem_score, roofline_frac / _MEMCPY_THRESHOLD) if max(sc_score, mem_score) < 1.0: return COMPUTE_BOUND return IDLE_STALL if sc_score >= mem_score else MEMORY_BOUND @@ -1007,11 +1038,14 @@ def autoresearch( guard). See :func:`autoresearch_v0`. When ``residuals`` (from :func:`gitm.optimizer.monitor.residuals`) are passed, - the search is repointed at the largest-residual op — the biggest gap vs the - predicted ceiling — rather than the whole trace. The loop already computes - these in its attribution phase, so it passes them straight through. + two things sharpen: the bottleneck class itself weighs the roofline-predicted + memory-bound fraction of matched kernel time (not just the memcpy-only + signal), and the search is repointed at the largest-residual op — the + biggest gap vs the predicted ceiling — rather than the whole trace. The loop + already computes these in its attribution phase, so it passes them straight + through. """ - bottleneck_class = classify_bottleneck(trace) + bottleneck_class = classify_bottleneck(trace, residuals) target = largest_residual(residuals) if residuals is not None else None return AutoresearchRun( bottleneck_class=bottleneck_class, diff --git a/gitm/optimizer/monitor.py b/gitm/optimizer/monitor.py index ef136e6..0a50575 100644 --- a/gitm/optimizer/monitor.py +++ b/gitm/optimizer/monitor.py @@ -28,6 +28,7 @@ class KernelResidual: r_mt: float | None # memory-traffic residual (None if bytes unavailable) t_obs_s: float | None = None t_pred_s: float | None = None + bound: str | None = None # the matched op's roofline bound: "compute" | "memory" @dataclass @@ -48,7 +49,10 @@ def residuals(trace: Trace, graph: Graph) -> Residuals: and compared to that op's roofline node — one representative node per op, since per-layer nodes share the same prediction. A kernel with no modeled op is unmodeled work and gets no residual; ``layer`` is unrecoverable from - name-based classification, so it's always ``None``. + name-based classification, so it's always ``None``. Each record also carries + the matched op's roofline ``bound`` ("compute" | "memory") so callers can + weight bottleneck classification by real arithmetic intensity instead of + guessing from wall-clock signals alone. """ obs = trace.kernels() pred = graph.nodes @@ -76,7 +80,7 @@ def residuals(trace: Trace, graph: Graph) -> Residuals: res.per_kernel.append( KernelResidual( op=pn.op, layer=None, r_kt=r_kt, r_mt=r_mt, - t_obs_s=t_obs, t_pred_s=t_pred, + t_obs_s=t_obs, t_pred_s=t_pred, bound=pn.prediction.bound, ) ) diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 65c6950..82163ac 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -590,7 +590,7 @@ def _unenactable(spec: Any) -> str | None: reject=_unenactable, ) else: - ar_run = AutoresearchRun(bottleneck_class=classify_bottleneck(trace), results=[]) + ar_run = AutoresearchRun(bottleneck_class=classify_bottleneck(trace, res), results=[]) ar_causal_evidence = ", ".join( f"{h.cause_op}→{h.effect_op} (p={h.p_value:.2g})" for h in hypotheses.top(2) diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 1a1abd6..4653c56 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -110,6 +110,48 @@ def test_classify_empty_trace_defaults_to_compute() -> None: assert classify_bottleneck(make_trace(events=[])) == "compute_bound" +# --- classify_bottleneck: roofline-weighted memory signal --------------------- + + +def test_roofline_memory_fraction_none_without_data() -> None: + from gitm.agents.autoresearch import _roofline_memory_fraction + + assert _roofline_memory_fraction(None) is None + assert _roofline_memory_fraction(Residuals()) is None + + +def test_roofline_memory_fraction_weights_by_observed_time() -> None: + from gitm.agents.autoresearch import _roofline_memory_fraction + + res = Residuals(per_kernel=[ + KernelResidual(op="a", layer=None, r_kt=0.0, r_mt=None, t_obs_s=3.0, bound="memory"), + KernelResidual(op="b", layer=None, r_kt=0.0, r_mt=None, t_obs_s=1.0, bound="compute"), + ]) + assert _roofline_memory_fraction(res) == 0.75 + + +def test_classify_catches_intrinsic_memory_boundedness_with_no_memcpy() -> None: + """Decode's real memory-boundedness is intra-kernel (low bytes/FLOP inside + the attention/GEMM kernel itself), not standalone memcpy ops — the memcpy- + only heuristic structurally can't see it. Two overlapping kernels with no + memcpy classify as compute_bound on the trace alone (see + test_classify_compute_bound_when_overlapped_and_no_memcpy); passing + residuals whose roofline prediction says those kernels are memory-bound + must flip the verdict.""" + events = [ + make_kernel("a", start_ns=0, end_ns=100, stream_id=0), + make_kernel("b", start_ns=50, end_ns=150, stream_id=1), + ] + trace = make_trace(events=events) + assert classify_bottleneck(trace) == "compute_bound" # unchanged, no residuals + + res = _residuals([("a", 0.0), ("b", 0.0)]) + for kr in res.per_kernel: + kr.bound = "memory" + kr.t_obs_s = 1.0 + assert classify_bottleneck(trace, res) == "memory_bound" + + # --- autoresearch_v0: apply / rollback --------------------------------------- diff --git a/tests/test_run_loop_workload.py b/tests/test_run_loop_workload.py index 78b3576..43cff96 100644 --- a/tests/test_run_loop_workload.py +++ b/tests/test_run_loop_workload.py @@ -133,8 +133,11 @@ def test_vllm_loop_runs_autoresearch(tmp_path: Path, monkeypatch): """The vllm path runs agentic autoresearch: it classifies the bottleneck, then *generates* non-catalog levers from the real EngineArgs surface and searches a value grid per knob (EngineArgsProposer), surfacing them in the summary + - report. The serialized same-stream kernels above classify as idle_stall; the - frozen fallback catalog yields three idle knobs across six value-grid points.""" + report. The serialized same-stream "paged_attention" kernels above are also + roofline-predicted memory-bound (attn_score_value at batch=1 is bandwidth- + bound) — a stronger signal than the serialization alone, so classification is + memory_bound; the frozen fallback catalog's two memory knobs (cpu_offload_gb, + preemption_mode) surface as candidates.""" import gitm.scheduler.loop as loop monkeypatch.setattr(loop, "capture", _fake_capture_with_kernels("paged_attention")) @@ -148,12 +151,14 @@ def test_vllm_loop_runs_autoresearch(tmp_path: Path, monkeypatch): workload="vllm-decode", budget="30s", scratch=str(tmp_path), workload_runner=lambda: {} ) s = result["summary"] - assert s["bottleneck_class"] == "idle_stall" - # Generated idle partial-prefill/DBO/KV-sharing knobs are filtered as noisy; - # catalog idle levers still run in the main candidate pass. - assert s["n_autoresearch"] == 0 + assert s["bottleneck_class"] == "memory_bound" + assert s["n_autoresearch"] == 3 # cpu_offload_gb x2 value-grid points + preemption_mode ar_json = (Path(result["run_dir"]) / "autoresearch.json").read_text(encoding="utf-8") + assert "cpu_offload_gb=" in ar_json + assert "preemption_mode=" in ar_json + # Denylisted knobs with unchecked prerequisites must never surface, regardless + # of which bottleneck class the run classifies as. assert "max_num_partial_prefills=" not in ar_json assert "long_prefill_token_threshold=" not in ar_json From da4c59443b3db8bb3f07a72ebd42657432048d86 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:46:16 -0700 Subject: [PATCH 11/24] autoresearch: replace the partial-prefill/dbo denylist with a live prerequisite check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit max_num_partial_prefills, max_long_partial_prefills, long_prefill_token_ threshold, and dbo_prefill_token_threshold were permanently excluded from the generated-knob surface after they caused live-apply failures ("Concurrent Partial Prefill is not supported..."). But they're real, sometimes-useful vLLM knobs gated by a live toggle, not permanently broken: the partial-prefill family requires enable_chunked_prefill, and dbo_prefill_token_threshold requires enable_dbo (vLLM docs). A blanket denylist forfeits them on every deployment, including ones where the prerequisite genuinely holds. Add KNOB_PREREQUISITES + unmet_prerequisite(engine, knob) to vllm_knobs.py: checks the live engine's actual flag (reusing get_knob) instead of assuming. No live engine -> conservative reject (can't verify), matching this module's other unknown-defaults-unsafe checks. Wire it into loop.py's existing reject-hook veto alongside the structural-knob-needs-restart check. kv_sharing_fast_prefill stays denylisted outright — it's WIP/no-op per vLLM docs regardless of config, not a prerequisite gap. Adds unmet_prerequisite unit tests and an autoresearch_v0 end-to-end test showing the same candidate rejected with enable_chunked_prefill off and allowed through with it on. --- gitm/agents/autoresearch.py | 10 ++---- gitm/optimizer/vllm_knobs.py | 38 ++++++++++++++++++++ gitm/scheduler/loop.py | 4 +-- tests/test_autoresearch.py | 52 +++++++++++++++++++++++++--- tests/test_vllm_knobs_and_restart.py | 50 +++++++++++++++++++++++++- 5 files changed, 140 insertions(+), 14 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index e27cf97..20ac2ba 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -488,13 +488,9 @@ def _field_kind_and_choices(annotation: object) -> tuple[str, tuple]: "trust_remote_code", "config_format", "download", - # Valid EngineArgs but poor standalone generated candidates: they require - # extra feature gates, are currently no-op/WIP, or are rejected by common vLLM - # runtime paths. Keep curated, reviewed catalog entries for these instead. - "partial_prefill", - "partial_prefills", - "long_prefill_token_threshold", - "dbo", + # WIP/no-op per vLLM docs ("no prefill optimization takes place with this + # flag enabled currently") — not a prerequisite gap, it just doesn't do + # anything yet. Nothing to check live; always exclude. "kv_sharing", ) diff --git a/gitm/optimizer/vllm_knobs.py b/gitm/optimizer/vllm_knobs.py index c5a887d..abdae4d 100644 --- a/gitm/optimizer/vllm_knobs.py +++ b/gitm/optimizer/vllm_knobs.py @@ -97,6 +97,8 @@ def is_env(self) -> bool: "VLLM_ATTENTION_BACKEND": KnobSpec( "VLLM_ATTENTION_BACKEND", "env:VLLM_ATTENTION_BACKEND", "structural" ), + # Prerequisite flags for the table below — not applied standalone, only read. + "enable_dbo": KnobSpec("enable_dbo", "scheduler_config.enable_dbo", "structural"), } @@ -185,3 +187,39 @@ def knob_kind(knob: str) -> KnobKind: """ spec = resolve_knob(knob) return spec.kind if spec is not None else "structural" + + +#: (name substring, prerequisite knob) — a generated candidate matching the +#: substring only does anything (or even applies) when the prerequisite is +#: already on. E.g. ``dbo_prefill_token_threshold`` requires ``--enable-dbo``; +#: ``max_num_partial_prefills``/``long_prefill_token_threshold`` require +#: ``--enable-chunked-prefill`` (vLLM docs). A blanket denylist would forfeit +#: these forever, including on deployments where the prerequisite genuinely +#: holds — check the live engine instead via :func:`unmet_prerequisite`. +KNOB_PREREQUISITES: tuple[tuple[str, str], ...] = ( + ("partial_prefill", "enable_chunked_prefill"), + ("long_prefill_token_threshold", "enable_chunked_prefill"), + ("dbo", "enable_dbo"), +) + + +def unmet_prerequisite(engine: Any | None, knob: str) -> str | None: + """None if ``knob`` has no prerequisite, or its prerequisite holds on ``engine``. + + Otherwise the rejection reason. With no live engine, a prerequisite-gated + knob can't be verified at all — conservative default is to reject rather + than assume it holds (mirrors this module's other unknown-defaults-unsafe + calls, e.g. :func:`knob_kind`). + """ + lname = knob.lower() + prereq = next((p for needle, p in KNOB_PREREQUISITES if needle in lname), None) + if prereq is None: + return None + if engine is None: + return f"prerequisite {prereq!r} unverifiable: no live engine" + try: + if get_knob(engine, prereq): + return None + return f"prerequisite {prereq!r} not enabled on this engine" + except AttributeError: + return f"prerequisite {prereq!r} unknown on this engine" diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 82163ac..3de7821 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -42,7 +42,7 @@ from gitm.optimizer.qualification import qualify from gitm.optimizer.report import Claim, build_provenance, write_report from gitm.optimizer.scheduler_attribution import scheduler_causes -from gitm.optimizer.vllm_knobs import knob_kind +from gitm.optimizer.vllm_knobs import knob_kind, unmet_prerequisite from gitm.planner.context import build_planner_context from gitm.planner.graph import predict_graph from gitm.safety.audit import AuditLog, _write_report @@ -578,7 +578,7 @@ def _unenactable(spec: Any) -> str | None: and knob_kind(spec.knob) == "structural" ): return "structural knob: needs engine restart, no restart_fn" - return None + return unmet_prerequisite(cfg.engine, spec.knob) ar_run = autoresearch( trace, diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 4653c56..d96a645 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -484,6 +484,46 @@ def test_apply_failure_surfaces_the_error_not_just_a_bare_none() -> None: assert all(r.apply_error and "distributed clash" in r.apply_error for r in results) +def test_unmet_prerequisite_vetoes_partial_prefill_without_chunked_prefill() -> None: + """End-to-end: max_num_partial_prefills is no longer denylisted (see + _is_tunable), but unmet_prerequisite wired as the reject hook still stops + it from being tried on an engine where enable_chunked_prefill is off — + the exact "Concurrent Partial Prefill is not supported" failure this was + built to prevent, without permanently forbidding the knob everywhere.""" + from gitm.optimizer.vllm_knobs import unmet_prerequisite + + class _Sched: + def __init__(self, enabled: bool): + self.chunked_prefill_enabled = enabled + + class _Engine: + def __init__(self, enabled: bool): + self.scheduler_config = _Sched(enabled) + + proposer = EngineArgsProposer( + knobs=[Knob("max_num_partial_prefills", "int", default=1)], catalog_knobs=set() + ) + + off = _Engine(enabled=False) + rejected = autoresearch_v0( + _trace(), "idle_stall", + applicator=DictApplicator({}, measure_fn=lambda s: 0.10), + proposer=proposer, + reject=lambda spec: unmet_prerequisite(off, spec.knob), + ) + assert rejected and all(not r.applicable for r in rejected) + assert all("enable_chunked_prefill" in (r.rejected_reason or "") for r in rejected) + + on = _Engine(enabled=True) + kept = autoresearch_v0( + _trace(), "idle_stall", + applicator=DictApplicator({}, measure_fn=lambda s: 0.10), + proposer=proposer, + reject=lambda spec: unmet_prerequisite(on, spec.knob), + ) + assert kept and all(r.applicable for r in kept) + + def test_apply_error_is_none_when_not_applicable_or_measured() -> None: # A candidate vetoed by the caller's reject hook never reaches apply — no # apply_error to report (the same shape as a plain gate rejection). @@ -596,11 +636,15 @@ def test_is_tunable_excludes_non_perf_engine_args() -> None: for name in ("max_num_seqs", "gpu_memory_utilization", "compilation_config", "cpu_offload_gb"): assert _is_tunable(name), f"{name} should be kept" - # Valid vLLM args, but not useful/supported as standalone generated candidates. + # kv_sharing_fast_prefill is WIP/no-op per vLLM docs regardless of engine + # config — nothing to check live, so it stays denylisted outright. + assert not _is_tunable("kv_sharing_fast_prefill") + # These are real, sometimes-useful knobs gated by a live prerequisite + # (enable_chunked_prefill / enable_dbo) — not denylisted anymore; they're + # tunable here and vetoed live via unmet_prerequisite (see test_vllm_knobs). for name in ("max_num_partial_prefills", "max_long_partial_prefills", - "long_prefill_token_threshold", "dbo_prefill_token_threshold", - "kv_sharing_fast_prefill"): - assert not _is_tunable(name), f"{name} should be excluded" + "long_prefill_token_threshold", "dbo_prefill_token_threshold"): + assert _is_tunable(name), f"{name} should be kept (prerequisite-gated, not denylisted)" def test_field_kind_and_choices_extracts_literal_enum() -> None: diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index 22cf881..d83b9f6 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -16,7 +16,13 @@ apply_intervention, ) from gitm.optimizer.scheduler_attribution import scheduler_causes -from gitm.optimizer.vllm_knobs import get_knob, knob_kind, resolve_knob, set_knob +from gitm.optimizer.vllm_knobs import ( + get_knob, + knob_kind, + resolve_knob, + set_knob, + unmet_prerequisite, +) from gitm.tracer.vllm_stats import SchedulerStatsSummary @@ -40,6 +46,48 @@ def test_knob_kind_classification(): assert knob_kind("totally_unknown_knob") == "structural" # safe default +class _SchedCfgFlags: + def __init__(self, *, chunked_prefill_enabled=False, enable_dbo=False): + self.chunked_prefill_enabled = chunked_prefill_enabled + self.enable_dbo = enable_dbo + + +class _EngineWithFlags: + def __init__(self, **flags): + self.scheduler_config = _SchedCfgFlags(**flags) + + +def test_unmet_prerequisite_none_for_knobs_without_one(): + assert unmet_prerequisite(_EngineWithFlags(), "max_num_seqs") is None + assert unmet_prerequisite(None, "max_num_seqs") is None # no engine, no prerequisite either + + +def test_unmet_prerequisite_rejects_without_live_engine(): + for knob in ("max_num_partial_prefills", "long_prefill_token_threshold", + "dbo_prefill_token_threshold"): + reason = unmet_prerequisite(None, knob) + assert reason is not None and "no live engine" in reason + + +def test_unmet_prerequisite_checks_the_live_flag(): + off = _EngineWithFlags(chunked_prefill_enabled=False, enable_dbo=False) + assert unmet_prerequisite(off, "max_num_partial_prefills") is not None + assert unmet_prerequisite(off, "long_prefill_token_threshold") is not None + assert unmet_prerequisite(off, "dbo_prefill_token_threshold") is not None + + on = _EngineWithFlags(chunked_prefill_enabled=True, enable_dbo=True) + assert unmet_prerequisite(on, "max_num_partial_prefills") is None + assert unmet_prerequisite(on, "long_prefill_token_threshold") is None + assert unmet_prerequisite(on, "dbo_prefill_token_threshold") is None + + +def test_unmet_prerequisite_unknown_flag_is_conservative(): + # An engine that doesn't expose the prerequisite at all (neither taxonomy + # path nor flat attr resolves) -> reject, don't assume it's satisfied. + reason = unmet_prerequisite(object(), "dbo_prefill_token_threshold") + assert reason is not None and "unknown on this engine" in reason + + class _SchedCfg: def __init__(self): self.max_num_seqs = 32 From b278c3e584a713c9ff10c07d5ca347b61208230e Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 15:55:45 -0700 Subject: [PATCH 12/24] autoresearch: joint (multi-knob) candidates, and rank generated candidates for real MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two independent fixes to the same limitation: InterventionSpec could only ever carry one knob=value pair, so autoresearch had no way to propose two things together, and predict_delta gave every generated candidate targeting the same op an identical expected_delta_mean (a flat 0.05 constant) — the "Predicted Δ" ranking column was blind within a run. 1. Joint candidates. InterventionSpec gains ``knobs: dict[str, Any]`` (additive, default {}) plus a ``knob_values`` property that returns ``knobs`` if set, else ``{knob: value}`` — one shape every applicator reads regardless of single vs joint. LiveEngineApplicator now applies a whole knob set atomically: hot-swaps together when every knob is "scheduling", otherwise routes the full set through restart_fn in one rebuild (restart_fn's contract changes from (engine, knob, value) to (engine, {knob: value, ...}) — updated in workloads.py's vLLM restart and every fake restart_fn in tests). DictApplicator/ConfigFileApplicator generalize the same way via a duck-type-safe _knob_values() helper (so bare test doubles exposing just .knob/.value still work). This directly resolves the prerequisite gap from the previous commit: unmet_prerequisite can only veto a standalone dependent-knob proposal when its prerequisite isn't already on — it can never turn the prerequisite on itself. EngineArgsProposer now also proposes {prerequisite: True, dependent: value} as ONE joint candidate (_joint_prerequisite_candidates, gated by the same bottleneck-class affinity as every other generated candidate) via a new GenerativeProposer._extra_candidates hook — a no-op in the workload-agnostic base class, overridden in the vLLM-specific binding. 2. Real ranking numbers. _affinity_strength(knob, keywords) counts how many affinity keywords actually match a knob's name — a real, computable per-candidate signal (an explicit Knob.classes tag counts as one strong match). expected_delta_mean is now min(0.05 * strength, expected_delta_hi) instead of a flat 0.05, so e.g. cpu_offload_gb (matches "cpu" + "offload") ranks ahead of preemption_mode (matches "preempt" only) within the same run — still an unproven v0 estimate, just one that varies with a real signal instead of tying every candidate. Updates the DRY-guard test that had hardcoded delta_mean equality across table/generated candidates (generated now legitimately varies). Adds tests for _knob_values/knob_values, the joint apply/restore/restart path, _affinity_strength/_delta_mean_for, and the joint proposer end-to-end (present only when the prerequisite flag is in the knob surface and affine to the bottleneck class, absent otherwise). --- gitm/agents/autoresearch.py | 122 ++++++++++++++++++++++++++- gitm/kernels/spec.py | 18 +++- gitm/optimizer/apply.py | 88 +++++++++++-------- gitm/workloads.py | 18 ++-- tests/test_apply_rollback.py | 2 +- tests/test_autoresearch.py | 122 ++++++++++++++++++++++++++- tests/test_knobs_via_restart.py | 10 +-- tests/test_vllm_knobs_and_restart.py | 10 +-- tests/test_workload_bootstrap.py | 2 +- 9 files changed, 331 insertions(+), 61 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 20ac2ba..c09928d 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -48,6 +48,7 @@ from gitm.kernels.spec import Applicability, InterventionSpec, SafetyGate from gitm.optimizer.apply import Applicator, apply_intervention from gitm.optimizer.monitor import Residuals, _serialized_fraction +from gitm.optimizer.vllm_knobs import KNOB_PREREQUISITES from gitm.tracer.schema import Trace if TYPE_CHECKING: @@ -267,6 +268,8 @@ def _candidate_spec( bottleneck_class: str, workload: str, source: str, + knobs: dict[str, object] | None = None, + delta_mean: float = _DELTA_MEAN, ) -> InterventionSpec: """Build a candidate spec with the fields every proposer forces. @@ -274,16 +277,24 @@ def _candidate_spec( applicability, and the moderate/rollback-gated safety posture — so the table and generative proposers can't drift apart on what a *candidate* is. Only the caller-varying parts (name, summary, knob/value, source) are parameters. + + ``knobs`` makes this a *joint* candidate (>1 knob=value pair applied and + rolled back together) — ``knob``/``value`` stay a display label in that + case; see :class:`gitm.kernels.spec.InterventionSpec`. ``delta_mean`` + overrides the flat default so a caller can scale confidence by a real, + computable signal (e.g. affinity-keyword strength) instead of every + candidate carrying the identical constant. """ return InterventionSpec( name=name, summary=summary, knob=knob, value=value, # int | float | str | bool + knobs=knobs or {}, applies_to_kernels=applies_to_kernels, # Proposed, not measured: an honest, modest range. The measured A/B is # what turns this into a real number. - expected_delta_mean=_DELTA_MEAN, + expected_delta_mean=min(delta_mean, _DELTA_HI), expected_delta_lo=_DELTA_LO, expected_delta_hi=_DELTA_HI, source=source, @@ -411,6 +422,31 @@ class vocabulary); otherwise fall back to a keyword-substring match on the return any(k in lname for k in keywords) +def _affinity_strength(knob: Knob, keywords: tuple[str, ...]) -> int: + """How many affinity keywords match ``knob``'s name — a real, computable + signal for how strongly it's implicated in the bottleneck class, not a + fabricated score. An explicit ``Knob.classes`` tag is one strong match (no + keyword multiplicity to count).""" + if knob.classes: + return 1 + lname = knob.name.lower() + return sum(1 for k in keywords if k in lname) + + +#: expected_delta_mean per unit of affinity strength — every generated candidate +#: previously carried the exact same constant (_DELTA_MEAN) regardless of how +#: strongly its name actually matched the bottleneck class, so predict_delta's +#: ranking was blind within a run (identical coverage, identical mean ⇒ +#: identical predicted_delta). Scaling by keyword-match count is still an +#: unproven v0 estimate — just one that varies with a real per-candidate signal +#: instead of a flat number. +_DELTA_MEAN_PER_MATCH = _DELTA_MEAN + + +def _delta_mean_for(strength: int) -> float: + return min(_DELTA_MEAN_PER_MATCH * max(strength, 1), _DELTA_HI) + + #: Multipliers applied to a numeric knob's default to derive its search points. _GRID_MULTIPLIERS = (0.5, 2.0, 4.0) @@ -739,6 +775,7 @@ def _spec( target_op: str | None, verb: str, source: str, + keywords: tuple[str, ...] = (), ) -> InterventionSpec: aim = f" (targeting {target_op})" if target_op else "" return _candidate_spec( @@ -750,8 +787,18 @@ def _spec( bottleneck_class=bottleneck_class, workload=self._workload, source=source, + delta_mean=_delta_mean_for(_affinity_strength(knob, keywords)), ) + def _extra_candidates( + self, bottleneck_class: str, target_op: str | None, keywords: tuple[str, ...] + ) -> list[InterventionSpec]: + """Hook for a workload-specific source to add candidates beyond the + per-knob value grid (e.g. joint prerequisite+dependent pairs). No-op + here — :class:`GenerativeProposer` stays workload-agnostic; the vLLM + binding (:class:`EngineArgsProposer`) overrides this.""" + return [] + class GenerativeProposer(_ProposerBase): """Search a workload's knob surface exhaustively, gated like any spec. @@ -794,16 +841,76 @@ def propose( target_op=target_op, verb="search", source="autoresearch-v0 (generated candidate; real workload knob, unproven delta)", + keywords=keywords, ) for knob in self._searchable() if _affine(knob, bottleneck_class, keywords) for value in _value_grid(knob) ] + out += self._extra_candidates(bottleneck_class, target_op, keywords) # Bound the per-class candidate count so a large config surface can't # flood the gate; the rollback gate still ranks and proves what survives. return out if self._max is None else out[: self._max] +def _joint_prerequisite_candidates( + knobs: list[Knob], + *, + bottleneck_class: str, + workload: str, + target_op: str | None, + keywords: tuple[str, ...] = (), +) -> list[InterventionSpec]: + """Pair a prerequisite-gated knob with its prerequisite, as ONE candidate. + + :func:`gitm.optimizer.vllm_knobs.unmet_prerequisite` can only veto a + standalone dependent-knob proposal when the prerequisite isn't already on + — it can never turn the prerequisite on itself (an ``InterventionSpec`` only + carries one knob). This is the other half: propose + ``{prerequisite: True, dependent: value}`` as a single joint spec, so the + dependent knob stays reachable (atomically tested, not just vetoed) even on + a deployment where the prerequisite isn't set yet. + + Gated the same way every other generated candidate is — ``_affine`` against + this call's bottleneck class — so a joint dbo/chunked-prefill candidate + only surfaces when the run is actually searching idle_stall, not every + class. Also requires the prerequisite itself to be a real, present, boolean + knob — a source that doesn't expose it (e.g. the offline fallback catalog) + yields nothing here, same as any other candidate needing a surface it + doesn't have. + """ + by_name = {k.name: k for k in knobs} + out: list[InterventionSpec] = [] + for needle, prereq_name in KNOB_PREREQUISITES: + prereq = by_name.get(prereq_name) + if prereq is None or prereq.kind != "bool": + continue + for knob in knobs: + if knob.name == prereq_name or needle not in knob.name.lower(): + continue + if not _affine(knob, bottleneck_class, keywords): + continue + aim = f" (targeting {target_op})" if target_op else "" + for value in _value_grid(knob): + pair = {prereq_name: True, knob.name: value} + label = ",".join(f"{k}={v}" for k, v in pair.items()) + out.append( + _candidate_spec( + name=f"autoresearch:{bottleneck_class}:{label}", + summary=f"enable {prereq_name} and set {knob.name}={value} " + f"for {bottleneck_class}{aim}", + knob=label, + value=None, + knobs=pair, + applies_to_kernels=[target_op] if target_op else [], + bottleneck_class=bottleneck_class, + workload=workload, + source="autoresearch-v0 (joint candidate: prerequisite + dependent knob)", + ) + ) + return out + + class EngineArgsProposer(GenerativeProposer): """vLLM binding of :class:`GenerativeProposer`: the EngineArgs surface + vllm-decode. @@ -814,6 +921,10 @@ class EngineArgsProposer(GenerativeProposer): ``gpu_count`` forwards to :class:`VLLMKnobSource` (autodetected when unset; ignored when ``knobs`` is given) so a 1-GPU box isn't handed a multi-GPU topology candidate whose engine build can only fail. + + Also proposes joint prerequisite+dependent candidates (see + :func:`_joint_prerequisite_candidates`) — the vLLM-specific extension of + :meth:`GenerativeProposer._extra_candidates`. """ def __init__( @@ -836,6 +947,14 @@ def __init__( max_candidates=max_candidates, ) + def _extra_candidates( + self, bottleneck_class: str, target_op: str | None, keywords: tuple[str, ...] + ) -> list[InterventionSpec]: + return _joint_prerequisite_candidates( + self._searchable(), bottleneck_class=bottleneck_class, + workload=self._workload, target_op=target_op, keywords=keywords, + ) + class FallbackProposer: """Try ``primary``; use ``secondary`` only when primary yields nothing. @@ -926,6 +1045,7 @@ def propose( target_op=target_op, verb="sample", source="autoresearch-v0 (stochastic sample; real workload knob, unproven delta)", + keywords=keywords, ) ) return out diff --git a/gitm/kernels/spec.py b/gitm/kernels/spec.py index d7cab42..302d658 100644 --- a/gitm/kernels/spec.py +++ b/gitm/kernels/spec.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel, ConfigDict, Field @@ -42,8 +42,14 @@ class InterventionSpec(BaseModel): name: str summary: str - knob: str # vLLM config key, e.g. "max_num_batched_tokens" - value: int | float | str | bool | None = None # value to set the knob to on apply + knob: str # vLLM config key, e.g. "max_num_batched_tokens" — or a display + # label ("k1=v1,k2=v2") for a joint candidate; see ``knobs`` below. + value: int | float | str | bool | None = None # value to set on apply (single-knob) + # A joint candidate: >1 knob=value pair applied/rolled back together as one + # atomic unit. Empty (the default) means "single knob" — use ``knob``/ + # ``value`` instead. Additive: every existing single-knob spec is + # unaffected (``knobs`` stays ``{}``); see ``knob_values``. + knobs: dict[str, Any] = Field(default_factory=dict) applies_to_kernels: list[str] = Field(default_factory=list) # substring match expected_delta_mean: float # signed, e.g. +0.08 = 8% improvement expected_delta_lo: float @@ -52,3 +58,9 @@ class InterventionSpec(BaseModel): applicability: Applicability = Field(default_factory=Applicability) safety: SafetyGate = Field(default_factory=SafetyGate) review: str | None = None # reviewer sign-off note (None until reviewed) + + @property + def knob_values(self) -> dict[str, Any]: + """The knob=value pairs this spec wants applied — the one shape every + applicator should read, whether the spec is single-knob or joint.""" + return dict(self.knobs) if self.knobs else {self.knob: self.value} diff --git a/gitm/optimizer/apply.py b/gitm/optimizer/apply.py index a1185cf..1849c98 100644 --- a/gitm/optimizer/apply.py +++ b/gitm/optimizer/apply.py @@ -84,7 +84,7 @@ def apply_intervention( _audit(audit, "revert", spec, cause=f"apply failed, restored: {exc}") return ApplyResult(False, rolled_back=True, measured_delta=None, error=f"apply failed, restored: {exc}") - _audit(audit, "apply", spec, cause="applied live", knob=spec.knob, value=spec.value) + _audit(audit, "apply", spec, cause="applied live", knobs=_knob_values(spec)) # Step 3: measure. A crash mid-measurement also rolls back. try: @@ -119,15 +119,25 @@ def _audit( pass +def _knob_values(spec: Any) -> dict[str, Any]: + """The knob=value pairs ``spec`` wants applied — single or joint. + + Duck-type-friendly: works for a real :class:`InterventionSpec` (via its + ``knobs``/``knob``/``value`` fields) or a bare test double exposing just + ``.knob``/``.value``, which several tests use for brevity. + """ + knobs = getattr(spec, "knobs", None) + return dict(knobs) if knobs else {spec.knob: spec.value} + + # --- reference applicators --------------------------------------------------- def _set_knob(config: dict, spec: InterventionSpec) -> None: - if spec.value is None: - raise ValueError( - f"intervention {spec.name!r} has no value to set on knob {spec.knob!r}" - ) - config[spec.knob] = spec.value + values = _knob_values(spec) + if not values or any(v is None for v in values.values()): + raise ValueError(f"intervention {spec.name!r} has no value(s) to set") + config.update(values) class DryRunApplicator: @@ -248,18 +258,21 @@ class StructuralKnobRequiresRestart(RuntimeError): class LiveEngineApplicator: - """Apply a knob to a live (vLLM) engine, gated by a real decode-throughput A/B. + """Apply a knob (or a joint set — see ``InterventionSpec.knobs``) to a live + (vLLM) engine, gated by a real decode-throughput A/B. Routes by knob taxonomy (:mod:`gitm.optimizer.vllm_knobs`): * **scheduling** knobs (``max_num_seqs``, ``max_num_batched_tokens``, ``scheduling_policy``) are *hot-swapped* in place — set on the live scheduler config, effective next step, restored by setting the old value. + A joint candidate whose knobs are ALL scheduling is hot-swapped as one set. * **structural** knobs (parallelism, dtype, quantization, block size, …) can - only take effect on a fresh engine, so they are routed through - ``restart_fn(engine, knob, value) -> new_engine``: the candidate engine - replaces the live one for the A/B, and restore swaps the original back - (shutting the candidate down best-effort). With no ``restart_fn`` a + only take effect on a fresh engine, so they (or a joint set containing + any structural knob) are routed through + ``restart_fn(engine, {knob: value, ...}) -> new_engine``: the candidate + engine replaces the live one for the A/B, and restore swaps the original + back (shutting the candidate down best-effort). With no ``restart_fn`` a structural knob raises :class:`StructuralKnobRequiresRestart`, which ``apply_intervention`` turns into a clean rollback — never a silent no-op. @@ -279,7 +292,7 @@ def __init__( engine: Any, *, throughput_fn: Callable[[Any], float], - restart_fn: Callable[[Any, str, Any], Any] | None = None, + restart_fn: Callable[[Any, dict[str, Any]], Any] | None = None, baseline_restart_fn: Callable[[Any], Any] | None = None, restart_mode: str = "parallel", getter: Callable[[Any, str], Any] | None = None, @@ -335,26 +348,31 @@ def snapshot(self) -> dict[str, Any]: return {"baseline_tps": self._baseline_tps} def apply(self, spec: InterventionSpec) -> None: - if spec.value is None: - raise ValueError(f"intervention {spec.name!r} has no value for knob {spec.knob!r}") - - if not self._force_restart and knob_kind(spec.knob) == "scheduling": - # Hot-swap in place. Record the restore point only AFTER a successful - # set: if the knob can't be located the setter raises, _prev stays - # None, and restore() is a no-op (nothing changed → nothing to undo). - try: - prev = self._getter(self.engine, spec.knob) - except AttributeError: - prev = None - self._setter(self.engine, spec.knob, spec.value) - self._prev = ("hotswap", spec.knob, prev) + values = _knob_values(spec) # single knob, or a joint candidate's set + if not values or any(v is None for v in values.values()): + raise ValueError(f"intervention {spec.name!r} has no value(s) to set") + + if not self._force_restart and all(knob_kind(k) == "scheduling" for k in values): + # Hot-swap every knob in place, atomically as a set. Record each + # restore point only AFTER its successful set: if a later knob's + # setter raises, restore() only undoes what was actually changed. + applied: dict[str, Any] = {} + self._prev = ("hotswap", applied) + for knob, value in values.items(): + try: + prev = self._getter(self.engine, knob) + except AttributeError: + prev = None + self._setter(self.engine, knob, value) + applied[knob] = prev return - # Structural knob — needs a restart to take effect. + # >=1 structural knob — the set can only take effect together via a + # rebuild, so the whole dict goes through restart_fn in one call. if self._restart_fn is None: raise StructuralKnobRequiresRestart( - f"knob {spec.knob!r} is structural (needs an engine restart); " - "no restart_fn supplied, so it cannot be applied live" + f"knob(s) {', '.join(values)} are structural (need an engine " + "restart); no restart_fn supplied, so they cannot be applied live" ) old_engine = self.engine if self._restart_mode == "serial": @@ -368,20 +386,20 @@ def restore_baseline() -> Any: self._shutdown(old_engine) self._prev = ("serial_restart", restore_baseline) try: - new_engine = self._restart_fn(old_engine, spec.knob, spec.value) + new_engine = self._restart_fn(old_engine, values) except Exception: self.engine = restore_baseline() self._prev = None raise else: - new_engine = self._restart_fn(old_engine, spec.knob, spec.value) + new_engine = self._restart_fn(old_engine, values) if new_engine is None: if self._restart_mode == "serial": _, restore_baseline = self._prev self.engine = restore_baseline() self._prev = None raise StructuralKnobRequiresRestart( - f"restart_fn produced no engine for structural knob {spec.knob!r}" + f"restart_fn produced no engine for knob(s) {', '.join(values)}" ) self.engine = new_engine self._activate(new_engine) @@ -393,8 +411,9 @@ def restore(self, snapshot: dict[str, Any]) -> None: return tag = self._prev[0] if tag == "hotswap": - _, knob, old = self._prev - self._setter(self.engine, knob, old) + _, applied = self._prev + for knob, old in applied.items(): + self._setter(self.engine, knob, old) elif tag == "restart": _, old_engine = self._prev self._shutdown(self.engine) # drop the candidate engine we built @@ -472,7 +491,8 @@ def measure(self, spec: InterventionSpec) -> float | None: significant = delta > noise_band via = "restart" if (self._prev and self._prev[0] == "restart") else "hot-swap" self.last_result = EngineABResult( - knob=spec.knob, value=spec.value, baseline_tps=baseline, + knob=spec.knob, value=(getattr(spec, "knobs", None) or spec.value), + baseline_tps=baseline, candidate_tps=candidate, speedup=speedup, kept=significant, via=via, baseline_std=self._baseline_std, candidate_std=cand_std, reps=self._reps, significant=significant, diff --git a/gitm/workloads.py b/gitm/workloads.py index e239492..2d9c8a3 100644 --- a/gitm/workloads.py +++ b/gitm/workloads.py @@ -701,22 +701,24 @@ def _throughput(eng: Any) -> float: sync_device() return toks / max(time.perf_counter() - t0, 1e-9) - def _restart(_old_engine: Any, knob: str, value: Any) -> Any: - """Rebuild a fresh vLLM engine with one structural knob changed. + def _restart(_old_engine: Any, knob_values: dict[str, Any]) -> Any: + """Rebuild a fresh vLLM engine with one or more structural knobs changed. The restart-apply path for structural levers that can't be hot-swapped on a running engine — most importantly ``kv_cache_dtype=fp8`` and ``quantization``, the real throughput/memory levers for decode. Structural knob names match the vLLM ``LLM`` kwargs (``kv_cache_dtype``, ``gpu_memory_utilization``, ``swap_space``, ``block_size``, - ``quantization``, …), so the change is a single kwarg. Returns the new - engine; ``LiveEngineApplicator`` swaps it in for the A/B and tears it down - on restore. Raises on an unsupported knob/value (no fp8 support on this - SKU, missing quantized checkpoint) → the candidate is rolled back cleanly, + ``quantization``, …), so the change is a kwargs update — one entry for a + single-knob candidate, more for a joint one (e.g. a prerequisite flag + turned on together with the knob it gates). Returns the new engine; + ``LiveEngineApplicator`` swaps it in for the A/B and tears it down on + restore. Raises on an unsupported knob/value (no fp8 support on this SKU, + missing quantized checkpoint) → the candidate is rolled back cleanly, never a silent no-op. """ kwargs = dict(getattr(_old_engine, "gitm_llm_kwargs", _base_kwargs)) - kwargs[knob] = value # the structural knob (wins if it IS gpu_memory_utilization) + kwargs.update(knob_values) # Give each restarted engine a fresh distributed port so V1 init does not # collide with any prior in-process engine state. os.environ["VLLM_PORT"] = str(_free_port()) @@ -724,7 +726,7 @@ def _restart(_old_engine: Any, knob: str, value: Any) -> Any: return _build_engine(kwargs) except Exception as exc: raise RuntimeError( - f"restart candidate for {knob!r} failed to build: {exc}" + f"restart candidate for {', '.join(knob_values)} failed to build: {exc}" ) from exc def _baseline_restart(old_engine: Any) -> Any: diff --git a/tests/test_apply_rollback.py b/tests/test_apply_rollback.py index 186495a..716761b 100644 --- a/tests/test_apply_rollback.py +++ b/tests/test_apply_rollback.py @@ -102,7 +102,7 @@ def test_apply_audits_kept_change(tmp_path): events = [e.event for e in log.entries()] assert events == ["apply"] # kept -> one apply, no revert - assert log.entries()[0].detail == {"knob": "block_size", "value": 16} + assert log.entries()[0].detail == {"knobs": {"block_size": 16}} def test_apply_audits_apply_then_revert_on_regression(tmp_path): diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index d96a645..861a0a7 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -311,10 +311,13 @@ def test_off_trace_residual_op_is_not_tagged() -> None: StochasticProposer, TableProposer, VLLMKnobSource, + _affinity_strength, _argparse_domains, _candidate_spec, + _delta_mean_for, _field_kind_and_choices, _is_tunable, + _joint_prerequisite_candidates, _knobs_from_engine_args, _requires_multi_gpu, _value_grid, @@ -777,15 +780,18 @@ def test_engineargs_proposer_accepts_gpu_count_offline() -> None: def test_candidate_specs_share_the_forced_fields() -> None: - """DRY guard: table and generated candidates agree on delta band, safety tier, - and applicability — the fields _candidate_spec centralizes.""" + """DRY guard: table and generated candidates agree on safety tier, delta + band bounds, and applicability — the fields _candidate_spec centralizes. + expected_delta_mean itself is excluded: generated candidates scale it by + affinity-keyword strength (see test_generated_delta_mean_scales_with_ + affinity_strength), so it legitimately differs candidate to candidate.""" table = propose("memory_bound")[0] generated = EngineArgsProposer( knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() ).propose("memory_bound")[0] for s in (table, generated): assert s.safety.tier == "moderate" - assert (s.expected_delta_mean, s.expected_delta_lo, s.expected_delta_hi) == (0.05, 0.0, 0.15) + assert (s.expected_delta_lo, s.expected_delta_hi) == (0.0, 0.15) assert s.applicability.workloads == ["vllm-decode"] @@ -858,6 +864,116 @@ def test_candidate_spec_helper_forces_safety_and_delta() -> None: assert spec.applicability.workloads == ["some-workload"] +# --- ranking: delta_mean scales with a real, computable signal --------------- + + +def test_affinity_strength_counts_keyword_matches() -> None: + keywords = ("cache", "swap", "offload", "cpu") + assert _affinity_strength(Knob("cpu_offload_gb", "int", default=0), keywords) == 2 + assert _affinity_strength(Knob("swap_space", "int", default=0), keywords) == 1 + assert _affinity_strength(Knob("unrelated_thing", "int", default=0), keywords) == 0 + # An explicit tag is one strong match regardless of name. + tagged = Knob("weird_lever", "int", default=0, classes=("idle_stall",)) + assert _affinity_strength(tagged, keywords) == 1 + + +def test_delta_mean_for_scales_and_caps_at_expected_delta_hi() -> None: + assert _delta_mean_for(0) == _delta_mean_for(1) # zero matches floors at 1 + assert abs(_delta_mean_for(2) - 2 * _delta_mean_for(1)) < 1e-9 + assert _delta_mean_for(100) == 0.15 # capped, never exceeds expected_delta_hi + + +def test_generated_delta_mean_scales_with_affinity_strength() -> None: + """The bug this fixes: every generated candidate carried the identical + expected_delta_mean regardless of how strongly its name actually matched + the bottleneck class, so predict_delta's ranking was blind within a run. + A knob matching 2 keywords must rank ahead of one matching 1.""" + p = EngineArgsProposer( + knobs=[ + Knob("cpu_offload_gb", "int", default=0), # matches "cpu" + "offload" + Knob("preemption_mode", "enum", default="recompute", + choices=("recompute", "swap")), # matches "preempt" only + ], + catalog_knobs=set(), + ) + specs = {s.knob: s for s in p.propose("memory_bound")} + assert specs["cpu_offload_gb"].expected_delta_mean > specs["preemption_mode"].expected_delta_mean + + +# --- joint candidates: prerequisite + dependent knob together ---------------- + + +def test_joint_prerequisite_candidate_pairs_flag_and_dependent_knob() -> None: + knobs = [ + Knob("enable_dbo", "bool", default=False), + Knob("dbo_prefill_token_threshold", "int", default=512), + ] + specs = _joint_prerequisite_candidates( + knobs, bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, + keywords=_CLASS_KEYWORDS["idle_stall"], + ) + assert specs # at least one value-grid point for the dependent knob + for s in specs: + assert set(s.knobs) == {"enable_dbo", "dbo_prefill_token_threshold"} + assert s.knobs["enable_dbo"] is True + assert s.value is None # display-only label lives in .knob + + +def test_joint_prerequisite_candidate_absent_without_the_flag_knob() -> None: + # No enable_dbo in the surface (e.g. offline fallback catalog) -> nothing. + knobs = [Knob("dbo_prefill_token_threshold", "int", default=512)] + assert _joint_prerequisite_candidates( + knobs, bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, + keywords=_CLASS_KEYWORDS["idle_stall"], + ) == [] + + +def test_joint_prerequisite_candidate_respects_bottleneck_class_affinity() -> None: + knobs = [ + Knob("enable_dbo", "bool", default=False), + Knob("dbo_prefill_token_threshold", "int", default=512), + ] + # compute_bound's keywords don't match "prefill" -> no joint candidate. + assert _joint_prerequisite_candidates( + knobs, bottleneck_class="compute_bound", workload="vllm-decode", target_op=None, + keywords=_CLASS_KEYWORDS["compute_bound"], + ) == [] + + +def test_engineargs_proposer_surfaces_joint_dbo_candidate() -> None: + """End-to-end through the proposer a real EngineArgs surface would yield.""" + p = EngineArgsProposer( + knobs=[ + Knob("enable_dbo", "bool", default=False), + Knob("dbo_prefill_token_threshold", "int", default=512), + ], + catalog_knobs=set(), + ) + specs = p.propose("idle_stall") + joint = [s for s in specs if len(s.knobs) > 1] + assert joint and all(set(s.knobs) == {"enable_dbo", "dbo_prefill_token_threshold"} + for s in joint) + + +def test_intervention_spec_knob_values_property() -> None: + from gitm.kernels.spec import InterventionSpec + + single = InterventionSpec( + name="n", summary="s", knob="max_num_seqs", value=256, + expected_delta_mean=0.05, expected_delta_lo=0.0, expected_delta_hi=0.1, + source="test", + ) + assert single.knob_values == {"max_num_seqs": 256} + + joint = InterventionSpec( + name="n", summary="s", knob="a=1,b=2", value=None, + knobs={"a": 1, "b": 2}, + expected_delta_mean=0.05, expected_delta_lo=0.0, expected_delta_hi=0.1, + source="test", + ) + assert joint.knob_values == {"a": 1, "b": 2} + + # --- stochastic (entropy-guided) proposer ------------------------------------ diff --git a/tests/test_knobs_via_restart.py b/tests/test_knobs_via_restart.py index ff20ed5..f2a606e 100644 --- a/tests/test_knobs_via_restart.py +++ b/tests/test_knobs_via_restart.py @@ -15,10 +15,10 @@ def test_force_restart_rebuilds_for_a_scheduling_knob(): - built: list[tuple[str, object]] = [] + built: list[dict] = [] - def restart(_engine, knob, value): - built.append((knob, value)) + def restart(_engine, knob_values): + built.append(dict(knob_values)) return SimpleNamespace(name="candidate") app = LiveEngineApplicator( @@ -29,7 +29,7 @@ def restart(_engine, knob, value): ) app.snapshot() app.apply(_SCHED_KNOB) - assert built == [("max_num_seqs", 256)] # rebuilt, not hot-swapped + assert built == [{"max_num_seqs": 256}] # rebuilt, not hot-swapped assert app._prev[0] == "restart" @@ -40,7 +40,7 @@ def test_default_hot_swaps_a_scheduling_knob(): throughput_fn=lambda _e: 100.0, getter=lambda _e, _k: None, setter=lambda _e, k, v: sets.append((k, v)), - restart_fn=lambda _e, _k, _v: SimpleNamespace(), + restart_fn=lambda _e, _kv: SimpleNamespace(), ) app.snapshot() app.apply(_SCHED_KNOB) diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index d83b9f6..7e9a220 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -158,8 +158,8 @@ def test_hotswap_scheduling_knob_kept(): def test_restart_apply_structural_knob_kept_and_swaps_engine(): e0 = _TpsEngine(100.0) - def restart_fn(old, knob, value): - assert knob == "block_size" and value == 16 + def restart_fn(old, knob_values): + assert knob_values == {"block_size": 16} return _TpsEngine(200.0) # the "restarted" engine is faster app = LiveEngineApplicator(e0, throughput_fn=_tps_of, restart_fn=restart_fn) @@ -194,10 +194,10 @@ def shutdown(self): restored = Engine(100.0) events: list[str] = [] - def restart_fn(old, knob, value): + def restart_fn(old, knob_values): assert old is baseline assert old.shutdown_called - events.append(f"restart:{knob}={value}") + events.append("restart:" + ",".join(f"{k}={v}" for k, v in knob_values.items())) return Engine(50.0) def baseline_restart_fn(old): @@ -224,7 +224,7 @@ def test_serial_restart_rebuilds_baseline_if_candidate_build_fails(): baseline = _TpsEngine(100.0) restored = _TpsEngine(100.0) - def restart_fn(_old, _knob, _value): + def restart_fn(_old, _knob_values): raise RuntimeError("candidate OOM") app = LiveEngineApplicator( diff --git a/tests/test_workload_bootstrap.py b/tests/test_workload_bootstrap.py index 4524bb5..bb0bd38 100644 --- a/tests/test_workload_bootstrap.py +++ b/tests/test_workload_bootstrap.py @@ -154,7 +154,7 @@ def generate(self, prompts, params): assert FakeLLM.instances[0].calls == 1 assert runner.engine.gitm_throughput_fn(runner.engine) > 0 baseline = runner.engine - restarted = runner.engine.gitm_restart_fn(runner.engine, "swap_space", 2) + restarted = runner.engine.gitm_restart_fn(runner.engine, {"swap_space": 2}) assert restarted.kwargs["gpu_memory_utilization"] == 0.45 assert restarted.kwargs["swap_space"] == 2 From 3e9a6f8f00244a5356bf81de7fa215447026e2ed Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 16:08:50 -0700 Subject: [PATCH 13/24] tests: compress redundant test functions from the last two commits Same coverage, fewer/denser test functions: merged near-duplicate tests that each set up similar fixtures to check one closely-related fact (the four unmet_prerequisite scenarios, the roofline-fraction cases, affinity_strength + delta_mean_for, the three joint-candidate scenarios plus the redundant EngineArgsProposer-level repeat of the same check), and trimmed docstrings down to what a reader actually needs. No assertions dropped. Net -48 lines across the two files. --- tests/test_autoresearch.py | 100 ++++++++------------------- tests/test_vllm_knobs_and_restart.py | 38 ++++------ 2 files changed, 45 insertions(+), 93 deletions(-) diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 861a0a7..4062859 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -113,31 +113,23 @@ def test_classify_empty_trace_defaults_to_compute() -> None: # --- classify_bottleneck: roofline-weighted memory signal --------------------- -def test_roofline_memory_fraction_none_without_data() -> None: +def test_roofline_memory_fraction() -> None: from gitm.agents.autoresearch import _roofline_memory_fraction assert _roofline_memory_fraction(None) is None - assert _roofline_memory_fraction(Residuals()) is None - - -def test_roofline_memory_fraction_weights_by_observed_time() -> None: - from gitm.agents.autoresearch import _roofline_memory_fraction - + assert _roofline_memory_fraction(Residuals()) is None # no per_kernel data res = Residuals(per_kernel=[ KernelResidual(op="a", layer=None, r_kt=0.0, r_mt=None, t_obs_s=3.0, bound="memory"), KernelResidual(op="b", layer=None, r_kt=0.0, r_mt=None, t_obs_s=1.0, bound="compute"), ]) - assert _roofline_memory_fraction(res) == 0.75 + assert _roofline_memory_fraction(res) == 0.75 # time-weighted, not a head count def test_classify_catches_intrinsic_memory_boundedness_with_no_memcpy() -> None: - """Decode's real memory-boundedness is intra-kernel (low bytes/FLOP inside - the attention/GEMM kernel itself), not standalone memcpy ops — the memcpy- - only heuristic structurally can't see it. Two overlapping kernels with no - memcpy classify as compute_bound on the trace alone (see - test_classify_compute_bound_when_overlapped_and_no_memcpy); passing - residuals whose roofline prediction says those kernels are memory-bound - must flip the verdict.""" + """Real memory-boundedness is intra-kernel (low bytes/FLOP inside the + attention/GEMM kernel), not standalone memcpy ops — passing residuals + whose roofline prediction says so must flip a memcpy-blind compute_bound + verdict (see test_classify_compute_bound_when_overlapped_and_no_memcpy).""" events = [ make_kernel("a", start_ns=0, end_ns=100, stream_id=0), make_kernel("b", start_ns=50, end_ns=150, stream_id=1), @@ -488,11 +480,9 @@ def test_apply_failure_surfaces_the_error_not_just_a_bare_none() -> None: def test_unmet_prerequisite_vetoes_partial_prefill_without_chunked_prefill() -> None: - """End-to-end: max_num_partial_prefills is no longer denylisted (see - _is_tunable), but unmet_prerequisite wired as the reject hook still stops - it from being tried on an engine where enable_chunked_prefill is off — - the exact "Concurrent Partial Prefill is not supported" failure this was - built to prevent, without permanently forbidding the knob everywhere.""" + """max_num_partial_prefills is no longer denylisted (see _is_tunable), but + unmet_prerequisite as the reject hook still stops it on an engine where + enable_chunked_prefill is off — without forbidding the knob everywhere.""" from gitm.optimizer.vllm_knobs import unmet_prerequisite class _Sched: @@ -867,27 +857,23 @@ def test_candidate_spec_helper_forces_safety_and_delta() -> None: # --- ranking: delta_mean scales with a real, computable signal --------------- -def test_affinity_strength_counts_keyword_matches() -> None: +def test_affinity_strength_and_delta_mean_for() -> None: keywords = ("cache", "swap", "offload", "cpu") assert _affinity_strength(Knob("cpu_offload_gb", "int", default=0), keywords) == 2 assert _affinity_strength(Knob("swap_space", "int", default=0), keywords) == 1 assert _affinity_strength(Knob("unrelated_thing", "int", default=0), keywords) == 0 # An explicit tag is one strong match regardless of name. - tagged = Knob("weird_lever", "int", default=0, classes=("idle_stall",)) - assert _affinity_strength(tagged, keywords) == 1 - + assert _affinity_strength(Knob("x", "int", default=0, classes=("idle_stall",)), keywords) == 1 -def test_delta_mean_for_scales_and_caps_at_expected_delta_hi() -> None: assert _delta_mean_for(0) == _delta_mean_for(1) # zero matches floors at 1 assert abs(_delta_mean_for(2) - 2 * _delta_mean_for(1)) < 1e-9 assert _delta_mean_for(100) == 0.15 # capped, never exceeds expected_delta_hi def test_generated_delta_mean_scales_with_affinity_strength() -> None: - """The bug this fixes: every generated candidate carried the identical - expected_delta_mean regardless of how strongly its name actually matched - the bottleneck class, so predict_delta's ranking was blind within a run. - A knob matching 2 keywords must rank ahead of one matching 1.""" + """Every generated candidate used to carry the identical expected_delta_mean + regardless of how strongly its name matched the class, so predict_delta's + ranking was blind within a run. A 2-keyword match must rank above a 1-match.""" p = EngineArgsProposer( knobs=[ Knob("cpu_offload_gb", "int", default=0), # matches "cpu" + "offload" @@ -903,13 +889,13 @@ def test_generated_delta_mean_scales_with_affinity_strength() -> None: # --- joint candidates: prerequisite + dependent knob together ---------------- -def test_joint_prerequisite_candidate_pairs_flag_and_dependent_knob() -> None: - knobs = [ +def test_joint_prerequisite_candidates() -> None: + dbo_knobs = [ Knob("enable_dbo", "bool", default=False), Knob("dbo_prefill_token_threshold", "int", default=512), ] specs = _joint_prerequisite_candidates( - knobs, bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, + dbo_knobs, bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, keywords=_CLASS_KEYWORDS["idle_stall"], ) assert specs # at least one value-grid point for the dependent knob @@ -918,60 +904,34 @@ def test_joint_prerequisite_candidate_pairs_flag_and_dependent_knob() -> None: assert s.knobs["enable_dbo"] is True assert s.value is None # display-only label lives in .knob - -def test_joint_prerequisite_candidate_absent_without_the_flag_knob() -> None: # No enable_dbo in the surface (e.g. offline fallback catalog) -> nothing. - knobs = [Knob("dbo_prefill_token_threshold", "int", default=512)] assert _joint_prerequisite_candidates( - knobs, bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, + [Knob("dbo_prefill_token_threshold", "int", default=512)], + bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, keywords=_CLASS_KEYWORDS["idle_stall"], ) == [] - -def test_joint_prerequisite_candidate_respects_bottleneck_class_affinity() -> None: - knobs = [ - Knob("enable_dbo", "bool", default=False), - Knob("dbo_prefill_token_threshold", "int", default=512), - ] # compute_bound's keywords don't match "prefill" -> no joint candidate. assert _joint_prerequisite_candidates( - knobs, bottleneck_class="compute_bound", workload="vllm-decode", target_op=None, + dbo_knobs, bottleneck_class="compute_bound", workload="vllm-decode", target_op=None, keywords=_CLASS_KEYWORDS["compute_bound"], ) == [] - -def test_engineargs_proposer_surfaces_joint_dbo_candidate() -> None: - """End-to-end through the proposer a real EngineArgs surface would yield.""" - p = EngineArgsProposer( - knobs=[ - Knob("enable_dbo", "bool", default=False), - Knob("dbo_prefill_token_threshold", "int", default=512), - ], - catalog_knobs=set(), - ) - specs = p.propose("idle_stall") - joint = [s for s in specs if len(s.knobs) > 1] - assert joint and all(set(s.knobs) == {"enable_dbo", "dbo_prefill_token_threshold"} - for s in joint) + # End-to-end through the proposer a real EngineArgs surface would yield. + joint = [s for s in EngineArgsProposer(knobs=dbo_knobs, catalog_knobs=set()).propose("idle_stall") + if len(s.knobs) > 1] + assert joint and all(set(s.knobs) == {"enable_dbo", "dbo_prefill_token_threshold"} for s in joint) def test_intervention_spec_knob_values_property() -> None: from gitm.kernels.spec import InterventionSpec - single = InterventionSpec( - name="n", summary="s", knob="max_num_seqs", value=256, - expected_delta_mean=0.05, expected_delta_lo=0.0, expected_delta_hi=0.1, - source="test", - ) - assert single.knob_values == {"max_num_seqs": 256} + def _spec(**kw): + return InterventionSpec(name="n", summary="s", expected_delta_mean=0.05, + expected_delta_lo=0.0, expected_delta_hi=0.1, source="test", **kw) - joint = InterventionSpec( - name="n", summary="s", knob="a=1,b=2", value=None, - knobs={"a": 1, "b": 2}, - expected_delta_mean=0.05, expected_delta_lo=0.0, expected_delta_hi=0.1, - source="test", - ) - assert joint.knob_values == {"a": 1, "b": 2} + assert _spec(knob="max_num_seqs", value=256).knob_values == {"max_num_seqs": 256} + assert _spec(knob="a=1,b=2", knobs={"a": 1, "b": 2}).knob_values == {"a": 1, "b": 2} # --- stochastic (entropy-guided) proposer ------------------------------------ diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index 7e9a220..8da5cdd 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -57,35 +57,27 @@ def __init__(self, **flags): self.scheduler_config = _SchedCfgFlags(**flags) -def test_unmet_prerequisite_none_for_knobs_without_one(): - assert unmet_prerequisite(_EngineWithFlags(), "max_num_seqs") is None - assert unmet_prerequisite(None, "max_num_seqs") is None # no engine, no prerequisite either - +def test_unmet_prerequisite(): + gated = ("max_num_partial_prefills", "long_prefill_token_threshold", + "dbo_prefill_token_threshold") -def test_unmet_prerequisite_rejects_without_live_engine(): - for knob in ("max_num_partial_prefills", "long_prefill_token_threshold", - "dbo_prefill_token_threshold"): - reason = unmet_prerequisite(None, knob) - assert reason is not None and "no live engine" in reason + # No prerequisite for this knob -> None regardless of engine. + assert unmet_prerequisite(_EngineWithFlags(), "max_num_seqs") is None + assert unmet_prerequisite(None, "max_num_seqs") is None + # No live engine -> can't verify a prerequisite-gated knob -> reject. + for knob in gated: + assert "no live engine" in unmet_prerequisite(None, knob) -def test_unmet_prerequisite_checks_the_live_flag(): + # Live flag off -> reject; on -> allowed. off = _EngineWithFlags(chunked_prefill_enabled=False, enable_dbo=False) - assert unmet_prerequisite(off, "max_num_partial_prefills") is not None - assert unmet_prerequisite(off, "long_prefill_token_threshold") is not None - assert unmet_prerequisite(off, "dbo_prefill_token_threshold") is not None - on = _EngineWithFlags(chunked_prefill_enabled=True, enable_dbo=True) - assert unmet_prerequisite(on, "max_num_partial_prefills") is None - assert unmet_prerequisite(on, "long_prefill_token_threshold") is None - assert unmet_prerequisite(on, "dbo_prefill_token_threshold") is None - + for knob in gated: + assert unmet_prerequisite(off, knob) is not None + assert unmet_prerequisite(on, knob) is None -def test_unmet_prerequisite_unknown_flag_is_conservative(): - # An engine that doesn't expose the prerequisite at all (neither taxonomy - # path nor flat attr resolves) -> reject, don't assume it's satisfied. - reason = unmet_prerequisite(object(), "dbo_prefill_token_threshold") - assert reason is not None and "unknown on this engine" in reason + # Engine exposes neither the taxonomy path nor a flat attr -> conservative reject. + assert "unknown on this engine" in unmet_prerequisite(object(), "dbo_prefill_token_threshold") class _SchedCfg: From 47072d2cbea712f1715c45587674747a1bfe2def Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 16:23:40 -0700 Subject: [PATCH 14/24] replay: unify predict_delta coverage onto classify_op, fix _op_present MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit predict_delta's coverage matched each spec's applies_to_kernels against raw kernel-name substrings a human wrote per catalog entry — a second, ad-hoc vocabulary duplicating (and drifting from) the canonical per-decode-step op vocabulary classify_op()/residuals() already use. Worse, an empty applies_to_kernels meant "matches every kernel" (100% coverage) rather than "targets nothing" — so 5 catalog entries that happened to leave it blank (max_num_batched_tokens, max_num_seqs, swap_space, gpu_memory_utilization, scheduling_policy) always won ranking by default, regardless of how much real upside a precisely -tagged lever like cuda_graphs_enable or quantization_awq had. If the tagged substrings didn't match real kernel names, those levers silently scored 0 and never got picked — which is exactly what happened. _applies() now checks classify_op(kernel_name) against the spec's tags first (the same canonical vocabulary residuals() uses), falling back to substring matching for tags classify_op doesn't cover (HFT's cudf_groupby_scan, edge's pointpillars/centerpoint — their own kernel-name vocabularies). An empty tag list is now 0 coverage, not 100%. Rebuilt every vLLM entry's applies_to_kernels in library.yaml against the canonical ops (qkv_proj, attn_score_value, attn_out_proj, mlp_gate_up, mlp_down, lm_head) via a shared &all_ops anchor for levers that genuinely span the whole decode step (scheduling/topology/graph- capture), a specific op subset for kernel-targeted levers (KV-cache, attention backend, quantization), and an honest empty list for levers predict_graph()'s decode-only model doesn't represent at all (prefix caching's prefill-side win, comm-path/orchestration levers) — 0 instead of a false 100%. Verified against a synthetic decode-shaped trace: cuda_graphs_enable/quantization_awq/speculative_decode now outrank the blank-list entries that used to dominate by default. Also fixes _op_present() in autoresearch.py: it checked whether the synthetic op label (e.g. "attn_score_value") was a literal substring of a real kernel name -- which is never true, since that label doesn't come from any real kernel's name. It now checks classify_op(k.name) == op, so autoresearch's generated candidates actually get scoped to their target op instead of silently falling back to unscoped every time. Adds tests for the classify_op-first/substring-fallback/empty-means-zero behavior in predict_delta, and for the _op_present fix. --- gitm/agents/autoresearch.py | 15 ++++++--- gitm/kernels/library.yaml | 61 ++++++++++++++++++++++++---------- gitm/optimizer/replay.py | 20 +++++++++++ tests/test_runtime_on_trace.py | 38 +++++++++++++++++++++ 4 files changed, 112 insertions(+), 22 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index c09928d..cc69410 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -47,6 +47,7 @@ from gitm.kernels.library import load_library from gitm.kernels.spec import Applicability, InterventionSpec, SafetyGate from gitm.optimizer.apply import Applicator, apply_intervention +from gitm.optimizer.deviation import classify_op from gitm.optimizer.monitor import Residuals, _serialized_fraction from gitm.optimizer.vllm_knobs import KNOB_PREREQUISITES from gitm.tracer.schema import Trace @@ -196,14 +197,18 @@ def largest_residual(res: Residuals) -> ResidualTarget | None: def _op_present(trace: Trace, op: str) -> bool: - """True if ``op`` names (a substring of) a real kernel in the trace. + """True if some kernel in the trace classifies to ``op``. Guards the ``applies_to_kernels`` tagging: the residual op label comes from - the predicted graph, so only tag a proposal with it when it actually matches - a captured kernel name — otherwise ``predict_delta`` coverage would be 0 and - the proposal would rank as worthless rather than untargeted. + the predicted graph (a synthetic name like ``attn_score_value``, never a + literal substring of a real kernel name — checking containment directly + against ``k.name`` would essentially never match), so classify by identity + the same way ``residuals()`` does, and only tag a proposal with the op when + it actually matches a captured kernel — otherwise ``predict_delta`` + coverage would be 0 and the proposal would rank as worthless rather than + untargeted. """ - return any(op in k.name for k in trace.kernels()) + return any(classify_op(k.name) == op for k in trace.kernels()) # --- candidate table -------------------------------------------------------- diff --git a/gitm/kernels/library.yaml b/gitm/kernels/library.yaml index ddd9a31..5885918 100644 --- a/gitm/kernels/library.yaml +++ b/gitm/kernels/library.yaml @@ -5,6 +5,25 @@ # library review. Expected deltas are estimates from the cited source, not yet # measured on our own workloads. The replay engine and the rollback-gated live # apply produce the measured numbers. +# +# applies_to_kernels scopes each vLLM entry to the canonical per-decode-step op +# vocabulary predict_graph()/classify_op() share (qkv_proj, attn_score_value, +# attn_out_proj, mlp_gate_up, mlp_down, lm_head) — NOT raw kernel-name +# substrings, which drift across attention backends/GPU vendors/torch.compile. +# A knob whose effect genuinely spans the whole decode step (scheduling/ +# admission/topology/graph-capture levers) is tagged with every op via the +# &all_ops anchor below; an empty list means "no modeled coverage" (0, not +# 100%) — e.g. a lever whose benefit is prefill-side or orchestration-level, +# which the decode-only predicted graph doesn't represent. Non-vLLM entries +# (edge/HFT) keep their own kernel-name vocabulary; classify_op doesn't cover +# them, so predict_delta falls back to substring matching for those. +_decode_step_ops: &all_ops + - qkv_proj + - attn_score_value + - attn_out_proj + - mlp_gate_up + - mlp_down + - lm_head interventions: # --- Tier 1: low-effort levers -------------------------------------------- @@ -12,7 +31,7 @@ interventions: summary: Set PagedAttention block size to 16 for short-context decode. knob: block_size value: 16 - applies_to_kernels: [paged_attention, attn_score_value] + applies_to_kernels: [attn_score_value] # KV-cache layout; reads live here expected_delta_mean: 0.05 expected_delta_lo: 0.02 expected_delta_hi: 0.10 @@ -31,7 +50,7 @@ interventions: summary: Store the KV cache in fp8 to roughly double cache capacity. knob: kv_cache_dtype value: fp8 - applies_to_kernels: [paged_attention, kv_cache_write, kv_cache_read] + applies_to_kernels: [attn_score_value] # KV-cache read/write traffic expected_delta_mean: 0.06 expected_delta_lo: 0.00 expected_delta_hi: 0.12 @@ -50,7 +69,7 @@ interventions: summary: Raise the per-step token budget to 8192 to fill decode batches. knob: max_num_batched_tokens value: 8192 - applies_to_kernels: [] + applies_to_kernels: *all_ops # batch-shape lever: touches every op's batch dim expected_delta_mean: 0.05 expected_delta_lo: 0.01 expected_delta_hi: 0.10 @@ -68,7 +87,7 @@ interventions: summary: Allow up to 256 concurrent sequences to raise decode concurrency. knob: max_num_seqs value: 256 - applies_to_kernels: [] + applies_to_kernels: *all_ops # concurrency lever: touches every op's batch dim expected_delta_mean: 0.04 expected_delta_lo: 0.00 expected_delta_hi: 0.09 @@ -86,7 +105,7 @@ interventions: summary: Give the block manager 4 GiB CPU swap to avoid preemption stalls. knob: swap_space value: 4 - applies_to_kernels: [] + applies_to_kernels: [attn_score_value] # preemption/swap is the KV-cache path expected_delta_mean: 0.03 expected_delta_lo: -0.01 expected_delta_hi: 0.07 @@ -103,7 +122,7 @@ interventions: summary: Raise gpu_memory_utilization to 0.92 to permit larger batch slots. knob: gpu_memory_utilization value: 0.92 - applies_to_kernels: [] + applies_to_kernels: *all_ops # raises the batch-slot ceiling for every op expected_delta_mean: 0.04 expected_delta_lo: 0.01 expected_delta_hi: 0.08 @@ -122,7 +141,7 @@ interventions: summary: Chunk prefill so it overlaps with decode instead of stalling it. knob: enable_chunked_prefill value: true - applies_to_kernels: [prefill, paged_attention] + applies_to_kernels: *all_ops # changes the prefill/decode mix every step runs expected_delta_mean: 0.08 expected_delta_lo: 0.02 expected_delta_hi: 0.15 @@ -139,7 +158,10 @@ interventions: summary: Cache shared prompt prefixes to skip repeated prefill work. knob: enable_prefix_caching value: true - applies_to_kernels: [prefill, kv_cache_write] + # No modeled coverage: the win is prefill-side (skipped prefill work), and + # predict_graph() only models one decode step — honestly 0 here, not a + # false 100%. The measured A/B is what actually judges this one. + applies_to_kernels: [] expected_delta_mean: 0.10 expected_delta_lo: 0.00 expected_delta_hi: 0.30 @@ -157,7 +179,7 @@ interventions: summary: Enable 5-token n-gram speculative decoding to cut decode steps. knob: num_speculative_tokens value: 5 - applies_to_kernels: [decode, paged_attention] + applies_to_kernels: *all_ops # changes the whole per-step decode cadence expected_delta_mean: 0.12 expected_delta_lo: -0.05 expected_delta_hi: 0.30 @@ -176,7 +198,7 @@ interventions: summary: Switch the scheduler to priority policy for mixed-SLO traffic. knob: scheduling_policy value: priority - applies_to_kernels: [] + applies_to_kernels: *all_ops # admission order changes every step's batch expected_delta_mean: 0.03 expected_delta_lo: -0.02 expected_delta_hi: 0.08 @@ -194,7 +216,7 @@ interventions: summary: Swap the attention backend to FlashInfer for fused decode attention. knob: VLLM_ATTENTION_BACKEND value: FLASHINFER - applies_to_kernels: [paged_attention, attn_score_value] + applies_to_kernels: [attn_score_value] # attention kernel swap only expected_delta_mean: 0.07 expected_delta_lo: -0.03 expected_delta_hi: 0.15 @@ -213,7 +235,8 @@ interventions: summary: Fall back to NCCL all-reduce when custom all-reduce stalls comm overlap. knob: disable_custom_all_reduce value: true - applies_to_kernels: [all_reduce, nccl] + # Comm-path change, not one of the modeled per-op decode kernels — honest 0. + applies_to_kernels: [] expected_delta_mean: 0.02 expected_delta_lo: -0.04 expected_delta_hi: 0.08 @@ -232,7 +255,7 @@ interventions: summary: Reshape to TP=2 to fit larger batches / longer context per replica. knob: tensor_parallel_size value: 2 - applies_to_kernels: [all_reduce, linear, paged_attention] + applies_to_kernels: *all_ops # every layer's op gets sharded expected_delta_mean: 0.00 expected_delta_lo: -0.15 expected_delta_hi: 0.20 @@ -252,7 +275,9 @@ interventions: summary: Load AWQ 4-bit weights to shrink the model and free KV-cache memory. knob: quantization value: awq - applies_to_kernels: [linear, gemm] + # Weight dtype/size shrinks — the weight-heavy projection ops, not + # attn_score_value (its cost is KV-cache bytes, not weight bytes). + applies_to_kernels: [qkv_proj, attn_out_proj, mlp_gate_up, mlp_down, lm_head] expected_delta_mean: 0.10 expected_delta_lo: -0.05 expected_delta_hi: 0.25 @@ -271,7 +296,7 @@ interventions: summary: Disable enforce_eager so CUDA graphs capture and replay the decode step. knob: enforce_eager value: false - applies_to_kernels: [decode, paged_attention, linear] + applies_to_kernels: *all_ops # graph capture wraps the whole decode step expected_delta_mean: 0.09 expected_delta_lo: 0.02 expected_delta_hi: 0.18 @@ -289,7 +314,7 @@ interventions: summary: Extend CUDA-graph capture to 8192 tokens so long contexts stay graphed. knob: max_seq_len_to_capture value: 8192 - applies_to_kernels: [decode, paged_attention] + applies_to_kernels: *all_ops # extends graph-capture scope for the whole step expected_delta_mean: 0.04 expected_delta_lo: -0.02 expected_delta_hi: 0.10 @@ -308,7 +333,7 @@ interventions: summary: Split layers across 2 pipeline stages to fit a model that won't shard on TP alone. knob: pipeline_parallel_size value: 2 - applies_to_kernels: [send_recv, linear] + applies_to_kernels: *all_ops # every layer's op is split across stages expected_delta_mean: -0.02 expected_delta_lo: -0.20 expected_delta_hi: 0.15 @@ -328,6 +353,8 @@ interventions: summary: Use the Ray executor backend for multi-node distributed decode. knob: distributed_executor_backend value: ray + # Multi-node orchestration overhead, not represented in the decode roofline + # model at all — honest 0. applies_to_kernels: [] expected_delta_mean: 0.00 expected_delta_lo: -0.10 diff --git a/gitm/optimizer/replay.py b/gitm/optimizer/replay.py index b64e4bc..99201fa 100644 --- a/gitm/optimizer/replay.py +++ b/gitm/optimizer/replay.py @@ -15,6 +15,7 @@ import yaml from gitm.kernels.spec import InterventionSpec +from gitm.optimizer.deviation import classify_op from gitm.tracer.schema import Trace @@ -35,7 +36,26 @@ def predict_delta(trace: Trace, spec: InterventionSpec) -> float: def _applies(spec: InterventionSpec, kernel_name: str) -> bool: + """Does ``kernel_name`` fall within ``spec``'s declared scope? + + Prefers op-identity classification (:func:`gitm.optimizer.deviation. + classify_op`) — the same canonical vLLM decode-step vocabulary + ``residuals()`` uses — over raw substring matching, so a tag like + ``attn_score_value`` means the same thing here as it does in residual + attribution, and both stay in sync when the name→op mapping is updated. + Falls back to substring matching for tags classify_op doesn't cover + (other workloads' own kernel-name vocabularies, e.g. HFT's + ``cudf_groupby_scan`` or edge's ``pointpillars``). + + An empty ``applies_to_kernels`` means "targets nothing measurable" (0 + coverage) — NOT "targets everything". A blank scope used to win every + ranking by default (100% coverage) regardless of whether the lever was + real; every candidate must now state its scope to be counted. + """ if not spec.applies_to_kernels: + return False + op = classify_op(kernel_name) + if op is not None and op in spec.applies_to_kernels: return True return any(pat in kernel_name for pat in spec.applies_to_kernels) diff --git a/tests/test_runtime_on_trace.py b/tests/test_runtime_on_trace.py index 56e52b4..780f071 100644 --- a/tests/test_runtime_on_trace.py +++ b/tests/test_runtime_on_trace.py @@ -182,6 +182,44 @@ def test_attribute_dr_ranks_pairs(): assert "doubly-robust ATE" in top.notes +# --- predict_delta coverage: unified onto classify_op ------------------------ + + +def test_predict_delta_coverage_uses_classify_op(): + from gitm.kernels.spec import InterventionSpec + from gitm.optimizer.replay import predict_delta + + def _spec(tags): + return InterventionSpec(name="n", summary="s", knob="k", value=1, + applies_to_kernels=tags, expected_delta_mean=0.10, + expected_delta_lo=0.0, expected_delta_hi=0.2, source="test") + + trace = _trace([ + _kernel("flash_attn_kernel", 0, 100), # classifies to attn_score_value + _kernel("triton_rms_norm_kernel", 100, 200), # unmodeled + ]) + # Canonical-op tag matches via classify_op, not literal substring. + assert predict_delta(trace, _spec(["attn_score_value"])) == pytest.approx(0.05) + # Unmatched op tag -> 0 coverage. + assert predict_delta(trace, _spec(["mlp_down"])) == 0.0 + # No tag at all -> 0 coverage (was 1.0; a blank scope no longer wins by default). + assert predict_delta(trace, _spec([])) == 0.0 + # A tag classify_op doesn't recognize still matches via substring fallback + # (other workloads' own kernel-name vocabularies, e.g. HFT/edge). + assert predict_delta(trace, _spec(["rms_norm"])) == pytest.approx(0.05) + + +def test_op_present_uses_classify_op_not_literal_substring(): + from gitm.agents.autoresearch import _op_present + + trace = _trace([_kernel("flash_attn_kernel", 0, 100)]) + # The op label is synthetic (from the predicted graph) and never literally + # appears in a real kernel name -- checking containment against the raw + # name would (and did) always be False. + assert _op_present(trace, "attn_score_value") is True + assert _op_present(trace, "mlp_down") is False + + # --- replay validation harness ---------------------------------------------- From 605dca877d4a5dc0873c53df235804bf91be5107 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 16:27:38 -0700 Subject: [PATCH 15/24] catalog: scale max_num_batched_tokens off the live engine, not a static 8192 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit max_num_batched_tokens_8192 hardcoded one absolute value for every deployment. There's no single right number here — it depends on model size, GPU, and workload shape, which is exactly why vLLM's own auto_tune.sh sweeps this knob per-deployment rather than hardcoding a constant. A tiny model already running at max_num_batched_tokens=512 and a 70B model running at 8192 don't want the same target. Add InterventionSpec.value_multiplier: when set, resolve_relative_value (vllm_knobs.py) reads the engine's CURRENT value for the knob and scales that instead of using the spec's literal value. Wired into loop.py's Phase 3 right after load_library, before ranking. value stays as the offline/predict-only fallback (no live engine, or the current value can't be read) — additive and backward compatible, every other catalog entry is untouched (value_multiplier defaults to None). Renamed max_num_batched_tokens_8192 -> max_num_batched_tokens_2x_current with value_multiplier: 2.0; value: 8192 remains as the offline fallback. No test referenced the old name or the literal 8192 value, so this is a clean rename. Adds resolve_relative_value unit tests: scales off the live current value, leaves specs with no multiplier untouched, and falls back to the static value with no engine or an unreadable current value. --- gitm/kernels/library.yaml | 10 ++++++--- gitm/kernels/spec.py | 8 +++++++ gitm/optimizer/vllm_knobs.py | 31 ++++++++++++++++++++++++++++ gitm/scheduler/loop.py | 6 ++++-- tests/test_vllm_knobs_and_restart.py | 31 ++++++++++++++++++++++++++++ 5 files changed, 81 insertions(+), 5 deletions(-) diff --git a/gitm/kernels/library.yaml b/gitm/kernels/library.yaml index 5885918..c64992c 100644 --- a/gitm/kernels/library.yaml +++ b/gitm/kernels/library.yaml @@ -65,10 +65,14 @@ interventions: notes: fp8 KV can shift logits slightly; gate on an accuracy sentinel before keep. review: null - - name: max_num_batched_tokens_8192 - summary: Raise the per-step token budget to 8192 to fill decode batches. + - name: max_num_batched_tokens_2x_current + summary: Double the per-step token budget (relative to whatever this engine + is already running) to fill decode batches. knob: max_num_batched_tokens - value: 8192 + value: 8192 # offline/predict-only fallback when there's no live engine to scale + value_multiplier: 2.0 # vLLM's own auto_tune.sh sweeps this per-deployment, + # not one static number — the right absolute value depends on model + # size/GPU/workload shape, so scale the engine's current setting instead. applies_to_kernels: *all_ops # batch-shape lever: touches every op's batch dim expected_delta_mean: 0.05 expected_delta_lo: 0.01 diff --git a/gitm/kernels/spec.py b/gitm/kernels/spec.py index 302d658..e5efce5 100644 --- a/gitm/kernels/spec.py +++ b/gitm/kernels/spec.py @@ -45,6 +45,14 @@ class InterventionSpec(BaseModel): knob: str # vLLM config key, e.g. "max_num_batched_tokens" — or a display # label ("k1=v1,k2=v2") for a joint candidate; see ``knobs`` below. value: int | float | str | bool | None = None # value to set on apply (single-knob) + # For a knob whose "right" absolute value depends on model size/GPU/ + # workload shape (e.g. max_num_batched_tokens): scale the engine's CURRENT + # value by this factor instead of hardcoding one number for every + # deployment. None (the default) means value is a static literal, applied + # as-is. See gitm.optimizer.vllm_knobs.resolve_relative_value — value + # stays the offline/predict-only fallback when there's no live engine to + # read a current setting from. + value_multiplier: float | None = None # A joint candidate: >1 knob=value pair applied/rolled back together as one # atomic unit. Empty (the default) means "single knob" — use ``knob``/ # ``value`` instead. Additive: every existing single-knob spec is diff --git a/gitm/optimizer/vllm_knobs.py b/gitm/optimizer/vllm_knobs.py index abdae4d..309b14f 100644 --- a/gitm/optimizer/vllm_knobs.py +++ b/gitm/optimizer/vllm_knobs.py @@ -24,6 +24,8 @@ from dataclasses import dataclass from typing import Any, Literal +from gitm.kernels.spec import InterventionSpec + KnobKind = Literal["scheduling", "structural"] # Prefixes a config object may hide behind, newest-first. ``""`` = the config is @@ -223,3 +225,32 @@ def unmet_prerequisite(engine: Any | None, knob: str) -> str | None: return f"prerequisite {prereq!r} not enabled on this engine" except AttributeError: return f"prerequisite {prereq!r} unknown on this engine" + + +def resolve_relative_value(spec: InterventionSpec, engine: Any | None) -> InterventionSpec: + """Scale a relative catalog lever's value off the engine's CURRENT setting. + + A knob like ``max_num_batched_tokens`` doesn't have one right absolute + value across every deployment — it depends on model size, GPU, and + workload shape (this is exactly why vLLM's own auto_tune.sh sweeps it + rather than hardcoding a number). A catalog entry with ``value_multiplier`` + set says "N times whatever this engine is already running at" instead of + one static guess. With a live engine, read its current value for the knob + and scale that. Falls back to the spec's static ``value`` when there's no + multiplier, no engine, or the current value can't be read/scaled + (offline/predict-only mode) — still a usable, honest estimate. + """ + if spec.value_multiplier is None or engine is None: + return spec + try: + current = get_knob(engine, spec.knob) + except AttributeError: + return spec + if not isinstance(current, int | float) or isinstance(current, bool) or current <= 0: + return spec + scaled = current * spec.value_multiplier + new_value = int(round(scaled)) if isinstance(current, int) else scaled + return spec.model_copy(update={ + "value": new_value, + "summary": f"{spec.summary} (scaled {spec.value_multiplier:g}x current {current} -> {new_value})", + }) diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 3de7821..69e77d2 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -42,7 +42,7 @@ from gitm.optimizer.qualification import qualify from gitm.optimizer.report import Claim, build_provenance, write_report from gitm.optimizer.scheduler_attribution import scheduler_causes -from gitm.optimizer.vllm_knobs import knob_kind, unmet_prerequisite +from gitm.optimizer.vllm_knobs import knob_kind, resolve_relative_value, unmet_prerequisite from gitm.planner.context import build_planner_context from gitm.planner.graph import predict_graph from gitm.safety.audit import AuditLog, _write_report @@ -466,7 +466,9 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # Phase 3 — library + counterfactual replay ranking pctx = build_planner_context(cfg.engine, workload = workload) - library = load_library(workload = workload) + # Relative levers (value_multiplier set) resolve against the live engine's + # current setting here, once, before ranking — see resolve_relative_value. + library = [resolve_relative_value(s, cfg.engine) for s in load_library(workload=workload)] policy = Policy(require_qualification_commit=qual.commit, skip_high_risk=not qual.commit) ranked = select_interventions(trace, library, policy, top_n=cfg.top_n_interventions, ctx=pctx.gate) (run_dir / "ranked_candidates.json").write_text( diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index 8da5cdd..7e92feb 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -20,6 +20,7 @@ get_knob, knob_kind, resolve_knob, + resolve_relative_value, set_knob, unmet_prerequisite, ) @@ -80,6 +81,36 @@ def test_unmet_prerequisite(): assert "unknown on this engine" in unmet_prerequisite(object(), "dbo_prefill_token_threshold") +def test_resolve_relative_value(): + relative = _spec("max_num_batched_tokens", 8192) + relative.value_multiplier = 2.0 + absolute = _spec("max_num_seqs", 256) # no multiplier -> untouched + + class _Sched: + def __init__(self, tokens): + self.max_num_batched_tokens = tokens + + class _Engine: + def __init__(self, tokens): + self.scheduler_config = _Sched(tokens) + + # No engine, or no multiplier -> the static YAML/literal value is unchanged. + assert resolve_relative_value(relative, None).value == 8192 + assert resolve_relative_value(absolute, _Engine(512)) is absolute + + # A live engine's current value scales instead of the hardcoded literal — + # a tiny model's 512 doubles to 1024, not always jumping to 8192. + small = resolve_relative_value(relative, _Engine(512)) + assert small.value == 1024 + assert "512 -> 1024" in small.summary + + big = resolve_relative_value(relative, _Engine(4096)) + assert big.value == 8192 + + # Can't read the current value -> falls back to the static value, not a crash. + assert resolve_relative_value(relative, object()).value == 8192 + + class _SchedCfg: def __init__(self): self.max_num_seqs = 32 From b549aa5730edf92808945c1099452375b8dd4405 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 16:35:20 -0700 Subject: [PATCH 16/24] catalog: sweep every relative lever instead of one guess, not just tokens MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the max_num_batched_tokens fix to the rest of the static-value entries with the same problem — max_num_seqs, gpu_memory_utilization, swap_space, max_seq_len_to_capture all hardcoded one absolute number regardless of model size/GPU/workload shape — and goes further: instead of a single value_multiplier guess, each entry now sweeps a small grid of multipliers (matching vLLM's own auto_tune.sh, which tries a LIST of candidate values per knob, and autoresearch's generated-knob value grid — same idea, now applied to reviewed catalog levers too). InterventionSpec gains value_multiplier_grid (list of multipliers, expands one entry into N candidates) and value_max/value_min (clamp the scaled result — needed for gpu_memory_utilization, a 0-1 fraction where the shared 0.5x/2x/4x ladder would be nonsensical: halving it wastes headroom, 4x is an invalid fraction above 1.0). expand_relative_candidates (vllm_knobs.py) resolves each grid point off the SAME current engine value via resolve_relative_value, dedupes candidates that collapse to the same value (e.g. swap_space starting at 0 — any multiplier of 0 stays 0, so the sweep collapses to one candidate at the static fallback instead of 3 duplicates), and collapses to a single predict-only candidate when there's no live engine to scale against. Wired into loop.py's Phase 3 in place of the single-resolve call. library.yaml: renamed the four entries to *_dynamic with a shared &sweep_grid anchor ([0.5, 2.0, 4.0] — the same ladder autoresearch's generated candidates already use); gpu_memory_utilization gets its own gentler grid ([0.95, 1.03, 1.10]) plus a 0.97 clamp. Updates the two integration tests that hardcoded max_num_seqs == 256: one engine starts at 32 (now legitimately ends at 128, not 256) and would have failed outright; the other happens to start at 64 (4x lands on 256 by coincidence) and would have passed for the wrong reason. Both now assert "kept and raised from the starting value" instead of pinning an exact literal that depends on where the sweep started. Adds unit tests for expand_relative_candidates (per-point scaling, dedup on collapse, offline fallback) and the value_max/value_min clamp. --- gitm/kernels/library.yaml | 57 ++++++++++++++++++--------- gitm/kernels/spec.py | 11 ++++++ gitm/optimizer/vllm_knobs.py | 39 ++++++++++++++++++ gitm/scheduler/loop.py | 17 ++++++-- tests/test_vllm_embodiment.py | 15 ++++--- tests/test_vllm_knobs_and_restart.py | 59 ++++++++++++++++++++++++++-- 6 files changed, 167 insertions(+), 31 deletions(-) diff --git a/gitm/kernels/library.yaml b/gitm/kernels/library.yaml index c64992c..9d3af01 100644 --- a/gitm/kernels/library.yaml +++ b/gitm/kernels/library.yaml @@ -25,6 +25,12 @@ _decode_step_ops: &all_ops - mlp_down - lm_head +# Shared multiplier ladder for relative levers that sweep — same ladder +# autoresearch's generated candidates already use (_GRID_MULTIPLIERS in +# gitm/agents/autoresearch.py), so there's one canonical "how wide a swing" +# convention instead of each entry inventing its own. +_relative_sweep: &sweep_grid [0.5, 2.0, 4.0] + interventions: # --- Tier 1: low-effort levers -------------------------------------------- - name: kv_cache_block_size_16 @@ -65,14 +71,14 @@ interventions: notes: fp8 KV can shift logits slightly; gate on an accuracy sentinel before keep. review: null - - name: max_num_batched_tokens_2x_current - summary: Double the per-step token budget (relative to whatever this engine - is already running) to fill decode batches. + - name: max_num_batched_tokens_dynamic + summary: Sweep the per-step token budget relative to whatever this engine + is already running (not one hardcoded number) to fill decode batches. knob: max_num_batched_tokens value: 8192 # offline/predict-only fallback when there's no live engine to scale - value_multiplier: 2.0 # vLLM's own auto_tune.sh sweeps this per-deployment, - # not one static number — the right absolute value depends on model - # size/GPU/workload shape, so scale the engine's current setting instead. + value_multiplier_grid: *sweep_grid # vLLM's own auto_tune.sh sweeps this + # per-deployment rather than hardcoding a number — the right absolute + # value depends on model size/GPU/workload shape. applies_to_kernels: *all_ops # batch-shape lever: touches every op's batch dim expected_delta_mean: 0.05 expected_delta_lo: 0.01 @@ -87,10 +93,12 @@ interventions: notes: Larger batches raise peak activation memory; watch for OOM. review: null - - name: max_num_seqs_256 - summary: Allow up to 256 concurrent sequences to raise decode concurrency. + - name: max_num_seqs_dynamic + summary: Sweep concurrent-sequence concurrency relative to whatever this + engine already allows (not one hardcoded 256) to raise decode concurrency. knob: max_num_seqs - value: 256 + value: 256 # offline/predict-only fallback when there's no live engine to scale + value_multiplier_grid: *sweep_grid applies_to_kernels: *all_ops # concurrency lever: touches every op's batch dim expected_delta_mean: 0.04 expected_delta_lo: 0.00 @@ -105,10 +113,14 @@ interventions: notes: Interacts with gpu_memory_utilization; tune one lever at a time. review: null - - name: swap_space_4gib - summary: Give the block manager 4 GiB CPU swap to avoid preemption stalls. + - name: swap_space_dynamic + summary: Sweep CPU swap relative to whatever this engine already has + configured (not one hardcoded 4 GiB) to avoid preemption stalls. knob: swap_space - value: 4 + value: 4 # offline/predict-only fallback; also the outcome when current + # swap_space is 0 (the common default) — any multiplier of 0 stays 0, + # so the sweep collapses to this one static candidate in that case. + value_multiplier_grid: *sweep_grid applies_to_kernels: [attn_score_value] # preemption/swap is the KV-cache path expected_delta_mean: 0.03 expected_delta_lo: -0.01 @@ -122,10 +134,17 @@ interventions: notes: Swap I/O can regress latency-bound workloads; keep only on a positive delta. review: null - - name: gpu_memory_utilization_092 - summary: Raise gpu_memory_utilization to 0.92 to permit larger batch slots. + - name: gpu_memory_utilization_dynamic + summary: Sweep gpu_memory_utilization relative to whatever this engine + already runs at (not one hardcoded 0.92) to permit larger batch slots. knob: gpu_memory_utilization - value: 0.92 + value: 0.92 # offline/predict-only fallback when there's no live engine to scale + # A fraction near 1.0, not a count — the *0.5/*4 ladder used elsewhere + # would be nonsensical here (halving it wastes memory headroom; 4x is an + # invalid fraction). A gentler ladder, clamped so it can never approach + # 1.0 and starve the engine of headroom entirely. + value_multiplier_grid: [0.95, 1.03, 1.10] + value_max: 0.97 applies_to_kernels: *all_ops # raises the batch-slot ceiling for every op expected_delta_mean: 0.04 expected_delta_lo: 0.01 @@ -314,10 +333,12 @@ interventions: notes: Graph capture raises warmup memory; watch for OOM on first capture. review: null - - name: max_seq_len_to_capture_8192 - summary: Extend CUDA-graph capture to 8192 tokens so long contexts stay graphed. + - name: max_seq_len_to_capture_dynamic + summary: Sweep CUDA-graph capture length relative to whatever this engine + already captures (not one hardcoded 8192) so long contexts stay graphed. knob: max_seq_len_to_capture - value: 8192 + value: 8192 # offline/predict-only fallback when there's no live engine to scale + value_multiplier_grid: *sweep_grid applies_to_kernels: *all_ops # extends graph-capture scope for the whole step expected_delta_mean: 0.04 expected_delta_lo: -0.02 diff --git a/gitm/kernels/spec.py b/gitm/kernels/spec.py index e5efce5..de7b357 100644 --- a/gitm/kernels/spec.py +++ b/gitm/kernels/spec.py @@ -53,6 +53,17 @@ class InterventionSpec(BaseModel): # stays the offline/predict-only fallback when there's no live engine to # read a current setting from. value_multiplier: float | None = None + # A sweep of multipliers instead of one — e.g. [0.5, 2.0, 4.0] — expands + # this single reviewed entry into one candidate per point (see + # gitm.optimizer.vllm_knobs.expand_relative_candidates), so the rollback + # gate picks whichever point actually wins instead of the catalog betting + # on a single guess. Empty (the default) means no sweep; value_multiplier + # (if set) resolves as one candidate as usual. + value_multiplier_grid: list[float] = Field(default_factory=list) + # Clamp the multiplier-scaled result (e.g. gpu_memory_utilization must stay + # a fraction below 1.0). None means no clamp. + value_max: float | None = None + value_min: float | None = None # A joint candidate: >1 knob=value pair applied/rolled back together as one # atomic unit. Empty (the default) means "single knob" — use ``knob``/ # ``value`` instead. Additive: every existing single-knob spec is diff --git a/gitm/optimizer/vllm_knobs.py b/gitm/optimizer/vllm_knobs.py index 309b14f..4ef8493 100644 --- a/gitm/optimizer/vllm_knobs.py +++ b/gitm/optimizer/vllm_knobs.py @@ -249,8 +249,47 @@ def resolve_relative_value(spec: InterventionSpec, engine: Any | None) -> Interv if not isinstance(current, int | float) or isinstance(current, bool) or current <= 0: return spec scaled = current * spec.value_multiplier + if spec.value_max is not None: + scaled = min(scaled, spec.value_max) + if spec.value_min is not None: + scaled = max(scaled, spec.value_min) new_value = int(round(scaled)) if isinstance(current, int) else scaled return spec.model_copy(update={ "value": new_value, "summary": f"{spec.summary} (scaled {spec.value_multiplier:g}x current {current} -> {new_value})", }) + + +def expand_relative_candidates(spec: InterventionSpec, engine: Any | None) -> list[InterventionSpec]: + """Sweep ``value_multiplier_grid`` into one resolved candidate per point. + + Mirrors vLLM's own auto_tune.sh (sweeps a list of candidate values per + knob rather than trying one guess) and autoresearch's generated-knob value + grid — the same idea applied to a reviewed catalog lever, so the rollback + gate picks whichever point actually wins instead of the catalog betting on + a single number. Each point is resolved off the SAME current engine value + via :func:`resolve_relative_value` and gets its own name/summary. + + Falls back to a single :func:`resolve_relative_value` call when there's no + grid (a plain ``value_multiplier`` spec, or a static spec with neither) — + additive, not a new contract every entry must opt into. With no live + engine there's nothing to scale relative to, so the sweep collapses to one + candidate at the static fallback value (not N duplicates of it); the same + collapse happens when every grid point resolves to the same value (e.g. + the knob's current value is 0, so any multiplier of it stays 0). + """ + if not spec.value_multiplier_grid: + return [resolve_relative_value(spec, engine)] + if engine is None: + return [spec.model_copy(update={"value_multiplier_grid": []})] + out: list[InterventionSpec] = [] + seen_values: set[Any] = set() + for m in spec.value_multiplier_grid: + variant = spec.model_copy(update={"value_multiplier": m, "value_multiplier_grid": []}) + resolved = resolve_relative_value(variant, engine) + if resolved.value in seen_values: + continue # collapsed to an already-queued value (e.g. current == 0) + seen_values.add(resolved.value) + suffix = f"x{m:g}".replace(".", "_").replace("-", "neg") + out.append(resolved.model_copy(update={"name": f"{spec.name}_{suffix}"})) + return out diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 69e77d2..4b064e8 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -42,7 +42,11 @@ from gitm.optimizer.qualification import qualify from gitm.optimizer.report import Claim, build_provenance, write_report from gitm.optimizer.scheduler_attribution import scheduler_causes -from gitm.optimizer.vllm_knobs import knob_kind, resolve_relative_value, unmet_prerequisite +from gitm.optimizer.vllm_knobs import ( + expand_relative_candidates, + knob_kind, + unmet_prerequisite, +) from gitm.planner.context import build_planner_context from gitm.planner.graph import predict_graph from gitm.safety.audit import AuditLog, _write_report @@ -466,9 +470,14 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # Phase 3 — library + counterfactual replay ranking pctx = build_planner_context(cfg.engine, workload = workload) - # Relative levers (value_multiplier set) resolve against the live engine's - # current setting here, once, before ranking — see resolve_relative_value. - library = [resolve_relative_value(s, cfg.engine) for s in load_library(workload=workload)] + # Relative levers (value_multiplier / value_multiplier_grid) resolve — and + # sweep, if a grid is set — against the live engine's current setting here, + # once, before ranking. See expand_relative_candidates. + library = [ + resolved + for s in load_library(workload=workload) + for resolved in expand_relative_candidates(s, cfg.engine) + ] policy = Policy(require_qualification_commit=qual.commit, skip_high_risk=not qual.commit) ranked = select_interventions(trace, library, policy, top_n=cfg.top_n_interventions, ctx=pctx.gate) (run_dir / "ranked_candidates.json").write_text( diff --git a/tests/test_vllm_embodiment.py b/tests/test_vllm_embodiment.py index 4fcd897..2deeafb 100644 --- a/tests/test_vllm_embodiment.py +++ b/tests/test_vllm_embodiment.py @@ -295,9 +295,10 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): cfg = LoopConfig( engine=engine, workload="vllm-decode", - # Non-expiring budget: max_num_seqs_256 ranks low, and this test asserts it - # is reached + kept. A short wall-clock budget would make that racy under - # load (the loop is budget-bounded by design), so give it plenty of room. + # Non-expiring budget: max_num_seqs_dynamic ranks low, and this test + # asserts it is reached + kept. A short wall-clock budget would make + # that racy under load (the loop is budget-bounded by design), so + # give it plenty of room. budget="24h", scratch=str(tmp_path), top_n_interventions=50, # evaluate the whole library, not just top 5 @@ -308,11 +309,15 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): assert summary["status"] == "ok" and summary["mode"] == "intervention" # The hot-swappable winning knob is kept with a measured positive delta. - assert engine.max_num_seqs == 256 + # max_num_seqs_dynamic now sweeps relative to the engine's starting value + # (32) rather than always jumping to one hardcoded 256 — assert it was + # raised and kept, not a specific literal that depends on where the sweep + # started. + assert engine.max_num_seqs > 32 import json claims = json.loads((Path(out["run_dir"]) / "ranked_candidates.json").read_text()) - assert any(c["name"] == "max_num_seqs_256" for c in claims) + assert any(c["name"].startswith("max_num_seqs_dynamic") for c in claims) # report carries the measured live A/B verdict (not a prediction). assert "live A/B" in out["report_md"] # Knobs the engine can't hot-swap were rolled back, not silently kept. diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index 7e92feb..6a88917 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -111,6 +111,54 @@ def __init__(self, tokens): assert resolve_relative_value(relative, object()).value == 8192 +def test_resolve_relative_value_respects_clamp(): + frac = _spec("gpu_memory_utilization", 0.92) + frac.value_multiplier = 1.10 + frac.value_max = 0.97 + + class _Engine: + cache_config = type("C", (), {"gpu_memory_utilization": 0.95})() + + # 0.95 * 1.10 = 1.045, which would be an invalid fraction -> clamped. + assert resolve_relative_value(frac, _Engine()).value == 0.97 + + +def test_expand_relative_candidates(): + from gitm.optimizer.vllm_knobs import expand_relative_candidates + + swept = _spec("max_num_seqs", 256) + swept.value_multiplier_grid = [0.5, 2.0, 4.0] + plain = _spec("block_size", 16) # no multiplier/grid -> untouched + + class _Sched: + def __init__(self, seqs): + self.max_num_seqs = seqs + + class _Engine: + def __init__(self, seqs): + self.scheduler_config = _Sched(seqs) + + # No grid at all -> single-item list, spec passed through as-is. + assert expand_relative_candidates(plain, _Engine(32)) == [plain] + + # Each grid point becomes its own candidate, scaled off the SAME current + # value, with a distinct name. + out = expand_relative_candidates(swept, _Engine(64)) + assert {s.value for s in out} == {32, 128, 256} + assert len({s.name for s in out}) == 3 + assert all(s.name.startswith("max_num_seqs") for s in out) + + # No live engine -> nothing to scale relative to; one candidate at the + # static fallback, not 3 duplicates of it. + offline = expand_relative_candidates(swept, None) + assert len(offline) == 1 and offline[0].value == 256 + + # Every grid point collapses to the same value (current == 0) -> dedup to + # one candidate, not 3 identical ones. + zeroed = expand_relative_candidates(swept, _Engine(0)) + assert len(zeroed) == 1 and zeroed[0].value == 256 # static fallback + + class _SchedCfg: def __init__(self): self.max_num_seqs = 32 @@ -381,8 +429,9 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): monkeypatch.setattr(loop, "sync_device", lambda: None) engine = _FullEngine() - # Non-expiring budget: max_num_seqs_256 ranks low and this asserts it's reached - # + kept; a short wall-clock would make that racy under load (loop is budget-bounded). + # Non-expiring budget: max_num_seqs_dynamic ranks low and this asserts it's + # reached + kept; a short wall-clock would make that racy under load (loop + # is budget-bounded). out = run_loop(LoopConfig(engine=engine, workload="vllm-decode", budget="24h", scratch=str(tmp_path), top_n_interventions=50)) @@ -393,8 +442,10 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): assert "under_filled_batch" in signals # And it reached the claim that it motivates (max_num_seqs), in the report. assert "scheduler[under_filled_batch]" in out["report_md"] - # The hot-swap won and was kept. - assert engine.scheduler_config.max_num_seqs == 256 + # The hot-swap won and was kept. max_num_seqs_dynamic now sweeps relative + # to the engine's starting value (64) rather than one hardcoded 256, so + # assert it was raised and kept rather than pin an exact literal. + assert engine.scheduler_config.max_num_seqs > 64 # Scheduler summary surfaced in the run summary (synchronous first sample). assert out["summary"]["scheduler_stats"] is not None From c2f278a1d38885435c7c36d2ec049c4a6b182cca Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 16:40:15 -0700 Subject: [PATCH 17/24] deviation: classify_op recognizes real vLLM attention/KV-cache kernels The 0.0% kernel_time residual wasn't a regression from the recent classify_op/predict_delta unification -- it was the original, never-confirmed vocabulary gap. Got a real kernel-name dump from an actual vLLM/L4/CUPTI trace and checked it against _OP_RULES directly: - FlashAttention's real kernel is `flash_fwd_splitkv_kernel` (the `flash` C++ namespace) -- the existing `flash_attn`/`flashattn` needles never matched it. ~6100 launches on that trace, previously all falling through to unmodeled. - vLLM's own KV-cache management kernels (`reshape_and_cache_flash_ kernel`, `_compute_slot_mapping_kernel`) weren't in the vocabulary at all. Another ~6650 launches. Added `flash_fwd`, `reshape_and_cache`, `slot_mapping` to the attn_score_value rule. Confirmed against the exact real kernel names. Also confirmed, and worth being honest about: ~35% of launches on that trace are bare `ampere_fp16_s16816gemm_*`/`cutlass_*gemm*` kernels -- generic, shape-tuned GEMMs PyTorch/cuBLAS reuses interchangeably across qkv_proj/attn_out_proj/mlp_gate_up/mlp_down/lm_head. The kernel name carries no information about which logical layer called it, so these will keep classifying to None regardless of vocabulary tweaks -- a structural limit of name-based matching (already flagged in this module's docstring), not something this fix (or any vocabulary addition) can close. Only attn_score_value gets real signal from this change; the other five ops stay unattributed on this stack. Adds a regression test using the exact real kernel names from that trace, including the two GEMM kernels that correctly stay unmodeled. --- gitm/optimizer/deviation.py | 24 +++++++++++++++++------- tests/test_deviation_alignment.py | 22 ++++++++++++++++++++++ 2 files changed, 39 insertions(+), 7 deletions(-) diff --git a/gitm/optimizer/deviation.py b/gitm/optimizer/deviation.py index 1b44c8d..879fe31 100644 --- a/gitm/optimizer/deviation.py +++ b/gitm/optimizer/deviation.py @@ -32,14 +32,24 @@ # when it maps to no modeled op (norms, activations, copies, launch overhead) — # those are unmodeled work and kept as departures. # -# The projection GEMMs (qkv/out/gate_up/down) are only distinguishable when the -# kernel name carries the projection tag (vLLM/Triton fused kernels usually do); a -# bare cuBLAS/cutlass GEMM with no tag falls through to None (unmodeled). -# Distinguishing those needs shape-matching (a follow-up) — but this already beats -# the old positional `i % len(pred)`, which aligned nothing under CUDA graphs. +# The projection GEMMs (qkv/out/gate_up/down/lm_head) are only distinguishable +# when the kernel name carries the projection tag (vLLM/Triton fused kernels +# usually do); a bare cuBLAS/cutlass GEMM (e.g. `ampere_fp16_s16816gemm_*`, +# `cutlass_*gemm*`) carries no such tag and falls through to None (unmodeled) +# — confirmed against a real vLLM/L4/CUPTI trace, where these generic GEMM +# kernels are ~35% of all kernel launches and are reused interchangeably +# across every projection. Distinguishing those needs shape-matching or +# launch-order instrumentation (a real follow-up, not a vocabulary fix) — but +# this already beats the old positional `i % len(pred)`, which aligned +# nothing under CUDA graphs. The attention/KV-cache rules below ARE confirmed +# against that same real trace: FlashAttention's actual kernel name is +# `flash_fwd_splitkv_kernel` (the `flash_attn`/`flashattn` needles alone miss +# it), and vLLM's own KV-cache write/bookkeeping kernels +# (`reshape_and_cache_flash_kernel`, `_compute_slot_mapping_kernel`) weren't +# covered at all. _OP_RULES: tuple[tuple[tuple[str, ...], str], ...] = ( - (("flash_attn", "flashattn", "paged_attention", "paged_attn", "fmha", "attention", - "attn_score"), "attn_score_value"), + (("flash_attn", "flashattn", "flash_fwd", "paged_attention", "paged_attn", "fmha", + "attention", "attn_score", "reshape_and_cache", "slot_mapping"), "attn_score_value"), (("qkv",), "qkv_proj"), (("o_proj", "out_proj", "attn_out"), "attn_out_proj"), (("gate_up", "gate_proj", "up_proj", "swiglu", "silu_and_mul"), "mlp_gate_up"), diff --git a/tests/test_deviation_alignment.py b/tests/test_deviation_alignment.py index e468904..772edc2 100644 --- a/tests/test_deviation_alignment.py +++ b/tests/test_deviation_alignment.py @@ -37,6 +37,28 @@ def test_classify_op(): assert classify_op("triton_rms_norm") is None # not a modeled op +def test_classify_op_matches_real_vllm_kernel_names(): + """Confirmed against a real vLLM decode trace (L4, CUPTI) — these exact + mangled kernel names came back from a live run. FlashAttention's real + kernel is flash_fwd_*, NOT flash_attn_* (that needle alone misses it); + vLLM's own KV-cache write/bookkeeping kernels weren't covered at all.""" + assert classify_op( + "_ZN5flash24flash_fwd_splitkv_kernelI23Flash_fwd_kernel_traitsILi64E" + "Li64ELi256ELi4ELb0ELb0EN7cutlass6half_tE19Flash_kernel_traitsILi64E" + ) == "attn_score_value" + assert classify_op( + "_ZN4vllm30reshape_and_cache_flash_kernelIttLNS_18Fp8KVCacheDataTypeE0EEE" + ) == "attn_score_value" + assert classify_op("_compute_slot_mapping_kernel") == "attn_score_value" + # The dominant real kernel type (~35% of launches on that trace) is a bare + # cuBLAS/cutlass GEMM shared across every projection — genuinely + # unattributable by name alone, not a bug to chase with more substrings. + assert classify_op("ampere_fp16_s16816gemm_fp16_128x128_ldg8_relu_f2f_stages_32x5_tn") is None + assert classify_op( + "_ZN7cutlass7Kernel2I66cutlass_80_tensorop_f16_s16816gemm_relu_f16_256x128_32x3_tn_align8EEE" + ) is None + + def test_in_band_op_not_kept_out_of_band_and_unmodeled_kept(): g = predict_graph() t_attn = next(n.prediction.t_pred_s for n in g.nodes if n.op == "attn_score_value") From dccc639d013d772104dfb573efd130ecdeb96139 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 17:05:43 -0700 Subject: [PATCH 18/24] comments: compress verbose docstrings from this session's changes Trimmed multi-paragraph docstrings/comments added across the recent commits (residual attribution, prerequisite veto, joint candidates, coverage unification, relative catalog values) down to what a reader actually needs, in apply.py, deviation.py, monitor.py, replay.py, vllm_knobs.py, spec.py, and loop.py. No behavior change. --- gitm/kernels/spec.py | 30 ++++++------------ gitm/optimizer/apply.py | 13 +++----- gitm/optimizer/deviation.py | 30 ++++++++---------- gitm/optimizer/monitor.py | 19 +++++------- gitm/optimizer/replay.py | 18 +++-------- gitm/optimizer/vllm_knobs.py | 59 ++++++++++++------------------------ gitm/scheduler/loop.py | 5 ++- 7 files changed, 61 insertions(+), 113 deletions(-) diff --git a/gitm/kernels/spec.py b/gitm/kernels/spec.py index de7b357..cfbd739 100644 --- a/gitm/kernels/spec.py +++ b/gitm/kernels/spec.py @@ -45,29 +45,19 @@ class InterventionSpec(BaseModel): knob: str # vLLM config key, e.g. "max_num_batched_tokens" — or a display # label ("k1=v1,k2=v2") for a joint candidate; see ``knobs`` below. value: int | float | str | bool | None = None # value to set on apply (single-knob) - # For a knob whose "right" absolute value depends on model size/GPU/ - # workload shape (e.g. max_num_batched_tokens): scale the engine's CURRENT - # value by this factor instead of hardcoding one number for every - # deployment. None (the default) means value is a static literal, applied - # as-is. See gitm.optimizer.vllm_knobs.resolve_relative_value — value - # stays the offline/predict-only fallback when there's no live engine to - # read a current setting from. + # Scale the engine's CURRENT value by this factor instead of hardcoding one + # number (see gitm.optimizer.vllm_knobs.resolve_relative_value). None (the + # default): value is a static literal. value is the offline/predict-only + # fallback when there's no live engine to read a current setting from. value_multiplier: float | None = None - # A sweep of multipliers instead of one — e.g. [0.5, 2.0, 4.0] — expands - # this single reviewed entry into one candidate per point (see - # gitm.optimizer.vllm_knobs.expand_relative_candidates), so the rollback - # gate picks whichever point actually wins instead of the catalog betting - # on a single guess. Empty (the default) means no sweep; value_multiplier - # (if set) resolves as one candidate as usual. + # A sweep of multipliers (e.g. [0.5, 2.0, 4.0]) instead of one — expands + # this entry into one candidate per point (expand_relative_candidates). + # Empty (default): no sweep. value_multiplier_grid: list[float] = Field(default_factory=list) - # Clamp the multiplier-scaled result (e.g. gpu_memory_utilization must stay - # a fraction below 1.0). None means no clamp. - value_max: float | None = None + value_max: float | None = None # clamp the scaled result (e.g. a fraction < 1.0) value_min: float | None = None - # A joint candidate: >1 knob=value pair applied/rolled back together as one - # atomic unit. Empty (the default) means "single knob" — use ``knob``/ - # ``value`` instead. Additive: every existing single-knob spec is - # unaffected (``knobs`` stays ``{}``); see ``knob_values``. + # >1 knob=value pair applied/rolled back together as one atomic unit. Empty + # (default) means single-knob — use knob/value instead. See knob_values. knobs: dict[str, Any] = Field(default_factory=dict) applies_to_kernels: list[str] = Field(default_factory=list) # substring match expected_delta_mean: float # signed, e.g. +0.08 = 8% improvement diff --git a/gitm/optimizer/apply.py b/gitm/optimizer/apply.py index 1849c98..01890f3 100644 --- a/gitm/optimizer/apply.py +++ b/gitm/optimizer/apply.py @@ -120,12 +120,8 @@ def _audit( def _knob_values(spec: Any) -> dict[str, Any]: - """The knob=value pairs ``spec`` wants applied — single or joint. - - Duck-type-friendly: works for a real :class:`InterventionSpec` (via its - ``knobs``/``knob``/``value`` fields) or a bare test double exposing just - ``.knob``/``.value``, which several tests use for brevity. - """ + """The knob=value pairs ``spec`` wants applied — single or joint. Duck-type + friendly: also works for a bare test double exposing just .knob/.value.""" knobs = getattr(spec, "knobs", None) return dict(knobs) if knobs else {spec.knob: spec.value} @@ -353,9 +349,8 @@ def apply(self, spec: InterventionSpec) -> None: raise ValueError(f"intervention {spec.name!r} has no value(s) to set") if not self._force_restart and all(knob_kind(k) == "scheduling" for k in values): - # Hot-swap every knob in place, atomically as a set. Record each - # restore point only AFTER its successful set: if a later knob's - # setter raises, restore() only undoes what was actually changed. + # Record each restore point only after its successful set, so a + # later setter raising only leaves what was actually changed to undo. applied: dict[str, Any] = {} self._prev = ("hotswap", applied) for knob, value in values.items(): diff --git a/gitm/optimizer/deviation.py b/gitm/optimizer/deviation.py index 879fe31..ef71fe1 100644 --- a/gitm/optimizer/deviation.py +++ b/gitm/optimizer/deviation.py @@ -28,25 +28,19 @@ from gitm.tracer.schema import KernelEvent, Trace # Ordered kernel-name → predicted-op rules (first match wins), case-insensitive -# substring. Maps a raw CUDA/Triton kernel name to a predicted-graph op, or None -# when it maps to no modeled op (norms, activations, copies, launch overhead) — -# those are unmodeled work and kept as departures. +# substring. None = no modeled op (norms, activations, copies, launch overhead) +# — unmodeled work, kept as departures. # -# The projection GEMMs (qkv/out/gate_up/down/lm_head) are only distinguishable -# when the kernel name carries the projection tag (vLLM/Triton fused kernels -# usually do); a bare cuBLAS/cutlass GEMM (e.g. `ampere_fp16_s16816gemm_*`, -# `cutlass_*gemm*`) carries no such tag and falls through to None (unmodeled) -# — confirmed against a real vLLM/L4/CUPTI trace, where these generic GEMM -# kernels are ~35% of all kernel launches and are reused interchangeably -# across every projection. Distinguishing those needs shape-matching or -# launch-order instrumentation (a real follow-up, not a vocabulary fix) — but -# this already beats the old positional `i % len(pred)`, which aligned -# nothing under CUDA graphs. The attention/KV-cache rules below ARE confirmed -# against that same real trace: FlashAttention's actual kernel name is -# `flash_fwd_splitkv_kernel` (the `flash_attn`/`flashattn` needles alone miss -# it), and vLLM's own KV-cache write/bookkeeping kernels -# (`reshape_and_cache_flash_kernel`, `_compute_slot_mapping_kernel`) weren't -# covered at all. +# The projection GEMMs (qkv/out/gate_up/down/lm_head) only classify when the +# kernel name carries a projection tag; a bare cuBLAS/cutlass GEMM (e.g. +# `ampere_fp16_s16816gemm_*`) carries none and stays unmodeled — confirmed +# against a real vLLM/L4/CUPTI trace, where these are ~35% of launches and +# reused across every projection. Needs shape-matching or launch-order +# instrumentation to fix, not a vocabulary tweak. The attention/KV-cache +# needles below ARE confirmed against that trace: FlashAttention's real +# kernel is `flash_fwd_splitkv_kernel` (`flash_attn` alone misses it), and +# vLLM's `reshape_and_cache_flash_kernel`/`_compute_slot_mapping_kernel` +# weren't covered before. _OP_RULES: tuple[tuple[tuple[str, ...], str], ...] = ( (("flash_attn", "flashattn", "flash_fwd", "paged_attention", "paged_attn", "fmha", "attention", "attn_score", "reshape_and_cache", "slot_mapping"), "attn_score_value"), diff --git a/gitm/optimizer/monitor.py b/gitm/optimizer/monitor.py index 0a50575..e23906d 100644 --- a/gitm/optimizer/monitor.py +++ b/gitm/optimizer/monitor.py @@ -42,17 +42,14 @@ class Residuals: def residuals(trace: Trace, graph: Graph) -> Residuals: """Pair observed kernels to predicted nodes by op identity, not position. - A real capture has orders of magnitude more kernels than the predicted - graph's ~5*n_layers+1 nodes; the old ordinal pairing matched a handful of - early kernels against unrelated ops, producing runaway r_kt ratios. Each - kernel is classified by name (:func:`gitm.optimizer.deviation.classify_op`) - and compared to that op's roofline node — one representative node per op, - since per-layer nodes share the same prediction. A kernel with no modeled - op is unmodeled work and gets no residual; ``layer`` is unrecoverable from - name-based classification, so it's always ``None``. Each record also carries - the matched op's roofline ``bound`` ("compute" | "memory") so callers can - weight bottleneck classification by real arithmetic intensity instead of - guessing from wall-clock signals alone. + The old ordinal pairing matched a handful of early kernels against + unrelated ops (orders of magnitude fewer predicted nodes than real + kernels), producing runaway r_kt ratios. Each kernel is now classified by + name (:func:`gitm.optimizer.deviation.classify_op`) against one + representative node per op (per-layer nodes share a prediction); + unmodeled kernels get no residual, and ``layer`` is always ``None`` + (unrecoverable from name-based classification). Carries the matched op's + roofline ``bound`` for bottleneck classification. """ obs = trace.kernels() pred = graph.nodes diff --git a/gitm/optimizer/replay.py b/gitm/optimizer/replay.py index 99201fa..bfd7871 100644 --- a/gitm/optimizer/replay.py +++ b/gitm/optimizer/replay.py @@ -38,19 +38,11 @@ def predict_delta(trace: Trace, spec: InterventionSpec) -> float: def _applies(spec: InterventionSpec, kernel_name: str) -> bool: """Does ``kernel_name`` fall within ``spec``'s declared scope? - Prefers op-identity classification (:func:`gitm.optimizer.deviation. - classify_op`) — the same canonical vLLM decode-step vocabulary - ``residuals()`` uses — over raw substring matching, so a tag like - ``attn_score_value`` means the same thing here as it does in residual - attribution, and both stay in sync when the name→op mapping is updated. - Falls back to substring matching for tags classify_op doesn't cover - (other workloads' own kernel-name vocabularies, e.g. HFT's - ``cudf_groupby_scan`` or edge's ``pointpillars``). - - An empty ``applies_to_kernels`` means "targets nothing measurable" (0 - coverage) — NOT "targets everything". A blank scope used to win every - ranking by default (100% coverage) regardless of whether the lever was - real; every candidate must now state its scope to be counted. + Prefers op-identity via :func:`gitm.optimizer.deviation.classify_op` (same + vocabulary ``residuals()`` uses), falling back to substring matching for + tags it doesn't cover (other workloads' own vocabularies, e.g. HFT's + ``cudf_groupby_scan``). An empty ``applies_to_kernels`` means 0 coverage, + not 100% — a blank scope no longer wins ranking by default. """ if not spec.applies_to_kernels: return False diff --git a/gitm/optimizer/vllm_knobs.py b/gitm/optimizer/vllm_knobs.py index 4ef8493..2de2dc7 100644 --- a/gitm/optimizer/vllm_knobs.py +++ b/gitm/optimizer/vllm_knobs.py @@ -191,13 +191,10 @@ def knob_kind(knob: str) -> KnobKind: return spec.kind if spec is not None else "structural" -#: (name substring, prerequisite knob) — a generated candidate matching the -#: substring only does anything (or even applies) when the prerequisite is -#: already on. E.g. ``dbo_prefill_token_threshold`` requires ``--enable-dbo``; -#: ``max_num_partial_prefills``/``long_prefill_token_threshold`` require -#: ``--enable-chunked-prefill`` (vLLM docs). A blanket denylist would forfeit -#: these forever, including on deployments where the prerequisite genuinely -#: holds — check the live engine instead via :func:`unmet_prerequisite`. +#: (name substring, prerequisite knob) — a candidate matching the substring +#: only applies when the prerequisite is on (e.g. dbo_prefill_token_threshold +#: needs --enable-dbo). Check the live engine via unmet_prerequisite rather +#: than denylisting forever, including on deployments where it genuinely holds. KNOB_PREREQUISITES: tuple[tuple[str, str], ...] = ( ("partial_prefill", "enable_chunked_prefill"), ("long_prefill_token_threshold", "enable_chunked_prefill"), @@ -206,13 +203,9 @@ def knob_kind(knob: str) -> KnobKind: def unmet_prerequisite(engine: Any | None, knob: str) -> str | None: - """None if ``knob`` has no prerequisite, or its prerequisite holds on ``engine``. - - Otherwise the rejection reason. With no live engine, a prerequisite-gated - knob can't be verified at all — conservative default is to reject rather - than assume it holds (mirrors this module's other unknown-defaults-unsafe - calls, e.g. :func:`knob_kind`). - """ + """None if ``knob`` has no prerequisite, or it holds on ``engine`` — else + the rejection reason. No engine -> can't verify -> reject (conservative, + like :func:`knob_kind`'s unknown-defaults-unsafe default).""" lname = knob.lower() prereq = next((p for needle, p in KNOB_PREREQUISITES if needle in lname), None) if prereq is None: @@ -230,15 +223,11 @@ def unmet_prerequisite(engine: Any | None, knob: str) -> str | None: def resolve_relative_value(spec: InterventionSpec, engine: Any | None) -> InterventionSpec: """Scale a relative catalog lever's value off the engine's CURRENT setting. - A knob like ``max_num_batched_tokens`` doesn't have one right absolute - value across every deployment — it depends on model size, GPU, and - workload shape (this is exactly why vLLM's own auto_tune.sh sweeps it - rather than hardcoding a number). A catalog entry with ``value_multiplier`` - set says "N times whatever this engine is already running at" instead of - one static guess. With a live engine, read its current value for the knob - and scale that. Falls back to the spec's static ``value`` when there's no - multiplier, no engine, or the current value can't be read/scaled - (offline/predict-only mode) — still a usable, honest estimate. + A knob like ``max_num_batched_tokens`` has no single right absolute value + across deployments (model size/GPU/workload shape vary) — vLLM's own + auto_tune.sh sweeps it rather than hardcoding a number. With a live engine, + read its current value and scale by ``value_multiplier``. Falls back to + the static ``value`` with no multiplier/engine/readable current value. """ if spec.value_multiplier is None or engine is None: return spec @@ -261,22 +250,14 @@ def resolve_relative_value(spec: InterventionSpec, engine: Any | None) -> Interv def expand_relative_candidates(spec: InterventionSpec, engine: Any | None) -> list[InterventionSpec]: - """Sweep ``value_multiplier_grid`` into one resolved candidate per point. - - Mirrors vLLM's own auto_tune.sh (sweeps a list of candidate values per - knob rather than trying one guess) and autoresearch's generated-knob value - grid — the same idea applied to a reviewed catalog lever, so the rollback - gate picks whichever point actually wins instead of the catalog betting on - a single number. Each point is resolved off the SAME current engine value - via :func:`resolve_relative_value` and gets its own name/summary. - - Falls back to a single :func:`resolve_relative_value` call when there's no - grid (a plain ``value_multiplier`` spec, or a static spec with neither) — - additive, not a new contract every entry must opt into. With no live - engine there's nothing to scale relative to, so the sweep collapses to one - candidate at the static fallback value (not N duplicates of it); the same - collapse happens when every grid point resolves to the same value (e.g. - the knob's current value is 0, so any multiplier of it stays 0). + """Sweep ``value_multiplier_grid`` into one resolved candidate per point — + same idea as vLLM's auto_tune.sh and autoresearch's value grid, applied to + a reviewed catalog lever. Each point resolves off the SAME current engine + value via :func:`resolve_relative_value`, with its own name. + + No grid -> a single :func:`resolve_relative_value` call. No live engine, + or every point collapsing to the same value (current == 0), -> one + candidate, not N duplicates. """ if not spec.value_multiplier_grid: return [resolve_relative_value(spec, engine)] diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 4b064e8..203ac73 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -470,9 +470,8 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # Phase 3 — library + counterfactual replay ranking pctx = build_planner_context(cfg.engine, workload = workload) - # Relative levers (value_multiplier / value_multiplier_grid) resolve — and - # sweep, if a grid is set — against the live engine's current setting here, - # once, before ranking. See expand_relative_candidates. + # Relative/swept levers resolve against the live engine here, once, before + # ranking. See expand_relative_candidates. library = [ resolved for s in load_library(workload=workload) From e28a160bcab5351bf72e58a7bef7fcdf7c0d6338 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 17:06:02 -0700 Subject: [PATCH 19/24] catalog: add async_scheduling, cpu_offload_gb, preemption_mode; retire from autoresearch's static table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three new curated library.yaml entries, all backed by real evidence: - async_scheduling: measured +1.7% (kept) on an actual L4/opt-125m trace via autoresearch's generative search. Overlaps CPU-side scheduling with GPU execution, directly targeting idle_stall (launch-bound gaps). Low risk, no numerics change. - cpu_offload_gb / preemption_mode=swap: promoted from autoresearch's static _RULES fallback table (memory_bound) to reviewed catalog entries with real applicability/safety/source metadata, instead of living only as a bare (knob, value, rationale) tuple. autoresearch proposes *outside* the catalog by design — cpu_offload_gb and preemption_mode are removed from _RULES["memory_bound"] (now empty, same as idle_stall already was) so the disjoint-from-catalog invariant (test_proposed_knobs_are_disjoint_from_catalog) holds. This also empties the offline/no-vLLM fallback's memory_bound coverage, since both of its only two candidates graduated — same gap idle_stall's fallback already had. Updated 6 tests that exercised the old memory_bound fallback path to use compute_bound (still has compilation_config) or to assert the now-honest empty result, rather than paper over the gap. Also compresses the comments/docstrings touched in autoresearch.py along the way. --- gitm/agents/autoresearch.py | 121 ++++++++++++-------------------- gitm/kernels/library.yaml | 103 ++++++++++++++++++--------- tests/test_apply_rollback.py | 4 +- tests/test_autoresearch.py | 56 +++++++-------- tests/test_run_loop_workload.py | 21 +++--- 5 files changed, 153 insertions(+), 152 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index cc69410..c632621 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -23,10 +23,9 @@ docs.vllm.ai); their expected deltas, however, are unproven estimates. The ``source`` field says so, and only the measured A/B keeps or discards them. -v0 classifies the bottleneck from trace telemetry (:func:`classify_bottleneck`) — -weighted by the roofline model's per-op compute/memory prediction when residuals -are available, not just wall-clock heuristics — and repoints the search at the -largest-residual op. Candidates come from one of +v0 classifies the bottleneck from trace telemetry (:func:`classify_bottleneck`, +weighted by the roofline model's per-op bound when residuals are available) and +repoints the search at the largest-residual op. Candidates come from one of two sources behind the :class:`Proposer` seam: the static per-class table (:func:`propose`) or :class:`GenerativeProposer`, which searches a workload's knob surface (supplied by a :class:`KnobSource` — vLLM's ``EngineArgs`` by default) at @@ -110,12 +109,9 @@ def _roofline_memory_fraction(residuals: Residuals | None) -> float | None: """Fraction of matched kernel time the roofline model predicts is memory-bound. - ``None`` when there's nothing to compute from (no residuals passed, or no - observed kernel matched a predicted op) — ``classify_bottleneck`` then falls - back to the memcpy-only signal. Unlike memcpy share, this catches the actual - decode bottleneck: low arithmetic intensity *inside* an attention/GEMM kernel - (roofline's ``t_memory > t_compute``), not just standalone copy ops — decode - is rarely dominated by explicit memcpys, so the memcpy signal alone misses it. + ``None`` with nothing to compute from — ``classify_bottleneck`` falls back + to the memcpy-only signal. Catches real arithmetic-intensity-bound decode + (roofline's ``t_memory > t_compute``), which explicit memcpy share misses. """ if residuals is None or not residuals.per_kernel: return None @@ -131,17 +127,12 @@ def _roofline_memory_fraction(residuals: Residuals | None) -> float | None: def classify_bottleneck(trace: Trace, residuals: Residuals | None = None) -> str: """Map a captured trace to one of ``idle_stall`` / ``memory_bound`` / ``compute_bound``. - Two signals, scored against a threshold each; the stronger wins, and neither - crossing its threshold defaults to compute bound (an empty trace too, having - no stall/movement signal at all): - - - serialized-concurrency fraction: poor kernel overlap ⇒ idle/scheduling gaps. - - memory pressure: the memcpy share of GPU-op time (data movement dominating), - widened by the roofline-predicted memory-bound fraction of matched kernel - time when ``residuals`` (from :func:`gitm.optimizer.monitor.residuals`) is - passed — real arithmetic intensity per op, not just standalone copies. - Whichever of the two reads higher memory pressure wins; without - ``residuals`` this is exactly the old memcpy-only heuristic. + Two signals, scored against a threshold each; the stronger wins, neither + crossing defaults to compute bound: serialized-concurrency fraction (poor + kernel overlap ⇒ idle/scheduling gaps), and memory pressure (memcpy share + of GPU-op time, widened by the roofline-predicted memory-bound fraction of + matched kernel time when ``residuals`` is passed). Without ``residuals`` + this is the memcpy-only heuristic. """ kernels = trace.kernels() if not kernels: @@ -199,14 +190,10 @@ def largest_residual(res: Residuals) -> ResidualTarget | None: def _op_present(trace: Trace, op: str) -> bool: """True if some kernel in the trace classifies to ``op``. - Guards the ``applies_to_kernels`` tagging: the residual op label comes from - the predicted graph (a synthetic name like ``attn_score_value``, never a - literal substring of a real kernel name — checking containment directly - against ``k.name`` would essentially never match), so classify by identity - the same way ``residuals()`` does, and only tag a proposal with the op when - it actually matches a captured kernel — otherwise ``predict_delta`` - coverage would be 0 and the proposal would rank as worthless rather than - untargeted. + The op label is synthetic (from the predicted graph), never a literal + substring of a real kernel name, so classify by identity like + ``residuals()`` does — otherwise an untargeted proposal would get tagged + anyway and rank as worthless (zero ``predict_delta`` coverage). """ return any(classify_op(k.name) == op for k in trace.kernels()) @@ -219,12 +206,9 @@ def _op_present(trace: Trace, op: str) -> bool: # rationales are plausibility arguments, not measured claims. _RULES: dict[str, list[tuple[str, object, str]]] = { "idle_stall": [], - "memory_bound": [ - ("cpu_offload_gb", 4, - "offload cold weights to host RAM to free HBM for a larger KV cache"), - ("preemption_mode", "swap", - "swap preempted KV blocks to host instead of recomputing them under memory pressure"), - ], + # cpu_offload_gb/preemption_mode moved to library.yaml (curated catalog); + # autoresearch only proposes outside the catalog. + "memory_bound": [], "compute_bound": [ ("compilation_config", 3, "raise torch.compile to level 3 for kernel fusion + piecewise CUDA graphs"), @@ -283,12 +267,10 @@ def _candidate_spec( and generative proposers can't drift apart on what a *candidate* is. Only the caller-varying parts (name, summary, knob/value, source) are parameters. - ``knobs`` makes this a *joint* candidate (>1 knob=value pair applied and - rolled back together) — ``knob``/``value`` stay a display label in that - case; see :class:`gitm.kernels.spec.InterventionSpec`. ``delta_mean`` - overrides the flat default so a caller can scale confidence by a real, - computable signal (e.g. affinity-keyword strength) instead of every - candidate carrying the identical constant. + ``knobs`` set makes this a *joint* candidate (see + :class:`gitm.kernels.spec.InterventionSpec`); ``knob``/``value`` become a + display label. ``delta_mean`` overrides the flat default so a caller can + scale confidence by a real per-candidate signal instead of one constant. """ return InterventionSpec( name=name, @@ -429,22 +411,17 @@ class vocabulary); otherwise fall back to a keyword-substring match on the def _affinity_strength(knob: Knob, keywords: tuple[str, ...]) -> int: """How many affinity keywords match ``knob``'s name — a real, computable - signal for how strongly it's implicated in the bottleneck class, not a - fabricated score. An explicit ``Knob.classes`` tag is one strong match (no - keyword multiplicity to count).""" + confidence signal, not a fabricated score. An explicit ``Knob.classes`` tag + counts as one strong match.""" if knob.classes: return 1 lname = knob.name.lower() return sum(1 for k in keywords if k in lname) -#: expected_delta_mean per unit of affinity strength — every generated candidate -#: previously carried the exact same constant (_DELTA_MEAN) regardless of how -#: strongly its name actually matched the bottleneck class, so predict_delta's -#: ranking was blind within a run (identical coverage, identical mean ⇒ -#: identical predicted_delta). Scaling by keyword-match count is still an -#: unproven v0 estimate — just one that varies with a real per-candidate signal -#: instead of a flat number. +#: expected_delta_mean per unit of affinity strength (see _affinity_strength) — +#: lets ranking vary per candidate instead of every generated candidate sharing +#: the identical _DELTA_MEAN constant. _DELTA_MEAN_PER_MATCH = _DELTA_MEAN @@ -799,9 +776,8 @@ def _extra_candidates( self, bottleneck_class: str, target_op: str | None, keywords: tuple[str, ...] ) -> list[InterventionSpec]: """Hook for a workload-specific source to add candidates beyond the - per-knob value grid (e.g. joint prerequisite+dependent pairs). No-op - here — :class:`GenerativeProposer` stays workload-agnostic; the vLLM - binding (:class:`EngineArgsProposer`) overrides this.""" + per-knob value grid. No-op here (stays workload-agnostic); overridden + by :class:`EngineArgsProposer`.""" return [] @@ -869,20 +845,14 @@ def _joint_prerequisite_candidates( """Pair a prerequisite-gated knob with its prerequisite, as ONE candidate. :func:`gitm.optimizer.vllm_knobs.unmet_prerequisite` can only veto a - standalone dependent-knob proposal when the prerequisite isn't already on - — it can never turn the prerequisite on itself (an ``InterventionSpec`` only - carries one knob). This is the other half: propose - ``{prerequisite: True, dependent: value}`` as a single joint spec, so the - dependent knob stays reachable (atomically tested, not just vetoed) even on - a deployment where the prerequisite isn't set yet. - - Gated the same way every other generated candidate is — ``_affine`` against - this call's bottleneck class — so a joint dbo/chunked-prefill candidate - only surfaces when the run is actually searching idle_stall, not every - class. Also requires the prerequisite itself to be a real, present, boolean - knob — a source that doesn't expose it (e.g. the offline fallback catalog) - yields nothing here, same as any other candidate needing a surface it - doesn't have. + standalone dependent-knob proposal when its prerequisite isn't already on + — it can't turn the prerequisite on itself. This proposes + ``{prerequisite: True, dependent: value}`` as a single joint spec instead, + so the dependent knob stays reachable rather than permanently vetoed. + + Gated by ``_affine`` like any other generated candidate, and requires the + prerequisite to be a real, present boolean knob — a source that doesn't + expose it (e.g. the offline fallback catalog) yields nothing here. """ by_name = {k.name: k for k in knobs} out: list[InterventionSpec] = [] @@ -927,9 +897,9 @@ class EngineArgsProposer(GenerativeProposer): ignored when ``knobs`` is given) so a 1-GPU box isn't handed a multi-GPU topology candidate whose engine build can only fail. - Also proposes joint prerequisite+dependent candidates (see - :func:`_joint_prerequisite_candidates`) — the vLLM-specific extension of - :meth:`GenerativeProposer._extra_candidates`. + Also proposes joint prerequisite+dependent candidates via + :func:`_joint_prerequisite_candidates` (the vLLM extension of + :meth:`GenerativeProposer._extra_candidates`). """ def __init__( @@ -1159,12 +1129,9 @@ def autoresearch( guard). See :func:`autoresearch_v0`. When ``residuals`` (from :func:`gitm.optimizer.monitor.residuals`) are passed, - two things sharpen: the bottleneck class itself weighs the roofline-predicted - memory-bound fraction of matched kernel time (not just the memcpy-only - signal), and the search is repointed at the largest-residual op — the - biggest gap vs the predicted ceiling — rather than the whole trace. The loop - already computes these in its attribution phase, so it passes them straight - through. + the bottleneck class weighs the roofline-predicted memory-bound fraction too, + and the search repoints at the largest-residual op rather than the whole + trace. The loop already computes these, so it passes them straight through. """ bottleneck_class = classify_bottleneck(trace, residuals) target = largest_residual(residuals) if residuals is not None else None diff --git a/gitm/kernels/library.yaml b/gitm/kernels/library.yaml index 9d3af01..1cc9c27 100644 --- a/gitm/kernels/library.yaml +++ b/gitm/kernels/library.yaml @@ -7,16 +7,11 @@ # apply produce the measured numbers. # # applies_to_kernels scopes each vLLM entry to the canonical per-decode-step op -# vocabulary predict_graph()/classify_op() share (qkv_proj, attn_score_value, -# attn_out_proj, mlp_gate_up, mlp_down, lm_head) — NOT raw kernel-name -# substrings, which drift across attention backends/GPU vendors/torch.compile. -# A knob whose effect genuinely spans the whole decode step (scheduling/ -# admission/topology/graph-capture levers) is tagged with every op via the -# &all_ops anchor below; an empty list means "no modeled coverage" (0, not -# 100%) — e.g. a lever whose benefit is prefill-side or orchestration-level, -# which the decode-only predicted graph doesn't represent. Non-vLLM entries -# (edge/HFT) keep their own kernel-name vocabulary; classify_op doesn't cover -# them, so predict_delta falls back to substring matching for those. +# vocabulary predict_graph()/classify_op() share — NOT raw kernel-name +# substrings, which drift across backends/vendors/torch.compile. Whole-step +# levers use &all_ops; empty means 0 coverage (not 100%) for levers the +# decode-only predicted graph can't represent (prefill-side, orchestration). +# Non-vLLM entries (edge/HFT) keep their own vocabulary via substring fallback. _decode_step_ops: &all_ops - qkv_proj - attn_score_value @@ -25,10 +20,8 @@ _decode_step_ops: &all_ops - mlp_down - lm_head -# Shared multiplier ladder for relative levers that sweep — same ladder -# autoresearch's generated candidates already use (_GRID_MULTIPLIERS in -# gitm/agents/autoresearch.py), so there's one canonical "how wide a swing" -# convention instead of each entry inventing its own. +# Shared multiplier ladder for sweeping relative levers — same as +# autoresearch's _GRID_MULTIPLIERS, one canonical convention. _relative_sweep: &sweep_grid [0.5, 2.0, 4.0] interventions: @@ -76,9 +69,7 @@ interventions: is already running (not one hardcoded number) to fill decode batches. knob: max_num_batched_tokens value: 8192 # offline/predict-only fallback when there's no live engine to scale - value_multiplier_grid: *sweep_grid # vLLM's own auto_tune.sh sweeps this - # per-deployment rather than hardcoding a number — the right absolute - # value depends on model size/GPU/workload shape. + value_multiplier_grid: *sweep_grid # right value depends on model/GPU/workload applies_to_kernels: *all_ops # batch-shape lever: touches every op's batch dim expected_delta_mean: 0.05 expected_delta_lo: 0.01 @@ -117,9 +108,8 @@ interventions: summary: Sweep CPU swap relative to whatever this engine already has configured (not one hardcoded 4 GiB) to avoid preemption stalls. knob: swap_space - value: 4 # offline/predict-only fallback; also the outcome when current - # swap_space is 0 (the common default) — any multiplier of 0 stays 0, - # so the sweep collapses to this one static candidate in that case. + value: 4 # fallback; also the outcome when current swap_space is 0 (common + # default) — 0 * any multiplier stays 0, so the sweep collapses to this. value_multiplier_grid: *sweep_grid applies_to_kernels: [attn_score_value] # preemption/swap is the KV-cache path expected_delta_mean: 0.03 @@ -139,10 +129,8 @@ interventions: already runs at (not one hardcoded 0.92) to permit larger batch slots. knob: gpu_memory_utilization value: 0.92 # offline/predict-only fallback when there's no live engine to scale - # A fraction near 1.0, not a count — the *0.5/*4 ladder used elsewhere - # would be nonsensical here (halving it wastes memory headroom; 4x is an - # invalid fraction). A gentler ladder, clamped so it can never approach - # 1.0 and starve the engine of headroom entirely. + # A fraction near 1.0, not a count — the shared *0.5/*4 ladder would be + # invalid here, so a gentler one, clamped below 1.0. value_multiplier_grid: [0.95, 1.03, 1.10] value_max: 0.97 applies_to_kernels: *all_ops # raises the batch-slot ceiling for every op @@ -159,7 +147,62 @@ interventions: notes: Watch for OOM during long-running decode; rollback to 0.85 on first OOM. review: null + - name: cpu_offload_gb_4 + summary: Offload cold weights to host RAM to free HBM for a larger KV cache. + knob: cpu_offload_gb + value: 4 + applies_to_kernels: [attn_score_value] # frees HBM for KV-cache headroom + expected_delta_mean: 0.04 + expected_delta_lo: -0.02 + expected_delta_hi: 0.10 + source: https://docs.vllm.ai/en/latest/serving/engine_args.html + applicability: + workloads: [vllm-decode] + other: Helps when weights, not KV cache, are the memory bottleneck; adds host<->device transfer overhead. + safety: + tier: moderate + requires_rollback_window_s: 120 + notes: Structural (needs restart); watch for transfer-latency regressions on small models. + review: null + + - name: preemption_mode_swap + summary: Swap preempted KV blocks to host instead of recomputing them under memory pressure. + knob: preemption_mode + value: swap + applies_to_kernels: [attn_score_value] # preemption is the KV-cache path + expected_delta_mean: 0.03 + expected_delta_lo: -0.01 + expected_delta_hi: 0.07 + source: https://docs.vllm.ai/en/latest/serving/engine_args.html + applicability: + workloads: [vllm-decode] + other: Only matters under real preemption pressure; neutral otherwise. + safety: + tier: low_risk + requires_rollback_window_s: 60 + notes: Swap I/O can regress latency-bound workloads; keep only on a positive delta. + review: null + # --- Tier 2: medium-effort levers ----------------------------------------- + - name: async_scheduling + summary: Overlap CPU-side request scheduling with GPU execution to close + idle/launch-bound gaps between decode steps. + knob: async_scheduling + value: true + applies_to_kernels: *all_ops # closes idle gaps around every op's launch + expected_delta_mean: 0.02 + expected_delta_lo: -0.02 + expected_delta_hi: 0.05 + source: https://docs.vllm.ai/en/latest/serving/engine_args.html + applicability: + workloads: [vllm-decode] + other: Targets idle_stall (scheduling/launch-bound gaps); measured ~+1.7% on a real L4/opt-125m trace. + safety: + tier: low_risk + requires_rollback_window_s: 60 + notes: Pure scheduling overlap; no numerics change. + review: null + - name: enable_chunked_prefill summary: Chunk prefill so it overlaps with decode instead of stalling it. knob: enable_chunked_prefill @@ -181,10 +224,7 @@ interventions: summary: Cache shared prompt prefixes to skip repeated prefill work. knob: enable_prefix_caching value: true - # No modeled coverage: the win is prefill-side (skipped prefill work), and - # predict_graph() only models one decode step — honestly 0 here, not a - # false 100%. The measured A/B is what actually judges this one. - applies_to_kernels: [] + applies_to_kernels: [] # win is prefill-side; the decode-only graph doesn't model it expected_delta_mean: 0.10 expected_delta_lo: 0.00 expected_delta_hi: 0.30 @@ -298,8 +338,7 @@ interventions: summary: Load AWQ 4-bit weights to shrink the model and free KV-cache memory. knob: quantization value: awq - # Weight dtype/size shrinks — the weight-heavy projection ops, not - # attn_score_value (its cost is KV-cache bytes, not weight bytes). + # Weight-heavy projection ops, not attn_score_value (KV-cache bytes, not weight bytes). applies_to_kernels: [qkv_proj, attn_out_proj, mlp_gate_up, mlp_down, lm_head] expected_delta_mean: 0.10 expected_delta_lo: -0.05 @@ -378,9 +417,7 @@ interventions: summary: Use the Ray executor backend for multi-node distributed decode. knob: distributed_executor_backend value: ray - # Multi-node orchestration overhead, not represented in the decode roofline - # model at all — honest 0. - applies_to_kernels: [] + applies_to_kernels: [] # multi-node orchestration overhead, not modeled at all expected_delta_mean: 0.00 expected_delta_lo: -0.10 expected_delta_hi: 0.05 diff --git a/tests/test_apply_rollback.py b/tests/test_apply_rollback.py index 716761b..e255696 100644 --- a/tests/test_apply_rollback.py +++ b/tests/test_apply_rollback.py @@ -165,11 +165,11 @@ def test_apply_from_file_with_config(tmp_path): # --- library now carries the 21 curated levers ------------------------------ -def test_library_has_18_validated_levers_with_values(): +def test_library_has_validated_levers_with_values(): from gitm.kernels import load_library specs = load_library() - assert len(specs) == 21 + assert len(specs) == 24 for s in specs: assert s.value is not None, f"{s.name} missing value" assert s.expected_delta_lo <= s.expected_delta_mean <= s.expected_delta_hi diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 4062859..93a4268 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -44,7 +44,7 @@ def test_public_api_names_all_resolve() -> None: def test_propose_known_class_returns_specs() -> None: - specs = propose("memory_bound") + specs = propose("compute_bound") assert specs, "known class with safe standalone rules should yield proposals" assert all(isinstance(s, InterventionSpec) for s in specs) # Never high-risk: unproven proposals stay at moderate and lean on the gate. @@ -126,9 +126,7 @@ def test_roofline_memory_fraction() -> None: def test_classify_catches_intrinsic_memory_boundedness_with_no_memcpy() -> None: - """Real memory-boundedness is intra-kernel (low bytes/FLOP inside the - attention/GEMM kernel), not standalone memcpy ops — passing residuals - whose roofline prediction says so must flip a memcpy-blind compute_bound + """Roofline-flagged memory-boundedness flips a memcpy-blind compute_bound verdict (see test_classify_compute_bound_when_overlapped_and_no_memcpy).""" events = [ make_kernel("a", start_ns=0, end_ns=100, stream_id=0), @@ -155,9 +153,12 @@ def test_applies_and_keeps_on_measured_win() -> None: config: dict = {} applicator = DictApplicator(config, measure_fn=lambda spec: 0.10) # +10% ⇒ keep - results = autoresearch_v0(_trace(), "memory_bound", applicator=applicator) + # compute_bound: memory_bound's static candidates (cpu_offload_gb, + # preemption_mode) graduated to the curated catalog, so the static table + # is empty for it now (see _RULES) — same as idle_stall already was. + results = autoresearch_v0(_trace(), "compute_bound", applicator=applicator) - assert results, "memory_bound has safe standalone proposals" + assert results, "compute_bound has safe standalone proposals" kept = [r for r in results if r.applicable and not r.rolled_back] assert kept, "a positive delta must be kept" for r in kept: @@ -169,7 +170,7 @@ def test_rolls_back_on_regression() -> None: config: dict = {} applicator = DictApplicator(config, measure_fn=lambda spec: -0.10) # slower ⇒ revert - results = autoresearch_v0(_trace(), "memory_bound", applicator=applicator) + results = autoresearch_v0(_trace(), "compute_bound", applicator=applicator) assert results assert all(r.applicable and r.rolled_back for r in results) @@ -407,15 +408,18 @@ def test_engineargs_proposer_scopes_specs_to_target_op() -> None: def test_engineargs_offline_fallback_runs_without_vllm() -> None: - """vLLM isn't importable in CI; the frozen catalog must still yield candidates - for every class, and every candidate stays outside the curated library.""" + """vLLM isn't importable in CI; the frozen fallback catalog still yields + candidates for compute_bound, and every candidate stays outside the + curated library. memory_bound's fallback knobs (cpu_offload_gb, + preemption_mode) graduated to the curated catalog, so the offline path + has nothing left for that class — same gap idle_stall already had.""" catalog = {s.knob for s in load_library()} p = EngineArgsProposer() # default knobs → _engine_arg_knobs() → frozen fallback - for cls in ("memory_bound", "compute_bound"): - specs = p.propose(cls) - assert specs, f"{cls} should yield fallback candidates offline" - assert all(s.safety.tier == "moderate" for s in specs) - assert all(s.knob not in catalog for s in specs) + assert p.propose("memory_bound") == [] + specs = p.propose("compute_bound") + assert specs, "compute_bound should yield fallback candidates offline" + assert all(s.safety.tier == "moderate" for s in specs) + assert all(s.knob not in catalog for s in specs) assert p.propose("idle_stall") == [] # disjoint from the library @@ -629,12 +633,10 @@ def test_is_tunable_excludes_non_perf_engine_args() -> None: for name in ("max_num_seqs", "gpu_memory_utilization", "compilation_config", "cpu_offload_gb"): assert _is_tunable(name), f"{name} should be kept" - # kv_sharing_fast_prefill is WIP/no-op per vLLM docs regardless of engine - # config — nothing to check live, so it stays denylisted outright. + # kv_sharing_fast_prefill is WIP/no-op per vLLM docs -> denylisted outright. assert not _is_tunable("kv_sharing_fast_prefill") - # These are real, sometimes-useful knobs gated by a live prerequisite - # (enable_chunked_prefill / enable_dbo) — not denylisted anymore; they're - # tunable here and vetoed live via unmet_prerequisite (see test_vllm_knobs). + # Prerequisite-gated (enable_chunked_prefill/enable_dbo), not denylisted — + # vetoed live via unmet_prerequisite instead (see test_vllm_knobs). for name in ("max_num_partial_prefills", "max_long_partial_prefills", "long_prefill_token_threshold", "dbo_prefill_token_threshold"): assert _is_tunable(name), f"{name} should be kept (prerequisite-gated, not denylisted)" @@ -766,16 +768,14 @@ def test_vllm_knob_source_gpu_count_override_is_accepted_offline() -> None: def test_engineargs_proposer_accepts_gpu_count_offline() -> None: # Same offline-safety guarantee at the EngineArgsProposer entry point. p = EngineArgsProposer(gpu_count=1) - assert p.propose("memory_bound") + assert p.propose("compute_bound") def test_candidate_specs_share_the_forced_fields() -> None: - """DRY guard: table and generated candidates agree on safety tier, delta - band bounds, and applicability — the fields _candidate_spec centralizes. - expected_delta_mean itself is excluded: generated candidates scale it by - affinity-keyword strength (see test_generated_delta_mean_scales_with_ - affinity_strength), so it legitimately differs candidate to candidate.""" - table = propose("memory_bound")[0] + """DRY guard: safety tier, delta bounds, applicability match. Excludes + expected_delta_mean, which now legitimately varies by affinity strength + (see test_generated_delta_mean_scales_with_affinity_strength).""" + table = propose("compute_bound")[0] generated = EngineArgsProposer( knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() ).propose("memory_bound")[0] @@ -871,9 +871,7 @@ def test_affinity_strength_and_delta_mean_for() -> None: def test_generated_delta_mean_scales_with_affinity_strength() -> None: - """Every generated candidate used to carry the identical expected_delta_mean - regardless of how strongly its name matched the class, so predict_delta's - ranking was blind within a run. A 2-keyword match must rank above a 1-match.""" + """A 2-keyword match must rank above a 1-match (previously identical).""" p = EngineArgsProposer( knobs=[ Knob("cpu_offload_gb", "int", default=0), # matches "cpu" + "offload" diff --git a/tests/test_run_loop_workload.py b/tests/test_run_loop_workload.py index 43cff96..8a0d101 100644 --- a/tests/test_run_loop_workload.py +++ b/tests/test_run_loop_workload.py @@ -130,14 +130,15 @@ def test_vllm_workload_still_uses_intervention_path(tmp_path: Path, monkeypatch) def test_vllm_loop_runs_autoresearch(tmp_path: Path, monkeypatch): - """The vllm path runs agentic autoresearch: it classifies the bottleneck, then - *generates* non-catalog levers from the real EngineArgs surface and searches a - value grid per knob (EngineArgsProposer), surfacing them in the summary + - report. The serialized same-stream "paged_attention" kernels above are also - roofline-predicted memory-bound (attn_score_value at batch=1 is bandwidth- - bound) — a stronger signal than the serialization alone, so classification is - memory_bound; the frozen fallback catalog's two memory knobs (cpu_offload_gb, - preemption_mode) surface as candidates.""" + """The vllm path runs agentic autoresearch: it classifies the bottleneck via + trace telemetry (the serialized same-stream "paged_attention" kernels are + also roofline-predicted memory-bound at batch=1, so classification is + memory_bound). With no vLLM installed, the frozen fallback catalog is used + — but its two former memory_bound candidates (cpu_offload_gb, + preemption_mode) have since graduated to the curated catalog (see + library.yaml), so there's honestly nothing left for autoresearch to + generate here: the catalog absorbing a proven lever is the intended + lifecycle, not a regression.""" import gitm.scheduler.loop as loop monkeypatch.setattr(loop, "capture", _fake_capture_with_kernels("paged_attention")) @@ -152,11 +153,9 @@ def test_vllm_loop_runs_autoresearch(tmp_path: Path, monkeypatch): ) s = result["summary"] assert s["bottleneck_class"] == "memory_bound" - assert s["n_autoresearch"] == 3 # cpu_offload_gb x2 value-grid points + preemption_mode + assert s["n_autoresearch"] == 0 ar_json = (Path(result["run_dir"]) / "autoresearch.json").read_text(encoding="utf-8") - assert "cpu_offload_gb=" in ar_json - assert "preemption_mode=" in ar_json # Denylisted knobs with unchecked prerequisites must never surface, regardless # of which bottleneck class the run classifies as. assert "max_num_partial_prefills=" not in ar_json From 491ee0012ac40b9f0ef91dfd8fc37dc71c2e2a47 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 17:17:32 -0700 Subject: [PATCH 20/24] docs: fix invalid UTF-8 byte in autoresearch.md breaking CI review bots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A markdown table cell had a raw Windows-1252 em-dash byte (0x97) instead of a proper UTF-8-encoded em-dash — every other em-dash in the file was already correctly encoded (0xE2 0x80 0x94), just this one stray byte from what was likely a copy-paste. Both the gemini-review and review CI checks read the full PR diff as UTF-8 and crashed with UnicodeDecodeError on this exact byte, which is why they were failing -- not any of the code-quality findings those bots had posted. --- docs/autoresearch.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/autoresearch.md b/docs/autoresearch.md index c5b3a81..e2db20e 100644 --- a/docs/autoresearch.md +++ b/docs/autoresearch.md @@ -150,7 +150,7 @@ docs.vllm.ai) that is **not** in `library.yaml`: | Class | Knob | Value | Rationale | |-------|------|-------|-----------| -| `idle_stall` | *(none by default)* | — | dependent prefill/DBO knobs are filtered unless represented by reviewed catalog entries | +| `idle_stall` | *(none by default)* | — | dependent prefill/DBO knobs are filtered unless represented by reviewed catalog entries | | `memory_bound` | `cpu_offload_gb` | 4 | offload cold weights to host RAM, freeing HBM for KV cache | | `memory_bound` | `preemption_mode` | `swap` | swap preempted KV blocks instead of recomputing under pressure | | `compute_bound` | `compilation_config` | 3 | torch.compile level 3: kernel fusion + piecewise CUDA graphs | From 4daa319abb86f15202eefb5d7b1aab06ab0a80fc Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Wed, 8 Jul 2026 17:21:14 -0700 Subject: [PATCH 21/24] autoresearch: joint candidates go first so max_candidates can't crowd them out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GenerativeProposer.propose() appended joint (prerequisite+dependent) candidates after the per-knob value grid, then sliced the combined list to max_candidates from the front. A large single-knob grid (the common case — real EngineArgs surfaces have many affine knobs) could fill the whole cap on its own, silently truncating every joint candidate off the end with no visibility into why. Put joint candidates first instead: they're deliberately targeted (unlocking a prerequisite-gated knob atomically), so they shouldn't lose a slot to the more generic grid search. Adds a regression test with a 50-knob filler surface confirming a joint candidate survives a max_candidates=10 cap that the filler grid alone would otherwise fill. --- gitm/agents/autoresearch.py | 7 ++++++- tests/test_autoresearch.py | 15 +++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index c632621..dc86074 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -828,7 +828,12 @@ def propose( if _affine(knob, bottleneck_class, keywords) for value in _value_grid(knob) ] - out += self._extra_candidates(bottleneck_class, target_op, keywords) + # Joint candidates (e.g. prerequisite+dependent pairs) go first: a large + # value-grid surface can easily fill the whole cap on its own, which + # would silently starve out the joint candidates the cap truncates + # from the end otherwise. + extra = self._extra_candidates(bottleneck_class, target_op, keywords) + out = extra + out # Bound the per-class candidate count so a large config surface can't # flood the gate; the rollback gate still ranks and proves what survives. return out if self._max is None else out[: self._max] diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 93a4268..5a5bbcc 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -921,6 +921,21 @@ def test_joint_prerequisite_candidates() -> None: assert joint and all(set(s.knobs) == {"enable_dbo", "dbo_prefill_token_threshold"} for s in joint) +def test_joint_candidates_survive_max_candidates_even_when_grid_alone_fills_it() -> None: + """Joint candidates go first, so a large single-knob value grid can't + silently truncate them off the end of the list.""" + dbo_knobs = [ + Knob("enable_dbo", "bool", default=False), + Knob("dbo_prefill_token_threshold", "int", default=512), + ] + # A single-knob grid big enough to fill the whole cap on its own. + filler = [Knob(f"prefill_filler_{i}", "bool", default=False) for i in range(50)] + p = EngineArgsProposer(knobs=dbo_knobs + filler, catalog_knobs=set(), max_candidates=10) + specs = p.propose("idle_stall") + assert len(specs) == 10 + assert any(len(s.knobs) > 1 for s in specs), "joint candidates were crowded out by the cap" + + def test_intervention_spec_knob_values_property() -> None: from gitm.kernels.spec import InterventionSpec From 574f5cdb0680ad45c26e14ce7b1b28b72a1c0d1b Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Thu, 9 Jul 2026 08:49:02 -0700 Subject: [PATCH 22/24] apply: audit revert calls now carry knobs=, matching apply MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All three apply_intervention() "revert" audit calls (apply-raises, measure-raises, regression rollback) omitted knobs=_knob_values(spec), so a rollback's audit entry recorded only the cause string, not which knob(s) were involved — the sibling "apply" entry already had this. Fixed all three call sites and extended test_apply_rollback.py to assert the knobs detail on every revert path, including a new test covering the apply-raises/measure-raises cases directly. autoresearch: explicit Knob.classes tag scores max, not flat 1 _affinity_strength() gave an explicit, authored Knob.classes tag a flat score of 1, while a knob matched by several affinity keywords incidentally could score higher (e.g. 2+) — inverting the intended confidence ordering. An authored tag is ground truth and must never be outranked by a coincidental multi-keyword name match, so it now scores max(len(keywords), 1), the highest score reachable for that call. Updated the corresponding assertion in test_autoresearch.py. Co-Authored-By: Claude Sonnet 5 --- gitm/agents/autoresearch.py | 7 +++++-- gitm/optimizer/apply.py | 8 +++++--- tests/test_apply_rollback.py | 35 +++++++++++++++++++++++++++++++++++ tests/test_autoresearch.py | 7 +++++-- 4 files changed, 50 insertions(+), 7 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index dc86074..28d5974 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -412,9 +412,12 @@ class vocabulary); otherwise fall back to a keyword-substring match on the def _affinity_strength(knob: Knob, keywords: tuple[str, ...]) -> int: """How many affinity keywords match ``knob``'s name — a real, computable confidence signal, not a fabricated score. An explicit ``Knob.classes`` tag - counts as one strong match.""" + is authored ground truth, so it counts as matching every keyword (the max + possible score for this class) — a knob whose name happens to contain + several keywords by coincidence must never outrank one a source explicitly + tagged.""" if knob.classes: - return 1 + return max(len(keywords), 1) lname = knob.name.lower() return sum(1 for k in keywords if k in lname) diff --git a/gitm/optimizer/apply.py b/gitm/optimizer/apply.py index 01890f3..d7d6e1c 100644 --- a/gitm/optimizer/apply.py +++ b/gitm/optimizer/apply.py @@ -81,7 +81,8 @@ def apply_intervention( applicator.apply(spec) except Exception as exc: applicator.restore(snapshot) - _audit(audit, "revert", spec, cause=f"apply failed, restored: {exc}") + _audit(audit, "revert", spec, cause=f"apply failed, restored: {exc}", + knobs=_knob_values(spec)) return ApplyResult(False, rolled_back=True, measured_delta=None, error=f"apply failed, restored: {exc}") _audit(audit, "apply", spec, cause="applied live", knobs=_knob_values(spec)) @@ -91,14 +92,15 @@ def apply_intervention( delta = applicator.measure(spec) except Exception as exc: applicator.restore(snapshot) - _audit(audit, "revert", spec, cause=f"measure failed, restored: {exc}") + _audit(audit, "revert", spec, cause=f"measure failed, restored: {exc}", + knobs=_knob_values(spec)) return ApplyResult(False, rolled_back=True, measured_delta=None, error=f"measure failed, restored: {exc}") # Step 4: keep-or-rollback on the regression threshold. if delta is not None and delta < min_keep_delta: applicator.restore(snapshot) - _audit(audit, "revert", spec, + _audit(audit, "revert", spec, knobs=_knob_values(spec), cause=f"regression {delta:+.3f} < keep threshold {min_keep_delta:+.3f}") return ApplyResult(True, rolled_back=True, measured_delta=delta, error=f"regression {delta:+.3f} < keep threshold " diff --git a/tests/test_apply_rollback.py b/tests/test_apply_rollback.py index e255696..1831d12 100644 --- a/tests/test_apply_rollback.py +++ b/tests/test_apply_rollback.py @@ -114,6 +114,41 @@ def test_apply_audits_apply_then_revert_on_regression(tmp_path): # A regressing change is applied then reverted — both must be on the trail. assert [e.event for e in log.entries()] == ["apply", "revert"] + # The revert entry names which knob(s) were involved, not just the cause — + # essential once a rollback can be a joint (multi-knob) candidate. + assert log.entries()[1].detail["knobs"] == {"block_size": 16} + + +def test_apply_audits_knobs_on_apply_and_measure_failures(tmp_path): + """The revert entry carries which knob(s) were attempted even when the + failure happens before or during measurement, not just on a regression.""" + from gitm.safety import AuditLog + + class _RaisesOnApply: + def snapshot(self): + return {} + + def apply(self, spec): + raise RuntimeError("bad value") + + def restore(self, snapshot): + pass + + def measure(self, spec): + return None + + class _RaisesOnMeasure(_RaisesOnApply): + def apply(self, spec): + pass + + def measure(self, spec): + raise RuntimeError("measurement crashed") + + for applicator in (_RaisesOnApply(), _RaisesOnMeasure()): + log = AuditLog(tmp_path / f"audit-{type(applicator).__name__}.jsonl") + apply_intervention(_spec(), applicator, audit=log) + assert log.entries()[-1].event == "revert" + assert log.entries()[-1].detail["knobs"] == {"block_size": 16} def test_apply_without_audit_writes_nothing(tmp_path): diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 5a5bbcc..094506d 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -862,8 +862,11 @@ def test_affinity_strength_and_delta_mean_for() -> None: assert _affinity_strength(Knob("cpu_offload_gb", "int", default=0), keywords) == 2 assert _affinity_strength(Knob("swap_space", "int", default=0), keywords) == 1 assert _affinity_strength(Knob("unrelated_thing", "int", default=0), keywords) == 0 - # An explicit tag is one strong match regardless of name. - assert _affinity_strength(Knob("x", "int", default=0, classes=("idle_stall",)), keywords) == 1 + # An explicit tag is authored ground truth: it scores as matching every + # keyword (the max), so it can never be outranked by a coincidental + # multi-keyword name match. + assert _affinity_strength(Knob("x", "int", default=0, classes=("idle_stall",)), keywords) == len(keywords) + assert _affinity_strength(Knob("x", "int", default=0, classes=("idle_stall",)), ()) == 1 # no keywords -> floor at 1 assert _delta_mean_for(0) == _delta_mean_for(1) # zero matches floors at 1 assert abs(_delta_mean_for(2) - 2 * _delta_mean_for(1)) < 1e-9 From e0d11bcd1fda1f0b444a0bbd0594db8ea6eb10a2 Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Thu, 9 Jul 2026 08:53:29 -0700 Subject: [PATCH 23/24] autoresearch: fix zero-default value grid and cache _searchable() _value_grid() substituted base=1 for an int/float knob whose real default is 0 (e.g. cpu_offload_gb), fabricating a grid like [0, 1, 2, 4] that doesn't represent a real multiplicative search around the knob's actual value. A zero default has no meaningful ratio-based grid, so it now returns [] instead of inventing one. Updated the test fixtures that relied on the old fabricated-grid behavior (cpu_offload_gb -> nonzero default) since they were exercising the value-grid path, not testing zero specifically. _searchable() re-ran self._source.knobs() (a fresh knob-surface parse, e.g. VLLMKnobSource's argparse introspection) on every call; propose() calls it twice per invocation (main loop + _extra_candidates), doubling that work for no reason since source/catalog are fixed for the proposer's lifetime. Now memoized per instance. Also: one-line comment documenting that classify_bottleneck's sc/mem tie resolves to idle_stall (was previously undocumented). Co-Authored-By: Claude Sonnet 5 --- gitm/agents/autoresearch.py | 25 ++++++++++++++++++------- tests/test_autoresearch.py | 22 +++++++++++----------- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/gitm/agents/autoresearch.py b/gitm/agents/autoresearch.py index 28d5974..8ba7d3c 100644 --- a/gitm/agents/autoresearch.py +++ b/gitm/agents/autoresearch.py @@ -152,7 +152,7 @@ def classify_bottleneck(trace: Trace, residuals: Residuals | None = None) -> str mem_score = max(mem_score, roofline_frac / _MEMCPY_THRESHOLD) if max(sc_score, mem_score) < 1.0: return COMPUTE_BOUND - return IDLE_STALL if sc_score >= mem_score else MEMORY_BOUND + return IDLE_STALL if sc_score >= mem_score else MEMORY_BOUND # ties favor idle_stall # "Repoint at the largest residual": aim the search at the single op whose @@ -451,8 +451,9 @@ def _value_grid(knob: Knob) -> list[object]: return [c for c in knob.choices if c != knob.default] if knob.kind in ("int", "float"): d = knob.default - base = d if isinstance(d, int | float) and d not in (0, False) else 1 - raw = [base * m for m in _GRID_MULTIPLIERS] + if not isinstance(d, int | float) or d in (0, False): + return [] # zero/non-numeric default: no meaningful multiplicative grid + raw = [d * m for m in _GRID_MULTIPLIERS] if knob.kind == "int": vals = sorted({int(round(x)) for x in raw if round(x) >= 1}) else: @@ -744,12 +745,22 @@ def __init__( if catalog_knobs is not None else {s.knob for s in load_library()} ) + self._searchable_cache: list[Knob] | None = None def _searchable(self) -> list[Knob]: - """Knobs worth searching: outside the catalog and with a non-empty grid.""" - return [ - k for k in self._source.knobs() if k.name not in self._catalog and _value_grid(k) - ] + """Knobs worth searching: outside the catalog and with a non-empty grid. + + Cached: ``propose()`` calls this twice per invocation (main loop + + ``_extra_candidates``), and ``self._source.knobs()`` re-parses the + knob surface (e.g. VLLMKnobSource re-runs argparse introspection) + every time — the source and catalog are fixed for this instance's + lifetime, so the result can't change between calls. + """ + if self._searchable_cache is None: + self._searchable_cache = [ + k for k in self._source.knobs() if k.name not in self._catalog and _value_grid(k) + ] + return self._searchable_cache def _spec( self, diff --git a/tests/test_autoresearch.py b/tests/test_autoresearch.py index 094506d..9ca52b3 100644 --- a/tests/test_autoresearch.py +++ b/tests/test_autoresearch.py @@ -356,7 +356,7 @@ def test_value_grid_explicit_grid_wins() -> None: def _stub_knobs() -> list[Knob]: return [ Knob("max_num_partial_prefills", "int", default=1), # idle_stall-affine - Knob("cpu_offload_gb", "int", default=0), # memory_bound-affine + Knob("cpu_offload_gb", "int", default=4), # memory_bound-affine Knob("block_size", "int", default=16), # memory-affine but IN the catalog ] @@ -400,7 +400,7 @@ def test_engineargs_proposer_unknown_class_is_empty() -> None: def test_engineargs_proposer_scopes_specs_to_target_op() -> None: p = EngineArgsProposer( - knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set() ) specs = p.propose("memory_bound", target_op="paged_attention") assert specs @@ -427,9 +427,9 @@ def test_engineargs_candidates_route_through_gate_and_rollback() -> None: """A generated candidate is kept on a measured win and reverted on a regression — the same gate the catalog goes through, nothing special for generated specs.""" proposer = EngineArgsProposer( - knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set() ) - grid = _value_grid(Knob("cpu_offload_gb", "int", default=0)) + grid = _value_grid(Knob("cpu_offload_gb", "int", default=4)) kept_cfg: dict = {} kept = autoresearch_v0( @@ -474,7 +474,7 @@ def test_apply_failure_surfaces_the_error_not_just_a_bare_none() -> None: — the same shape as 'measured and lost'. apply_error must distinguish the two so a report can say *why* it failed instead of a bare unexplained '-'.""" proposer = EngineArgsProposer( - knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set() ) results = autoresearch_v0( _trace(), "memory_bound", applicator=_FailingApplicator(), proposer=proposer @@ -537,7 +537,7 @@ def test_apply_error_is_none_when_not_applicable_or_measured() -> None: kept = autoresearch_v0( _trace(), "memory_bound", applicator=DictApplicator({}, measure_fn=lambda s: 0.10), - proposer=EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set()), + proposer=EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set()), ) assert kept and all(r.apply_error is None for r in kept) @@ -546,7 +546,7 @@ def test_fallback_proposer_uses_table_only_when_primary_is_empty() -> None: table = TableProposer() # Primary has no compute_bound knob here → falls back to the table's compute lever. fb = FallbackProposer( - EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set()), + EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set()), table, ) # memory_bound: primary yields cpu_offload_gb candidates → table NOT consulted. @@ -609,7 +609,7 @@ def test_knob_explicit_classes_override_keyword_affinity() -> None: def test_engineargs_proposer_is_a_vllm_bound_generative_proposer() -> None: """EngineArgsProposer is just GenerativeProposer bound to the vLLM surface.""" - p = EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set()) + p = EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set()) assert isinstance(p, GenerativeProposer) specs = p.propose("memory_bound") assert specs and all(s.applicability.workloads == ["vllm-decode"] for s in specs) @@ -777,7 +777,7 @@ def test_candidate_specs_share_the_forced_fields() -> None: (see test_generated_delta_mean_scales_with_affinity_strength).""" table = propose("compute_bound")[0] generated = EngineArgsProposer( - knobs=[Knob("cpu_offload_gb", "int", default=0)], catalog_knobs=set() + knobs=[Knob("cpu_offload_gb", "int", default=4)], catalog_knobs=set() ).propose("memory_bound")[0] for s in (table, generated): assert s.safety.tier == "moderate" @@ -859,7 +859,7 @@ def test_candidate_spec_helper_forces_safety_and_delta() -> None: def test_affinity_strength_and_delta_mean_for() -> None: keywords = ("cache", "swap", "offload", "cpu") - assert _affinity_strength(Knob("cpu_offload_gb", "int", default=0), keywords) == 2 + assert _affinity_strength(Knob("cpu_offload_gb", "int", default=4), keywords) == 2 assert _affinity_strength(Knob("swap_space", "int", default=0), keywords) == 1 assert _affinity_strength(Knob("unrelated_thing", "int", default=0), keywords) == 0 # An explicit tag is authored ground truth: it scores as matching every @@ -877,7 +877,7 @@ def test_generated_delta_mean_scales_with_affinity_strength() -> None: """A 2-keyword match must rank above a 1-match (previously identical).""" p = EngineArgsProposer( knobs=[ - Knob("cpu_offload_gb", "int", default=0), # matches "cpu" + "offload" + Knob("cpu_offload_gb", "int", default=4), # matches "cpu" + "offload" Knob("preemption_mode", "enum", default="recompute", choices=("recompute", "swap")), # matches "preempt" only ], From de4e0e902f25bced4471901fc242b17c96f230dd Mon Sep 17 00:00:00 2001 From: Nicholas Lawrence Date: Thu, 9 Jul 2026 08:57:45 -0700 Subject: [PATCH 24/24] apply: LiveEngineApplicator shutdown releases IPC memory like the vllm-decode path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _shutdown() called gc.collect() + torch.cuda.empty_cache() but not torch.cuda.ipc_collect(), unlike gitm/workloads.py's _vllm_decode_factory shutdown, which already calls all three. Brought the two in line so a restart-apply releases GPU memory as completely as a fresh workload boot. docs: memory_bound fallback-table rows were stale after the cpu_offload_gb/preemption_mode catalog graduation _RULES["memory_bound"] is empty in code (those two knobs moved to library.yaml), matching how the idle_stall row already reads "(none by default)" — but the memory_bound rows still listed them as if they were still _RULES entries. Updated to match. Co-Authored-By: Claude Sonnet 5 --- docs/autoresearch.md | 3 +- gitm/optimizer/apply.py | 17 +++++----- gitm/optimizer/vllm_knobs.py | 48 +++++++++++++++++++++++----- gitm/scheduler/loop.py | 33 ++++++++++--------- tests/test_knobs_via_restart.py | 36 ++++++++++++--------- tests/test_vllm_embodiment.py | 30 +++++++++-------- tests/test_vllm_knobs_and_restart.py | 47 +++++++++++++++++---------- tests/test_vllm_stress.py | 45 ++++++++++++++++---------- 8 files changed, 164 insertions(+), 95 deletions(-) diff --git a/docs/autoresearch.md b/docs/autoresearch.md index e2db20e..c5adf26 100644 --- a/docs/autoresearch.md +++ b/docs/autoresearch.md @@ -151,8 +151,7 @@ docs.vllm.ai) that is **not** in `library.yaml`: | Class | Knob | Value | Rationale | |-------|------|-------|-----------| | `idle_stall` | *(none by default)* | — | dependent prefill/DBO knobs are filtered unless represented by reviewed catalog entries | -| `memory_bound` | `cpu_offload_gb` | 4 | offload cold weights to host RAM, freeing HBM for KV cache | -| `memory_bound` | `preemption_mode` | `swap` | swap preempted KV blocks instead of recomputing under pressure | +| `memory_bound` | *(none by default)* | — | `cpu_offload_gb` and `preemption_mode` graduated to `library.yaml` (see below) | | `compute_bound` | `compilation_config` | 3 | torch.compile level 3: kernel fusion + piecewise CUDA graphs | The **rationales are plausibility arguments, not measured claims**. Each proposal diff --git a/gitm/optimizer/apply.py b/gitm/optimizer/apply.py index d7d6e1c..9c9704c 100644 --- a/gitm/optimizer/apply.py +++ b/gitm/optimizer/apply.py @@ -261,10 +261,10 @@ class LiveEngineApplicator: Routes by knob taxonomy (:mod:`gitm.optimizer.vllm_knobs`): - * **scheduling** knobs (``max_num_seqs``, ``max_num_batched_tokens``, - ``scheduling_policy``) are *hot-swapped* in place — set on the live - scheduler config, effective next step, restored by setting the old value. - A joint candidate whose knobs are ALL scheduling is hot-swapped as one set. + * **scheduling** knobs are rare deployment-specific live controls. A joint + candidate whose knobs are ALL scheduling is hot-swapped as one set. + The curated vLLM EngineArg map intentionally treats scheduler-looking + EngineArgs as structural because vLLM reads them at construction time. * **structural** knobs (parallelism, dtype, quantization, block size, …) can only take effect on a fresh engine, so they (or a joint set containing any structural knob) are routed through @@ -308,9 +308,8 @@ def __init__( self._getter = getter or get_knob self._setter = setter or set_knob self._reps = max(1, reps) - # force_restart: apply scheduling knobs via the restart path too, for - # engines that don't honor a live scheduler-config mutation (vLLM V1 reads - # scheduler_config at construction). Requires a restart_fn. + # force_restart is kept for custom deployments that still classify a + # knob as scheduling but want to measure it through the restart path. self._force_restart = force_restart self._baseline_tps: float | None = None self._baseline_std: float = 0.0 @@ -462,10 +461,10 @@ def _shutdown(engine: Any) -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() + torch.cuda.ipc_collect() except Exception: pass - def measure(self, spec: InterventionSpec) -> float | None: baseline = self._baseline_tps if self._baseline_tps is not None else self._bench_stats()[0] # A non-positive baseline means the A/B has no valid reference — an idle @@ -486,7 +485,7 @@ def measure(self, spec: InterventionSpec) -> float | None: # anything within noise. reps=1 → std=0 → band=0 → keep iff delta>0. noise_band = (self._baseline_std + cand_std) / baseline significant = delta > noise_band - via = "restart" if (self._prev and self._prev[0] == "restart") else "hot-swap" + via = "restart" if (self._prev and self._prev[0] in {"restart", "serial_restart"}) else "hot-swap" self.last_result = EngineABResult( knob=spec.knob, value=(getattr(spec, "knobs", None) or spec.value), baseline_tps=baseline, diff --git a/gitm/optimizer/vllm_knobs.py b/gitm/optimizer/vllm_knobs.py index 2de2dc7..6b0f70c 100644 --- a/gitm/optimizer/vllm_knobs.py +++ b/gitm/optimizer/vllm_knobs.py @@ -53,16 +53,17 @@ def is_env(self) -> bool: return self.path.startswith("env:") -# The curated map. Only the three scheduler knobs vLLM honors mid-run are -# "scheduling"; everything else is fixed at construction → "structural". +# The curated map. vLLM EngineArgs are treated as construction-time knobs in the +# real in-process engine path, even when a Python config object exposes the same +# field. Route them through restart so the A/B measures a rebuilt engine, not a +# config mutation the scheduler/cache/model may ignore. _KNOBS: dict[str, KnobSpec] = { - # --- scheduling: live-settable on the scheduler, effective next step ------- - "max_num_seqs": KnobSpec("max_num_seqs", "scheduler_config.max_num_seqs", "scheduling"), + # --- structural: fixed at engine construction, need a restart to change ---- + "max_num_seqs": KnobSpec("max_num_seqs", "scheduler_config.max_num_seqs", "structural"), "max_num_batched_tokens": KnobSpec( - "max_num_batched_tokens", "scheduler_config.max_num_batched_tokens", "scheduling" + "max_num_batched_tokens", "scheduler_config.max_num_batched_tokens", "structural" ), - "scheduling_policy": KnobSpec("scheduling_policy", "scheduler_config.policy", "scheduling"), - # --- structural: fixed at engine construction, need a restart to change ---- + "scheduling_policy": KnobSpec("scheduling_policy", "scheduler_config.policy", "structural"), "block_size": KnobSpec("block_size", "cache_config.block_size", "structural"), "kv_cache_dtype": KnobSpec("kv_cache_dtype", "cache_config.cache_dtype", "structural"), "gpu_memory_utilization": KnobSpec( @@ -75,6 +76,27 @@ def is_env(self) -> bool: "enable_chunked_prefill": KnobSpec( "enable_chunked_prefill", "scheduler_config.chunked_prefill_enabled", "structural" ), + "async_scheduling": KnobSpec( + "async_scheduling", "scheduler_config.async_scheduling", "structural" + ), + "scheduler_delay_factor": KnobSpec( + "scheduler_delay_factor", "scheduler_config.scheduler_delay_factor", "structural" + ), + "max_num_partial_prefills": KnobSpec( + "max_num_partial_prefills", "scheduler_config.max_num_partial_prefills", "structural" + ), + "max_long_partial_prefills": KnobSpec( + "max_long_partial_prefills", "scheduler_config.max_long_partial_prefills", "structural" + ), + "long_prefill_token_threshold": KnobSpec( + "long_prefill_token_threshold", "scheduler_config.long_prefill_token_threshold", "structural" + ), + "dbo_decode_token_threshold": KnobSpec( + "dbo_decode_token_threshold", "scheduler_config.dbo_decode_token_threshold", "structural" + ), + "dbo_prefill_token_threshold": KnobSpec( + "dbo_prefill_token_threshold", "scheduler_config.dbo_prefill_token_threshold", "structural" + ), "num_speculative_tokens": KnobSpec( "num_speculative_tokens", "speculative_config.num_speculative_tokens", "structural" ), @@ -82,7 +104,17 @@ def is_env(self) -> bool: "max_seq_len_to_capture": KnobSpec( "max_seq_len_to_capture", "model_config.max_seq_len_to_capture", "structural" ), + "max_model_len": KnobSpec("max_model_len", "model_config.max_model_len", "structural"), + "disable_sliding_window": KnobSpec( + "disable_sliding_window", "model_config.disable_sliding_window", "structural" + ), "quantization": KnobSpec("quantization", "model_config.quantization", "structural"), + "calculate_kv_scales": KnobSpec( + "calculate_kv_scales", "cache_config.calculate_kv_scales", "structural" + ), + "kv_sharing_fast_prefill": KnobSpec( + "kv_sharing_fast_prefill", "cache_config.kv_sharing_fast_prefill", "structural" + ), "tensor_parallel_size": KnobSpec( "tensor_parallel_size", "parallel_config.tensor_parallel_size", "structural" ), @@ -153,7 +185,7 @@ def get_knob(engine: Any, knob: str) -> Any: def set_knob(engine: Any, knob: str, value: Any) -> None: - """Set ``knob`` on the live engine via its taxonomy path (scheduling knobs). + """Set ``knob`` on the live engine via its taxonomy path (live knobs). Raises ``AttributeError`` when the knob can't be located — the caller (:class:`~gitm.optimizer.apply.LiveEngineApplicator`) turns that into a diff --git a/gitm/scheduler/loop.py b/gitm/scheduler/loop.py index 203ac73..fd9fa3e 100644 --- a/gitm/scheduler/loop.py +++ b/gitm/scheduler/loop.py @@ -228,11 +228,14 @@ def _agg_kt_residual(res: Any) -> float: return _clamp_pct(value) -def _ar_target_residual(ar_run: AutoresearchRun) -> float: - """The residual autoresearch's claims should report: the largest-residual - op's mean r_kt that the search targeted (:func:`gitm.agents.autoresearch. - largest_residual`), same for every claim in the phase. 0.0 with no target.""" - return _clamp_pct(ar_run.target.residual) if ar_run.target is not None else 0.0 +def _ar_target_residual(ar_run: AutoresearchRun, fallback: float = 0.0) -> float: + """Residual for autoresearch claims. + + Prefer the largest-residual op that autoresearch targeted; when there is no + target, fall back to the run-level kernel-time residual so generated claims + do not all display a misleading +0.0% gap. + """ + return _clamp_pct(ar_run.target.residual) if ar_run.target is not None else fallback def run_loop(cfg: LoopConfig) -> dict[str, Any]: @@ -495,13 +498,13 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: # Phase 4 — apply with rollback gates. # With a live engine attached, each candidate runs the rollback-gated decode- - # throughput A/B (LiveEngineApplicator): snapshot baseline tps → hot-swap the - # knob → measure candidate tps → keep only on a non-negative delta, else - # restore. Scheduling knobs are hot-swapped; structural knobs are routed - # through ``engine.gitm_restart_fn`` (if the deployment provides one) or - # rolled back rather than silently no-op'd. With no engine it's predict-only - # (DryRunApplicator): candidates land in the report as unverified - # (measured_delta=None), never claimed as won. + # throughput A/B (LiveEngineApplicator): snapshot baseline tps, apply the + # candidate, measure candidate tps, keep only on a non-negative delta, else + # restore. vLLM EngineArgs are routed through ``engine.gitm_restart_fn`` + # (if the deployment provides one) because the real engine reads them at + # construction time. With no engine it is predict-only (DryRunApplicator): + # candidates land in the report as unverified (measured_delta=None), never + # claimed as won. live_restart_fn = getattr(cfg.engine, "gitm_restart_fn", None) if cfg.engine else None if cfg.engine is not None: applicator: Applicator = LiveEngineApplicator( @@ -511,8 +514,8 @@ def run_loop(cfg: LoopConfig) -> dict[str, Any]: baseline_restart_fn=getattr(cfg.engine, "gitm_baseline_restart_fn", None), restart_mode=os.environ.get("GITM_RESTART_MODE", "parallel"), reps=int(os.environ.get("GITM_AB_REPS", "1")), - # GITM_KNOBS_VIA_RESTART=1 applies scheduling knobs via engine rebuild - # too — for V1, which ignores a live scheduler-config mutation. + # Compatibility escape hatch for custom scheduling-classified knobs + # that should still be measured through engine rebuild. force_restart=os.environ.get("GITM_KNOBS_VIA_RESTART") == "1", ) else: @@ -605,7 +608,7 @@ def _unenactable(spec: Any) -> str | None: ar_causal_evidence = ", ".join( f"{h.cause_op}→{h.effect_op} (p={h.p_value:.2g})" for h in hypotheses.top(2) ) or "no strong causal signal" - ar_residual = _ar_target_residual(ar_run) + ar_residual = _ar_target_residual(ar_run, kt_residual) for r in ar_run.results: if not r.applicable: rejected.append(f"{r.spec.name} ({r.rejected_reason})") diff --git a/tests/test_knobs_via_restart.py b/tests/test_knobs_via_restart.py index f2a606e..2261825 100644 --- a/tests/test_knobs_via_restart.py +++ b/tests/test_knobs_via_restart.py @@ -1,8 +1,8 @@ -"""GITM_KNOBS_VIA_RESTART routes scheduling knobs through the rebuild. +"""EngineArg knobs route through restart by default. -If an engine ignores a live scheduler-config mutation (vLLM V1 reads it at -construction), force_restart makes the applicator apply scheduling knobs via the -restart path — a fresh engine with the knob baked in — instead of hot-swapping. +vLLM V1 reads many scheduler-looking EngineArgs at construction time. The +applicator should rebuild a fresh engine with those values baked in instead of +mutating scheduler_config fields that the live engine may ignore. """ from __future__ import annotations @@ -11,10 +11,10 @@ from gitm.optimizer.apply import LiveEngineApplicator -_SCHED_KNOB = SimpleNamespace(knob="max_num_seqs", value=256) # a scheduling knob +_ENGINE_ARG_KNOB = SimpleNamespace(knob="max_num_seqs", value=256) -def test_force_restart_rebuilds_for_a_scheduling_knob(): +def test_engine_arg_rebuilds_by_default(): built: list[dict] = [] def restart(_engine, knob_values): @@ -25,24 +25,30 @@ def restart(_engine, knob_values): SimpleNamespace(name="baseline"), throughput_fn=lambda _e: 100.0, restart_fn=restart, - force_restart=True, ) app.snapshot() - app.apply(_SCHED_KNOB) - assert built == [{"max_num_seqs": 256}] # rebuilt, not hot-swapped + app.apply(_ENGINE_ARG_KNOB) + assert built == [{"max_num_seqs": 256}] assert app._prev[0] == "restart" -def test_default_hot_swaps_a_scheduling_knob(): +def test_unknown_knob_defaults_to_restart_not_hotswap(): + built: list[dict] = [] sets: list[tuple[str, object]] = [] + + def restart(_engine, knob_values): + built.append(dict(knob_values)) + return SimpleNamespace(name="candidate") + app = LiveEngineApplicator( - SimpleNamespace(), + SimpleNamespace(name="baseline"), throughput_fn=lambda _e: 100.0, getter=lambda _e, _k: None, setter=lambda _e, k, v: sets.append((k, v)), - restart_fn=lambda _e, _kv: SimpleNamespace(), + restart_fn=restart, ) app.snapshot() - app.apply(_SCHED_KNOB) - assert sets == [("max_num_seqs", 256)] # hot-swapped in place - assert app._prev[0] == "hotswap" + app.apply(SimpleNamespace(knob="future_engine_arg", value=1)) + assert sets == [] + assert built == [{"future_engine_arg": 1}] + assert app._prev[0] == "restart" diff --git a/tests/test_vllm_embodiment.py b/tests/test_vllm_embodiment.py index 2deeafb..ba7d91e 100644 --- a/tests/test_vllm_embodiment.py +++ b/tests/test_vllm_embodiment.py @@ -259,16 +259,21 @@ class _ModelConfig: class _LiveEngine: - """Duck-typed vLLM engine with one hot-swappable knob and a deterministic - throughput probe, so the rollback-gated A/B has a real (noise-free) signal.""" + """Duck-typed vLLM engine with a restartable max_num_seqs lever.""" - def __init__(self): + def __init__(self, max_num_seqs: int = 32): self.model_config = _ModelConfig() - self.max_num_seqs = 32 # the one knob this fake engine can hot-swap - # Decode throughput scales with batch width — so raising max_num_seqs is a - # measurable win and is kept; knobs with no matching attribute raise on - # apply and are rolled back (the "structural knob can't hot-swap" path). + self.max_num_seqs = max_num_seqs + # Decode throughput scales with batch width, so raising max_num_seqs is a + # measurable win and is kept. Other structural knobs raise from restart + # and are rolled back. self.gitm_throughput_fn = lambda e: float(e.max_num_seqs) + self.gitm_restart_fn = self._restart + + def _restart(self, _old_engine, knob_values): + if set(knob_values) != {"max_num_seqs"}: + raise AttributeError("unsupported fake-engine restart knob") + return _LiveEngine(max_num_seqs=int(knob_values["max_num_seqs"])) def test_run_loop_live_engine_keeps_winning_knob_and_rolls_back_unswappable( @@ -308,12 +313,11 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): summary = out["summary"] assert summary["status"] == "ok" and summary["mode"] == "intervention" - # The hot-swappable winning knob is kept with a measured positive delta. - # max_num_seqs_dynamic now sweeps relative to the engine's starting value - # (32) rather than always jumping to one hardcoded 256 — assert it was - # raised and kept, not a specific literal that depends on where the sweep - # started. - assert engine.max_num_seqs > 32 + # The restartable winning knob is kept with a measured positive delta. + # The original fake engine is not mutated; the live applicator swaps in a + # rebuilt candidate engine, just like the vLLM runner does. + assert "via restart" in out["report_md"] + assert "max_num_seqs" in out["report_md"] import json claims = json.loads((Path(out["run_dir"]) / "ranked_candidates.json").read_text()) diff --git a/tests/test_vllm_knobs_and_restart.py b/tests/test_vllm_knobs_and_restart.py index 6a88917..a4bbd44 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -39,9 +39,12 @@ def _spec(knob: str, value) -> InterventionSpec: # knob taxonomy # # --------------------------------------------------------------------------- # def test_knob_kind_classification(): - assert knob_kind("max_num_seqs") == "scheduling" - assert knob_kind("max_num_batched_tokens") == "scheduling" - assert knob_kind("scheduling_policy") == "scheduling" + assert knob_kind("max_num_seqs") == "structural" + assert knob_kind("max_num_batched_tokens") == "structural" + assert knob_kind("scheduling_policy") == "structural" + assert knob_kind("async_scheduling") == "structural" + assert knob_kind("dbo_prefill_token_threshold") == "structural" + assert knob_kind("kv_sharing_fast_prefill") == "structural" assert knob_kind("tensor_parallel_size") == "structural" assert knob_kind("block_size") == "structural" assert knob_kind("totally_unknown_knob") == "structural" # safe default @@ -201,7 +204,7 @@ def test_unresolvable_knob_raises(): # --------------------------------------------------------------------------- # -# LiveEngineApplicator: hot-swap vs restart-apply # +# LiveEngineApplicator: restart-apply # # --------------------------------------------------------------------------- # class _TpsEngine: """Engine whose throughput is a settable number; restart yields a new one.""" @@ -215,15 +218,22 @@ def _tps_of(e): return e._tps -def test_hotswap_scheduling_knob_kept(): +def test_engine_arg_knob_uses_restart_not_hotswap(): e = _TpsEngine(100.0) - # throughput rises with max_num_seqs so the hot-swap is a measurable win. - app = LiveEngineApplicator(e, throughput_fn=lambda x: float(x.scheduler_config.max_num_seqs)) + built: list[dict[str, object]] = [] + + def restart_fn(old, knob_values): + assert old is e + built.append(dict(knob_values)) + return _TpsEngine(200.0) + + app = LiveEngineApplicator(e, throughput_fn=_tps_of, restart_fn=restart_fn) res = apply_intervention(_spec("max_num_seqs", 256), app, min_keep_delta=0.0) assert res.applied and not res.rolled_back - assert e.scheduler_config.max_num_seqs == 256 - assert app.last_result.via == "hot-swap" - assert app.last_result.speedup == pytest.approx(256 / 32) + assert built == [{"max_num_seqs": 256}] + assert app.engine is not e and app.engine._tps == 200.0 + assert app.last_result.via == "restart" + assert app.last_result.speedup == pytest.approx(2.0) def test_restart_apply_structural_knob_kept_and_swaps_engine(): @@ -289,6 +299,7 @@ def baseline_restart_fn(old): assert res.applied and res.rolled_back assert events == ["restart:block_size=16", "restore-baseline"] assert app.engine is restored + assert app.last_result.via == "restart" def test_serial_restart_rebuilds_baseline_if_candidate_build_fails(): @@ -397,17 +408,22 @@ class _SchedCfg64: class _FullEngine: - """Fake engine exposing scheduler stats AND a hot-swappable max_num_seqs.""" + """Fake engine exposing scheduler stats and a restartable max_num_seqs.""" - def __init__(self): + def __init__(self, max_num_seqs: int = 64): self.model_config = _ModelCfgBf16() self.scheduler = [_LowOccScheduler()] self.scheduler_config = _SchedCfg64() + self.scheduler_config.max_num_seqs = max_num_seqs self.gitm_throughput_fn = lambda e: float(e.scheduler_config.max_num_seqs) + self.gitm_restart_fn = self._restart def get_num_unfinished_requests(self): return 2 + def _restart(self, _old_engine, knob_values): + return _FullEngine(max_num_seqs=int(knob_values.get("max_num_seqs", 64))) + def test_run_loop_scheduler_stats_feed_attribution_and_claims(tmp_path, monkeypatch): import json @@ -442,10 +458,9 @@ def fake_capture(out_path, *, workload_id="w", fingerprint="f", run_id=None): assert "under_filled_batch" in signals # And it reached the claim that it motivates (max_num_seqs), in the report. assert "scheduler[under_filled_batch]" in out["report_md"] - # The hot-swap won and was kept. max_num_seqs_dynamic now sweeps relative - # to the engine's starting value (64) rather than one hardcoded 256, so - # assert it was raised and kept rather than pin an exact literal. - assert engine.scheduler_config.max_num_seqs > 64 + # The restart path won and was kept for the generated max_num_seqs lever. + assert "via restart" in out["report_md"] + assert "max_num_seqs" in out["report_md"] # Scheduler summary surfaced in the run summary (synchronous first sample). assert out["summary"]["scheduler_stats"] is not None diff --git a/tests/test_vllm_stress.py b/tests/test_vllm_stress.py index aa0a7ae..adde9cf 100644 --- a/tests/test_vllm_stress.py +++ b/tests/test_vllm_stress.py @@ -36,38 +36,49 @@ def __init__(self): class _Engine: - def __init__(self): + def __init__(self, max_num_seqs: int = 32): self.scheduler_config = _SchedCfg() + self.scheduler_config.max_num_seqs = max_num_seqs + + def restart(self, _old_engine, knob_values): + if set(knob_values) != {"max_num_seqs"}: + raise AttributeError("unsupported fake-engine restart knob") + return _Engine(max_num_seqs=int(knob_values["max_num_seqs"])) def test_failed_apply_clears_previous_ab_result(): - """A structural-knob candidate with no restart hook must not surface the - *previous* hot-swap's A/B result (snapshot() resets last_result).""" + """A failed restart candidate must not surface the previous A/B result.""" e = _Engine() - app = LiveEngineApplicator(e, throughput_fn=lambda x: float(x.scheduler_config.max_num_seqs)) + app = LiveEngineApplicator( + e, + throughput_fn=lambda x: float(x.scheduler_config.max_num_seqs), + restart_fn=e.restart, + ) r1 = apply_intervention(_spec("max_num_seqs", 256), app, min_keep_delta=0.0) - assert r1.applied and app.last_result is not None # hot-swap measured + assert r1.applied and app.last_result is not None - # Next: a structural knob with no restart_fn → apply raises → rolled back, - # measure never runs → last_result must be reset to None, not stale. + # Next: an unsupported structural knob raises during restart. measure never + # runs, so last_result must be reset to None, not stale from r1. r2 = apply_intervention(_spec("quantization", "awq"), app, min_keep_delta=0.0) assert r2.rolled_back and app.last_result is None def test_many_candidates_restore_is_clean(): - """Hot-swap → rollback → hot-swap across many candidates leaves the engine at - the last *kept* value, never a half-applied state.""" + """Restart, rollback, restart leaves the applicator at the last kept engine.""" e = _Engine() - # Throughput rises with the knob, so higher is kept, lower is rolled back. - app = LiveEngineApplicator(e, throughput_fn=lambda x: float(x.scheduler_config.max_num_seqs)) + app = LiveEngineApplicator( + e, + throughput_fn=lambda x: float(x.scheduler_config.max_num_seqs), + restart_fn=e.restart, + ) apply_intervention(_spec("max_num_seqs", 256), app, min_keep_delta=0.0) # 32->256 kept - assert e.scheduler_config.max_num_seqs == 256 - apply_intervention(_spec("max_num_seqs", 64), app, min_keep_delta=0.0) # 256->64 slower, rolled back - assert e.scheduler_config.max_num_seqs == 256 # restored + assert app.engine.scheduler_config.max_num_seqs == 256 + apply_intervention(_spec("max_num_seqs", 64), app, min_keep_delta=0.0) # 256->64 slower + assert app.engine.scheduler_config.max_num_seqs == 256 # restored apply_intervention(_spec("max_num_seqs", 512), app, min_keep_delta=0.0) # 256->512 kept - assert e.scheduler_config.max_num_seqs == 512 + assert app.engine.scheduler_config.max_num_seqs == 512 # --------------------------------------------------------------------------- # @@ -213,8 +224,8 @@ def test_measure_rolls_back_on_nonpositive_baseline(): class _Idle: scheduler_config = _SchedCfg() - # Baseline probe returns 0 (idle engine) → measure() raises → rollback. - app = LiveEngineApplicator(_Idle(), throughput_fn=lambda e: 0.0) + # Baseline probe returns 0 (idle engine), so measure() raises after apply. + app = LiveEngineApplicator(_Idle(), throughput_fn=lambda e: 0.0, restart_fn=lambda _e, _kv: _Idle()) res = apply_intervention(_spec("max_num_seqs", 256), app, min_keep_delta=0.0) assert not res.applied and res.rolled_back