Refactor pipeline execution functionality to allow for multiple backends by emiehling · Pull Request #22 · generative-computing/steerability

emiehling · 2026-08-20T15:41:08Z

Add aisteer360/backends/ with a common backend interface to support (in-process) Hugging Face backend (backends/huggingface/) and a vLLM backend (backends/vllm/, offline engine and server); vLLM server functionality supported by merged vLLM-Hook changes
Implement single backend per pipeline and a deterministic acquire/release lifecycle with compatibility check before load/inference
Add offline fitting and capture sessions so activation methods are fit through the backend
Introduce an intervention representation for state controls
Add the constrained_decoding output control (automaton-based)
Unify the evaluation generation path, make the Benchmark backend-aware, and route judge metrics and Perplexity through the backend (evaluation/metrics/backend_utils.py)
RAD / SASA fixes
Support for reasoning models
Restructure gating into structured gates and relocate routing / extract pipeline internals into modules
Update docs, README, notebooks
Rerun notebooks
Add/modify tests for backend, session lifecycle, vLLM engine/plugin/release/serve, intervention IR/lowering/export, constrained-decoding, RAD, verbosity, and backend-execution
Bump version to 0.4.0
Closes Vllm support #17

Introduce a backend seam between SteeringPipeline and the runtime that executes it. The seam defines BackendSpec, capability atoms, a phase-keyed requirement language, SupportReport, ModelLayout, the SteeringSession protocol, work units, PreparedPrompt, GenerationParams, and a lazy backend registry. Pipeline generation and scoring route through backend sessions rather than a private in-process loop. Ship three backends behind the seam: - HFBackend / ExclusiveSession for the in-process Hugging Face path, kept byte-identical to the prior generation path apart from a versioned stop-string and finish-reason change recorded in the changelog. - VLLMBackend over the offline engine, with strict parameter rendering and per-item seed derivation. - VLLMServeBackend over token-id completions, with retries and PartialBatchError. Controls declare what they need through BaseControl.requirements() and a serializer-derived spec surface; the pipeline negotiates the intersection of advertised and required capability kinds and selects entries accordingly. Sessions lower steering specs, salt captures, stage structural artifacts with provenance, refuse scoring items that carry constraints, and remap scoring. Static vLLM capability tables import without vllm installed and raise only on execution; add the vllm optional extra. A pipeline-level check() reports support before any model or engine work. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Extend the backend seam to activation capture and constrained generation, so fitting and guided decoding run on any backend that advertises the capability. Capture and fitting: - VLLMOfflineSession.capture serves the capture surface with a fresh random salt per request, right-padded assembly, negotiated capture kinds, and engine-fact refusals; capture_hidden bridges fitters onto session.capture while the in-process path stays byte-identical. - Estimators, ConditionPointSelector, fit_probe, ProbeSet, and ActivationStats thread a session through. Data-fitted CAA, DirectionalAblation, and AngularSteering steer on any hidden-capture backend; ActAdd and ITI fitting stay in-process. ProbeSet.read scores through a capture session at the layer-input boundary. - RoutedDecoding rolls out over sessions with its probe pass on session.capture. - SteeringVector and probes gain provenance metadata (model, config, and chat-template fingerprints, stamped at fit); entry selection warns on a served-model fingerprint mismatch. Constrained decoding: - ConstraintSource (json_schema, regex, grammar, choice) is the portable constraint form. ConstrainedDecoding renders it per backend: a client-compiled xgrammar automaton driving ConstraintProcessor in-process (new guided extra), and the engine's native structured-output parameters on vLLM. - Add the GUIDED_DECODING atom and a ConstraintKinds set advertised by both vLLM backends. Automaton-object configurations stay in-process only. Sessions refuse scoring items carrying constraints and more than one constraint per request. Document the resulting backend compatibility matrix and update AGENTS.md for the execution layer. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Replace the imperative hook wiring of residual-stream state controls with a declarative intervention representation that the backend seam can lower to either in-process hooks or a served plugin. - Define Intervention, TokenScope, Condition, and WireForm, and give each component a self-description. Controls emit interventions from sources; lower_interventions turns them into wire forms. - Add build_hooks and split the taxonomy into InterventionControl (declare an intervention template) and HookControl (own get_hooks for mechanisms other than the residual stream). Rewrite the seven residual controls declaratively. - Carry hooks through a single session-owned registration path (SteeredSession), delete the No-op hook shims, and rewire the decoding driver against it. - Cache steer-time lowering eagerly, add a stage_artifacts protocol, and delete the old exporter. - Stage backend artifacts through a registry with a serve PUT route, and rewrite the driver rollout anchor on spec backends. - Reshape the execution module into contracts, payloads, and backend layers with ModelFacts, unify the activation collectors, and add a processor-spec rung. Close out with a batch of correctness fixes: opener tie-break, o_proj layer-0 wire form, probe-ordered condition layers, follower gate lowering, capability-gated entry collection, rollout guards for uneven batches, last_k, and conditional gates, and the custom-estimator steer need. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Bring the evaluation stack to the core's construction standard and route every benchmark generation through SteeringPipeline.generate, so exactly one prompt path exists. - UseCase declares extra constructor parameters as class-level annotations (a bare annotation is required, a class-attribute default is optional); unknown kwargs and missing required parameters raise TypeError, and schema-invalid rows raise ValueError with an index prefix. Mutable class-level defaults are copied per instance. - Delete chat_generate_model, render_inference_prompts, and chat_generate_pipeline. Every generation routes through pipeline.generate(messages=...) (or text= for template-less tokenizers), so message-level input controls fire and the pipeline owns templating, tokenization, and padding. Runtime-override columns resolve against the prompt rows, so retry and prompt expansion stay aligned. - Config identity becomes a canonical digest over the materialized pipeline (control classes and their full parameters), stable across processes, with the baseline unified on the literal "baseline". Checkpoints are a versioned envelope carrying identity metadata; resume is trial-granular and accepts only a current-format envelope whose identity matches. - Benchmark trials are reproducible via a benchmark-level seed derived per (config, trial) and threaded through gen_kwargs and use-case RNG. Backends pass through to the pipelines, with a pre-flight check() over every sweep point before any model or engine work. A shared-base fingerprint tripwire warns and reloads a clean base on detected mutation. Behavior change: benchmark results for adapt_messages controls (such as FewShot) differ because the message phase now fires; runtime-override values align per row under retry and expansion. Swap the vLLM engine test model to JackFram/llama-68m, whose head_dim clears FlexAttention's minimum. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Execute LLMJudgeMetric and Perplexity through Backend and SteeringSession rather than a private HF-only loop. Judges are configured by model reference and backend and gain a declarative class-level authoring surface, shedding the legacy configuration surface. Add backend_utils.resolve_metric_backend with a spec-keyed cache. Sampling gen_kwargs default to non-greedy. RewardScore is left on its existing path as a noted follow-up. This changes how judge metrics are authored and configured; the changelog records the migration. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Document backend usage in pipelines, installation, and quickstart; cover RUNTIME_KWARGS_SCHEMA in the tutorials; and add navigation and index entries for ActivationAdapter, AngularSteering, DirectionalAblation, and ConstrainedDecoding. Recompile the affected notebooks. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Collapse the two-backend model into a single backend per pipeline and give engine-owning backends a deterministic shutdown. - Add Backend.release(), SteeringPipeline.release_backends(), and context-manager support. VLLMBackend.release() tears down the engine and distributed state idempotently. The benchmark releases each config's backends after its trials, and a failed steer() releases the backends it constructed. - SteeringPipeline keeps a single backend (backend= replaces inference_backend=; steer_backend= is removed) plus a fit= venue policy. Each control declares its steer step's model access on a ladder (facts < rollouts < capture < module), and check() returns a deterministic steer plan alongside the generate and score verdicts. - On engine backends, stage-venued steps run on a temporary in-process model that is freed (weakref-verified) before the engine boots, with exported artifacts as the handoff. The benchmark moves to backend= and fit= with a single checkpoint identity gate. Signed-off-by: Erik Miehling <emiehling@gmail.com>

