diff --git a/docs/autoresearch.md b/docs/autoresearch.md new file mode 100644 index 0000000..c5adf26 --- /dev/null +++ b/docs/autoresearch.md @@ -0,0 +1,261 @@ +# 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 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. +- **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. +- **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` | *(none by default)* | — | dependent prefill/DBO knobs are filtered unless represented by reviewed catalog entries | +| `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 +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`, `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 + +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. +- **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`. +- **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. +- **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) 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) + 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..8ba7d3c --- /dev/null +++ b/gitm/agents/autoresearch.py @@ -0,0 +1,1171 @@ +"""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 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 +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.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 + +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 _roofline_memory_fraction(residuals: Residuals | None) -> float | None: + """Fraction of matched kernel time the roofline model predicts is memory-bound. + + ``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 + 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``. + + 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: + 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 + 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 # ties favor idle_stall + + +# "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.""" + + 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 some kernel in the trace classifies to ``op``. + + 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()) + + +# --- 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": [], + # 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"), + ], +} + + +@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 + # 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 +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, + knobs: dict[str, object] | None = None, + delta_mean: float = _DELTA_MEAN, +) -> 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. + + ``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, + 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=min(delta_mean, _DELTA_HI), + 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("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) + + +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 + 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 max(len(keywords), 1) + lname = knob.name.lower() + return sum(1 for k in keywords if k in lname) + + +#: 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 + + +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) + + +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 + 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: + 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", + # 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", +) + + +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) + + +#: 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. + + ``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). + + 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 argparse + + parser = engine_args_cls.add_cli_args(argparse.ArgumentParser()) # type: ignore[attr-defined] + except Exception: + 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, *, 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`), 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 + 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(*, 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, + 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, gpu_count=gpu_count) or list(_FALLBACK_KNOBS) + except Exception: + return list(_FALLBACK_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). + + ``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(gpu_count=self._gpu_count) + + +@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()} + ) + self._searchable_cache: list[Knob] | None = None + + def _searchable(self) -> list[Knob]: + """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, + bottleneck_class: str, + knob: Knob, + value: object, + *, + target_op: str | None, + verb: str, + source: str, + keywords: tuple[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, + 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. No-op here (stays workload-agnostic); overridden + by :class:`EngineArgsProposer`.""" + return [] + + +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)", + keywords=keywords, + ) + for knob in self._searchable() + if _affine(knob, bottleneck_class, keywords) + for value in _value_grid(knob) + ] + # 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] + + +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 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] = [] + 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. + + 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. + ``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 via + :func:`_joint_prerequisite_candidates` (the vLLM extension of + :meth:`GenerativeProposer._extra_candidates`). + """ + + def __init__( + self, + *, + 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(gpu_count=gpu_count) + ) + super().__init__( + source, + workload="vllm-decode", + catalog_knobs=catalog_knobs, + 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. + + 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)", + keywords=keywords, + ) + ) + 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, + apply_error=applied.error if applied else None, + ) + ) + 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 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 + 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/kernels/library.yaml b/gitm/kernels/library.yaml index ddd9a31..1cc9c27 100644 --- a/gitm/kernels/library.yaml +++ b/gitm/kernels/library.yaml @@ -5,6 +5,24 @@ # 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 — 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 + - attn_out_proj + - mlp_gate_up + - mlp_down + - lm_head + +# 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: # --- Tier 1: low-effort levers -------------------------------------------- @@ -12,7 +30,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 +49,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 @@ -46,11 +64,13 @@ 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_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 - applies_to_kernels: [] + value: 8192 # offline/predict-only fallback when there's no live engine to scale + 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 expected_delta_hi: 0.10 @@ -64,11 +84,13 @@ 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 - applies_to_kernels: [] + 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 expected_delta_hi: 0.09 @@ -82,11 +104,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 - applies_to_kernels: [] + 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 expected_delta_lo: -0.01 expected_delta_hi: 0.07 @@ -99,11 +124,16 @@ 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 - applies_to_kernels: [] + value: 0.92 # offline/predict-only fallback when there's no live engine to scale + # 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 expected_delta_mean: 0.04 expected_delta_lo: 0.01 expected_delta_hi: 0.08 @@ -117,12 +147,67 @@ 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 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 +224,7 @@ interventions: summary: Cache shared prompt prefixes to skip repeated prefill work. knob: enable_prefix_caching value: true - applies_to_kernels: [prefill, kv_cache_write] + 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 @@ -157,7 +242,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 +261,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 +279,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 +298,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 +318,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 +338,8 @@ 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-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 expected_delta_hi: 0.25 @@ -271,7 +358,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 @@ -285,11 +372,13 @@ 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 - applies_to_kernels: [decode, paged_attention] + 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 expected_delta_hi: 0.10 @@ -308,7 +397,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,7 +417,7 @@ interventions: summary: Use the Ray executor backend for multi-node distributed decode. knob: distributed_executor_backend value: ray - 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/gitm/kernels/spec.py b/gitm/kernels/spec.py index d7cab42..cfbd739 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,23 @@ 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) + # 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 (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) + value_max: float | None = None # clamp the scaled result (e.g. a fraction < 1.0) + value_min: float | None = None + # >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 expected_delta_lo: float @@ -52,3 +67,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 8c24bbe..9c9704c 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 @@ -80,24 +81,26 @@ 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", 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: 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 " @@ -118,15 +121,21 @@ def _audit( pass +def _knob_values(spec: Any) -> dict[str, Any]: + """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} + + # --- 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: @@ -247,18 +256,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. + * **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 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. @@ -278,7 +290,9 @@ 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, setter: Callable[[Any, str, Any], None] | None = None, reps: int = 1, @@ -286,17 +300,21 @@ 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) - # 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 - # 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 @@ -327,54 +345,104 @@ 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): + # 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(): + 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 - 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, values) + except Exception: + self.engine = restore_baseline() + self._prev = None + raise + else: + 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._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: 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 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 +454,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() + 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] @@ -408,9 +485,10 @@ 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=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/optimizer/deviation.py b/gitm/optimizer/deviation.py index 1b44c8d..ef71fe1 100644 --- a/gitm/optimizer/deviation.py +++ b/gitm/optimizer/deviation.py @@ -28,18 +28,22 @@ 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) 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) 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", "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/gitm/optimizer/monitor.py b/gitm/optimizer/monitor.py index 9b7cfed..e23906d 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 @@ -25,6 +26,9 @@ 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 + bound: str | None = None # the matched op's roofline bound: "compute" | "memory" @dataclass @@ -36,20 +40,30 @@ 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. + + 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 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 @@ -60,7 +74,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=None, r_kt=r_kt, r_mt=r_mt, + t_obs_s=t_obs, t_pred_s=t_pred, bound=pn.prediction.bound, + ) + ) res.serialized_concurrency_fraction = _serialized_fraction(obs) return res diff --git a/gitm/optimizer/replay.py b/gitm/optimizer/replay.py index b64e4bc..bfd7871 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,18 @@ 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 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 + 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/gitm/optimizer/vllm_knobs.py b/gitm/optimizer/vllm_knobs.py index c5a887d..6b0f70c 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 @@ -51,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( @@ -73,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" ), @@ -80,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" ), @@ -97,6 +131,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"), } @@ -149,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 @@ -185,3 +221,88 @@ 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 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"), + ("dbo", "enable_dbo"), +) + + +def unmet_prerequisite(engine: Any | None, knob: str) -> str | None: + """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: + 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" + + +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`` 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 + 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 + 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 — + 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)] + 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 a3fed2e..fd9fa3e 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 ( @@ -34,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 +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 @@ -191,16 +203,39 @@ def _model_spec_from_engine(engine: Any): return None +def _clamp_pct(value: float) -> float: + """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 — 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. - """ - kts = sorted(kr.r_kt for kr in res.per_kernel) - if not kts: + """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 - 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 _clamp_pct(value) + + +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]: @@ -438,7 +473,13 @@ 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/swept levers resolve against the live engine 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( @@ -457,22 +498,24 @@ 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( 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. + # 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: @@ -536,6 +579,94 @@ 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 unmet_prerequisite(cfg.engine, spec.knob) + + 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, res), 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" + 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})") + 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=ar_residual, + causal_evidence=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, + "apply_error": r.apply_error, + } + for r in ar_run.results + ], + }, + indent=2, + ) + ) + # Phase 5 — stabilize + write report provenance = build_provenance( workload_id=workload, @@ -572,6 +703,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"), } @@ -853,7 +986,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, @@ -992,7 +1125,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/gitm/workloads.py b/gitm/workloads.py index 236abd7..d730936 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. @@ -573,22 +579,110 @@ 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: _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) + 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} @@ -607,49 +701,51 @@ 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(_base_kwargs) # inherit the mem cap so both engines coexist - 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). + kwargs = dict(getattr(_old_engine, "gitm_llm_kwargs", _base_kwargs)) + 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()) 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 {', '.join(knob_values)} 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) + # Expose the live engine + its A/B hooks so the loop can (a) sample scheduler # stats and (b) run the Phase-4 decode-throughput A/B on it. ``run.engine`` is # picked up as ``cfg.engine``; the loop reads ``gitm_throughput_fn`` / # ``gitm_restart_fn`` off the engine (gitm.scheduler.loop Phase 4). The restart # hook is what lets structural knobs (fp8 KV cache, quantization) be *measured* - # via an engine rebuild instead of rejected. Mirrors the ``.applicator`` - # convention the hft/edge/openfold factories use. + # via an engine rebuild instead of rejected; ``gitm_baseline_restart_fn`` + # rebuilds the baseline for ``restart_mode="serial"``. Mirrors the + # ``.applicator`` convention the hft/edge/openfold factories use. run.engine = llm run.workload_id = "vllm-decode" 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_apply_rollback.py b/tests/test_apply_rollback.py index 186495a..1831d12 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): @@ -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): @@ -165,11 +200,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 new file mode 100644 index 0000000..9ca52b3 --- /dev/null +++ b/tests/test_autoresearch.py @@ -0,0 +1,1041 @@ +"""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("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. + 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" + + +# --- classify_bottleneck: roofline-weighted memory signal --------------------- + + +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 # 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 # time-weighted, not a head count + + +def test_classify_catches_intrinsic_memory_boundedness_with_no_memcpy() -> None: + """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), + 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 --------------------------------------- + + +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 + + # 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, "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: + 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(), "compute_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 == [] + + + +# --- 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 == [] + + # The idle generated partial-prefill knobs are filtered as unsupported/noisy. + + + +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, + _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, + _visible_gpu_count, +) + + +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=4), # 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=4)], 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 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 + 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 + + +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=4)], catalog_knobs=set() + ) + grid = _value_grid(Knob("cpu_offload_gb", "int", default=4)) + + 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" + + +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=4)], 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_unmet_prerequisite_vetoes_partial_prefill_without_chunked_prefill() -> None: + """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: + 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). + 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=4)], 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. + fb = FallbackProposer( + EngineArgsProposer(knobs=[Knob("cpu_offload_gb", "int", default=4)], 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=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) + + +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) + 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: + # 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 standalone optimization knobs pass. + 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 -> denylisted outright. + assert not _is_tunable("kv_sharing_fast_prefill") + # 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)" + + +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", ()) + + +# --- 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 + tensor_parallel_size: int = 1 # multi-GPU topology → skipped on a 1-GPU box + + @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=[]) + 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, gpu_count=1)} + 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, 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, 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 + + +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) == {} + + +# --- 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("compute_bound") + + +def test_candidate_specs_share_the_forced_fields() -> None: + """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=4)], catalog_knobs=set() + ).propose("memory_bound")[0] + for s in (table, generated): + assert s.safety.tier == "moderate" + assert (s.expected_delta_lo, s.expected_delta_hi) == (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"] + + +# --- ranking: delta_mean scales with a real, computable signal --------------- + + +def test_affinity_strength_and_delta_mean_for() -> None: + keywords = ("cache", "swap", "offload", "cpu") + 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 + # 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 + assert _delta_mean_for(100) == 0.15 # capped, never exceeds expected_delta_hi + + +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=4), # 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_candidates() -> None: + dbo_knobs = [ + Knob("enable_dbo", "bool", default=False), + Knob("dbo_prefill_token_threshold", "int", default=512), + ] + specs = _joint_prerequisite_candidates( + 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 + 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 + + # No enable_dbo in the surface (e.g. offline fallback catalog) -> nothing. + assert _joint_prerequisite_candidates( + [Knob("dbo_prefill_token_threshold", "int", default=512)], + bottleneck_class="idle_stall", workload="vllm-decode", target_op=None, + keywords=_CLASS_KEYWORDS["idle_stall"], + ) == [] + + # compute_bound's keywords don't match "prefill" -> no joint candidate. + assert _joint_prerequisite_candidates( + dbo_knobs, bottleneck_class="compute_bound", workload="vllm-decode", target_op=None, + keywords=_CLASS_KEYWORDS["compute_bound"], + ) == [] + + # 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_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 + + 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) + + 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 ------------------------------------ + + +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_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") 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_knobs_via_restart.py b/tests/test_knobs_via_restart.py index ff20ed5..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,38 +11,44 @@ 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(): - built: list[tuple[str, object]] = [] +def test_engine_arg_rebuilds_by_default(): + 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( 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, _k, _v: 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_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 diff --git a/tests/test_run_loop_workload.py b/tests/test_run_loop_workload.py index 5affde5..da36de8 100644 --- a/tests/test_run_loop_workload.py +++ b/tests/test_run_loop_workload.py @@ -129,6 +129,42 @@ 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 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")) + 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"] == "memory_bound" + assert s["n_autoresearch"] == 0 + + ar_json = (Path(result["run_dir"]) / "autoresearch.json").read_text(encoding="utf-8") + # 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 + + # 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 @@ -215,5 +251,5 @@ def __init__(self, max_tokens: int = 0, temperature: float = 0.0): assert runner()["generated_tokens"] == 3 * 5 assert engine.gitm_throughput_fn(engine) > 0 # The restart hook rebuilds a fresh engine with the structural knob applied. - rebuilt = engine.gitm_restart_fn(engine, "kv_cache_dtype", "fp8") + rebuilt = engine.gitm_restart_fn(engine, {"kv_cache_dtype": "fp8"}) assert rebuilt is not engine and rebuilt.kwargs["kv_cache_dtype"] == "fp8" diff --git a/tests/test_runtime_on_trace.py b/tests/test_runtime_on_trace.py index f5fec49..780f071 100644 --- a/tests/test_runtime_on_trace.py +++ b/tests/test_runtime_on_trace.py @@ -49,6 +49,43 @@ 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(): + """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 + + 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 ------------------------------------------------------ @@ -145,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 ---------------------------------------------- diff --git a/tests/test_vllm_embodiment.py b/tests/test_vllm_embodiment.py index 4fcd897..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( @@ -295,9 +300,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 @@ -307,12 +313,15 @@ 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. - assert engine.max_num_seqs == 256 + # 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()) - 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 0de7c96..a4bbd44 100644 --- a/tests/test_vllm_knobs_and_restart.py +++ b/tests/test_vllm_knobs_and_restart.py @@ -16,7 +16,14 @@ 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, + resolve_relative_value, + set_knob, + unmet_prerequisite, +) from gitm.tracer.vllm_stats import SchedulerStatsSummary @@ -32,14 +39,129 @@ 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 +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(): + gated = ("max_num_partial_prefills", "long_prefill_token_threshold", + "dbo_prefill_token_threshold") + + # 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) + + # Live flag off -> reject; on -> allowed. + off = _EngineWithFlags(chunked_prefill_enabled=False, enable_dbo=False) + on = _EngineWithFlags(chunked_prefill_enabled=True, enable_dbo=True) + for knob in gated: + assert unmet_prerequisite(off, knob) is not None + assert unmet_prerequisite(on, knob) is None + + # 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") + + +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 + + +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 @@ -82,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.""" @@ -96,22 +218,29 @@ 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(): 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) @@ -132,6 +261,68 @@ 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_values): + assert old is baseline + assert old.shutdown_called + events.append("restart:" + ",".join(f"{k}={v}" for k, v in knob_values.items())) + 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 + assert app.last_result.via == "restart" + + +def test_serial_restart_rebuilds_baseline_if_candidate_build_fails(): + baseline = _TpsEngine(100.0) + restored = _TpsEngine(100.0) + + def restart_fn(_old, _knob_values): + 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 @@ -217,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 @@ -249,8 +445,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)) @@ -261,7 +458,51 @@ 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 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 + +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) + + +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 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 diff --git a/tests/test_workload_bootstrap.py b/tests/test_workload_bootstrap.py index d197983..bb0bd38 100644 --- a/tests/test_workload_bootstrap.py +++ b/tests/test_workload_bootstrap.py @@ -96,3 +96,75 @@ 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) + 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