…tracts Make two contracts uniform across the Hugging Face and vLLM backends. JSON-schema constraints compile compact (any_whitespace=False) on every backend. Capture zeroes rows outside the attention mask on every surface, and the final layer_output boundary is re-captured pre-norm via a forward hook. vLLM constraint kwargs go through render_constraint_sampling_args. Add engine-gated constraint parity and capture-boundary tests. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Support reasoning models that emit a thinking span before their answer, and evaluate on the answer alone. - Add per-call chat_template_kwargs passthrough and a shared thinking splitter (split_thinking, DEFAULT_THINK_TAGS). generate_on_pipeline returns the decoded text, records, and thinking; batch_retry_generate gains return_thinking and think_tags. The commonsense_mcqa, instruction_following, and truthful_qa use cases retain a thinking column and score the answer. - Add shared-filesystem artifact-visibility verification to VLLMServeBackend: when discovery advertises an artifact registry root, stage_artifacts HEADs the server's artifact route to confirm visibility and raises a configuration error naming both roots otherwise; older servers skip the probe. - Consolidate the CAA notebooks, promoting the executed vLLM-serve walkthrough into the main notebook and dropping the standalone variant and tracked run artifacts. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Pin isort settings (black profile, line length 120) in pyproject.toml for pre-commit, editors, and manual runs, skipping the worktree directory. Run isort over the package and tests, merging fragmented imports and restoring sort order, and apply the pre-commit whitespace and end-of-file fixers. Follow with two typing corrections that leave runtime behavior unchanged: annotate the pipeline and benchmark model and tokenizer holders while keeping the factories for call sites, collapse the dataclasses import, suppress the model and tokenizer fields from SteeringPipeline's repr, and fix the generate() overloads so the positional text overloads pin attention_mask to None and two catch-all overloads accept return_output. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Move per-call machinery out of SteeringPipeline into focused modules and defer all I/O to steer(). - specs.py holds only the intervention representation; fit-time vocabulary moves to fit_specs.py and the wire compiler to lowering.py, with ScopeKindLiteral renamed ScopeKind. - The prompt front-end (resolve_generate_source, resolve_text_prompt, resolve_messages_prompt, resolve_token_prompt, prepare_inputs) moves to core/utils/generation.py as module functions with explicit keyword parameters, and the warn-once flags collapse into one holder. - Per-call payload assembly (state-entry collection, state-control lowering, rollout entries, output-control collectors, decoding-driver resolution, stack composition, scoring processors) moves to core/utils/assembly.py, keeping the id(control)-keyed semantics. - Stage-free verification and the capture smoke test move to core/execution/staging.py. - Add generate_text, generate_messages, and generate_tokens mirroring generate()'s keyword sources, with the two reserved gen_kwargs keys promoted to named parameters and per-source return types restored. - Construction performs no I/O: steer() acquires the model and tokenizer, preloaded objects are injectable via model= and tokenizer=, and misconfiguration raises at construction. The dataclass gains identity equality and a weakref slot; lazy_init is retained, deprecated, and inert. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Factor gating into one model in the state-control common library: a Gate is an operation on Evidence decided by a Rule, with affine, cosine, projected-cosine, and callable readouts. Cosine and projected-cosine gating, including CAST, lower to the served plugin's structured gate spec. Delete the previous gate, scorer, and condition classes, and rename DirectionalAblationTransform to ProjectionTransform. Relocate routing into the routed-decoding package with renamed types (ProbePredicate to Predicate, Rule to Route, RoutingRules to Router, Readout to ProbeReadings), leaving probes as measurement-only. Signed-off-by: Erik Miehling <emiehling@gmail.com>

Consolidate late structural and behavioral fixes toward the branch tip. - Fix ActAdd's extraction boundary and positional semantics: the single-pair estimator now reads the layer-input boundary where ActAdd injects (rather than layer_output), drops the fabricated BOS row, and records the extraction location, which Intervention.bind checks against the intervention boundary. AdditiveTransform's mode becomes an explicit positional flag, removing the alignment=0 decode leak, and ActAdd counts real token positions. The demonstration notebook moves to Qwen2.5-1.5B with seeded generations. - Enforce one candidate per prompt on the decoded text= and messages= return, raising a ValueError that points to return_output=True or input_ids= when num_return_sequences or n exceeds one, and document the contract. - Remove the thinking_intervention preset; the capability survives as a PhasedDecoding configuration documented in the driver docstrings. - Split backends/huggingface.py and backends/vllm.py into packages by role, with __init__ re-exporting the public surface and vllm imports kept function-local; no behavior, signature, or naming changes. - Rename the per-category _common/ component libraries to common/ across code, tests, docs, navigation, and notebooks, since it is a documented public composition surface; no runtime behavior change. - Align docs prose with the style guidelines, dissolve the backends reference page into the concepts pages, add vLLM serving sections to the trl and activation_adapter notebooks, rerun notebooks, and remove tracked test logs. Signed-off-by: Erik Miehling <emiehling@gmail.com>

…erbosity utility Rework the RAD control around a required beta and reward_model_id, dropping the legacy path and caching the reward value; converge SASA and the subspace-margin value onto the shared core probes substrate, and add fisher probe fitting with unpaired data and chunked feature extraction. Read candidate hidden states at the raw final-layer boundary and pin fit/apply boundary consistency along with the SASA wv_path compatibility matrix. Add a verbosity utility and wire it through the Hugging Face session. Rework and rerun the RAD and SASA notebook demos with fresh executed outputs.

Recovers examples/notebooks/algorithms/trl.ipynb, which was deleted in 8fe2970. Restored from 0afe8e8, the last version retaining the vLLM serving sections (the intervening a3a98b7 rerun left the file empty). Signed-off-by: Erik Miehling <emiehling@gmail.com>

ingelise

Hi @emiehling,

I have some notes on code as comments here, mainly distilled from the AI code review.

I had a few test failures on CCC - logs mentioned to you earlier.

For the AI assisted code review:

Reproduction environment

python 3.11 · venv at /workspace/.venv-rev (ephemeral)
torch 2.13.0+cpu · transformers 4.56.2 (inside the declared >=4.52,<5.0 range)
scikit-learn, pandas, pytest installed; vllm / vllm_hook_plugins / xgrammar NOT installed
trl 1.10.0 and peft 0.20.0 were OUT of range (pins are trl<0.28, peft<0.16) —
  this accounts for 3 of the 18 suite failures and nothing else

Suite command used:

python -m pytest tests/ -q --no-header -p no:cacheprovider \
  --ignore=tests/controls/test_ppo_wrapper.py --ignore=tests/controls/test_trl_release.py \
  --ignore=tests/core/test_registry.py --ignore=tests/core/test_spec_hook_equivalence.py \
  --ignore=tests/controls/test_intervention_export.py --ignore=tests/core/test_intervention_lowering.py
# -> 18 failed, 2179 passed, 727 skipped in 298s

You might consider documenting the breaking changes somewhere as well. For example common to _common.

I will continue running through the notebooks and will flag if I come across any errors.

ingelise · 2026-08-21T14:37:28Z

+        """One `model.generate` pass over all items (identical entries, one shared seed)."""
+        model = self.model
+        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
+        input_ids, attention_mask = self._stack_prompt_rows(rows)


@emiehling, There is an issue with padding here where batched generation right-pads prompts on the HF backend, which you can solve the same way you have in _score_batched by adding
input_ids, attention_mask = to_left_pad(input_ids, attention_mask)

ingelise · 2026-08-21T14:40:21Z

+        engine_kwargs = dict(spec.get_option("engine_kwargs", default={}) or {})
+        # default to a compact grammar so json constraints match the in-process automaton
+        # (disable_any_whitespace needs an explicit backend); caller kwargs win
+        engine_kwargs.setdefault("structured_outputs_config", {"disable_any_whitespace": True, "backend": "xgrammar"})


This can fail in older vllm . Your Pyproject toml has vllm>=0.8.5,<1.0.0 . Maybe you can increase it there?

ingelise · 2026-08-21T14:41:53Z

+
+BackendConfig = "BackendSpec | str | Backend | None"
+
+_METRIC_BACKENDS: dict[BackendSpec, Backend] = {}


One of the AI assisted code review suggestions:

3. Metric backends are cached process-globally with no release path evaluation/metrics/backend_utils.py:16 — _METRIC_BACKENDS is a module-level dict with no eviction and no public release API. Two consequences: - A judge or Perplexity configured with backend="vllm" boots an engine that is never shut down. Benchmark correctly calls pipeline.release_backends() between configs (benchmark.py:514), but nothing reaches the metric cache. - Because _METRIC_BACKENDS and SteeringPipeline._backends are separate dicts, a vLLM judge alongside a vLLM pipeline yields two live engines in one process — which VLLMBackend.release() explicitly documents as unsupported: "The distributed-state teardown is process-global, so release assumes no other live vLLM engine in the process." Releasing the pipeline's engine calls destroy_model_parallel() / destroy_distributed_environment() and pulls the floor out from under the still-cached judge engine. The tests already reach into the private dict (backend_utils._METRIC_BACKENDS.clear() in test_perplexity.py:188, test_base_judge.py:247), which is a good signal the public surface is missing. A release_metric_backends() plus a documented note on the single-engine constraint would close it.```

ingelise · 2026-08-21T16:14:45Z

@@ -63,3 +63,26 @@ def test_config_changes_digest(self, saved_model_dir):
        torch.manual_seed(0)
        b = tiny_llama(num_layers=2)
        assert model_fingerprint(a) != model_fingerprint(b)
+


Should this file and others use vllm_hook_plugins importorskip?

tests/internals/test_fingerprint.py

tests/core/test_spec_hook_equivalence.py (collection error)

tests/core/test_vllm_serve_backend.py (8 failures)

tests/core/test_steering_pipeline.py

tests/core/test_intervention_lowering.py (collection error)

tests/controls/test_intervention_export.py (collection error)

tests/core/test_declarative_phases.py

tests/core/test_driver_rollout_anchor.py

tests/controls/test_intervention_ir.py

ingelise · 2026-08-21T16:20:48Z

+        if spec.get_option("hook_plugin"):
+            os.environ["VLLM_HOOK_WORKER"] = "unified"
+        try:
+            self._llm = LLM(model=model_ref, **engine_kwargs)


VLLMBackend.init boots the engine here, then resolves tokenizer (208), layout (209), and discovery (214). If any of those have errors the engine will be live but the instance never return,
therefore _backend_for never records it in self._backends and steer()'s except: release_backends() cannot reach it.

You could wrap everything after LLM(...) in a try/except that calls self.release() before re-raising.

ingelise · 2026-08-21T16:23:48Z

+            ) from error
+        return adapter_name
+
+    def _verify_fingerprints(self, tokenizer_source: str) -> None:


tokenizer_source is not used here.

ingelise · 2026-08-21T16:25:46Z

+        self.spec = spec
+        self._released = False
+        require("vllm")
+        import os


This is imported at top of file already

emiehling added 14 commits August 19, 2026 06:44

emiehling requested a review from ingelise August 20, 2026 15:41

emiehling force-pushed the backend-layer branch from 569df2b to d239cca Compare August 20, 2026 19:07

emiehling force-pushed the backend-layer branch from d239cca to d691e30 Compare August 20, 2026 22:28

ingelise requested changes Aug 21, 2026

View reviewed changes

Navigation Menu

Uh oh!

Refactor pipeline execution functionality to allow for multiple backends - #22

Refactor pipeline execution functionality to allow for multiple backends#22
emiehling wants to merge 15 commits into
generative-computing:mainfrom
emiehling:backend-layer

emiehling commented Aug 20, 2026 •

edited

Loading

Uh oh!

ingelise left a comment

Uh oh!

ingelise Aug 21, 2026

Uh oh!

ingelise Aug 21, 2026

Uh oh!

ingelise Aug 21, 2026

Uh oh!

ingelise Aug 21, 2026

Uh oh!

ingelise Aug 21, 2026

Uh oh!

ingelise Aug 21, 2026

Uh oh!

ingelise Aug 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants


		BackendConfig = "BackendSpec \| str \| Backend \| None"

		_METRIC_BACKENDS: dict[BackendSpec, Backend] = {}

Uh oh!

Conversation

emiehling commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ingelise left a comment

Choose a reason for hiding this comment

Reproduction environment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

ingelise Aug 21, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

emiehling commented Aug 20, 2026 •

edited

Loading