Backends

steerability.backends

Backend implementations of the execution seam.

Each package implements the Backend and SteeringSession protocols from steerability.algorithms.core.execution for one backend family. Specs resolve to these classes through steerability.algorithms.core.execution.backend; nothing in steerability.algorithms imports this package at module level.

ExclusiveSession

The in-process session: direct model access, hook scopes, and the default decode loop.

Exposes .model for components whose requirements include Capability.IN_PROCESS_TORCH. The default decode delegates to model.generate. Items execute serially, each under its own hook registrations, which preserves in-process semantics for every entry combination.

Source code in steerability/backends/huggingface/session.py

class ExclusiveSession:
    """The in-process session: direct model access, hook scopes, and the default decode loop.

    Exposes `.model` for components whose requirements include `Capability.IN_PROCESS_TORCH`.
    The default decode delegates to `model.generate`. Items execute serially, each under its own
    hook registrations, which preserves in-process semantics for every entry combination.
    """

    def __init__(self, backend: "HFBackend") -> None:
        self._backend = backend
        self._closed = False
        self._generate_count = 0

    @property
    def closed(self) -> bool:
        """Whether the session has been closed."""
        return self._closed

    def close(self) -> None:
        """Close the session; further use raises `RuntimeError`."""
        self._closed = True

    def __enter__(self) -> "ExclusiveSession":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("This session is closed; open a new session on the backend.")

    @property
    def model(self) -> PreTrainedModel:
        """The live model.

        Raises:
            RuntimeError: If the session is closed or no model is available yet.
        """
        self._ensure_open()
        model = self._backend._model_provider()
        if model is None:
            raise RuntimeError("No model is available on this session.")
        return model

    @property
    def tokenizer(self) -> PreTrainedTokenizerBase | None:
        """The tokenizer, or None when the adopting caller has not resolved one yet."""
        self._ensure_open()
        return self._backend._tokenizer_provider()

    @property
    def layout(self) -> ModelFacts:
        """Structural facts derived from the loaded model, computed on every access so weight
        edits and model replacements are always reflected.

        `num_layers` comes from the resolved decoder layer list; `hidden_size`,
        `num_attention_heads`, and `head_dim` come from the text config (`text_config(model)`,
        the text sub-config on composite multimodal models), with a
        `hidden_size // num_attention_heads` fallback for `head_dim`; `dtype` from the model;
        `model_fingerprint` from the weight/config fingerprint; and `model_type` from the
        composite config, so a multimodal checkpoint keeps its wrapper `model_type`.
        """
        model = self.model

        from steerability.algorithms.core.internals.fingerprint import model_fingerprint

        _, layer_names = get_model_layer_list(model)
        text_cfg = text_config(model)
        hidden_size = text_cfg.hidden_size
        num_heads = getattr(text_cfg, "num_attention_heads", None)
        head_dim = getattr(text_cfg, "head_dim", None)
        if head_dim is None and num_heads:
            head_dim = hidden_size // num_heads
        return ModelFacts(
            num_layers=len(layer_names),
            hidden_size=hidden_size,
            num_attention_heads=num_heads,
            head_dim=head_dim,
            dtype=str(model.dtype).removeprefix("torch."),
            model_fingerprint=model_fingerprint(model),
            model_type=getattr(model.config, "model_type", None),
            model_ref=getattr(model, "name_or_path", None),
        )

    def _resolve_prompt_tensors(self, prompt: PreparedPrompt) -> tuple[torch.Tensor, torch.Tensor]:
        """Token ids and attention mask for one prompt, on the model device."""
        resolved = prompt.resolve_token_ids(self.tokenizer)
        device = self.model.device
        input_ids = resolved.token_ids.to(device)
        attention_mask = resolved.attention_mask
        if attention_mask is None:
            tokenizer = self.tokenizer
            if tokenizer is not None and tokenizer.pad_token_id is not None:
                attention_mask = infer_attention_mask_from_ids(input_ids, tokenizer.pad_token_id)
            else:
                attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        attention_mask = attention_mask.to(dtype=input_ids.dtype, device=device)
        return input_ids, attention_mask

    def _compose_entry_stacks(
        self, output_entries, extra_processors=(), extra_criteria=(),
    ) -> tuple[LogitsProcessorList, StoppingCriteriaList]:
        """Compose the items' stack entries, appending caller extras after entry contributions."""
        processors: list = []
        criteria: list = []
        for entry in output_entries:
            if not isinstance(entry, StackEntry):
                raise UnsupportedOperationError(
                    f"{type(entry).__name__} requires an engine-hosted processor path; the "
                    "in-process session consumes StackEntry contributions."
                )
            processors.extend(entry.logits_processors)
            criteria.extend(entry.stopping_criteria)
        processors.extend(extra_processors)
        criteria.extend(extra_criteria)
        return LogitsProcessorList(processors), StoppingCriteriaList(criteria)

    def _register_state_entries(self, model: PreTrainedModel, state_entries) -> list:
        """Attach each entry's hook specifications to `model`, returning removable handles.

        The session is the single registrar of state hooks: entries are the only carriage, and
        registration lives strictly inside the session's execution of work. Pre and forward
        hooks register with `with_kwargs=True`; backward hooks register as full backward hooks.
        If registration fails partway, handles already attached are removed before re-raising.
        """
        handles: list = []
        try:
            for entry in state_entries:
                if not isinstance(entry, HookEntry):
                    raise UnsupportedOperationError(
                        f"{type(entry).__name__} requires an intervention-capable backend; the "
                        "in-process session consumes HookEntry contributions."
                    )
                for phase in ("pre", "forward", "backward"):
                    for spec in entry.hooks.get(phase, []):
                        module = model.get_submodule(spec["module"])
                        if phase == "pre":
                            handle = module.register_forward_pre_hook(spec["hook_func"], with_kwargs=True)
                        elif phase == "forward":
                            handle = module.register_forward_hook(spec["hook_func"], with_kwargs=True)
                        else:
                            handle = module.register_full_backward_hook(spec["hook_func"])
                        handles.append(handle)
        except Exception:
            for handle in handles:
                handle.remove()
            raise
        return handles

    @contextlib.contextmanager
    def entries_applied(self, state_entries):
        """Apply state entries to the live model for the duration of the context.

        Used by the pipeline around a client-side decoding driver's `decode`, so every forward
        the driver issues on the live model, including rollouts through this session and
        auxiliary scoring passes, runs under the generation's hooks. The session owns
        registration; hooks are removed when the context exits, even on error.
        """
        handles = self._register_state_entries(self.model, state_entries)
        try:
            yield self
        finally:
            for handle in handles:
                handle.remove()

    def _seeded(self, seed: int | None):
        """A context that snapshots and restores RNG state around a seeded decode.

        The CPU generator is always covered. On CUDA models every CUDA device generator is
        covered, since sharded models may sample on a device other than the first parameter's.
        On MPS models the MPS generator is covered.
        """
        if seed is None:
            return contextlib.nullcontext()
        device = self.model.device
        if device.type == "cuda":
            return torch.random.fork_rng(devices=list(range(torch.cuda.device_count())))
        if device.type == "mps":
            return _mps_rng_fork()
        return torch.random.fork_rng(devices=[])

    def _apply_seed(self, seed: int) -> None:
        """Seed the generators covered by `_seeded` for this decode."""
        torch.default_generator.manual_seed(seed)
        device = self.model.device
        if device.type == "cuda":
            torch.cuda.manual_seed_all(seed)
        elif device.type == "mps":
            torch.mps.manual_seed(seed)

    def _item_seeds(self, items: Sequence[GenerationItem], params: GenerationParams) -> list[int | None]:
        """Effective per-item seeds.

        An item's own seed is always honored. Otherwise, under `params.seed` with
        `seed_scope="item"`, each item derives its own seed from this call's operation id and its
        index; with `seed_scope="dispatch"`, every item takes the dispatch seed, derived at index 0
        so a single-item dispatch decodes identically under either scope. A dispatch mixing explicit
        and absent item seeds falls back to per-item derivation. Unseeded calls yield None.
        """
        operation_id = f"generate-{self._generate_count}"
        if (
            params.seed is not None
            and params.seed_scope == "dispatch"
            and all(item.seed is None for item in items)
        ):
            return [derive_item_seed(params.seed, operation_id, 0)] * len(items)
        seeds: list[int | None] = []
        for index, item in enumerate(items):
            if item.seed is not None:
                seeds.append(item.seed)
            elif params.seed is not None:
                seeds.append(derive_item_seed(params.seed, operation_id, index))
            else:
                seeds.append(None)
        return seeds

    def _report_serial_fallback(self, items: Sequence[GenerationItem]) -> None:
        """Log once per backend and reason why a multi-item dispatch decodes one item at a time.

        Reached only when `batchable` is False, so identical entries imply distinct seeds.
        """
        if not self._entries_identical(items):
            key, detail = "entries", (
                "items carry distinct state or output entries (row-scoped runtime kwargs or per-row hooks)"
            )
        else:
            key, detail = "seeds", (
                "items carry distinct seeds (seed_scope='item'); pass seed_scope='dispatch' to generate(), or "
                "set ProviderOptions.seed_scope on the evaluation provider, to decode the dispatch in one pass"
            )
        if self._backend.report_once(f"serial_fallback:{key}"):
            logger.info("Multi-item generate (%d items) decodes serially: %s.", len(items), detail)

    @staticmethod
    def _entries_identical(items: Sequence[GenerationItem | ScoringItem]) -> bool:
        """True when every item carries the same state and output entry objects."""
        first = items[0]
        for item in items[1:]:
            if len(item.state_entries) != len(first.state_entries) or any(
                a is not b for a, b in zip(item.state_entries, first.state_entries)
            ):
                return False
            if len(item.output_entries) != len(first.output_entries) or any(
                a is not b for a, b in zip(item.output_entries, first.output_entries)
            ):
                return False
        return True

    def _stack_prompt_rows(self, rows: list[tuple[torch.Tensor, torch.Tensor]]):
        """Stack resolved single-row prompts into one right-padded batch.

        The batched generate and score paths both left-pack the result via `to_left_pad` before
        the forward pass.
        """
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or 0
        max_len = max(ids.size(1) for ids, _ in rows)
        device = rows[0][0].device
        input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
        attention_mask = torch.zeros((len(rows), max_len), dtype=rows[0][1].dtype, device=device)
        for row, (ids, mask) in enumerate(rows):
            length = ids.size(1)
            input_ids[row, :length] = ids[0]
            attention_mask[row, :length] = mask[0]
        return input_ids, attention_mask

    def _classify(self, new_tokens: torch.Tensor, gen_kwargs: dict, params: GenerationParams) -> list[str | None]:
        """Per-row finish reasons under the pinned precedence, from the composed stop rules."""
        tokenizer = self.tokenizer
        return infer_finish_reasons(
            new_tokens,
            gen_kwargs,
            eos_token_id=getattr(tokenizer, "eos_token_id", None),
            pad_token_id=getattr(tokenizer, "pad_token_id", None),
            stop_strings=params.stop_strings,
            stop_token_ids=params.stop_token_ids,
            tokenizer=tokenizer,
        )

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item, each under its own hook registrations.

        Items sharing identical state entries, identical output entries, and identical-or-absent
        effective seeds execute in one batched `model.generate` pass (right-padded to a common
        prompt length, then left-packed together); otherwise items decode serially. Caller-supplied `logits_processor` and
        `stopping_criteria` entries in `params.extra` append after the items' own contributions,
        and the normalized stop fields compose as stop rules anchored at the prompt length. A
        seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the
        covered generator state (CPU, plus the model device's) is restored afterwards; when
        `params.seed` is set and an item carries no seed of its own, the item's seed derives per
        index, so multi-item fan-outs sample distinct streams.

        Args:
            items: The generation items.
            params: Normalized generation parameters shared by all items.

        Returns:
            One `ItemResult` per item, in item order. Each result's `finish_reasons` carries one
            reason per candidate with the precedence stop, then eos, then length, then None.
        """
        self._ensure_open()
        model = self.model
        tokenizer = self.tokenizer
        gen_kwargs = render_hf_gen_kwargs(params)
        # default pad_token_id per call to avoid the transformers open-end-generation warning
        # without mutating the model's generation config; caller kwargs and a model-configured
        # value both take precedence over the tokenizer fallback
        if (
            "pad_token_id" not in gen_kwargs
            and model.generation_config.pad_token_id is None
            and getattr(tokenizer, "pad_token_id", None) is not None
        ):
            gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
        user_processors = tuple(gen_kwargs.pop("logits_processor", None) or ())
        user_criteria = tuple(gen_kwargs.pop("stopping_criteria", None) or ())

        if not items:
            return []
        seeds = self._item_seeds(items, params)
        self._generate_count += 1

        batchable = (
            len(items) > 1
            and self._entries_identical(items)
            and len(set(seeds)) == 1
        )
        if batchable:
            return self._generate_batched(
                items, params, gen_kwargs, user_processors, user_criteria, seeds[0],
            )
        if len(items) > 1:
            self._report_serial_fallback(items)

        results: list[ItemResult] = []
        for index, item in enumerate(items):
            input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
            processors, criteria = self._compose_entry_stacks(
                item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
            )
            criteria.extend(compose_stop_criteria(params, input_ids.size(1), tokenizer))
            stacks = stack_generate_kwargs(processors, criteria)
            handles = self._register_state_entries(model, item.state_entries)
            try:
                seed = seeds[index]
                with self._seeded(seed):
                    if seed is not None:
                        self._apply_seed(seed)
                    full_ids = model.generate(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        **stacks,
                        **gen_kwargs,
                    )
            finally:
                for handle in handles:
                    handle.remove()

            new_tokens = full_ids[:, input_ids.size(1):]
            reasons = self._classify(new_tokens, gen_kwargs, params)
            results.append(ItemResult(
                index=index,
                output=Output(
                    output_ids=new_tokens,
                    adapted_input_ids=input_ids,
                    finish_reason=reasons[0],
                    finish_reasons=tuple(reasons),
                ),
            ))
        return results

    def _generate_batched(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
        gen_kwargs: dict,
        user_processors: tuple,
        user_criteria: tuple,
        seed: int | None,
    ) -> list[ItemResult]:
        """One `model.generate` pass over all items (identical entries, one shared seed).

        The stacked prompts left-pack (pad positions move before the real tokens), so every row's
        first generated token is predicted from its last real token rather than from a trailing pad.
        """
        model = self.model
        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
        input_ids, attention_mask = self._stack_prompt_rows(rows)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
        processors, criteria = self._compose_entry_stacks(
            items[0].output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
        )
        criteria.extend(compose_stop_criteria(params, input_ids.size(1), self.tokenizer))
        stacks = stack_generate_kwargs(processors, criteria)
        handles = self._register_state_entries(model, items[0].state_entries)
        try:
            with self._seeded(seed):
                if seed is not None:
                    self._apply_seed(seed)
                full_ids = model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **stacks,
                    **gen_kwargs,
                )
        finally:
            for handle in handles:
                handle.remove()

        prompt_len = input_ids.size(1)
        new_tokens = full_ids[:, prompt_len:]
        num_candidates = params.n or 1
        results: list[ItemResult] = []
        for index in range(len(items)):
            item_rows = new_tokens[index * num_candidates:(index + 1) * num_candidates]
            reasons = self._classify(item_rows, gen_kwargs, params)
            results.append(ItemResult(
                index=index,
                output=Output(
                    output_ids=item_rows,
                    adapted_input_ids=input_ids[index:index + 1],
                    finish_reason=reasons[0],
                    finish_reasons=tuple(reasons),
                ),
            ))
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens.

        For each item, the prompt left-packs (pad positions move before the real tokens) so the
        reference follows the prompt's last real token, then prompt and reference concatenate
        into one causal forward pass under the item's hook registrations; the item's logits
        processors replay position-by-position with the same `(prefix_ids, scores)` view they
        receive during generation. Items sharing identical state and output entries score in one
        batched forward pass (prompts right-padded to a common length, then left-packed
        together); otherwise items score serially. Stopping criteria never apply. Decoder-only
        models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: `params.extra` passes through as forward keyword arguments.

        Returns:
            Log probabilities of shape `[num_items, ref_len]`.

        Raises:
            ValueError: If items carry differing reference lengths.
            UnsupportedOperationError: If the model is an encoder-decoder model.
        """
        self._ensure_open()
        model = self.model
        if getattr(model.config, "is_encoder_decoder", False):
            raise UnsupportedOperationError(
                "Session scoring supports decoder-only models; encoder-decoder scoring runs "
                "through SteeringPipeline.compute_logprobs."
            )
        device = model.device
        forward_kwargs = dict(params.extra)

        if not items:
            return torch.zeros((0, 0), device=device, dtype=torch.float32)
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")

        if len(items) > 1 and self._entries_identical(items):
            return self._score_batched(items, forward_kwargs)

        all_logprobs: list[torch.Tensor] = []
        for item in items:
            input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
            input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
            ref_output_ids = item.ref_output_ids
            if ref_output_ids.ndim == 1:
                ref_output_ids = ref_output_ids.unsqueeze(0)
            ref_output_ids = ref_output_ids.to(device)
            ref_len = ref_output_ids.size(1)
            if ref_len == 0:
                all_logprobs.append(torch.zeros((1, 0), device=device, dtype=torch.float32))
                continue

            processors, _ = self._compose_entry_stacks(item.output_entries)
            handles = self._register_state_entries(model, item.state_entries)
            try:
                with torch.no_grad():
                    combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                    combined_mask = torch.cat([
                        attention_mask,
                        torch.ones(1, ref_len, device=device, dtype=attention_mask.dtype),
                    ], dim=1)
                    outputs = model(
                        input_ids=combined_ids,
                        attention_mask=combined_mask,
                        **forward_kwargs,
                    )
                    input_len = input_ids.size(1)
                    logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                    if len(processors):
                        for t in range(logits.size(1)):
                            prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                            logits[:, t, :] = processors(prefix, logits[:, t, :])
                    logprobs = torch.log_softmax(logits, dim=-1)
                    all_logprobs.append(
                        logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
                    )
            finally:
                for handle in handles:
                    handle.remove()

        return torch.cat(all_logprobs, dim=0)

    def _score_batched(self, items: Sequence[ScoringItem], forward_kwargs: dict) -> torch.Tensor:
        """One causal forward pass over all items (identical entries)."""
        model = self.model
        device = model.device
        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
        input_ids, attention_mask = self._stack_prompt_rows(rows)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)

        refs: list[torch.Tensor] = []
        for item in items:
            ref = item.ref_output_ids
            if ref.ndim == 1:
                ref = ref.unsqueeze(0)
            refs.append(ref.to(device))
        ref_output_ids = torch.cat(refs, dim=0)
        ref_len = ref_output_ids.size(1)
        if ref_len == 0:
            return torch.zeros((len(items), 0), device=device, dtype=torch.float32)

        processors, _ = self._compose_entry_stacks(items[0].output_entries)
        handles = self._register_state_entries(model, items[0].state_entries)
        try:
            with torch.no_grad():
                combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                combined_mask = torch.cat([
                    attention_mask,
                    torch.ones(len(items), ref_len, device=device, dtype=attention_mask.dtype),
                ], dim=1)
                outputs = model(
                    input_ids=combined_ids,
                    attention_mask=combined_mask,
                    **forward_kwargs,
                )
                input_len = input_ids.size(1)
                logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                if len(processors):
                    for t in range(logits.size(1)):
                        prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                        logits[:, t, :] = processors(prefix, logits[:, t, :])
                logprobs = torch.log_softmax(logits, dim=-1)
                return logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
        finally:
            for handle in handles:
                handle.remove()

    def capture(
        self,
        prompts: list[PreparedPrompt],
        layers: list[int],
        mode: Literal["all_tokens", "last_token"],
        location: Literal["layer_output", "layer_input"] = "layer_output",
    ) -> CaptureResult:
        """Capture residual-stream hidden states for `prompts` at `layers`.

        Prompts resolve to token ids, right-pad into one batch, and run through the shared
        layerwise extraction. In `"all_tokens"` mode each layer's tensor is `[N, T, H]` with rows
        outside the attention mask zeroed; in `"last_token"` mode the last real (non-pad) position
        of each row is selected, giving `[N, H]`.

        Args:
            prompts: The prompts to capture.
            layers: 0-based layer ids to keep.
            mode: `"all_tokens"` or `"last_token"`.
            location: `"layer_output"` (a layer's output boundary) or `"layer_input"` (the
                boundary a forward pre-hook observes).

        Returns:
            The captured tensors and the batch attention mask, on CPU.

        Raises:
            ValueError: If `mode` is unknown or a requested layer id is out of range.
        """
        self._ensure_open()
        if mode not in ("all_tokens", "last_token"):
            raise ValueError(f"Unknown capture mode {mode!r}; modes are 'all_tokens', 'last_token'.")
        if not prompts:
            raise ValueError("capture() requires at least one prompt.")

        from steerability.algorithms.core.internals.capture import layerwise_tokenwise_hidden
        from steerability.algorithms.core.internals.pooling import aggregate_condition_hidden

        model = self.model
        device = model.device
        tokenizer = self.tokenizer
        pad_token_id = getattr(tokenizer, "pad_token_id", None) or 0

        rows = [self._resolve_prompt_tensors(prompt) for prompt in prompts]
        max_len = max(ids.size(1) for ids, _ in rows)
        input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
        attention_mask = torch.zeros((len(rows), max_len), dtype=torch.long, device=device)
        for row, (ids, mask) in enumerate(rows):
            length = ids.size(1)
            input_ids[row, :length] = ids[0]
            attention_mask[row, :length] = mask[0]

        enc = {"input_ids": input_ids, "attention_mask": attention_mask}
        hidden = layerwise_tokenwise_hidden(
            model, enc, batch_size=_CAPTURE_BATCH_SIZE, location=location,
        )
        missing = [layer for layer in layers if layer not in hidden]
        if missing:
            raise ValueError(
                f"Requested layer ids {missing} are out of range; the model has {len(hidden)} layers."
            )

        mask_cpu = attention_mask.cpu()
        selected = {layer: hidden[layer] for layer in layers}
        if mode == "last_token":
            selected = {
                layer: aggregate_condition_hidden(tensor, "last", mask_cpu)
                for layer, tensor in selected.items()
            }
        return CaptureResult(
            hidden=selected, attention_mask=mask_cpu, mode=mode, location=location,
        )

closed property

Whether the session has been closed.

layout property

Structural facts derived from the loaded model, computed on every access so weight edits and model replacements are always reflected.

num_layers comes from the resolved decoder layer list; hidden_size, num_attention_heads, and head_dim come from the text config (text_config(model), the text sub-config on composite multimodal models), with a hidden_size // num_attention_heads fallback for head_dim; dtype from the model; model_fingerprint from the weight/config fingerprint; and model_type from the composite config, so a multimodal checkpoint keeps its wrapper model_type.

model property

The live model.

Raises:

Type	Description
RuntimeError	If the session is closed or no model is available yet.

tokenizer property

The tokenizer, or None when the adopting caller has not resolved one yet.

capture(prompts, layers, mode, location='layer_output')

Capture residual-stream hidden states for prompts at layers.

Prompts resolve to token ids, right-pad into one batch, and run through the shared layerwise extraction. In "all_tokens" mode each layer's tensor is [N, T, H] with rows outside the attention mask zeroed; in "last_token" mode the last real (non-pad) position of each row is selected, giving [N, H].

Parameters:

Name	Type	Description	Default
prompts	list[PreparedPrompt]	The prompts to capture.	required
layers	list[int]	0-based layer ids to keep.	required
mode	Literal['all_tokens', 'last_token']	"all_tokens" or "last_token".	required
location	Literal['layer_output', 'layer_input']	"layer_output" (a layer's output boundary) or "layer_input" (the boundary a forward pre-hook observes).	'layer_output'

Returns:

Type	Description
CaptureResult	The captured tensors and the batch attention mask, on CPU.

Raises:

Type	Description
ValueError	If mode is unknown or a requested layer id is out of range.

Source code in steerability/backends/huggingface/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Capture residual-stream hidden states for `prompts` at `layers`.

    Prompts resolve to token ids, right-pad into one batch, and run through the shared
    layerwise extraction. In `"all_tokens"` mode each layer's tensor is `[N, T, H]` with rows
    outside the attention mask zeroed; in `"last_token"` mode the last real (non-pad) position
    of each row is selected, giving `[N, H]`.

    Args:
        prompts: The prompts to capture.
        layers: 0-based layer ids to keep.
        mode: `"all_tokens"` or `"last_token"`.
        location: `"layer_output"` (a layer's output boundary) or `"layer_input"` (the
            boundary a forward pre-hook observes).

    Returns:
        The captured tensors and the batch attention mask, on CPU.

    Raises:
        ValueError: If `mode` is unknown or a requested layer id is out of range.
    """
    self._ensure_open()
    if mode not in ("all_tokens", "last_token"):
        raise ValueError(f"Unknown capture mode {mode!r}; modes are 'all_tokens', 'last_token'.")
    if not prompts:
        raise ValueError("capture() requires at least one prompt.")

    from steerability.algorithms.core.internals.capture import layerwise_tokenwise_hidden
    from steerability.algorithms.core.internals.pooling import aggregate_condition_hidden

    model = self.model
    device = model.device
    tokenizer = self.tokenizer
    pad_token_id = getattr(tokenizer, "pad_token_id", None) or 0

    rows = [self._resolve_prompt_tensors(prompt) for prompt in prompts]
    max_len = max(ids.size(1) for ids, _ in rows)
    input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
    attention_mask = torch.zeros((len(rows), max_len), dtype=torch.long, device=device)
    for row, (ids, mask) in enumerate(rows):
        length = ids.size(1)
        input_ids[row, :length] = ids[0]
        attention_mask[row, :length] = mask[0]

    enc = {"input_ids": input_ids, "attention_mask": attention_mask}
    hidden = layerwise_tokenwise_hidden(
        model, enc, batch_size=_CAPTURE_BATCH_SIZE, location=location,
    )
    missing = [layer for layer in layers if layer not in hidden]
    if missing:
        raise ValueError(
            f"Requested layer ids {missing} are out of range; the model has {len(hidden)} layers."
        )

    mask_cpu = attention_mask.cpu()
    selected = {layer: hidden[layer] for layer in layers}
    if mode == "last_token":
        selected = {
            layer: aggregate_condition_hidden(tensor, "last", mask_cpu)
            for layer, tensor in selected.items()
        }
    return CaptureResult(
        hidden=selected, attention_mask=mask_cpu, mode=mode, location=location,
    )

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/huggingface/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

entries_applied(state_entries)

Apply state entries to the live model for the duration of the context.

Used by the pipeline around a client-side decoding driver's decode, so every forward the driver issues on the live model, including rollouts through this session and auxiliary scoring passes, runs under the generation's hooks. The session owns registration; hooks are removed when the context exits, even on error.

Source code in steerability/backends/huggingface/session.py

@contextlib.contextmanager
def entries_applied(self, state_entries):
    """Apply state entries to the live model for the duration of the context.

    Used by the pipeline around a client-side decoding driver's `decode`, so every forward
    the driver issues on the live model, including rollouts through this session and
    auxiliary scoring passes, runs under the generation's hooks. The session owns
    registration; hooks are removed when the context exits, even on error.
    """
    handles = self._register_state_entries(self.model, state_entries)
    try:
        yield self
    finally:
        for handle in handles:
            handle.remove()

generate(items, params)

Generate one result per item, each under its own hook registrations.

Items sharing identical state entries, identical output entries, and identical-or-absent effective seeds execute in one batched model.generate pass (right-padded to a common prompt length, then left-packed together); otherwise items decode serially. Caller-supplied logits_processor and stopping_criteria entries in params.extra append after the items' own contributions, and the normalized stop fields compose as stop rules anchored at the prompt length. A seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the covered generator state (CPU, plus the model device's) is restored afterwards; when params.seed is set and an item carries no seed of its own, the item's seed derives per index, so multi-item fan-outs sample distinct streams.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items.	required
params	GenerationParams	Normalized generation parameters shared by all items.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order. Each result's finish_reasons carries one
list[ItemResult]	reason per candidate with the precedence stop, then eos, then length, then None.

Source code in steerability/backends/huggingface/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item, each under its own hook registrations.

    Items sharing identical state entries, identical output entries, and identical-or-absent
    effective seeds execute in one batched `model.generate` pass (right-padded to a common
    prompt length, then left-packed together); otherwise items decode serially. Caller-supplied `logits_processor` and
    `stopping_criteria` entries in `params.extra` append after the items' own contributions,
    and the normalized stop fields compose as stop rules anchored at the prompt length. A
    seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the
    covered generator state (CPU, plus the model device's) is restored afterwards; when
    `params.seed` is set and an item carries no seed of its own, the item's seed derives per
    index, so multi-item fan-outs sample distinct streams.

    Args:
        items: The generation items.
        params: Normalized generation parameters shared by all items.

    Returns:
        One `ItemResult` per item, in item order. Each result's `finish_reasons` carries one
        reason per candidate with the precedence stop, then eos, then length, then None.
    """
    self._ensure_open()
    model = self.model
    tokenizer = self.tokenizer
    gen_kwargs = render_hf_gen_kwargs(params)
    # default pad_token_id per call to avoid the transformers open-end-generation warning
    # without mutating the model's generation config; caller kwargs and a model-configured
    # value both take precedence over the tokenizer fallback
    if (
        "pad_token_id" not in gen_kwargs
        and model.generation_config.pad_token_id is None
        and getattr(tokenizer, "pad_token_id", None) is not None
    ):
        gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
    user_processors = tuple(gen_kwargs.pop("logits_processor", None) or ())
    user_criteria = tuple(gen_kwargs.pop("stopping_criteria", None) or ())

    if not items:
        return []
    seeds = self._item_seeds(items, params)
    self._generate_count += 1

    batchable = (
        len(items) > 1
        and self._entries_identical(items)
        and len(set(seeds)) == 1
    )
    if batchable:
        return self._generate_batched(
            items, params, gen_kwargs, user_processors, user_criteria, seeds[0],
        )
    if len(items) > 1:
        self._report_serial_fallback(items)

    results: list[ItemResult] = []
    for index, item in enumerate(items):
        input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
        processors, criteria = self._compose_entry_stacks(
            item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
        )
        criteria.extend(compose_stop_criteria(params, input_ids.size(1), tokenizer))
        stacks = stack_generate_kwargs(processors, criteria)
        handles = self._register_state_entries(model, item.state_entries)
        try:
            seed = seeds[index]
            with self._seeded(seed):
                if seed is not None:
                    self._apply_seed(seed)
                full_ids = model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **stacks,
                    **gen_kwargs,
                )
        finally:
            for handle in handles:
                handle.remove()

        new_tokens = full_ids[:, input_ids.size(1):]
        reasons = self._classify(new_tokens, gen_kwargs, params)
        results.append(ItemResult(
            index=index,
            output=Output(
                output_ids=new_tokens,
                adapted_input_ids=input_ids,
                finish_reason=reasons[0],
                finish_reasons=tuple(reasons),
            ),
        ))
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens.

For each item, the prompt left-packs (pad positions move before the real tokens) so the reference follows the prompt's last real token, then prompt and reference concatenate into one causal forward pass under the item's hook registrations; the item's logits processors replay position-by-position with the same (prefix_ids, scores) view they receive during generation. Items sharing identical state and output entries score in one batched forward pass (prompts right-padded to a common length, then left-packed together); otherwise items score serially. Stopping criteria never apply. Decoder-only models only; the pipeline's compute_logprobs serves encoder-decoder models in-process.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	params.extra passes through as forward keyword arguments.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len].

Raises:

Type	Description
ValueError	If items carry differing reference lengths.
UnsupportedOperationError	If the model is an encoder-decoder model.

Source code in steerability/backends/huggingface/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens.

    For each item, the prompt left-packs (pad positions move before the real tokens) so the
    reference follows the prompt's last real token, then prompt and reference concatenate
    into one causal forward pass under the item's hook registrations; the item's logits
    processors replay position-by-position with the same `(prefix_ids, scores)` view they
    receive during generation. Items sharing identical state and output entries score in one
    batched forward pass (prompts right-padded to a common length, then left-packed
    together); otherwise items score serially. Stopping criteria never apply. Decoder-only
    models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: `params.extra` passes through as forward keyword arguments.

    Returns:
        Log probabilities of shape `[num_items, ref_len]`.

    Raises:
        ValueError: If items carry differing reference lengths.
        UnsupportedOperationError: If the model is an encoder-decoder model.
    """
    self._ensure_open()
    model = self.model
    if getattr(model.config, "is_encoder_decoder", False):
        raise UnsupportedOperationError(
            "Session scoring supports decoder-only models; encoder-decoder scoring runs "
            "through SteeringPipeline.compute_logprobs."
        )
    device = model.device
    forward_kwargs = dict(params.extra)

    if not items:
        return torch.zeros((0, 0), device=device, dtype=torch.float32)
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")

    if len(items) > 1 and self._entries_identical(items):
        return self._score_batched(items, forward_kwargs)

    all_logprobs: list[torch.Tensor] = []
    for item in items:
        input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
        ref_output_ids = item.ref_output_ids
        if ref_output_ids.ndim == 1:
            ref_output_ids = ref_output_ids.unsqueeze(0)
        ref_output_ids = ref_output_ids.to(device)
        ref_len = ref_output_ids.size(1)
        if ref_len == 0:
            all_logprobs.append(torch.zeros((1, 0), device=device, dtype=torch.float32))
            continue

        processors, _ = self._compose_entry_stacks(item.output_entries)
        handles = self._register_state_entries(model, item.state_entries)
        try:
            with torch.no_grad():
                combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                combined_mask = torch.cat([
                    attention_mask,
                    torch.ones(1, ref_len, device=device, dtype=attention_mask.dtype),
                ], dim=1)
                outputs = model(
                    input_ids=combined_ids,
                    attention_mask=combined_mask,
                    **forward_kwargs,
                )
                input_len = input_ids.size(1)
                logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                if len(processors):
                    for t in range(logits.size(1)):
                        prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                        logits[:, t, :] = processors(prefix, logits[:, t, :])
                logprobs = torch.log_softmax(logits, dim=-1)
                all_logprobs.append(
                    logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
                )
        finally:
            for handle in handles:
                handle.remove()

    return torch.cat(all_logprobs, dim=0)

HFBackend

Bases: Backend

The in-process Hugging Face backend.

Owns a loaded model and tokenizer, either loaded from a spec or adopted from a caller that already holds them. At most one session may be open per backend at a time, so the backend runs one generation at a time.

Source code in steerability/backends/huggingface/backend.py

class HFBackend(Backend):
    """The in-process Hugging Face backend.

    Owns a loaded model and tokenizer, either loaded from a spec or adopted from a caller that
    already holds them. At most one session may be open per backend at a time, so the backend
    runs one generation at a time.
    """

    def __init__(
        self,
        spec: BackendSpec,
        *,
        model_provider: Callable[[], PreTrainedModel | None] | None = None,
        tokenizer_provider: Callable[[], object | None] | None = None,
    ) -> None:
        """Construct the backend, loading the model from `spec` unless providers are given.

        Loading reads the options `hf_model_kwargs`, `device_map`, `tokenizer_name_or_path`,
        and `trust_remote_code`. Option values must be plain data, since spec canonicalization
        renders live objects (e.g. a quantization config instance) as strings that
        `from_pretrained` cannot consume. A `device_map` key inside `hf_model_kwargs` is used
        when the spec carries no top-level `device_map` option.

        Args:
            spec: The backend spec.
            model_provider: Callable returning the adopted model; used with
                `tokenizer_provider` instead of loading.
            tokenizer_provider: Callable returning the adopted tokenizer.

        Raises:
            ValueError: If `spec.kind` is not `"huggingface"`, or no model reference is
                available to load from.
        """
        if spec.kind != "huggingface":
            raise ValueError(f"HFBackend requires a 'huggingface' spec; got kind {spec.kind!r}.")
        self.spec = spec
        self._open_session: ExclusiveSession | None = None
        self._reported_once: set[str] = set()

        if model_provider is not None:
            self._model_provider = model_provider
            self._tokenizer_provider = tokenizer_provider or (lambda: None)
            return

        if spec.model is None:
            raise ValueError(
                "HFBackend needs a model reference on the spec, or model_provider/"
                "tokenizer_provider for an already-loaded model."
            )
        hf_model_kwargs = dict(spec.get_option("hf_model_kwargs", default={}))
        device_map = spec.get_option("device_map", default=hf_model_kwargs.pop("device_map", "auto"))
        model = AutoModelForCausalLM.from_pretrained(
            spec.model,
            device_map=device_map,
            **hf_model_kwargs,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            spec.get_option("tokenizer_name_or_path") or spec.model,
            trust_remote_code=bool(spec.get_option("trust_remote_code", default=False)),
        )
        tokenizer = ensure_pad_token(tokenizer)
        self._model_provider = lambda: model
        self._tokenizer_provider = lambda: tokenizer

    @classmethod
    def adopt(
        cls,
        spec: BackendSpec,
        model_provider: Callable[[], PreTrainedModel | None],
        tokenizer_provider: Callable[[], object | None],
    ) -> "HFBackend":
        """Wrap an already-loaded model and tokenizer without loading anything.

        Providers are read on every access, so a caller whose model is replaced mid-steer (a
        structural control returning a new model) always exposes the current one to sessions.

        Args:
            spec: The backend spec identifying this configuration.
            model_provider: Callable returning the current model (may return None before one
                exists).
            tokenizer_provider: Callable returning the current tokenizer.

        Returns:
            The adopting backend.
        """
        return cls(spec, model_provider=model_provider, tokenizer_provider=tokenizer_provider)

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The static Hugging Face capability advertisement (spec-independent)."""
        return HF_CAPABILITIES

    def open_session(self) -> "ExclusiveSession":
        """Open the backend's one exclusive session.

        Returns:
            The session, usable as a context manager.

        Raises:
            RuntimeError: If an exclusive session is already open on this backend.
        """
        if self._open_session is not None and not self._open_session.closed:
            raise RuntimeError(
                "An exclusive session is already open on this backend; close it before opening "
                "another."
            )
        self._open_session = ExclusiveSession(self)
        return self._open_session

    def report_once(self, key: str) -> bool:
        """Return True the first time `key` is seen on this backend, False afterwards.

        Sessions are opened per pipeline call, so a per-session flag would fire on every call;
        this guard lives on the backend so a repeated informational readout fires once.
        """
        if key in self._reported_once:
            return False
        self._reported_once.add(key)
        return True

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

adopt(spec, model_provider, tokenizer_provider) classmethod

Wrap an already-loaded model and tokenizer without loading anything.

Providers are read on every access, so a caller whose model is replaced mid-steer (a structural control returning a new model) always exposes the current one to sessions.

Parameters:

Name	Type	Description	Default
spec	BackendSpec	The backend spec identifying this configuration.	required
model_provider	Callable[[], PreTrainedModel | None]	Callable returning the current model (may return None before one exists).	required
tokenizer_provider	Callable[[], object | None]	Callable returning the current tokenizer.	required

Returns:

Type	Description
HFBackend	The adopting backend.

Source code in steerability/backends/huggingface/backend.py

@classmethod
def adopt(
    cls,
    spec: BackendSpec,
    model_provider: Callable[[], PreTrainedModel | None],
    tokenizer_provider: Callable[[], object | None],
) -> "HFBackend":
    """Wrap an already-loaded model and tokenizer without loading anything.

    Providers are read on every access, so a caller whose model is replaced mid-steer (a
    structural control returning a new model) always exposes the current one to sessions.

    Args:
        spec: The backend spec identifying this configuration.
        model_provider: Callable returning the current model (may return None before one
            exists).
        tokenizer_provider: Callable returning the current tokenizer.

    Returns:
        The adopting backend.
    """
    return cls(spec, model_provider=model_provider, tokenizer_provider=tokenizer_provider)

capabilities_for_spec(spec) classmethod

The static Hugging Face capability advertisement (spec-independent).

Source code in steerability/backends/huggingface/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The static Hugging Face capability advertisement (spec-independent)."""
    return HF_CAPABILITIES

open_session()

Open the backend's one exclusive session.

Returns:

Type	Description
ExclusiveSession	The session, usable as a context manager.

Raises:

Type	Description
RuntimeError	If an exclusive session is already open on this backend.

Source code in steerability/backends/huggingface/backend.py

def open_session(self) -> "ExclusiveSession":
    """Open the backend's one exclusive session.

    Returns:
        The session, usable as a context manager.

    Raises:
        RuntimeError: If an exclusive session is already open on this backend.
    """
    if self._open_session is not None and not self._open_session.closed:
        raise RuntimeError(
            "An exclusive session is already open on this backend; close it before opening "
            "another."
        )
    self._open_session = ExclusiveSession(self)
    return self._open_session

release()

Free the resources this backend owns beyond what the caller owns.

The default is a no-op. Engine-owning backends override this to shut their engines down deterministically. Release is idempotent.

Source code in steerability/algorithms/core/execution/backend.py

def release(self) -> None:
    """Free the resources this backend owns beyond what the caller owns.

    The default is a no-op. Engine-owning backends override this to shut their engines down
    deterministically. Release is idempotent.
    """
    return None

report_once(key)

Return True the first time key is seen on this backend, False afterwards.

Sessions are opened per pipeline call, so a per-session flag would fire on every call; this guard lives on the backend so a repeated informational readout fires once.

Source code in steerability/backends/huggingface/backend.py

def report_once(self, key: str) -> bool:
    """Return True the first time `key` is seen on this backend, False afterwards.

    Sessions are opened per pipeline call, so a per-session flag would fire on every call;
    this guard lives on the backend so a repeated informational readout fires once.
    """
    if key in self._reported_once:
        return False
    self._reported_once.add(key)
    return True

stage_artifacts(payloads)

Make each content-addressed artifact available to the execution side.

Called by the pipeline at the end of steer() with the tensor payloads of every lowered intervention spec, keyed by content-addressed artifact id. Staging is idempotent: an artifact that already exists at the destination is success. The in-process backend keeps live tensors and needs no staging; engine-backed backends write into the registry the serving engine reads.

Parameters:

Name	Type	Description	Default
payloads	dict[str, dict[str, Tensor]]	Mapping from artifact id to a name-to-tensor mapping.	required

Source code in steerability/algorithms/core/execution/backend.py

def stage_artifacts(self, payloads: dict[str, dict[str, torch.Tensor]]) -> None:
    """Make each content-addressed artifact available to the execution side.

    Called by the pipeline at the end of `steer()` with the tensor payloads of every
    lowered intervention spec, keyed by content-addressed artifact id. Staging is
    idempotent: an artifact that already exists at the destination is success. The
    in-process backend keeps live tensors and needs no staging; engine-backed backends
    write into the registry the serving engine reads.

    Args:
        payloads: Mapping from artifact id to a name-to-tensor mapping.
    """
    return None

huggingface

The in-process Hugging Face backend and its exclusive session.

backend.py holds HFBackend and its static capability advertisement; session.py holds ExclusiveSession (direct model access, hook scopes, and the default decode loop) and the generation-parameter rendering helpers. Constructing HFBackend loads a model and tokenizer in the current process.

HF_CAPABILITIES = BackendCapabilities(atoms=frozenset({Capability.IN_PROCESS_TORCH, Capability.HIDDEN_CAPTURE, Capability.BEAM_PROPOSALS}), capture_kinds=CaptureKinds(kinds=frozenset({'residual'}), locations=frozenset({'layer_output', 'layer_input'}), modes=frozenset({'all_tokens', 'last_token'}))) module-attribute

ExclusiveSession

The in-process session: direct model access, hook scopes, and the default decode loop.

Exposes .model for components whose requirements include Capability.IN_PROCESS_TORCH. The default decode delegates to model.generate. Items execute serially, each under its own hook registrations, which preserves in-process semantics for every entry combination.

Source code in steerability/backends/huggingface/session.py

class ExclusiveSession:
    """The in-process session: direct model access, hook scopes, and the default decode loop.

    Exposes `.model` for components whose requirements include `Capability.IN_PROCESS_TORCH`.
    The default decode delegates to `model.generate`. Items execute serially, each under its own
    hook registrations, which preserves in-process semantics for every entry combination.
    """

    def __init__(self, backend: "HFBackend") -> None:
        self._backend = backend
        self._closed = False
        self._generate_count = 0

    @property
    def closed(self) -> bool:
        """Whether the session has been closed."""
        return self._closed

    def close(self) -> None:
        """Close the session; further use raises `RuntimeError`."""
        self._closed = True

    def __enter__(self) -> "ExclusiveSession":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("This session is closed; open a new session on the backend.")

    @property
    def model(self) -> PreTrainedModel:
        """The live model.

        Raises:
            RuntimeError: If the session is closed or no model is available yet.
        """
        self._ensure_open()
        model = self._backend._model_provider()
        if model is None:
            raise RuntimeError("No model is available on this session.")
        return model

    @property
    def tokenizer(self) -> PreTrainedTokenizerBase | None:
        """The tokenizer, or None when the adopting caller has not resolved one yet."""
        self._ensure_open()
        return self._backend._tokenizer_provider()

    @property
    def layout(self) -> ModelFacts:
        """Structural facts derived from the loaded model, computed on every access so weight
        edits and model replacements are always reflected.

        `num_layers` comes from the resolved decoder layer list; `hidden_size`,
        `num_attention_heads`, and `head_dim` come from the text config (`text_config(model)`,
        the text sub-config on composite multimodal models), with a
        `hidden_size // num_attention_heads` fallback for `head_dim`; `dtype` from the model;
        `model_fingerprint` from the weight/config fingerprint; and `model_type` from the
        composite config, so a multimodal checkpoint keeps its wrapper `model_type`.
        """
        model = self.model

        from steerability.algorithms.core.internals.fingerprint import model_fingerprint

        _, layer_names = get_model_layer_list(model)
        text_cfg = text_config(model)
        hidden_size = text_cfg.hidden_size
        num_heads = getattr(text_cfg, "num_attention_heads", None)
        head_dim = getattr(text_cfg, "head_dim", None)
        if head_dim is None and num_heads:
            head_dim = hidden_size // num_heads
        return ModelFacts(
            num_layers=len(layer_names),
            hidden_size=hidden_size,
            num_attention_heads=num_heads,
            head_dim=head_dim,
            dtype=str(model.dtype).removeprefix("torch."),
            model_fingerprint=model_fingerprint(model),
            model_type=getattr(model.config, "model_type", None),
            model_ref=getattr(model, "name_or_path", None),
        )

    def _resolve_prompt_tensors(self, prompt: PreparedPrompt) -> tuple[torch.Tensor, torch.Tensor]:
        """Token ids and attention mask for one prompt, on the model device."""
        resolved = prompt.resolve_token_ids(self.tokenizer)
        device = self.model.device
        input_ids = resolved.token_ids.to(device)
        attention_mask = resolved.attention_mask
        if attention_mask is None:
            tokenizer = self.tokenizer
            if tokenizer is not None and tokenizer.pad_token_id is not None:
                attention_mask = infer_attention_mask_from_ids(input_ids, tokenizer.pad_token_id)
            else:
                attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        attention_mask = attention_mask.to(dtype=input_ids.dtype, device=device)
        return input_ids, attention_mask

    def _compose_entry_stacks(
        self, output_entries, extra_processors=(), extra_criteria=(),
    ) -> tuple[LogitsProcessorList, StoppingCriteriaList]:
        """Compose the items' stack entries, appending caller extras after entry contributions."""
        processors: list = []
        criteria: list = []
        for entry in output_entries:
            if not isinstance(entry, StackEntry):
                raise UnsupportedOperationError(
                    f"{type(entry).__name__} requires an engine-hosted processor path; the "
                    "in-process session consumes StackEntry contributions."
                )
            processors.extend(entry.logits_processors)
            criteria.extend(entry.stopping_criteria)
        processors.extend(extra_processors)
        criteria.extend(extra_criteria)
        return LogitsProcessorList(processors), StoppingCriteriaList(criteria)

    def _register_state_entries(self, model: PreTrainedModel, state_entries) -> list:
        """Attach each entry's hook specifications to `model`, returning removable handles.

        The session is the single registrar of state hooks: entries are the only carriage, and
        registration lives strictly inside the session's execution of work. Pre and forward
        hooks register with `with_kwargs=True`; backward hooks register as full backward hooks.
        If registration fails partway, handles already attached are removed before re-raising.
        """
        handles: list = []
        try:
            for entry in state_entries:
                if not isinstance(entry, HookEntry):
                    raise UnsupportedOperationError(
                        f"{type(entry).__name__} requires an intervention-capable backend; the "
                        "in-process session consumes HookEntry contributions."
                    )
                for phase in ("pre", "forward", "backward"):
                    for spec in entry.hooks.get(phase, []):
                        module = model.get_submodule(spec["module"])
                        if phase == "pre":
                            handle = module.register_forward_pre_hook(spec["hook_func"], with_kwargs=True)
                        elif phase == "forward":
                            handle = module.register_forward_hook(spec["hook_func"], with_kwargs=True)
                        else:
                            handle = module.register_full_backward_hook(spec["hook_func"])
                        handles.append(handle)
        except Exception:
            for handle in handles:
                handle.remove()
            raise
        return handles

    @contextlib.contextmanager
    def entries_applied(self, state_entries):
        """Apply state entries to the live model for the duration of the context.

        Used by the pipeline around a client-side decoding driver's `decode`, so every forward
        the driver issues on the live model, including rollouts through this session and
        auxiliary scoring passes, runs under the generation's hooks. The session owns
        registration; hooks are removed when the context exits, even on error.
        """
        handles = self._register_state_entries(self.model, state_entries)
        try:
            yield self
        finally:
            for handle in handles:
                handle.remove()

    def _seeded(self, seed: int | None):
        """A context that snapshots and restores RNG state around a seeded decode.

        The CPU generator is always covered. On CUDA models every CUDA device generator is
        covered, since sharded models may sample on a device other than the first parameter's.
        On MPS models the MPS generator is covered.
        """
        if seed is None:
            return contextlib.nullcontext()
        device = self.model.device
        if device.type == "cuda":
            return torch.random.fork_rng(devices=list(range(torch.cuda.device_count())))
        if device.type == "mps":
            return _mps_rng_fork()
        return torch.random.fork_rng(devices=[])

    def _apply_seed(self, seed: int) -> None:
        """Seed the generators covered by `_seeded` for this decode."""
        torch.default_generator.manual_seed(seed)
        device = self.model.device
        if device.type == "cuda":
            torch.cuda.manual_seed_all(seed)
        elif device.type == "mps":
            torch.mps.manual_seed(seed)

    def _item_seeds(self, items: Sequence[GenerationItem], params: GenerationParams) -> list[int | None]:
        """Effective per-item seeds.

        An item's own seed is always honored. Otherwise, under `params.seed` with
        `seed_scope="item"`, each item derives its own seed from this call's operation id and its
        index; with `seed_scope="dispatch"`, every item takes the dispatch seed, derived at index 0
        so a single-item dispatch decodes identically under either scope. A dispatch mixing explicit
        and absent item seeds falls back to per-item derivation. Unseeded calls yield None.
        """
        operation_id = f"generate-{self._generate_count}"
        if (
            params.seed is not None
            and params.seed_scope == "dispatch"
            and all(item.seed is None for item in items)
        ):
            return [derive_item_seed(params.seed, operation_id, 0)] * len(items)
        seeds: list[int | None] = []
        for index, item in enumerate(items):
            if item.seed is not None:
                seeds.append(item.seed)
            elif params.seed is not None:
                seeds.append(derive_item_seed(params.seed, operation_id, index))
            else:
                seeds.append(None)
        return seeds

    def _report_serial_fallback(self, items: Sequence[GenerationItem]) -> None:
        """Log once per backend and reason why a multi-item dispatch decodes one item at a time.

        Reached only when `batchable` is False, so identical entries imply distinct seeds.
        """
        if not self._entries_identical(items):
            key, detail = "entries", (
                "items carry distinct state or output entries (row-scoped runtime kwargs or per-row hooks)"
            )
        else:
            key, detail = "seeds", (
                "items carry distinct seeds (seed_scope='item'); pass seed_scope='dispatch' to generate(), or "
                "set ProviderOptions.seed_scope on the evaluation provider, to decode the dispatch in one pass"
            )
        if self._backend.report_once(f"serial_fallback:{key}"):
            logger.info("Multi-item generate (%d items) decodes serially: %s.", len(items), detail)

    @staticmethod
    def _entries_identical(items: Sequence[GenerationItem | ScoringItem]) -> bool:
        """True when every item carries the same state and output entry objects."""
        first = items[0]
        for item in items[1:]:
            if len(item.state_entries) != len(first.state_entries) or any(
                a is not b for a, b in zip(item.state_entries, first.state_entries)
            ):
                return False
            if len(item.output_entries) != len(first.output_entries) or any(
                a is not b for a, b in zip(item.output_entries, first.output_entries)
            ):
                return False
        return True

    def _stack_prompt_rows(self, rows: list[tuple[torch.Tensor, torch.Tensor]]):
        """Stack resolved single-row prompts into one right-padded batch.

        The batched generate and score paths both left-pack the result via `to_left_pad` before
        the forward pass.
        """
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or 0
        max_len = max(ids.size(1) for ids, _ in rows)
        device = rows[0][0].device
        input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
        attention_mask = torch.zeros((len(rows), max_len), dtype=rows[0][1].dtype, device=device)
        for row, (ids, mask) in enumerate(rows):
            length = ids.size(1)
            input_ids[row, :length] = ids[0]
            attention_mask[row, :length] = mask[0]
        return input_ids, attention_mask

    def _classify(self, new_tokens: torch.Tensor, gen_kwargs: dict, params: GenerationParams) -> list[str | None]:
        """Per-row finish reasons under the pinned precedence, from the composed stop rules."""
        tokenizer = self.tokenizer
        return infer_finish_reasons(
            new_tokens,
            gen_kwargs,
            eos_token_id=getattr(tokenizer, "eos_token_id", None),
            pad_token_id=getattr(tokenizer, "pad_token_id", None),
            stop_strings=params.stop_strings,
            stop_token_ids=params.stop_token_ids,
            tokenizer=tokenizer,
        )

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item, each under its own hook registrations.

        Items sharing identical state entries, identical output entries, and identical-or-absent
        effective seeds execute in one batched `model.generate` pass (right-padded to a common
        prompt length, then left-packed together); otherwise items decode serially. Caller-supplied `logits_processor` and
        `stopping_criteria` entries in `params.extra` append after the items' own contributions,
        and the normalized stop fields compose as stop rules anchored at the prompt length. A
        seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the
        covered generator state (CPU, plus the model device's) is restored afterwards; when
        `params.seed` is set and an item carries no seed of its own, the item's seed derives per
        index, so multi-item fan-outs sample distinct streams.

        Args:
            items: The generation items.
            params: Normalized generation parameters shared by all items.

        Returns:
            One `ItemResult` per item, in item order. Each result's `finish_reasons` carries one
            reason per candidate with the precedence stop, then eos, then length, then None.
        """
        self._ensure_open()
        model = self.model
        tokenizer = self.tokenizer
        gen_kwargs = render_hf_gen_kwargs(params)
        # default pad_token_id per call to avoid the transformers open-end-generation warning
        # without mutating the model's generation config; caller kwargs and a model-configured
        # value both take precedence over the tokenizer fallback
        if (
            "pad_token_id" not in gen_kwargs
            and model.generation_config.pad_token_id is None
            and getattr(tokenizer, "pad_token_id", None) is not None
        ):
            gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
        user_processors = tuple(gen_kwargs.pop("logits_processor", None) or ())
        user_criteria = tuple(gen_kwargs.pop("stopping_criteria", None) or ())

        if not items:
            return []
        seeds = self._item_seeds(items, params)
        self._generate_count += 1

        batchable = (
            len(items) > 1
            and self._entries_identical(items)
            and len(set(seeds)) == 1
        )
        if batchable:
            return self._generate_batched(
                items, params, gen_kwargs, user_processors, user_criteria, seeds[0],
            )
        if len(items) > 1:
            self._report_serial_fallback(items)

        results: list[ItemResult] = []
        for index, item in enumerate(items):
            input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
            processors, criteria = self._compose_entry_stacks(
                item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
            )
            criteria.extend(compose_stop_criteria(params, input_ids.size(1), tokenizer))
            stacks = stack_generate_kwargs(processors, criteria)
            handles = self._register_state_entries(model, item.state_entries)
            try:
                seed = seeds[index]
                with self._seeded(seed):
                    if seed is not None:
                        self._apply_seed(seed)
                    full_ids = model.generate(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        **stacks,
                        **gen_kwargs,
                    )
            finally:
                for handle in handles:
                    handle.remove()

            new_tokens = full_ids[:, input_ids.size(1):]
            reasons = self._classify(new_tokens, gen_kwargs, params)
            results.append(ItemResult(
                index=index,
                output=Output(
                    output_ids=new_tokens,
                    adapted_input_ids=input_ids,
                    finish_reason=reasons[0],
                    finish_reasons=tuple(reasons),
                ),
            ))
        return results

    def _generate_batched(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
        gen_kwargs: dict,
        user_processors: tuple,
        user_criteria: tuple,
        seed: int | None,
    ) -> list[ItemResult]:
        """One `model.generate` pass over all items (identical entries, one shared seed).

        The stacked prompts left-pack (pad positions move before the real tokens), so every row's
        first generated token is predicted from its last real token rather than from a trailing pad.
        """
        model = self.model
        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
        input_ids, attention_mask = self._stack_prompt_rows(rows)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
        processors, criteria = self._compose_entry_stacks(
            items[0].output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
        )
        criteria.extend(compose_stop_criteria(params, input_ids.size(1), self.tokenizer))
        stacks = stack_generate_kwargs(processors, criteria)
        handles = self._register_state_entries(model, items[0].state_entries)
        try:
            with self._seeded(seed):
                if seed is not None:
                    self._apply_seed(seed)
                full_ids = model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **stacks,
                    **gen_kwargs,
                )
        finally:
            for handle in handles:
                handle.remove()

        prompt_len = input_ids.size(1)
        new_tokens = full_ids[:, prompt_len:]
        num_candidates = params.n or 1
        results: list[ItemResult] = []
        for index in range(len(items)):
            item_rows = new_tokens[index * num_candidates:(index + 1) * num_candidates]
            reasons = self._classify(item_rows, gen_kwargs, params)
            results.append(ItemResult(
                index=index,
                output=Output(
                    output_ids=item_rows,
                    adapted_input_ids=input_ids[index:index + 1],
                    finish_reason=reasons[0],
                    finish_reasons=tuple(reasons),
                ),
            ))
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens.

        For each item, the prompt left-packs (pad positions move before the real tokens) so the
        reference follows the prompt's last real token, then prompt and reference concatenate
        into one causal forward pass under the item's hook registrations; the item's logits
        processors replay position-by-position with the same `(prefix_ids, scores)` view they
        receive during generation. Items sharing identical state and output entries score in one
        batched forward pass (prompts right-padded to a common length, then left-packed
        together); otherwise items score serially. Stopping criteria never apply. Decoder-only
        models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: `params.extra` passes through as forward keyword arguments.

        Returns:
            Log probabilities of shape `[num_items, ref_len]`.

        Raises:
            ValueError: If items carry differing reference lengths.
            UnsupportedOperationError: If the model is an encoder-decoder model.
        """
        self._ensure_open()
        model = self.model
        if getattr(model.config, "is_encoder_decoder", False):
            raise UnsupportedOperationError(
                "Session scoring supports decoder-only models; encoder-decoder scoring runs "
                "through SteeringPipeline.compute_logprobs."
            )
        device = model.device
        forward_kwargs = dict(params.extra)

        if not items:
            return torch.zeros((0, 0), device=device, dtype=torch.float32)
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")

        if len(items) > 1 and self._entries_identical(items):
            return self._score_batched(items, forward_kwargs)

        all_logprobs: list[torch.Tensor] = []
        for item in items:
            input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
            input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
            ref_output_ids = item.ref_output_ids
            if ref_output_ids.ndim == 1:
                ref_output_ids = ref_output_ids.unsqueeze(0)
            ref_output_ids = ref_output_ids.to(device)
            ref_len = ref_output_ids.size(1)
            if ref_len == 0:
                all_logprobs.append(torch.zeros((1, 0), device=device, dtype=torch.float32))
                continue

            processors, _ = self._compose_entry_stacks(item.output_entries)
            handles = self._register_state_entries(model, item.state_entries)
            try:
                with torch.no_grad():
                    combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                    combined_mask = torch.cat([
                        attention_mask,
                        torch.ones(1, ref_len, device=device, dtype=attention_mask.dtype),
                    ], dim=1)
                    outputs = model(
                        input_ids=combined_ids,
                        attention_mask=combined_mask,
                        **forward_kwargs,
                    )
                    input_len = input_ids.size(1)
                    logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                    if len(processors):
                        for t in range(logits.size(1)):
                            prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                            logits[:, t, :] = processors(prefix, logits[:, t, :])
                    logprobs = torch.log_softmax(logits, dim=-1)
                    all_logprobs.append(
                        logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
                    )
            finally:
                for handle in handles:
                    handle.remove()

        return torch.cat(all_logprobs, dim=0)

    def _score_batched(self, items: Sequence[ScoringItem], forward_kwargs: dict) -> torch.Tensor:
        """One causal forward pass over all items (identical entries)."""
        model = self.model
        device = model.device
        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
        input_ids, attention_mask = self._stack_prompt_rows(rows)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)

        refs: list[torch.Tensor] = []
        for item in items:
            ref = item.ref_output_ids
            if ref.ndim == 1:
                ref = ref.unsqueeze(0)
            refs.append(ref.to(device))
        ref_output_ids = torch.cat(refs, dim=0)
        ref_len = ref_output_ids.size(1)
        if ref_len == 0:
            return torch.zeros((len(items), 0), device=device, dtype=torch.float32)

        processors, _ = self._compose_entry_stacks(items[0].output_entries)
        handles = self._register_state_entries(model, items[0].state_entries)
        try:
            with torch.no_grad():
                combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                combined_mask = torch.cat([
                    attention_mask,
                    torch.ones(len(items), ref_len, device=device, dtype=attention_mask.dtype),
                ], dim=1)
                outputs = model(
                    input_ids=combined_ids,
                    attention_mask=combined_mask,
                    **forward_kwargs,
                )
                input_len = input_ids.size(1)
                logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                if len(processors):
                    for t in range(logits.size(1)):
                        prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                        logits[:, t, :] = processors(prefix, logits[:, t, :])
                logprobs = torch.log_softmax(logits, dim=-1)
                return logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
        finally:
            for handle in handles:
                handle.remove()

    def capture(
        self,
        prompts: list[PreparedPrompt],
        layers: list[int],
        mode: Literal["all_tokens", "last_token"],
        location: Literal["layer_output", "layer_input"] = "layer_output",
    ) -> CaptureResult:
        """Capture residual-stream hidden states for `prompts` at `layers`.

        Prompts resolve to token ids, right-pad into one batch, and run through the shared
        layerwise extraction. In `"all_tokens"` mode each layer's tensor is `[N, T, H]` with rows
        outside the attention mask zeroed; in `"last_token"` mode the last real (non-pad) position
        of each row is selected, giving `[N, H]`.

        Args:
            prompts: The prompts to capture.
            layers: 0-based layer ids to keep.
            mode: `"all_tokens"` or `"last_token"`.
            location: `"layer_output"` (a layer's output boundary) or `"layer_input"` (the
                boundary a forward pre-hook observes).

        Returns:
            The captured tensors and the batch attention mask, on CPU.

        Raises:
            ValueError: If `mode` is unknown or a requested layer id is out of range.
        """
        self._ensure_open()
        if mode not in ("all_tokens", "last_token"):
            raise ValueError(f"Unknown capture mode {mode!r}; modes are 'all_tokens', 'last_token'.")
        if not prompts:
            raise ValueError("capture() requires at least one prompt.")

        from steerability.algorithms.core.internals.capture import layerwise_tokenwise_hidden
        from steerability.algorithms.core.internals.pooling import aggregate_condition_hidden

        model = self.model
        device = model.device
        tokenizer = self.tokenizer
        pad_token_id = getattr(tokenizer, "pad_token_id", None) or 0

        rows = [self._resolve_prompt_tensors(prompt) for prompt in prompts]
        max_len = max(ids.size(1) for ids, _ in rows)
        input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
        attention_mask = torch.zeros((len(rows), max_len), dtype=torch.long, device=device)
        for row, (ids, mask) in enumerate(rows):
            length = ids.size(1)
            input_ids[row, :length] = ids[0]
            attention_mask[row, :length] = mask[0]

        enc = {"input_ids": input_ids, "attention_mask": attention_mask}
        hidden = layerwise_tokenwise_hidden(
            model, enc, batch_size=_CAPTURE_BATCH_SIZE, location=location,
        )
        missing = [layer for layer in layers if layer not in hidden]
        if missing:
            raise ValueError(
                f"Requested layer ids {missing} are out of range; the model has {len(hidden)} layers."
            )

        mask_cpu = attention_mask.cpu()
        selected = {layer: hidden[layer] for layer in layers}
        if mode == "last_token":
            selected = {
                layer: aggregate_condition_hidden(tensor, "last", mask_cpu)
                for layer, tensor in selected.items()
            }
        return CaptureResult(
            hidden=selected, attention_mask=mask_cpu, mode=mode, location=location,
        )

closed property

Whether the session has been closed.

layout property

Structural facts derived from the loaded model, computed on every access so weight edits and model replacements are always reflected.

num_layers comes from the resolved decoder layer list; hidden_size, num_attention_heads, and head_dim come from the text config (text_config(model), the text sub-config on composite multimodal models), with a hidden_size // num_attention_heads fallback for head_dim; dtype from the model; model_fingerprint from the weight/config fingerprint; and model_type from the composite config, so a multimodal checkpoint keeps its wrapper model_type.

model property

The live model.

Raises:

Type	Description
RuntimeError	If the session is closed or no model is available yet.

tokenizer property

The tokenizer, or None when the adopting caller has not resolved one yet.

capture(prompts, layers, mode, location='layer_output')

Capture residual-stream hidden states for prompts at layers.

Prompts resolve to token ids, right-pad into one batch, and run through the shared layerwise extraction. In "all_tokens" mode each layer's tensor is [N, T, H] with rows outside the attention mask zeroed; in "last_token" mode the last real (non-pad) position of each row is selected, giving [N, H].

Parameters:

Name	Type	Description	Default
prompts	list[PreparedPrompt]	The prompts to capture.	required
layers	list[int]	0-based layer ids to keep.	required
mode	Literal['all_tokens', 'last_token']	"all_tokens" or "last_token".	required
location	Literal['layer_output', 'layer_input']	"layer_output" (a layer's output boundary) or "layer_input" (the boundary a forward pre-hook observes).	'layer_output'

Returns:

Type	Description
CaptureResult	The captured tensors and the batch attention mask, on CPU.

Raises:

Type	Description
ValueError	If mode is unknown or a requested layer id is out of range.

Source code in steerability/backends/huggingface/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Capture residual-stream hidden states for `prompts` at `layers`.

    Prompts resolve to token ids, right-pad into one batch, and run through the shared
    layerwise extraction. In `"all_tokens"` mode each layer's tensor is `[N, T, H]` with rows
    outside the attention mask zeroed; in `"last_token"` mode the last real (non-pad) position
    of each row is selected, giving `[N, H]`.

    Args:
        prompts: The prompts to capture.
        layers: 0-based layer ids to keep.
        mode: `"all_tokens"` or `"last_token"`.
        location: `"layer_output"` (a layer's output boundary) or `"layer_input"` (the
            boundary a forward pre-hook observes).

    Returns:
        The captured tensors and the batch attention mask, on CPU.

    Raises:
        ValueError: If `mode` is unknown or a requested layer id is out of range.
    """
    self._ensure_open()
    if mode not in ("all_tokens", "last_token"):
        raise ValueError(f"Unknown capture mode {mode!r}; modes are 'all_tokens', 'last_token'.")
    if not prompts:
        raise ValueError("capture() requires at least one prompt.")

    from steerability.algorithms.core.internals.capture import layerwise_tokenwise_hidden
    from steerability.algorithms.core.internals.pooling import aggregate_condition_hidden

    model = self.model
    device = model.device
    tokenizer = self.tokenizer
    pad_token_id = getattr(tokenizer, "pad_token_id", None) or 0

    rows = [self._resolve_prompt_tensors(prompt) for prompt in prompts]
    max_len = max(ids.size(1) for ids, _ in rows)
    input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
    attention_mask = torch.zeros((len(rows), max_len), dtype=torch.long, device=device)
    for row, (ids, mask) in enumerate(rows):
        length = ids.size(1)
        input_ids[row, :length] = ids[0]
        attention_mask[row, :length] = mask[0]

    enc = {"input_ids": input_ids, "attention_mask": attention_mask}
    hidden = layerwise_tokenwise_hidden(
        model, enc, batch_size=_CAPTURE_BATCH_SIZE, location=location,
    )
    missing = [layer for layer in layers if layer not in hidden]
    if missing:
        raise ValueError(
            f"Requested layer ids {missing} are out of range; the model has {len(hidden)} layers."
        )

    mask_cpu = attention_mask.cpu()
    selected = {layer: hidden[layer] for layer in layers}
    if mode == "last_token":
        selected = {
            layer: aggregate_condition_hidden(tensor, "last", mask_cpu)
            for layer, tensor in selected.items()
        }
    return CaptureResult(
        hidden=selected, attention_mask=mask_cpu, mode=mode, location=location,
    )

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/huggingface/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

entries_applied(state_entries)

Apply state entries to the live model for the duration of the context.

Used by the pipeline around a client-side decoding driver's decode, so every forward the driver issues on the live model, including rollouts through this session and auxiliary scoring passes, runs under the generation's hooks. The session owns registration; hooks are removed when the context exits, even on error.

Source code in steerability/backends/huggingface/session.py

@contextlib.contextmanager
def entries_applied(self, state_entries):
    """Apply state entries to the live model for the duration of the context.

    Used by the pipeline around a client-side decoding driver's `decode`, so every forward
    the driver issues on the live model, including rollouts through this session and
    auxiliary scoring passes, runs under the generation's hooks. The session owns
    registration; hooks are removed when the context exits, even on error.
    """
    handles = self._register_state_entries(self.model, state_entries)
    try:
        yield self
    finally:
        for handle in handles:
            handle.remove()

generate(items, params)

Generate one result per item, each under its own hook registrations.

Items sharing identical state entries, identical output entries, and identical-or-absent effective seeds execute in one batched model.generate pass (right-padded to a common prompt length, then left-packed together); otherwise items decode serially. Caller-supplied logits_processor and stopping_criteria entries in params.extra append after the items' own contributions, and the normalized stop fields compose as stop rules anchored at the prompt length. A seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the covered generator state (CPU, plus the model device's) is restored afterwards; when params.seed is set and an item carries no seed of its own, the item's seed derives per index, so multi-item fan-outs sample distinct streams.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items.	required
params	GenerationParams	Normalized generation parameters shared by all items.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order. Each result's finish_reasons carries one
list[ItemResult]	reason per candidate with the precedence stop, then eos, then length, then None.

Source code in steerability/backends/huggingface/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item, each under its own hook registrations.

    Items sharing identical state entries, identical output entries, and identical-or-absent
    effective seeds execute in one batched `model.generate` pass (right-padded to a common
    prompt length, then left-packed together); otherwise items decode serially. Caller-supplied `logits_processor` and
    `stopping_criteria` entries in `params.extra` append after the items' own contributions,
    and the normalized stop fields compose as stop rules anchored at the prompt length. A
    seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the
    covered generator state (CPU, plus the model device's) is restored afterwards; when
    `params.seed` is set and an item carries no seed of its own, the item's seed derives per
    index, so multi-item fan-outs sample distinct streams.

    Args:
        items: The generation items.
        params: Normalized generation parameters shared by all items.

    Returns:
        One `ItemResult` per item, in item order. Each result's `finish_reasons` carries one
        reason per candidate with the precedence stop, then eos, then length, then None.
    """
    self._ensure_open()
    model = self.model
    tokenizer = self.tokenizer
    gen_kwargs = render_hf_gen_kwargs(params)
    # default pad_token_id per call to avoid the transformers open-end-generation warning
    # without mutating the model's generation config; caller kwargs and a model-configured
    # value both take precedence over the tokenizer fallback
    if (
        "pad_token_id" not in gen_kwargs
        and model.generation_config.pad_token_id is None
        and getattr(tokenizer, "pad_token_id", None) is not None
    ):
        gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
    user_processors = tuple(gen_kwargs.pop("logits_processor", None) or ())
    user_criteria = tuple(gen_kwargs.pop("stopping_criteria", None) or ())

    if not items:
        return []
    seeds = self._item_seeds(items, params)
    self._generate_count += 1

    batchable = (
        len(items) > 1
        and self._entries_identical(items)
        and len(set(seeds)) == 1
    )
    if batchable:
        return self._generate_batched(
            items, params, gen_kwargs, user_processors, user_criteria, seeds[0],
        )
    if len(items) > 1:
        self._report_serial_fallback(items)

    results: list[ItemResult] = []
    for index, item in enumerate(items):
        input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
        processors, criteria = self._compose_entry_stacks(
            item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
        )
        criteria.extend(compose_stop_criteria(params, input_ids.size(1), tokenizer))
        stacks = stack_generate_kwargs(processors, criteria)
        handles = self._register_state_entries(model, item.state_entries)
        try:
            seed = seeds[index]
            with self._seeded(seed):
                if seed is not None:
                    self._apply_seed(seed)
                full_ids = model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **stacks,
                    **gen_kwargs,
                )
        finally:
            for handle in handles:
                handle.remove()

        new_tokens = full_ids[:, input_ids.size(1):]
        reasons = self._classify(new_tokens, gen_kwargs, params)
        results.append(ItemResult(
            index=index,
            output=Output(
                output_ids=new_tokens,
                adapted_input_ids=input_ids,
                finish_reason=reasons[0],
                finish_reasons=tuple(reasons),
            ),
        ))
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens.

For each item, the prompt left-packs (pad positions move before the real tokens) so the reference follows the prompt's last real token, then prompt and reference concatenate into one causal forward pass under the item's hook registrations; the item's logits processors replay position-by-position with the same (prefix_ids, scores) view they receive during generation. Items sharing identical state and output entries score in one batched forward pass (prompts right-padded to a common length, then left-packed together); otherwise items score serially. Stopping criteria never apply. Decoder-only models only; the pipeline's compute_logprobs serves encoder-decoder models in-process.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	params.extra passes through as forward keyword arguments.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len].

Raises:

Type	Description
ValueError	If items carry differing reference lengths.
UnsupportedOperationError	If the model is an encoder-decoder model.

Source code in steerability/backends/huggingface/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens.

    For each item, the prompt left-packs (pad positions move before the real tokens) so the
    reference follows the prompt's last real token, then prompt and reference concatenate
    into one causal forward pass under the item's hook registrations; the item's logits
    processors replay position-by-position with the same `(prefix_ids, scores)` view they
    receive during generation. Items sharing identical state and output entries score in one
    batched forward pass (prompts right-padded to a common length, then left-packed
    together); otherwise items score serially. Stopping criteria never apply. Decoder-only
    models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: `params.extra` passes through as forward keyword arguments.

    Returns:
        Log probabilities of shape `[num_items, ref_len]`.

    Raises:
        ValueError: If items carry differing reference lengths.
        UnsupportedOperationError: If the model is an encoder-decoder model.
    """
    self._ensure_open()
    model = self.model
    if getattr(model.config, "is_encoder_decoder", False):
        raise UnsupportedOperationError(
            "Session scoring supports decoder-only models; encoder-decoder scoring runs "
            "through SteeringPipeline.compute_logprobs."
        )
    device = model.device
    forward_kwargs = dict(params.extra)

    if not items:
        return torch.zeros((0, 0), device=device, dtype=torch.float32)
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")

    if len(items) > 1 and self._entries_identical(items):
        return self._score_batched(items, forward_kwargs)

    all_logprobs: list[torch.Tensor] = []
    for item in items:
        input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
        ref_output_ids = item.ref_output_ids
        if ref_output_ids.ndim == 1:
            ref_output_ids = ref_output_ids.unsqueeze(0)
        ref_output_ids = ref_output_ids.to(device)
        ref_len = ref_output_ids.size(1)
        if ref_len == 0:
            all_logprobs.append(torch.zeros((1, 0), device=device, dtype=torch.float32))
            continue

        processors, _ = self._compose_entry_stacks(item.output_entries)
        handles = self._register_state_entries(model, item.state_entries)
        try:
            with torch.no_grad():
                combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                combined_mask = torch.cat([
                    attention_mask,
                    torch.ones(1, ref_len, device=device, dtype=attention_mask.dtype),
                ], dim=1)
                outputs = model(
                    input_ids=combined_ids,
                    attention_mask=combined_mask,
                    **forward_kwargs,
                )
                input_len = input_ids.size(1)
                logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                if len(processors):
                    for t in range(logits.size(1)):
                        prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                        logits[:, t, :] = processors(prefix, logits[:, t, :])
                logprobs = torch.log_softmax(logits, dim=-1)
                all_logprobs.append(
                    logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
                )
        finally:
            for handle in handles:
                handle.remove()

    return torch.cat(all_logprobs, dim=0)

HFBackend

Bases: Backend

The in-process Hugging Face backend.

Owns a loaded model and tokenizer, either loaded from a spec or adopted from a caller that already holds them. At most one session may be open per backend at a time, so the backend runs one generation at a time.

Source code in steerability/backends/huggingface/backend.py

class HFBackend(Backend):
    """The in-process Hugging Face backend.

    Owns a loaded model and tokenizer, either loaded from a spec or adopted from a caller that
    already holds them. At most one session may be open per backend at a time, so the backend
    runs one generation at a time.
    """

    def __init__(
        self,
        spec: BackendSpec,
        *,
        model_provider: Callable[[], PreTrainedModel | None] | None = None,
        tokenizer_provider: Callable[[], object | None] | None = None,
    ) -> None:
        """Construct the backend, loading the model from `spec` unless providers are given.

        Loading reads the options `hf_model_kwargs`, `device_map`, `tokenizer_name_or_path`,
        and `trust_remote_code`. Option values must be plain data, since spec canonicalization
        renders live objects (e.g. a quantization config instance) as strings that
        `from_pretrained` cannot consume. A `device_map` key inside `hf_model_kwargs` is used
        when the spec carries no top-level `device_map` option.

        Args:
            spec: The backend spec.
            model_provider: Callable returning the adopted model; used with
                `tokenizer_provider` instead of loading.
            tokenizer_provider: Callable returning the adopted tokenizer.

        Raises:
            ValueError: If `spec.kind` is not `"huggingface"`, or no model reference is
                available to load from.
        """
        if spec.kind != "huggingface":
            raise ValueError(f"HFBackend requires a 'huggingface' spec; got kind {spec.kind!r}.")
        self.spec = spec
        self._open_session: ExclusiveSession | None = None
        self._reported_once: set[str] = set()

        if model_provider is not None:
            self._model_provider = model_provider
            self._tokenizer_provider = tokenizer_provider or (lambda: None)
            return

        if spec.model is None:
            raise ValueError(
                "HFBackend needs a model reference on the spec, or model_provider/"
                "tokenizer_provider for an already-loaded model."
            )
        hf_model_kwargs = dict(spec.get_option("hf_model_kwargs", default={}))
        device_map = spec.get_option("device_map", default=hf_model_kwargs.pop("device_map", "auto"))
        model = AutoModelForCausalLM.from_pretrained(
            spec.model,
            device_map=device_map,
            **hf_model_kwargs,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            spec.get_option("tokenizer_name_or_path") or spec.model,
            trust_remote_code=bool(spec.get_option("trust_remote_code", default=False)),
        )
        tokenizer = ensure_pad_token(tokenizer)
        self._model_provider = lambda: model
        self._tokenizer_provider = lambda: tokenizer

    @classmethod
    def adopt(
        cls,
        spec: BackendSpec,
        model_provider: Callable[[], PreTrainedModel | None],
        tokenizer_provider: Callable[[], object | None],
    ) -> "HFBackend":
        """Wrap an already-loaded model and tokenizer without loading anything.

        Providers are read on every access, so a caller whose model is replaced mid-steer (a
        structural control returning a new model) always exposes the current one to sessions.

        Args:
            spec: The backend spec identifying this configuration.
            model_provider: Callable returning the current model (may return None before one
                exists).
            tokenizer_provider: Callable returning the current tokenizer.

        Returns:
            The adopting backend.
        """
        return cls(spec, model_provider=model_provider, tokenizer_provider=tokenizer_provider)

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The static Hugging Face capability advertisement (spec-independent)."""
        return HF_CAPABILITIES

    def open_session(self) -> "ExclusiveSession":
        """Open the backend's one exclusive session.

        Returns:
            The session, usable as a context manager.

        Raises:
            RuntimeError: If an exclusive session is already open on this backend.
        """
        if self._open_session is not None and not self._open_session.closed:
            raise RuntimeError(
                "An exclusive session is already open on this backend; close it before opening "
                "another."
            )
        self._open_session = ExclusiveSession(self)
        return self._open_session

    def report_once(self, key: str) -> bool:
        """Return True the first time `key` is seen on this backend, False afterwards.

        Sessions are opened per pipeline call, so a per-session flag would fire on every call;
        this guard lives on the backend so a repeated informational readout fires once.
        """
        if key in self._reported_once:
            return False
        self._reported_once.add(key)
        return True

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

adopt(spec, model_provider, tokenizer_provider) classmethod

Wrap an already-loaded model and tokenizer without loading anything.

Providers are read on every access, so a caller whose model is replaced mid-steer (a structural control returning a new model) always exposes the current one to sessions.

Parameters:

Name	Type	Description	Default
spec	BackendSpec	The backend spec identifying this configuration.	required
model_provider	Callable[[], PreTrainedModel | None]	Callable returning the current model (may return None before one exists).	required
tokenizer_provider	Callable[[], object | None]	Callable returning the current tokenizer.	required

Returns:

Type	Description
HFBackend	The adopting backend.

Source code in steerability/backends/huggingface/backend.py

@classmethod
def adopt(
    cls,
    spec: BackendSpec,
    model_provider: Callable[[], PreTrainedModel | None],
    tokenizer_provider: Callable[[], object | None],
) -> "HFBackend":
    """Wrap an already-loaded model and tokenizer without loading anything.

    Providers are read on every access, so a caller whose model is replaced mid-steer (a
    structural control returning a new model) always exposes the current one to sessions.

    Args:
        spec: The backend spec identifying this configuration.
        model_provider: Callable returning the current model (may return None before one
            exists).
        tokenizer_provider: Callable returning the current tokenizer.

    Returns:
        The adopting backend.
    """
    return cls(spec, model_provider=model_provider, tokenizer_provider=tokenizer_provider)

capabilities_for_spec(spec) classmethod

The static Hugging Face capability advertisement (spec-independent).

Source code in steerability/backends/huggingface/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The static Hugging Face capability advertisement (spec-independent)."""
    return HF_CAPABILITIES

open_session()

Open the backend's one exclusive session.

Returns:

Type	Description
ExclusiveSession	The session, usable as a context manager.

Raises:

Type	Description
RuntimeError	If an exclusive session is already open on this backend.

Source code in steerability/backends/huggingface/backend.py

def open_session(self) -> "ExclusiveSession":
    """Open the backend's one exclusive session.

    Returns:
        The session, usable as a context manager.

    Raises:
        RuntimeError: If an exclusive session is already open on this backend.
    """
    if self._open_session is not None and not self._open_session.closed:
        raise RuntimeError(
            "An exclusive session is already open on this backend; close it before opening "
            "another."
        )
    self._open_session = ExclusiveSession(self)
    return self._open_session

release()

Free the resources this backend owns beyond what the caller owns.

The default is a no-op. Engine-owning backends override this to shut their engines down deterministically. Release is idempotent.

Source code in steerability/algorithms/core/execution/backend.py

def release(self) -> None:
    """Free the resources this backend owns beyond what the caller owns.

    The default is a no-op. Engine-owning backends override this to shut their engines down
    deterministically. Release is idempotent.
    """
    return None

report_once(key)

Return True the first time key is seen on this backend, False afterwards.

Sessions are opened per pipeline call, so a per-session flag would fire on every call; this guard lives on the backend so a repeated informational readout fires once.

Source code in steerability/backends/huggingface/backend.py

def report_once(self, key: str) -> bool:
    """Return True the first time `key` is seen on this backend, False afterwards.

    Sessions are opened per pipeline call, so a per-session flag would fire on every call;
    this guard lives on the backend so a repeated informational readout fires once.
    """
    if key in self._reported_once:
        return False
    self._reported_once.add(key)
    return True

stage_artifacts(payloads)

Make each content-addressed artifact available to the execution side.

Called by the pipeline at the end of steer() with the tensor payloads of every lowered intervention spec, keyed by content-addressed artifact id. Staging is idempotent: an artifact that already exists at the destination is success. The in-process backend keeps live tensors and needs no staging; engine-backed backends write into the registry the serving engine reads.

Parameters:

Name	Type	Description	Default
payloads	dict[str, dict[str, Tensor]]	Mapping from artifact id to a name-to-tensor mapping.	required

Source code in steerability/algorithms/core/execution/backend.py

def stage_artifacts(self, payloads: dict[str, dict[str, torch.Tensor]]) -> None:
    """Make each content-addressed artifact available to the execution side.

    Called by the pipeline at the end of `steer()` with the tensor payloads of every
    lowered intervention spec, keyed by content-addressed artifact id. Staging is
    idempotent: an artifact that already exists at the destination is success. The
    in-process backend keeps live tensors and needs no staging; engine-backed backends
    write into the registry the serving engine reads.

    Args:
        payloads: Mapping from artifact id to a name-to-tensor mapping.
    """
    return None

compose_stop_criteria(params, prompt_len, tokenizer)

The stop criteria implied by the normalized stop fields, anchored at prompt_len.

Parameters:

Name	Type	Description	Default
params	GenerationParams	The normalized parameters carrying stop_strings and stop_token_ids.	required
prompt_len	int	Prompt length the substring criteria decode past.	required
tokenizer	PreTrainedTokenizerBase | None	Tokenizer for substring decoding; required when stop strings are set.	required

Returns:

Type	Description
list	The composed criteria (possibly empty).

Raises:

Type	Description
ValueError	If stop strings are set and no tokenizer is available.

Source code in steerability/backends/huggingface/session.py

def compose_stop_criteria(
    params: GenerationParams, prompt_len: int, tokenizer: PreTrainedTokenizerBase | None
) -> list:
    """The stop criteria implied by the normalized stop fields, anchored at `prompt_len`.

    Args:
        params: The normalized parameters carrying `stop_strings` and `stop_token_ids`.
        prompt_len: Prompt length the substring criteria decode past.
        tokenizer: Tokenizer for substring decoding; required when stop strings are set.

    Returns:
        The composed criteria (possibly empty).

    Raises:
        ValueError: If stop strings are set and no tokenizer is available.
    """
    criteria: list = []
    if params.stop_strings:
        if tokenizer is None:
            raise ValueError("stop_strings require a tokenizer on the session.")
        for text in params.stop_strings:
            criteria.append(StopOnSubstring(tokenizer, text, prompt_len))
    if params.stop_token_ids:
        criteria.append(StopOnTokens(params.stop_token_ids))
    return criteria

render_hf_gen_kwargs(params)

Render normalized generation parameters onto model.generate keyword arguments.

Keys in params.extra pass through untouched; a normalized field always takes precedence over a same-named extra key. seed is not rendered here, since the session applies it as a fork_rng-scoped manual_seed around the item's decode; the stop fields are not rendered either, since the session composes them as prompt-anchored stop criteria (compose_stop_criteria).

Parameters:

Name	Type	Description	Default
params	GenerationParams	The normalized parameters.	required

Returns:

Type	Description
dict	Keyword arguments for model.generate.

Source code in steerability/backends/huggingface/session.py

def render_hf_gen_kwargs(params: GenerationParams) -> dict:
    """Render normalized generation parameters onto `model.generate` keyword arguments.

    Keys in `params.extra` pass through untouched; a normalized field always takes precedence
    over a same-named extra key. `seed` is not rendered here, since the session applies it as a
    `fork_rng`-scoped `manual_seed` around the item's decode; the stop fields are not rendered
    either, since the session composes them as prompt-anchored stop criteria
    (`compose_stop_criteria`).

    Args:
        params: The normalized parameters.

    Returns:
        Keyword arguments for `model.generate`.
    """
    gen_kwargs = dict(params.extra)
    if params.max_new_tokens is not None:
        gen_kwargs["max_new_tokens"] = params.max_new_tokens
    if params.min_new_tokens is not None:
        gen_kwargs["min_new_tokens"] = params.min_new_tokens
    if params.temperature is not None:
        gen_kwargs["temperature"] = params.temperature
    if params.top_p is not None:
        gen_kwargs["top_p"] = params.top_p
    if params.top_k is not None:
        gen_kwargs["top_k"] = params.top_k
    if params.repetition_penalty is not None:
        gen_kwargs["repetition_penalty"] = params.repetition_penalty
    if params.greedy is not None:
        gen_kwargs["do_sample"] = not params.greedy
    if params.n is not None:
        gen_kwargs["num_return_sequences"] = params.n
    return gen_kwargs

backend

The in-process Hugging Face backend and its capability advertisement.

HF_CAPABILITIES = BackendCapabilities(atoms=frozenset({Capability.IN_PROCESS_TORCH, Capability.HIDDEN_CAPTURE, Capability.BEAM_PROPOSALS}), capture_kinds=CaptureKinds(kinds=frozenset({'residual'}), locations=frozenset({'layer_output', 'layer_input'}), modes=frozenset({'all_tokens', 'last_token'}))) module-attribute

HFBackend

Bases: Backend

The in-process Hugging Face backend.

Owns a loaded model and tokenizer, either loaded from a spec or adopted from a caller that already holds them. At most one session may be open per backend at a time, so the backend runs one generation at a time.

Source code in steerability/backends/huggingface/backend.py

class HFBackend(Backend):
    """The in-process Hugging Face backend.

    Owns a loaded model and tokenizer, either loaded from a spec or adopted from a caller that
    already holds them. At most one session may be open per backend at a time, so the backend
    runs one generation at a time.
    """

    def __init__(
        self,
        spec: BackendSpec,
        *,
        model_provider: Callable[[], PreTrainedModel | None] | None = None,
        tokenizer_provider: Callable[[], object | None] | None = None,
    ) -> None:
        """Construct the backend, loading the model from `spec` unless providers are given.

        Loading reads the options `hf_model_kwargs`, `device_map`, `tokenizer_name_or_path`,
        and `trust_remote_code`. Option values must be plain data, since spec canonicalization
        renders live objects (e.g. a quantization config instance) as strings that
        `from_pretrained` cannot consume. A `device_map` key inside `hf_model_kwargs` is used
        when the spec carries no top-level `device_map` option.

        Args:
            spec: The backend spec.
            model_provider: Callable returning the adopted model; used with
                `tokenizer_provider` instead of loading.
            tokenizer_provider: Callable returning the adopted tokenizer.

        Raises:
            ValueError: If `spec.kind` is not `"huggingface"`, or no model reference is
                available to load from.
        """
        if spec.kind != "huggingface":
            raise ValueError(f"HFBackend requires a 'huggingface' spec; got kind {spec.kind!r}.")
        self.spec = spec
        self._open_session: ExclusiveSession | None = None
        self._reported_once: set[str] = set()

        if model_provider is not None:
            self._model_provider = model_provider
            self._tokenizer_provider = tokenizer_provider or (lambda: None)
            return

        if spec.model is None:
            raise ValueError(
                "HFBackend needs a model reference on the spec, or model_provider/"
                "tokenizer_provider for an already-loaded model."
            )
        hf_model_kwargs = dict(spec.get_option("hf_model_kwargs", default={}))
        device_map = spec.get_option("device_map", default=hf_model_kwargs.pop("device_map", "auto"))
        model = AutoModelForCausalLM.from_pretrained(
            spec.model,
            device_map=device_map,
            **hf_model_kwargs,
        )
        tokenizer = AutoTokenizer.from_pretrained(
            spec.get_option("tokenizer_name_or_path") or spec.model,
            trust_remote_code=bool(spec.get_option("trust_remote_code", default=False)),
        )
        tokenizer = ensure_pad_token(tokenizer)
        self._model_provider = lambda: model
        self._tokenizer_provider = lambda: tokenizer

    @classmethod
    def adopt(
        cls,
        spec: BackendSpec,
        model_provider: Callable[[], PreTrainedModel | None],
        tokenizer_provider: Callable[[], object | None],
    ) -> "HFBackend":
        """Wrap an already-loaded model and tokenizer without loading anything.

        Providers are read on every access, so a caller whose model is replaced mid-steer (a
        structural control returning a new model) always exposes the current one to sessions.

        Args:
            spec: The backend spec identifying this configuration.
            model_provider: Callable returning the current model (may return None before one
                exists).
            tokenizer_provider: Callable returning the current tokenizer.

        Returns:
            The adopting backend.
        """
        return cls(spec, model_provider=model_provider, tokenizer_provider=tokenizer_provider)

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The static Hugging Face capability advertisement (spec-independent)."""
        return HF_CAPABILITIES

    def open_session(self) -> "ExclusiveSession":
        """Open the backend's one exclusive session.

        Returns:
            The session, usable as a context manager.

        Raises:
            RuntimeError: If an exclusive session is already open on this backend.
        """
        if self._open_session is not None and not self._open_session.closed:
            raise RuntimeError(
                "An exclusive session is already open on this backend; close it before opening "
                "another."
            )
        self._open_session = ExclusiveSession(self)
        return self._open_session

    def report_once(self, key: str) -> bool:
        """Return True the first time `key` is seen on this backend, False afterwards.

        Sessions are opened per pipeline call, so a per-session flag would fire on every call;
        this guard lives on the backend so a repeated informational readout fires once.
        """
        if key in self._reported_once:
            return False
        self._reported_once.add(key)
        return True

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

adopt(spec, model_provider, tokenizer_provider) classmethod

Wrap an already-loaded model and tokenizer without loading anything.

Providers are read on every access, so a caller whose model is replaced mid-steer (a structural control returning a new model) always exposes the current one to sessions.

Parameters:

Name	Type	Description	Default
spec	BackendSpec	The backend spec identifying this configuration.	required
model_provider	Callable[[], PreTrainedModel | None]	Callable returning the current model (may return None before one exists).	required
tokenizer_provider	Callable[[], object | None]	Callable returning the current tokenizer.	required

Returns:

Type	Description
HFBackend	The adopting backend.

Source code in steerability/backends/huggingface/backend.py

@classmethod
def adopt(
    cls,
    spec: BackendSpec,
    model_provider: Callable[[], PreTrainedModel | None],
    tokenizer_provider: Callable[[], object | None],
) -> "HFBackend":
    """Wrap an already-loaded model and tokenizer without loading anything.

    Providers are read on every access, so a caller whose model is replaced mid-steer (a
    structural control returning a new model) always exposes the current one to sessions.

    Args:
        spec: The backend spec identifying this configuration.
        model_provider: Callable returning the current model (may return None before one
            exists).
        tokenizer_provider: Callable returning the current tokenizer.

    Returns:
        The adopting backend.
    """
    return cls(spec, model_provider=model_provider, tokenizer_provider=tokenizer_provider)

capabilities_for_spec(spec) classmethod

The static Hugging Face capability advertisement (spec-independent).

Source code in steerability/backends/huggingface/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The static Hugging Face capability advertisement (spec-independent)."""
    return HF_CAPABILITIES

open_session()

Open the backend's one exclusive session.

Returns:

Type	Description
ExclusiveSession	The session, usable as a context manager.

Raises:

Type	Description
RuntimeError	If an exclusive session is already open on this backend.

Source code in steerability/backends/huggingface/backend.py

def open_session(self) -> "ExclusiveSession":
    """Open the backend's one exclusive session.

    Returns:
        The session, usable as a context manager.

    Raises:
        RuntimeError: If an exclusive session is already open on this backend.
    """
    if self._open_session is not None and not self._open_session.closed:
        raise RuntimeError(
            "An exclusive session is already open on this backend; close it before opening "
            "another."
        )
    self._open_session = ExclusiveSession(self)
    return self._open_session

release()

Free the resources this backend owns beyond what the caller owns.

The default is a no-op. Engine-owning backends override this to shut their engines down deterministically. Release is idempotent.

Source code in steerability/algorithms/core/execution/backend.py

def release(self) -> None:
    """Free the resources this backend owns beyond what the caller owns.

    The default is a no-op. Engine-owning backends override this to shut their engines down
    deterministically. Release is idempotent.
    """
    return None

report_once(key)

Return True the first time key is seen on this backend, False afterwards.

Sessions are opened per pipeline call, so a per-session flag would fire on every call; this guard lives on the backend so a repeated informational readout fires once.

Source code in steerability/backends/huggingface/backend.py

def report_once(self, key: str) -> bool:
    """Return True the first time `key` is seen on this backend, False afterwards.

    Sessions are opened per pipeline call, so a per-session flag would fire on every call;
    this guard lives on the backend so a repeated informational readout fires once.
    """
    if key in self._reported_once:
        return False
    self._reported_once.add(key)
    return True

stage_artifacts(payloads)

Make each content-addressed artifact available to the execution side.

Called by the pipeline at the end of steer() with the tensor payloads of every lowered intervention spec, keyed by content-addressed artifact id. Staging is idempotent: an artifact that already exists at the destination is success. The in-process backend keeps live tensors and needs no staging; engine-backed backends write into the registry the serving engine reads.

Parameters:

Name	Type	Description	Default
payloads	dict[str, dict[str, Tensor]]	Mapping from artifact id to a name-to-tensor mapping.	required

Source code in steerability/algorithms/core/execution/backend.py

def stage_artifacts(self, payloads: dict[str, dict[str, torch.Tensor]]) -> None:
    """Make each content-addressed artifact available to the execution side.

    Called by the pipeline at the end of `steer()` with the tensor payloads of every
    lowered intervention spec, keyed by content-addressed artifact id. Staging is
    idempotent: an artifact that already exists at the destination is success. The
    in-process backend keeps live tensors and needs no staging; engine-backed backends
    write into the registry the serving engine reads.

    Args:
        payloads: Mapping from artifact id to a name-to-tensor mapping.
    """
    return None

session

The in-process exclusive session: direct model access, hook scopes, and the default decode loop.

logger = logging.getLogger(__name__) module-attribute

ExclusiveSession

The in-process session: direct model access, hook scopes, and the default decode loop.

Exposes .model for components whose requirements include Capability.IN_PROCESS_TORCH. The default decode delegates to model.generate. Items execute serially, each under its own hook registrations, which preserves in-process semantics for every entry combination.

Source code in steerability/backends/huggingface/session.py

class ExclusiveSession:
    """The in-process session: direct model access, hook scopes, and the default decode loop.

    Exposes `.model` for components whose requirements include `Capability.IN_PROCESS_TORCH`.
    The default decode delegates to `model.generate`. Items execute serially, each under its own
    hook registrations, which preserves in-process semantics for every entry combination.
    """

    def __init__(self, backend: "HFBackend") -> None:
        self._backend = backend
        self._closed = False
        self._generate_count = 0

    @property
    def closed(self) -> bool:
        """Whether the session has been closed."""
        return self._closed

    def close(self) -> None:
        """Close the session; further use raises `RuntimeError`."""
        self._closed = True

    def __enter__(self) -> "ExclusiveSession":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("This session is closed; open a new session on the backend.")

    @property
    def model(self) -> PreTrainedModel:
        """The live model.

        Raises:
            RuntimeError: If the session is closed or no model is available yet.
        """
        self._ensure_open()
        model = self._backend._model_provider()
        if model is None:
            raise RuntimeError("No model is available on this session.")
        return model

    @property
    def tokenizer(self) -> PreTrainedTokenizerBase | None:
        """The tokenizer, or None when the adopting caller has not resolved one yet."""
        self._ensure_open()
        return self._backend._tokenizer_provider()

    @property
    def layout(self) -> ModelFacts:
        """Structural facts derived from the loaded model, computed on every access so weight
        edits and model replacements are always reflected.

        `num_layers` comes from the resolved decoder layer list; `hidden_size`,
        `num_attention_heads`, and `head_dim` come from the text config (`text_config(model)`,
        the text sub-config on composite multimodal models), with a
        `hidden_size // num_attention_heads` fallback for `head_dim`; `dtype` from the model;
        `model_fingerprint` from the weight/config fingerprint; and `model_type` from the
        composite config, so a multimodal checkpoint keeps its wrapper `model_type`.
        """
        model = self.model

        from steerability.algorithms.core.internals.fingerprint import model_fingerprint

        _, layer_names = get_model_layer_list(model)
        text_cfg = text_config(model)
        hidden_size = text_cfg.hidden_size
        num_heads = getattr(text_cfg, "num_attention_heads", None)
        head_dim = getattr(text_cfg, "head_dim", None)
        if head_dim is None and num_heads:
            head_dim = hidden_size // num_heads
        return ModelFacts(
            num_layers=len(layer_names),
            hidden_size=hidden_size,
            num_attention_heads=num_heads,
            head_dim=head_dim,
            dtype=str(model.dtype).removeprefix("torch."),
            model_fingerprint=model_fingerprint(model),
            model_type=getattr(model.config, "model_type", None),
            model_ref=getattr(model, "name_or_path", None),
        )

    def _resolve_prompt_tensors(self, prompt: PreparedPrompt) -> tuple[torch.Tensor, torch.Tensor]:
        """Token ids and attention mask for one prompt, on the model device."""
        resolved = prompt.resolve_token_ids(self.tokenizer)
        device = self.model.device
        input_ids = resolved.token_ids.to(device)
        attention_mask = resolved.attention_mask
        if attention_mask is None:
            tokenizer = self.tokenizer
            if tokenizer is not None and tokenizer.pad_token_id is not None:
                attention_mask = infer_attention_mask_from_ids(input_ids, tokenizer.pad_token_id)
            else:
                attention_mask = torch.ones_like(input_ids, dtype=torch.long)
        attention_mask = attention_mask.to(dtype=input_ids.dtype, device=device)
        return input_ids, attention_mask

    def _compose_entry_stacks(
        self, output_entries, extra_processors=(), extra_criteria=(),
    ) -> tuple[LogitsProcessorList, StoppingCriteriaList]:
        """Compose the items' stack entries, appending caller extras after entry contributions."""
        processors: list = []
        criteria: list = []
        for entry in output_entries:
            if not isinstance(entry, StackEntry):
                raise UnsupportedOperationError(
                    f"{type(entry).__name__} requires an engine-hosted processor path; the "
                    "in-process session consumes StackEntry contributions."
                )
            processors.extend(entry.logits_processors)
            criteria.extend(entry.stopping_criteria)
        processors.extend(extra_processors)
        criteria.extend(extra_criteria)
        return LogitsProcessorList(processors), StoppingCriteriaList(criteria)

    def _register_state_entries(self, model: PreTrainedModel, state_entries) -> list:
        """Attach each entry's hook specifications to `model`, returning removable handles.

        The session is the single registrar of state hooks: entries are the only carriage, and
        registration lives strictly inside the session's execution of work. Pre and forward
        hooks register with `with_kwargs=True`; backward hooks register as full backward hooks.
        If registration fails partway, handles already attached are removed before re-raising.
        """
        handles: list = []
        try:
            for entry in state_entries:
                if not isinstance(entry, HookEntry):
                    raise UnsupportedOperationError(
                        f"{type(entry).__name__} requires an intervention-capable backend; the "
                        "in-process session consumes HookEntry contributions."
                    )
                for phase in ("pre", "forward", "backward"):
                    for spec in entry.hooks.get(phase, []):
                        module = model.get_submodule(spec["module"])
                        if phase == "pre":
                            handle = module.register_forward_pre_hook(spec["hook_func"], with_kwargs=True)
                        elif phase == "forward":
                            handle = module.register_forward_hook(spec["hook_func"], with_kwargs=True)
                        else:
                            handle = module.register_full_backward_hook(spec["hook_func"])
                        handles.append(handle)
        except Exception:
            for handle in handles:
                handle.remove()
            raise
        return handles

    @contextlib.contextmanager
    def entries_applied(self, state_entries):
        """Apply state entries to the live model for the duration of the context.

        Used by the pipeline around a client-side decoding driver's `decode`, so every forward
        the driver issues on the live model, including rollouts through this session and
        auxiliary scoring passes, runs under the generation's hooks. The session owns
        registration; hooks are removed when the context exits, even on error.
        """
        handles = self._register_state_entries(self.model, state_entries)
        try:
            yield self
        finally:
            for handle in handles:
                handle.remove()

    def _seeded(self, seed: int | None):
        """A context that snapshots and restores RNG state around a seeded decode.

        The CPU generator is always covered. On CUDA models every CUDA device generator is
        covered, since sharded models may sample on a device other than the first parameter's.
        On MPS models the MPS generator is covered.
        """
        if seed is None:
            return contextlib.nullcontext()
        device = self.model.device
        if device.type == "cuda":
            return torch.random.fork_rng(devices=list(range(torch.cuda.device_count())))
        if device.type == "mps":
            return _mps_rng_fork()
        return torch.random.fork_rng(devices=[])

    def _apply_seed(self, seed: int) -> None:
        """Seed the generators covered by `_seeded` for this decode."""
        torch.default_generator.manual_seed(seed)
        device = self.model.device
        if device.type == "cuda":
            torch.cuda.manual_seed_all(seed)
        elif device.type == "mps":
            torch.mps.manual_seed(seed)

    def _item_seeds(self, items: Sequence[GenerationItem], params: GenerationParams) -> list[int | None]:
        """Effective per-item seeds.

        An item's own seed is always honored. Otherwise, under `params.seed` with
        `seed_scope="item"`, each item derives its own seed from this call's operation id and its
        index; with `seed_scope="dispatch"`, every item takes the dispatch seed, derived at index 0
        so a single-item dispatch decodes identically under either scope. A dispatch mixing explicit
        and absent item seeds falls back to per-item derivation. Unseeded calls yield None.
        """
        operation_id = f"generate-{self._generate_count}"
        if (
            params.seed is not None
            and params.seed_scope == "dispatch"
            and all(item.seed is None for item in items)
        ):
            return [derive_item_seed(params.seed, operation_id, 0)] * len(items)
        seeds: list[int | None] = []
        for index, item in enumerate(items):
            if item.seed is not None:
                seeds.append(item.seed)
            elif params.seed is not None:
                seeds.append(derive_item_seed(params.seed, operation_id, index))
            else:
                seeds.append(None)
        return seeds

    def _report_serial_fallback(self, items: Sequence[GenerationItem]) -> None:
        """Log once per backend and reason why a multi-item dispatch decodes one item at a time.

        Reached only when `batchable` is False, so identical entries imply distinct seeds.
        """
        if not self._entries_identical(items):
            key, detail = "entries", (
                "items carry distinct state or output entries (row-scoped runtime kwargs or per-row hooks)"
            )
        else:
            key, detail = "seeds", (
                "items carry distinct seeds (seed_scope='item'); pass seed_scope='dispatch' to generate(), or "
                "set ProviderOptions.seed_scope on the evaluation provider, to decode the dispatch in one pass"
            )
        if self._backend.report_once(f"serial_fallback:{key}"):
            logger.info("Multi-item generate (%d items) decodes serially: %s.", len(items), detail)

    @staticmethod
    def _entries_identical(items: Sequence[GenerationItem | ScoringItem]) -> bool:
        """True when every item carries the same state and output entry objects."""
        first = items[0]
        for item in items[1:]:
            if len(item.state_entries) != len(first.state_entries) or any(
                a is not b for a, b in zip(item.state_entries, first.state_entries)
            ):
                return False
            if len(item.output_entries) != len(first.output_entries) or any(
                a is not b for a, b in zip(item.output_entries, first.output_entries)
            ):
                return False
        return True

    def _stack_prompt_rows(self, rows: list[tuple[torch.Tensor, torch.Tensor]]):
        """Stack resolved single-row prompts into one right-padded batch.

        The batched generate and score paths both left-pack the result via `to_left_pad` before
        the forward pass.
        """
        pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or 0
        max_len = max(ids.size(1) for ids, _ in rows)
        device = rows[0][0].device
        input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
        attention_mask = torch.zeros((len(rows), max_len), dtype=rows[0][1].dtype, device=device)
        for row, (ids, mask) in enumerate(rows):
            length = ids.size(1)
            input_ids[row, :length] = ids[0]
            attention_mask[row, :length] = mask[0]
        return input_ids, attention_mask

    def _classify(self, new_tokens: torch.Tensor, gen_kwargs: dict, params: GenerationParams) -> list[str | None]:
        """Per-row finish reasons under the pinned precedence, from the composed stop rules."""
        tokenizer = self.tokenizer
        return infer_finish_reasons(
            new_tokens,
            gen_kwargs,
            eos_token_id=getattr(tokenizer, "eos_token_id", None),
            pad_token_id=getattr(tokenizer, "pad_token_id", None),
            stop_strings=params.stop_strings,
            stop_token_ids=params.stop_token_ids,
            tokenizer=tokenizer,
        )

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item, each under its own hook registrations.

        Items sharing identical state entries, identical output entries, and identical-or-absent
        effective seeds execute in one batched `model.generate` pass (right-padded to a common
        prompt length, then left-packed together); otherwise items decode serially. Caller-supplied `logits_processor` and
        `stopping_criteria` entries in `params.extra` append after the items' own contributions,
        and the normalized stop fields compose as stop rules anchored at the prompt length. A
        seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the
        covered generator state (CPU, plus the model device's) is restored afterwards; when
        `params.seed` is set and an item carries no seed of its own, the item's seed derives per
        index, so multi-item fan-outs sample distinct streams.

        Args:
            items: The generation items.
            params: Normalized generation parameters shared by all items.

        Returns:
            One `ItemResult` per item, in item order. Each result's `finish_reasons` carries one
            reason per candidate with the precedence stop, then eos, then length, then None.
        """
        self._ensure_open()
        model = self.model
        tokenizer = self.tokenizer
        gen_kwargs = render_hf_gen_kwargs(params)
        # default pad_token_id per call to avoid the transformers open-end-generation warning
        # without mutating the model's generation config; caller kwargs and a model-configured
        # value both take precedence over the tokenizer fallback
        if (
            "pad_token_id" not in gen_kwargs
            and model.generation_config.pad_token_id is None
            and getattr(tokenizer, "pad_token_id", None) is not None
        ):
            gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
        user_processors = tuple(gen_kwargs.pop("logits_processor", None) or ())
        user_criteria = tuple(gen_kwargs.pop("stopping_criteria", None) or ())

        if not items:
            return []
        seeds = self._item_seeds(items, params)
        self._generate_count += 1

        batchable = (
            len(items) > 1
            and self._entries_identical(items)
            and len(set(seeds)) == 1
        )
        if batchable:
            return self._generate_batched(
                items, params, gen_kwargs, user_processors, user_criteria, seeds[0],
            )
        if len(items) > 1:
            self._report_serial_fallback(items)

        results: list[ItemResult] = []
        for index, item in enumerate(items):
            input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
            processors, criteria = self._compose_entry_stacks(
                item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
            )
            criteria.extend(compose_stop_criteria(params, input_ids.size(1), tokenizer))
            stacks = stack_generate_kwargs(processors, criteria)
            handles = self._register_state_entries(model, item.state_entries)
            try:
                seed = seeds[index]
                with self._seeded(seed):
                    if seed is not None:
                        self._apply_seed(seed)
                    full_ids = model.generate(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        **stacks,
                        **gen_kwargs,
                    )
            finally:
                for handle in handles:
                    handle.remove()

            new_tokens = full_ids[:, input_ids.size(1):]
            reasons = self._classify(new_tokens, gen_kwargs, params)
            results.append(ItemResult(
                index=index,
                output=Output(
                    output_ids=new_tokens,
                    adapted_input_ids=input_ids,
                    finish_reason=reasons[0],
                    finish_reasons=tuple(reasons),
                ),
            ))
        return results

    def _generate_batched(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
        gen_kwargs: dict,
        user_processors: tuple,
        user_criteria: tuple,
        seed: int | None,
    ) -> list[ItemResult]:
        """One `model.generate` pass over all items (identical entries, one shared seed).

        The stacked prompts left-pack (pad positions move before the real tokens), so every row's
        first generated token is predicted from its last real token rather than from a trailing pad.
        """
        model = self.model
        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
        input_ids, attention_mask = self._stack_prompt_rows(rows)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
        processors, criteria = self._compose_entry_stacks(
            items[0].output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
        )
        criteria.extend(compose_stop_criteria(params, input_ids.size(1), self.tokenizer))
        stacks = stack_generate_kwargs(processors, criteria)
        handles = self._register_state_entries(model, items[0].state_entries)
        try:
            with self._seeded(seed):
                if seed is not None:
                    self._apply_seed(seed)
                full_ids = model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **stacks,
                    **gen_kwargs,
                )
        finally:
            for handle in handles:
                handle.remove()

        prompt_len = input_ids.size(1)
        new_tokens = full_ids[:, prompt_len:]
        num_candidates = params.n or 1
        results: list[ItemResult] = []
        for index in range(len(items)):
            item_rows = new_tokens[index * num_candidates:(index + 1) * num_candidates]
            reasons = self._classify(item_rows, gen_kwargs, params)
            results.append(ItemResult(
                index=index,
                output=Output(
                    output_ids=item_rows,
                    adapted_input_ids=input_ids[index:index + 1],
                    finish_reason=reasons[0],
                    finish_reasons=tuple(reasons),
                ),
            ))
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens.

        For each item, the prompt left-packs (pad positions move before the real tokens) so the
        reference follows the prompt's last real token, then prompt and reference concatenate
        into one causal forward pass under the item's hook registrations; the item's logits
        processors replay position-by-position with the same `(prefix_ids, scores)` view they
        receive during generation. Items sharing identical state and output entries score in one
        batched forward pass (prompts right-padded to a common length, then left-packed
        together); otherwise items score serially. Stopping criteria never apply. Decoder-only
        models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: `params.extra` passes through as forward keyword arguments.

        Returns:
            Log probabilities of shape `[num_items, ref_len]`.

        Raises:
            ValueError: If items carry differing reference lengths.
            UnsupportedOperationError: If the model is an encoder-decoder model.
        """
        self._ensure_open()
        model = self.model
        if getattr(model.config, "is_encoder_decoder", False):
            raise UnsupportedOperationError(
                "Session scoring supports decoder-only models; encoder-decoder scoring runs "
                "through SteeringPipeline.compute_logprobs."
            )
        device = model.device
        forward_kwargs = dict(params.extra)

        if not items:
            return torch.zeros((0, 0), device=device, dtype=torch.float32)
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")

        if len(items) > 1 and self._entries_identical(items):
            return self._score_batched(items, forward_kwargs)

        all_logprobs: list[torch.Tensor] = []
        for item in items:
            input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
            input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
            ref_output_ids = item.ref_output_ids
            if ref_output_ids.ndim == 1:
                ref_output_ids = ref_output_ids.unsqueeze(0)
            ref_output_ids = ref_output_ids.to(device)
            ref_len = ref_output_ids.size(1)
            if ref_len == 0:
                all_logprobs.append(torch.zeros((1, 0), device=device, dtype=torch.float32))
                continue

            processors, _ = self._compose_entry_stacks(item.output_entries)
            handles = self._register_state_entries(model, item.state_entries)
            try:
                with torch.no_grad():
                    combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                    combined_mask = torch.cat([
                        attention_mask,
                        torch.ones(1, ref_len, device=device, dtype=attention_mask.dtype),
                    ], dim=1)
                    outputs = model(
                        input_ids=combined_ids,
                        attention_mask=combined_mask,
                        **forward_kwargs,
                    )
                    input_len = input_ids.size(1)
                    logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                    if len(processors):
                        for t in range(logits.size(1)):
                            prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                            logits[:, t, :] = processors(prefix, logits[:, t, :])
                    logprobs = torch.log_softmax(logits, dim=-1)
                    all_logprobs.append(
                        logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
                    )
            finally:
                for handle in handles:
                    handle.remove()

        return torch.cat(all_logprobs, dim=0)

    def _score_batched(self, items: Sequence[ScoringItem], forward_kwargs: dict) -> torch.Tensor:
        """One causal forward pass over all items (identical entries)."""
        model = self.model
        device = model.device
        rows = [self._resolve_prompt_tensors(item.prompt) for item in items]
        input_ids, attention_mask = self._stack_prompt_rows(rows)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)

        refs: list[torch.Tensor] = []
        for item in items:
            ref = item.ref_output_ids
            if ref.ndim == 1:
                ref = ref.unsqueeze(0)
            refs.append(ref.to(device))
        ref_output_ids = torch.cat(refs, dim=0)
        ref_len = ref_output_ids.size(1)
        if ref_len == 0:
            return torch.zeros((len(items), 0), device=device, dtype=torch.float32)

        processors, _ = self._compose_entry_stacks(items[0].output_entries)
        handles = self._register_state_entries(model, items[0].state_entries)
        try:
            with torch.no_grad():
                combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                combined_mask = torch.cat([
                    attention_mask,
                    torch.ones(len(items), ref_len, device=device, dtype=attention_mask.dtype),
                ], dim=1)
                outputs = model(
                    input_ids=combined_ids,
                    attention_mask=combined_mask,
                    **forward_kwargs,
                )
                input_len = input_ids.size(1)
                logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                if len(processors):
                    for t in range(logits.size(1)):
                        prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                        logits[:, t, :] = processors(prefix, logits[:, t, :])
                logprobs = torch.log_softmax(logits, dim=-1)
                return logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
        finally:
            for handle in handles:
                handle.remove()

    def capture(
        self,
        prompts: list[PreparedPrompt],
        layers: list[int],
        mode: Literal["all_tokens", "last_token"],
        location: Literal["layer_output", "layer_input"] = "layer_output",
    ) -> CaptureResult:
        """Capture residual-stream hidden states for `prompts` at `layers`.

        Prompts resolve to token ids, right-pad into one batch, and run through the shared
        layerwise extraction. In `"all_tokens"` mode each layer's tensor is `[N, T, H]` with rows
        outside the attention mask zeroed; in `"last_token"` mode the last real (non-pad) position
        of each row is selected, giving `[N, H]`.

        Args:
            prompts: The prompts to capture.
            layers: 0-based layer ids to keep.
            mode: `"all_tokens"` or `"last_token"`.
            location: `"layer_output"` (a layer's output boundary) or `"layer_input"` (the
                boundary a forward pre-hook observes).

        Returns:
            The captured tensors and the batch attention mask, on CPU.

        Raises:
            ValueError: If `mode` is unknown or a requested layer id is out of range.
        """
        self._ensure_open()
        if mode not in ("all_tokens", "last_token"):
            raise ValueError(f"Unknown capture mode {mode!r}; modes are 'all_tokens', 'last_token'.")
        if not prompts:
            raise ValueError("capture() requires at least one prompt.")

        from steerability.algorithms.core.internals.capture import layerwise_tokenwise_hidden
        from steerability.algorithms.core.internals.pooling import aggregate_condition_hidden

        model = self.model
        device = model.device
        tokenizer = self.tokenizer
        pad_token_id = getattr(tokenizer, "pad_token_id", None) or 0

        rows = [self._resolve_prompt_tensors(prompt) for prompt in prompts]
        max_len = max(ids.size(1) for ids, _ in rows)
        input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
        attention_mask = torch.zeros((len(rows), max_len), dtype=torch.long, device=device)
        for row, (ids, mask) in enumerate(rows):
            length = ids.size(1)
            input_ids[row, :length] = ids[0]
            attention_mask[row, :length] = mask[0]

        enc = {"input_ids": input_ids, "attention_mask": attention_mask}
        hidden = layerwise_tokenwise_hidden(
            model, enc, batch_size=_CAPTURE_BATCH_SIZE, location=location,
        )
        missing = [layer for layer in layers if layer not in hidden]
        if missing:
            raise ValueError(
                f"Requested layer ids {missing} are out of range; the model has {len(hidden)} layers."
            )

        mask_cpu = attention_mask.cpu()
        selected = {layer: hidden[layer] for layer in layers}
        if mode == "last_token":
            selected = {
                layer: aggregate_condition_hidden(tensor, "last", mask_cpu)
                for layer, tensor in selected.items()
            }
        return CaptureResult(
            hidden=selected, attention_mask=mask_cpu, mode=mode, location=location,
        )

closed property

Whether the session has been closed.

layout property

Structural facts derived from the loaded model, computed on every access so weight edits and model replacements are always reflected.

num_layers comes from the resolved decoder layer list; hidden_size, num_attention_heads, and head_dim come from the text config (text_config(model), the text sub-config on composite multimodal models), with a hidden_size // num_attention_heads fallback for head_dim; dtype from the model; model_fingerprint from the weight/config fingerprint; and model_type from the composite config, so a multimodal checkpoint keeps its wrapper model_type.

model property

The live model.

Raises:

Type	Description
RuntimeError	If the session is closed or no model is available yet.

tokenizer property

The tokenizer, or None when the adopting caller has not resolved one yet.

capture(prompts, layers, mode, location='layer_output')

Capture residual-stream hidden states for prompts at layers.

Prompts resolve to token ids, right-pad into one batch, and run through the shared layerwise extraction. In "all_tokens" mode each layer's tensor is [N, T, H] with rows outside the attention mask zeroed; in "last_token" mode the last real (non-pad) position of each row is selected, giving [N, H].

Parameters:

Name	Type	Description	Default
prompts	list[PreparedPrompt]	The prompts to capture.	required
layers	list[int]	0-based layer ids to keep.	required
mode	Literal['all_tokens', 'last_token']	"all_tokens" or "last_token".	required
location	Literal['layer_output', 'layer_input']	"layer_output" (a layer's output boundary) or "layer_input" (the boundary a forward pre-hook observes).	'layer_output'

Returns:

Type	Description
CaptureResult	The captured tensors and the batch attention mask, on CPU.

Raises:

Type	Description
ValueError	If mode is unknown or a requested layer id is out of range.

Source code in steerability/backends/huggingface/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Capture residual-stream hidden states for `prompts` at `layers`.

    Prompts resolve to token ids, right-pad into one batch, and run through the shared
    layerwise extraction. In `"all_tokens"` mode each layer's tensor is `[N, T, H]` with rows
    outside the attention mask zeroed; in `"last_token"` mode the last real (non-pad) position
    of each row is selected, giving `[N, H]`.

    Args:
        prompts: The prompts to capture.
        layers: 0-based layer ids to keep.
        mode: `"all_tokens"` or `"last_token"`.
        location: `"layer_output"` (a layer's output boundary) or `"layer_input"` (the
            boundary a forward pre-hook observes).

    Returns:
        The captured tensors and the batch attention mask, on CPU.

    Raises:
        ValueError: If `mode` is unknown or a requested layer id is out of range.
    """
    self._ensure_open()
    if mode not in ("all_tokens", "last_token"):
        raise ValueError(f"Unknown capture mode {mode!r}; modes are 'all_tokens', 'last_token'.")
    if not prompts:
        raise ValueError("capture() requires at least one prompt.")

    from steerability.algorithms.core.internals.capture import layerwise_tokenwise_hidden
    from steerability.algorithms.core.internals.pooling import aggregate_condition_hidden

    model = self.model
    device = model.device
    tokenizer = self.tokenizer
    pad_token_id = getattr(tokenizer, "pad_token_id", None) or 0

    rows = [self._resolve_prompt_tensors(prompt) for prompt in prompts]
    max_len = max(ids.size(1) for ids, _ in rows)
    input_ids = torch.full((len(rows), max_len), pad_token_id, dtype=torch.long, device=device)
    attention_mask = torch.zeros((len(rows), max_len), dtype=torch.long, device=device)
    for row, (ids, mask) in enumerate(rows):
        length = ids.size(1)
        input_ids[row, :length] = ids[0]
        attention_mask[row, :length] = mask[0]

    enc = {"input_ids": input_ids, "attention_mask": attention_mask}
    hidden = layerwise_tokenwise_hidden(
        model, enc, batch_size=_CAPTURE_BATCH_SIZE, location=location,
    )
    missing = [layer for layer in layers if layer not in hidden]
    if missing:
        raise ValueError(
            f"Requested layer ids {missing} are out of range; the model has {len(hidden)} layers."
        )

    mask_cpu = attention_mask.cpu()
    selected = {layer: hidden[layer] for layer in layers}
    if mode == "last_token":
        selected = {
            layer: aggregate_condition_hidden(tensor, "last", mask_cpu)
            for layer, tensor in selected.items()
        }
    return CaptureResult(
        hidden=selected, attention_mask=mask_cpu, mode=mode, location=location,
    )

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/huggingface/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

entries_applied(state_entries)

Apply state entries to the live model for the duration of the context.

Used by the pipeline around a client-side decoding driver's decode, so every forward the driver issues on the live model, including rollouts through this session and auxiliary scoring passes, runs under the generation's hooks. The session owns registration; hooks are removed when the context exits, even on error.

Source code in steerability/backends/huggingface/session.py

@contextlib.contextmanager
def entries_applied(self, state_entries):
    """Apply state entries to the live model for the duration of the context.

    Used by the pipeline around a client-side decoding driver's `decode`, so every forward
    the driver issues on the live model, including rollouts through this session and
    auxiliary scoring passes, runs under the generation's hooks. The session owns
    registration; hooks are removed when the context exits, even on error.
    """
    handles = self._register_state_entries(self.model, state_entries)
    try:
        yield self
    finally:
        for handle in handles:
            handle.remove()

generate(items, params)

Generate one result per item, each under its own hook registrations.

Items sharing identical state entries, identical output entries, and identical-or-absent effective seeds execute in one batched model.generate pass (right-padded to a common prompt length, then left-packed together); otherwise items decode serially. Caller-supplied logits_processor and stopping_criteria entries in params.extra append after the items' own contributions, and the normalized stop fields compose as stop rules anchored at the prompt length. A seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the covered generator state (CPU, plus the model device's) is restored afterwards; when params.seed is set and an item carries no seed of its own, the item's seed derives per index, so multi-item fan-outs sample distinct streams.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items.	required
params	GenerationParams	Normalized generation parameters shared by all items.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order. Each result's finish_reasons carries one
list[ItemResult]	reason per candidate with the precedence stop, then eos, then length, then None.

Source code in steerability/backends/huggingface/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item, each under its own hook registrations.

    Items sharing identical state entries, identical output entries, and identical-or-absent
    effective seeds execute in one batched `model.generate` pass (right-padded to a common
    prompt length, then left-packed together); otherwise items decode serially. Caller-supplied `logits_processor` and
    `stopping_criteria` entries in `params.extra` append after the items' own contributions,
    and the normalized stop fields compose as stop rules anchored at the prompt length. A
    seeded item decodes inside a seeded RNG fork, so seeded runs are reproducible and the
    covered generator state (CPU, plus the model device's) is restored afterwards; when
    `params.seed` is set and an item carries no seed of its own, the item's seed derives per
    index, so multi-item fan-outs sample distinct streams.

    Args:
        items: The generation items.
        params: Normalized generation parameters shared by all items.

    Returns:
        One `ItemResult` per item, in item order. Each result's `finish_reasons` carries one
        reason per candidate with the precedence stop, then eos, then length, then None.
    """
    self._ensure_open()
    model = self.model
    tokenizer = self.tokenizer
    gen_kwargs = render_hf_gen_kwargs(params)
    # default pad_token_id per call to avoid the transformers open-end-generation warning
    # without mutating the model's generation config; caller kwargs and a model-configured
    # value both take precedence over the tokenizer fallback
    if (
        "pad_token_id" not in gen_kwargs
        and model.generation_config.pad_token_id is None
        and getattr(tokenizer, "pad_token_id", None) is not None
    ):
        gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
    user_processors = tuple(gen_kwargs.pop("logits_processor", None) or ())
    user_criteria = tuple(gen_kwargs.pop("stopping_criteria", None) or ())

    if not items:
        return []
    seeds = self._item_seeds(items, params)
    self._generate_count += 1

    batchable = (
        len(items) > 1
        and self._entries_identical(items)
        and len(set(seeds)) == 1
    )
    if batchable:
        return self._generate_batched(
            items, params, gen_kwargs, user_processors, user_criteria, seeds[0],
        )
    if len(items) > 1:
        self._report_serial_fallback(items)

    results: list[ItemResult] = []
    for index, item in enumerate(items):
        input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
        processors, criteria = self._compose_entry_stacks(
            item.output_entries, extra_processors=user_processors, extra_criteria=user_criteria,
        )
        criteria.extend(compose_stop_criteria(params, input_ids.size(1), tokenizer))
        stacks = stack_generate_kwargs(processors, criteria)
        handles = self._register_state_entries(model, item.state_entries)
        try:
            seed = seeds[index]
            with self._seeded(seed):
                if seed is not None:
                    self._apply_seed(seed)
                full_ids = model.generate(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    **stacks,
                    **gen_kwargs,
                )
        finally:
            for handle in handles:
                handle.remove()

        new_tokens = full_ids[:, input_ids.size(1):]
        reasons = self._classify(new_tokens, gen_kwargs, params)
        results.append(ItemResult(
            index=index,
            output=Output(
                output_ids=new_tokens,
                adapted_input_ids=input_ids,
                finish_reason=reasons[0],
                finish_reasons=tuple(reasons),
            ),
        ))
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens.

For each item, the prompt left-packs (pad positions move before the real tokens) so the reference follows the prompt's last real token, then prompt and reference concatenate into one causal forward pass under the item's hook registrations; the item's logits processors replay position-by-position with the same (prefix_ids, scores) view they receive during generation. Items sharing identical state and output entries score in one batched forward pass (prompts right-padded to a common length, then left-packed together); otherwise items score serially. Stopping criteria never apply. Decoder-only models only; the pipeline's compute_logprobs serves encoder-decoder models in-process.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	params.extra passes through as forward keyword arguments.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len].

Raises:

Type	Description
ValueError	If items carry differing reference lengths.
UnsupportedOperationError	If the model is an encoder-decoder model.

Source code in steerability/backends/huggingface/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens.

    For each item, the prompt left-packs (pad positions move before the real tokens) so the
    reference follows the prompt's last real token, then prompt and reference concatenate
    into one causal forward pass under the item's hook registrations; the item's logits
    processors replay position-by-position with the same `(prefix_ids, scores)` view they
    receive during generation. Items sharing identical state and output entries score in one
    batched forward pass (prompts right-padded to a common length, then left-packed
    together); otherwise items score serially. Stopping criteria never apply. Decoder-only
    models only; the pipeline's `compute_logprobs` serves encoder-decoder models in-process.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: `params.extra` passes through as forward keyword arguments.

    Returns:
        Log probabilities of shape `[num_items, ref_len]`.

    Raises:
        ValueError: If items carry differing reference lengths.
        UnsupportedOperationError: If the model is an encoder-decoder model.
    """
    self._ensure_open()
    model = self.model
    if getattr(model.config, "is_encoder_decoder", False):
        raise UnsupportedOperationError(
            "Session scoring supports decoder-only models; encoder-decoder scoring runs "
            "through SteeringPipeline.compute_logprobs."
        )
    device = model.device
    forward_kwargs = dict(params.extra)

    if not items:
        return torch.zeros((0, 0), device=device, dtype=torch.float32)
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")

    if len(items) > 1 and self._entries_identical(items):
        return self._score_batched(items, forward_kwargs)

    all_logprobs: list[torch.Tensor] = []
    for item in items:
        input_ids, attention_mask = self._resolve_prompt_tensors(item.prompt)
        input_ids, attention_mask = to_left_pad(input_ids, attention_mask)
        ref_output_ids = item.ref_output_ids
        if ref_output_ids.ndim == 1:
            ref_output_ids = ref_output_ids.unsqueeze(0)
        ref_output_ids = ref_output_ids.to(device)
        ref_len = ref_output_ids.size(1)
        if ref_len == 0:
            all_logprobs.append(torch.zeros((1, 0), device=device, dtype=torch.float32))
            continue

        processors, _ = self._compose_entry_stacks(item.output_entries)
        handles = self._register_state_entries(model, item.state_entries)
        try:
            with torch.no_grad():
                combined_ids = torch.cat([input_ids, ref_output_ids], dim=1)
                combined_mask = torch.cat([
                    attention_mask,
                    torch.ones(1, ref_len, device=device, dtype=attention_mask.dtype),
                ], dim=1)
                outputs = model(
                    input_ids=combined_ids,
                    attention_mask=combined_mask,
                    **forward_kwargs,
                )
                input_len = input_ids.size(1)
                logits = outputs.logits[:, input_len - 1: input_len + ref_len - 1, :]
                if len(processors):
                    for t in range(logits.size(1)):
                        prefix = torch.cat([input_ids, ref_output_ids[:, :t]], dim=1)
                        logits[:, t, :] = processors(prefix, logits[:, t, :])
                logprobs = torch.log_softmax(logits, dim=-1)
                all_logprobs.append(
                    logprobs.gather(dim=-1, index=ref_output_ids.unsqueeze(-1)).squeeze(-1)
                )
        finally:
            for handle in handles:
                handle.remove()

    return torch.cat(all_logprobs, dim=0)

compose_stop_criteria(params, prompt_len, tokenizer)

The stop criteria implied by the normalized stop fields, anchored at prompt_len.

Parameters:

Name	Type	Description	Default
params	GenerationParams	The normalized parameters carrying stop_strings and stop_token_ids.	required
prompt_len	int	Prompt length the substring criteria decode past.	required
tokenizer	PreTrainedTokenizerBase | None	Tokenizer for substring decoding; required when stop strings are set.	required

Returns:

Type	Description
list	The composed criteria (possibly empty).

Raises:

Type	Description
ValueError	If stop strings are set and no tokenizer is available.

Source code in steerability/backends/huggingface/session.py

def compose_stop_criteria(
    params: GenerationParams, prompt_len: int, tokenizer: PreTrainedTokenizerBase | None
) -> list:
    """The stop criteria implied by the normalized stop fields, anchored at `prompt_len`.

    Args:
        params: The normalized parameters carrying `stop_strings` and `stop_token_ids`.
        prompt_len: Prompt length the substring criteria decode past.
        tokenizer: Tokenizer for substring decoding; required when stop strings are set.

    Returns:
        The composed criteria (possibly empty).

    Raises:
        ValueError: If stop strings are set and no tokenizer is available.
    """
    criteria: list = []
    if params.stop_strings:
        if tokenizer is None:
            raise ValueError("stop_strings require a tokenizer on the session.")
        for text in params.stop_strings:
            criteria.append(StopOnSubstring(tokenizer, text, prompt_len))
    if params.stop_token_ids:
        criteria.append(StopOnTokens(params.stop_token_ids))
    return criteria

render_hf_gen_kwargs(params)

Render normalized generation parameters onto model.generate keyword arguments.

Keys in params.extra pass through untouched; a normalized field always takes precedence over a same-named extra key. seed is not rendered here, since the session applies it as a fork_rng-scoped manual_seed around the item's decode; the stop fields are not rendered either, since the session composes them as prompt-anchored stop criteria (compose_stop_criteria).

Parameters:

Name	Type	Description	Default
params	GenerationParams	The normalized parameters.	required

Returns:

Type	Description
dict	Keyword arguments for model.generate.

Source code in steerability/backends/huggingface/session.py

def render_hf_gen_kwargs(params: GenerationParams) -> dict:
    """Render normalized generation parameters onto `model.generate` keyword arguments.

    Keys in `params.extra` pass through untouched; a normalized field always takes precedence
    over a same-named extra key. `seed` is not rendered here, since the session applies it as a
    `fork_rng`-scoped `manual_seed` around the item's decode; the stop fields are not rendered
    either, since the session composes them as prompt-anchored stop criteria
    (`compose_stop_criteria`).

    Args:
        params: The normalized parameters.

    Returns:
        Keyword arguments for `model.generate`.
    """
    gen_kwargs = dict(params.extra)
    if params.max_new_tokens is not None:
        gen_kwargs["max_new_tokens"] = params.max_new_tokens
    if params.min_new_tokens is not None:
        gen_kwargs["min_new_tokens"] = params.min_new_tokens
    if params.temperature is not None:
        gen_kwargs["temperature"] = params.temperature
    if params.top_p is not None:
        gen_kwargs["top_p"] = params.top_p
    if params.top_k is not None:
        gen_kwargs["top_k"] = params.top_k
    if params.repetition_penalty is not None:
        gen_kwargs["repetition_penalty"] = params.repetition_penalty
    if params.greedy is not None:
        gen_kwargs["do_sample"] = not params.greedy
    if params.n is not None:
        gen_kwargs["num_return_sequences"] = params.n
    return gen_kwargs

vllm

The vLLM backends: the offline engine ("vllm") and the OpenAI-compatible server ("vllm-serve").

This package imports cleanly without vLLM installed. The capability tables (capabilities.py) are static data used by check(); the strict parameter-rendering table and the request/response mapping helpers (rendering.py) are plain functions; the two backend classes and their backend-side helpers live in backend.py, and the request sessions in session.py. Constructing VLLMBackend requires the vllm optional dependency (it boots an engine); VLLMServeBackend needs only a reachable vLLM server.

VLLM_BASELINE_CAPABILITIES = BackendCapabilities(atoms=frozenset({Capability.SERVE_CHECKPOINT, Capability.SERVE_LORA, Capability.GUIDED_DECODING}), constraint_kinds=_VLLM_CONSTRAINT_KINDS) module-attribute

VLLMBackend

Bases: Backend

The offline vLLM engine backend.

Boots one engine per backend instance from the spec (engine_kwargs option forwarded to vllm.LLM); requires the vllm optional dependency. A CheckpointArtifact overrides the served model reference and a LoRAArtifact attaches as a LoRA request on every generation. When the spec declares hook_plugin, the unified worker is selected via VLLM_HOOK_WORKER=unified and the discovery payload is fetched once and cached by spec hash. Capability advertisement is available through capabilities_for_spec without constructing the backend.

Source code in steerability/backends/vllm/backend.py

class VLLMBackend(Backend):
    """The offline vLLM engine backend.

    Boots one engine per backend instance from the spec (`engine_kwargs` option forwarded to
    `vllm.LLM`); requires the `vllm` optional dependency. A `CheckpointArtifact` overrides the
    served model reference and a `LoRAArtifact` attaches as a LoRA request on every generation.
    When the spec declares `hook_plugin`, the unified worker is selected via
    `VLLM_HOOK_WORKER=unified` and the discovery payload is fetched once and cached by spec
    hash. Capability advertisement is available through `capabilities_for_spec` without
    constructing the backend.
    """

    def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> None:
        if spec.kind != "vllm":
            raise ValueError(f"VLLMBackend requires a 'vllm' spec; got kind {spec.kind!r}.")
        self.spec = spec
        self._released = False
        require("vllm")

        from vllm import LLM

        checkpoint, lora = _split_artifacts(artifacts)
        model_ref = checkpoint.path if checkpoint is not None else spec.model
        if model_ref is None:
            raise ValueError("VLLMBackend needs a model reference on the spec or a checkpoint artifact.")
        trust_remote_code = bool(spec.get_option("trust_remote_code", default=False))
        _reject_encoder_decoder(model_ref, trust_remote_code)

        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"})
        if lora is not None:
            engine_kwargs.setdefault("enable_lora", True)
        if trust_remote_code:
            engine_kwargs.setdefault("trust_remote_code", True)
        if spec.get_option("hook_plugin"):
            # worker hooks do not run under CUDA-graph replay; spec construction rejects an
            # explicit False, so this only fills the default
            engine_kwargs.setdefault("enforce_eager", True)

        # the boot environment is scoped to this engine's construction and restored afterwards,
        # so a later plugin-free engine in the same process is unaffected
        forced, defaults = engine_boot_environment(bool(spec.get_option("hook_plugin")))
        with engine_environment(forced, defaults):
            self._llm = LLM(model=model_ref, **engine_kwargs)

        # the pipeline records a backend only when the constructor returns, so a failure from
        # here on must release the engine before propagating
        try:
            self._lora_request = None
            if lora is not None:
                from vllm.lora.request import LoRARequest

                self._lora_request = LoRARequest("steered", 1, lora.path)

            tokenizer_source = (
                spec.get_option("tokenizer_name_or_path")
                or model_ref
            )
            self.tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code)
            self._layout = _config_layout(model_ref, trust_remote_code)
            self._plain_salt = uuid.uuid4().hex
            self._artifact_uploader = _ArtifactUploader(spec.get_option("artifact_dir"))
            self._discovery: dict | None = None
            if spec.get_option("hook_plugin"):
                self._discovery = self._fetch_discovery()
        except Exception:
            self.release()
            raise

    def stage_artifacts(self, payloads) -> None:
        """Write each content-addressed artifact into the plugin registry the engine reads.

        The offline engine shares the process's filesystem, so staging is a registry write
        (idempotent, verified against the content address).
        """
        self._artifact_uploader.upload_payloads(payloads)

    def _fetch_discovery(self) -> dict | None:
        cached = _DISCOVERY_CACHE.get(self.spec.spec_hash)
        if cached is not None:
            return cached
        payload = None
        for target in (self._llm, getattr(self._llm, "llm_engine", None)):
            rpc = getattr(target, "collective_rpc", None)
            if callable(rpc):
                try:
                    replies = rpc("hook_capabilities")
                except Exception as error:
                    logger.warning("vLLM-Hook discovery failed: %s", error)
                    return None
                payload = next((reply for reply in replies if reply), None)
                break
        if payload is None:
            logger.warning(
                "vLLM-Hook discovery returned no payload; is VLLM_HOOK_WORKER=unified active?"
            )
            return None
        _DISCOVERY_CACHE[self.spec.spec_hash] = payload
        _reconcile_discovery(self.spec, self.capabilities_for_spec(self.spec), payload)
        return payload

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The capability advertisement implied by `spec`."""
        return _vllm_capabilities(spec, offline=True)

    def open_session(self) -> "VLLMOfflineSession":
        """Open a request session over the shared engine.

        Raises:
            RuntimeError: If the backend has been released.
        """
        self._require_llm()
        return VLLMOfflineSession(self)

    def _require_llm(self):
        """The live engine, or a `RuntimeError` when the backend has been released."""
        if self._llm is None:
            raise RuntimeError(
                "This VLLMBackend was released; construct a new backend (or a new "
                "SteeringPipeline operation, which reconstructs backends automatically)."
            )
        return self._llm

    def release(self) -> None:
        """Shut the engine down explicitly and mark the backend unusable.

        Release is idempotent; after it, a new backend must be constructed. The distributed-state
        teardown is process-global, so release assumes no other live vLLM engine in the process.
        Ray-based executors are out of scope. Engine-touching calls on any still-open session raise
        after release.
        """
        if self._released:
            return
        self._released = True
        llm = self._llm
        self._llm = None
        self._lora_request = None

        for resolve in (
            lambda: getattr(llm, "shutdown", None),
            lambda: getattr(getattr(llm, "llm_engine", None), "shutdown", None),
            lambda: getattr(
                getattr(getattr(llm, "llm_engine", None), "engine_core", None), "shutdown", None
            ),
        ):
            shutdown = resolve()
            if callable(shutdown):
                try:
                    shutdown()
                except Exception:
                    logger.warning("vLLM engine shutdown hop failed; continuing.", exc_info=True)
                break

        del llm
        gc.collect()

        try:
            from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel

            destroy_model_parallel()
            destroy_distributed_environment()
        except Exception:
            logger.warning("vLLM distributed-state teardown failed; continuing.", exc_info=True)

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code) instance-attribute

capabilities_for_spec(spec) classmethod

The capability advertisement implied by spec.

Source code in steerability/backends/vllm/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The capability advertisement implied by `spec`."""
    return _vllm_capabilities(spec, offline=True)

open_session()

Open a request session over the shared engine.

Raises:

Type	Description
RuntimeError	If the backend has been released.

Source code in steerability/backends/vllm/backend.py

def open_session(self) -> "VLLMOfflineSession":
    """Open a request session over the shared engine.

    Raises:
        RuntimeError: If the backend has been released.
    """
    self._require_llm()
    return VLLMOfflineSession(self)

release()

Shut the engine down explicitly and mark the backend unusable.

Release is idempotent; after it, a new backend must be constructed. The distributed-state teardown is process-global, so release assumes no other live vLLM engine in the process. Ray-based executors are out of scope. Engine-touching calls on any still-open session raise after release.

Source code in steerability/backends/vllm/backend.py

def release(self) -> None:
    """Shut the engine down explicitly and mark the backend unusable.

    Release is idempotent; after it, a new backend must be constructed. The distributed-state
    teardown is process-global, so release assumes no other live vLLM engine in the process.
    Ray-based executors are out of scope. Engine-touching calls on any still-open session raise
    after release.
    """
    if self._released:
        return
    self._released = True
    llm = self._llm
    self._llm = None
    self._lora_request = None

    for resolve in (
        lambda: getattr(llm, "shutdown", None),
        lambda: getattr(getattr(llm, "llm_engine", None), "shutdown", None),
        lambda: getattr(
            getattr(getattr(llm, "llm_engine", None), "engine_core", None), "shutdown", None
        ),
    ):
        shutdown = resolve()
        if callable(shutdown):
            try:
                shutdown()
            except Exception:
                logger.warning("vLLM engine shutdown hop failed; continuing.", exc_info=True)
            break

    del llm
    gc.collect()

    try:
        from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel

        destroy_model_parallel()
        destroy_distributed_environment()
    except Exception:
        logger.warning("vLLM distributed-state teardown failed; continuing.", exc_info=True)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

stage_artifacts(payloads)

Write each content-addressed artifact into the plugin registry the engine reads.

The offline engine shares the process's filesystem, so staging is a registry write (idempotent, verified against the content address).

Source code in steerability/backends/vllm/backend.py

def stage_artifacts(self, payloads) -> None:
    """Write each content-addressed artifact into the plugin registry the engine reads.

    The offline engine shares the process's filesystem, so staging is a registry write
    (idempotent, verified against the content address).
    """
    self._artifact_uploader.upload_payloads(payloads)

VLLMOfflineSession

Bases: _RequestSessionBase

Request session over the offline engine.

Token-id prompts submit as TokensPrompts in one engine call with per-item sampling parameters; the engine schedules the batch internally, so no client-side fan-out is needed.

Source code in steerability/backends/vllm/session.py

class VLLMOfflineSession(_RequestSessionBase):
    """Request session over the offline engine.

    Token-id prompts submit as `TokensPrompt`s in one engine call with per-item sampling
    parameters; the engine schedules the batch internally, so no client-side fan-out is needed.
    """

    def capture(
        self,
        prompts: list[PreparedPrompt],
        layers: list[int],
        mode: Literal["all_tokens", "last_token"],
        location: Literal["layer_output", "layer_input"] = "layer_output",
    ) -> CaptureResult:
        """Hidden-state capture over the plugin's capture surface.

        One request per prompt carries a `capture` spec and a fresh random `cache_salt`
        (a prefix-cache hit skips forward passes, so capture cannot tolerate reused salts) with
        `max_tokens=1`; the surplus decode position is truncated by the plugin. Per-layer
        tensors are stacked and right-padded to the batch's longest prompt.

        Args:
            prompts: The prompts to capture over.
            layers: 0-based decoder-layer indices to capture.
            mode: `"all_tokens"` for every prompt position, `"last_token"` for the final real
                position per row.
            location: The residual-stream boundary, `"layer_output"` or `"layer_input"`.

        Returns:
            The capture result: `[N, T, H]` per layer for `"all_tokens"` or `[N, H]` for
            `"last_token"`, on CPU in the engine's native dtype, with the derived `[N, T]`
            attention mask.

        Raises:
            UnsupportedOperationError: If the spec declares no `hook_plugin`, the negotiated
                capture kinds lack the requested mode or location, or the engine facts refuse
                capture (speculative decoding, non-eager execution).
            ValueError: If `prompts` is empty, a layer id is out of range, or the engine
                returned no capture payload.
        """
        self._ensure_open()
        backend = self._backend
        if not backend.spec.get_option("hook_plugin"):
            raise UnsupportedOperationError(
                "Hidden-state capture requires the vLLM-Hook plugin; declare hook_plugin=True "
                "on the vllm backend spec, or run capture on the huggingface backend."
            )
        capture_kinds = backend.capture_kinds
        required = CaptureKinds(
            kinds=frozenset({"residual"}),
            locations=frozenset({location}),
            modes=frozenset({mode}),
        )
        if capture_kinds is None or not capture_kinds.contains(required):
            raise UnsupportedOperationError(
                f"The serving backend does not advertise capture mode {mode!r} at location "
                f"{location!r}; update the server's vllm_hook_plugins or run capture on the "
                "huggingface backend."
            )
        _refuse_by_engine_facts(backend._discovery, "capture")
        if not prompts:
            raise ValueError("capture() requires at least one prompt.")
        num_layers = self.layout.num_layers
        missing = sorted(int(layer) for layer in layers if not 0 <= int(layer) < num_layers)
        if missing:
            raise ValueError(
                f"Requested layer ids {missing} are out of range; the model has {num_layers} layers."
            )

        from vllm import SamplingParams, TokensPrompt

        layer_ids = [int(layer) for layer in layers]
        # the client's validator and assembly expect full prompt coverage and pool the last real
        # position themselves, so every wire capture requests all_tokens
        wire_mode = "all_tokens" if mode == "last_token" else mode
        capture_spec = {"layers": layer_ids, "mode": wire_mode, "location": location}
        engine_prompts = []
        prompt_lens: list[int] = []
        for prompt in prompts:
            resolved = prompt.resolve_token_ids(self.tokenizer)
            ids = resolved.token_ids[0]
            if resolved.attention_mask is not None:
                ids = ids[resolved.attention_mask[0].bool()]
            ids = ids.tolist()
            prompt_lens.append(len(ids))
            engine_prompt = TokensPrompt(prompt_token_ids=ids)
            engine_prompt["cache_salt"] = uuid.uuid4().hex
            engine_prompts.append(engine_prompt)
        sampling = SamplingParams(max_tokens=1, temperature=0.0, extra_args={"capture": capture_spec})

        request_outputs = self._backend._require_llm().generate(engine_prompts, sampling, use_tqdm=False)

        rows_per_layer: dict[int, list[torch.Tensor]] = {layer: [] for layer in layer_ids}
        for index, request_output in enumerate(request_outputs):
            payload = getattr(request_output, "captures", None)
            if payload is None:
                raise ValueError(
                    "The engine returned no capture payload; is the vLLM-Hook unified worker "
                    "active on this engine?"
                )
            manifest_json, data = payload
            manifest = json.loads(manifest_json)
            tensors = _load_safetensors_bytes(data)
            for layer in layer_ids:
                stacked = tensors.get(f"layer_{layer}")
                if stacked is None or stacked.size(0) < prompt_lens[index]:
                    raise ValueError(
                        f"The capture payload covers layer {layer} at "
                        f"{0 if stacked is None else stacked.size(0)} of {prompt_lens[index]} "
                        f"prompt positions for prompt {index}; positions recorded: "
                        f"{manifest.get('positions', {}).get(str(layer))}."
                    )
                rows_per_layer[layer].append(stacked[: prompt_lens[index]])

        max_len = max(prompt_lens)
        attention_mask = torch.zeros(len(prompts), max_len, dtype=torch.long)
        for index, length in enumerate(prompt_lens):
            attention_mask[index, :length] = 1

        hidden: dict[int, torch.Tensor] = {}
        for layer, rows in rows_per_layer.items():
            if mode == "last_token":
                hidden[layer] = torch.stack([row[-1] for row in rows])
            else:
                padded = torch.zeros(len(rows), max_len, rows[0].size(-1), dtype=rows[0].dtype)
                for index, row in enumerate(rows):
                    padded[index, : row.size(0)] = row
                hidden[layer] = padded
        return CaptureResult(hidden=hidden, attention_mask=attention_mask, mode=mode, location=location)

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item through the engine.

        Args:
            items: The generation items; state entries lower as intervention specs on
                plugin-active backends, and no client-side hooks or live processors execute
                here.
            params: Normalized generation parameters shared by all items; unmapped `extra` keys
                raise.

        Returns:
            One `ItemResult` per item, in item order.
        """
        self._ensure_open()
        if not items:
            return []
        item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm")
        base_args = render_vllm_sampling_args(params)

        from vllm import SamplingParams, TokensPrompt

        prompts = []
        sampling = []
        prompt_ids_per_item: list[list[int]] = []
        for index, item in enumerate(items):
            ids = self._resolve_item_ids(item)
            prompt_ids_per_item.append(ids)
            args = dict(base_args)
            seed = self._item_seed(item, params, index)
            if seed is not None:
                args["seed"] = seed
            if item_constraints[index] is not None:
                field, value = render_guided_decoding_field(item_constraints[index])
                args.update(render_constraint_sampling_args(field, value))
            if item_specs[index] is not None:
                args["extra_args"] = {"intervention_spec": item_specs[index].to_wire()}
            prompt = TokensPrompt(prompt_token_ids=ids)
            if item_salts is not None:
                prompt["cache_salt"] = item_salts[index]
            prompts.append(prompt)
            sampling.append(SamplingParams(**args))
        self._generate_count += 1

        generate_kwargs: dict[str, Any] = {"use_tqdm": False}
        if self._backend._lora_request is not None:
            generate_kwargs["lora_request"] = self._backend._lora_request
        request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)

        results: list[ItemResult] = []
        for index, request_output in enumerate(request_outputs):
            candidates = [
                (
                    list(candidate.token_ids),
                    map_vllm_finish_reason(
                        candidate.finish_reason, getattr(candidate, "stop_reason", None),
                    ),
                )
                for candidate in request_output.outputs
            ]
            results.append(self._pack_output(index, prompt_ids_per_item[index], candidates))
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

        Each item's prompt and reference concatenate into one token-id prompt submitted with
        `prompt_logprobs=0`, and the reference positions' log-probabilities are read back.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: Must carry no `extra` keys; forward keyword arguments have no remote
                rendering.

        Returns:
            Log probabilities of shape `[num_items, ref_len]` on CPU.

        Raises:
            ValueError: If items carry differing reference lengths or `params.extra` is
                non-empty.
        """
        self._ensure_open()
        if params.extra:
            raise ValueError(
                f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
                "scoring accepts no forward keyword arguments."
            )
        if not items:
            return torch.zeros((0, 0), dtype=torch.float32)
        item_specs, _, item_salts = self._prepare_spec_submission(
            items, "vllm", allow_constraints=False,
        )
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
        ref_len = ref_lens.pop()
        if ref_len == 0:
            return torch.zeros((len(items), 0), dtype=torch.float32)

        from vllm import SamplingParams, TokensPrompt

        prompts = []
        sampling = []
        ref_ids_per_item: list[list[int]] = []
        for index, item in enumerate(items):
            prompt_ids = self._resolve_item_ids(item)
            ref_ids = item.ref_output_ids.reshape(-1).tolist()
            ref_ids_per_item.append(ref_ids)
            prompt = TokensPrompt(prompt_token_ids=[*prompt_ids, *ref_ids])
            args: dict[str, Any] = {"max_tokens": 1, "temperature": 0.0, "prompt_logprobs": 0}
            if item_specs[index] is not None:
                scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids))
                args["extra_args"] = {"intervention_spec": scoring_spec.to_wire()}
                if item_salts is not None:
                    item_salts[index] = scoring_spec.salt()
            if item_salts is not None:
                prompt["cache_salt"] = item_salts[index]
            prompts.append(prompt)
            sampling.append(SamplingParams(**args))

        generate_kwargs: dict[str, Any] = {"use_tqdm": False}
        if self._backend._lora_request is not None:
            generate_kwargs["lora_request"] = self._backend._lora_request
        request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)
        rows = [
            extract_ref_logprobs(request_output.prompt_logprobs, ref_ids)
            for request_output, ref_ids in zip(request_outputs, ref_ids_per_item)
        ]
        return torch.tensor(rows, dtype=torch.float32)

closed property

Whether the session has been closed.

layout property

Structural facts from the model config (client-side).

Raises:

Type	Description
RuntimeError	If the model config could not be resolved.

tokenizer property

The backend's client-side tokenizer.

capture(prompts, layers, mode, location='layer_output')

Hidden-state capture over the plugin's capture surface.

One request per prompt carries a capture spec and a fresh random cache_salt (a prefix-cache hit skips forward passes, so capture cannot tolerate reused salts) with max_tokens=1; the surplus decode position is truncated by the plugin. Per-layer tensors are stacked and right-padded to the batch's longest prompt.

Parameters:

Name	Type	Description	Default
prompts	list[PreparedPrompt]	The prompts to capture over.	required
layers	list[int]	0-based decoder-layer indices to capture.	required
mode	Literal['all_tokens', 'last_token']	"all_tokens" for every prompt position, "last_token" for the final real position per row.	required
location	Literal['layer_output', 'layer_input']	The residual-stream boundary, "layer_output" or "layer_input".	'layer_output'

Returns:

Type	Description
CaptureResult	The capture result: [N, T, H] per layer for "all_tokens" or [N, H] for
CaptureResult	"last_token", on CPU in the engine's native dtype, with the derived [N, T]
CaptureResult	attention mask.

Raises:

Type	Description
UnsupportedOperationError	If the spec declares no hook_plugin, the negotiated capture kinds lack the requested mode or location, or the engine facts refuse capture (speculative decoding, non-eager execution).
ValueError	If prompts is empty, a layer id is out of range, or the engine returned no capture payload.

Source code in steerability/backends/vllm/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Hidden-state capture over the plugin's capture surface.

    One request per prompt carries a `capture` spec and a fresh random `cache_salt`
    (a prefix-cache hit skips forward passes, so capture cannot tolerate reused salts) with
    `max_tokens=1`; the surplus decode position is truncated by the plugin. Per-layer
    tensors are stacked and right-padded to the batch's longest prompt.

    Args:
        prompts: The prompts to capture over.
        layers: 0-based decoder-layer indices to capture.
        mode: `"all_tokens"` for every prompt position, `"last_token"` for the final real
            position per row.
        location: The residual-stream boundary, `"layer_output"` or `"layer_input"`.

    Returns:
        The capture result: `[N, T, H]` per layer for `"all_tokens"` or `[N, H]` for
        `"last_token"`, on CPU in the engine's native dtype, with the derived `[N, T]`
        attention mask.

    Raises:
        UnsupportedOperationError: If the spec declares no `hook_plugin`, the negotiated
            capture kinds lack the requested mode or location, or the engine facts refuse
            capture (speculative decoding, non-eager execution).
        ValueError: If `prompts` is empty, a layer id is out of range, or the engine
            returned no capture payload.
    """
    self._ensure_open()
    backend = self._backend
    if not backend.spec.get_option("hook_plugin"):
        raise UnsupportedOperationError(
            "Hidden-state capture requires the vLLM-Hook plugin; declare hook_plugin=True "
            "on the vllm backend spec, or run capture on the huggingface backend."
        )
    capture_kinds = backend.capture_kinds
    required = CaptureKinds(
        kinds=frozenset({"residual"}),
        locations=frozenset({location}),
        modes=frozenset({mode}),
    )
    if capture_kinds is None or not capture_kinds.contains(required):
        raise UnsupportedOperationError(
            f"The serving backend does not advertise capture mode {mode!r} at location "
            f"{location!r}; update the server's vllm_hook_plugins or run capture on the "
            "huggingface backend."
        )
    _refuse_by_engine_facts(backend._discovery, "capture")
    if not prompts:
        raise ValueError("capture() requires at least one prompt.")
    num_layers = self.layout.num_layers
    missing = sorted(int(layer) for layer in layers if not 0 <= int(layer) < num_layers)
    if missing:
        raise ValueError(
            f"Requested layer ids {missing} are out of range; the model has {num_layers} layers."
        )

    from vllm import SamplingParams, TokensPrompt

    layer_ids = [int(layer) for layer in layers]
    # the client's validator and assembly expect full prompt coverage and pool the last real
    # position themselves, so every wire capture requests all_tokens
    wire_mode = "all_tokens" if mode == "last_token" else mode
    capture_spec = {"layers": layer_ids, "mode": wire_mode, "location": location}
    engine_prompts = []
    prompt_lens: list[int] = []
    for prompt in prompts:
        resolved = prompt.resolve_token_ids(self.tokenizer)
        ids = resolved.token_ids[0]
        if resolved.attention_mask is not None:
            ids = ids[resolved.attention_mask[0].bool()]
        ids = ids.tolist()
        prompt_lens.append(len(ids))
        engine_prompt = TokensPrompt(prompt_token_ids=ids)
        engine_prompt["cache_salt"] = uuid.uuid4().hex
        engine_prompts.append(engine_prompt)
    sampling = SamplingParams(max_tokens=1, temperature=0.0, extra_args={"capture": capture_spec})

    request_outputs = self._backend._require_llm().generate(engine_prompts, sampling, use_tqdm=False)

    rows_per_layer: dict[int, list[torch.Tensor]] = {layer: [] for layer in layer_ids}
    for index, request_output in enumerate(request_outputs):
        payload = getattr(request_output, "captures", None)
        if payload is None:
            raise ValueError(
                "The engine returned no capture payload; is the vLLM-Hook unified worker "
                "active on this engine?"
            )
        manifest_json, data = payload
        manifest = json.loads(manifest_json)
        tensors = _load_safetensors_bytes(data)
        for layer in layer_ids:
            stacked = tensors.get(f"layer_{layer}")
            if stacked is None or stacked.size(0) < prompt_lens[index]:
                raise ValueError(
                    f"The capture payload covers layer {layer} at "
                    f"{0 if stacked is None else stacked.size(0)} of {prompt_lens[index]} "
                    f"prompt positions for prompt {index}; positions recorded: "
                    f"{manifest.get('positions', {}).get(str(layer))}."
                )
            rows_per_layer[layer].append(stacked[: prompt_lens[index]])

    max_len = max(prompt_lens)
    attention_mask = torch.zeros(len(prompts), max_len, dtype=torch.long)
    for index, length in enumerate(prompt_lens):
        attention_mask[index, :length] = 1

    hidden: dict[int, torch.Tensor] = {}
    for layer, rows in rows_per_layer.items():
        if mode == "last_token":
            hidden[layer] = torch.stack([row[-1] for row in rows])
        else:
            padded = torch.zeros(len(rows), max_len, rows[0].size(-1), dtype=rows[0].dtype)
            for index, row in enumerate(rows):
                padded[index, : row.size(0)] = row
            hidden[layer] = padded
    return CaptureResult(hidden=hidden, attention_mask=attention_mask, mode=mode, location=location)

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/vllm/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

generate(items, params)

Generate one result per item through the engine.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items; state entries lower as intervention specs on plugin-active backends, and no client-side hooks or live processors execute here.	required
params	GenerationParams	Normalized generation parameters shared by all items; unmapped extra keys raise.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order.

Source code in steerability/backends/vllm/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item through the engine.

    Args:
        items: The generation items; state entries lower as intervention specs on
            plugin-active backends, and no client-side hooks or live processors execute
            here.
        params: Normalized generation parameters shared by all items; unmapped `extra` keys
            raise.

    Returns:
        One `ItemResult` per item, in item order.
    """
    self._ensure_open()
    if not items:
        return []
    item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm")
    base_args = render_vllm_sampling_args(params)

    from vllm import SamplingParams, TokensPrompt

    prompts = []
    sampling = []
    prompt_ids_per_item: list[list[int]] = []
    for index, item in enumerate(items):
        ids = self._resolve_item_ids(item)
        prompt_ids_per_item.append(ids)
        args = dict(base_args)
        seed = self._item_seed(item, params, index)
        if seed is not None:
            args["seed"] = seed
        if item_constraints[index] is not None:
            field, value = render_guided_decoding_field(item_constraints[index])
            args.update(render_constraint_sampling_args(field, value))
        if item_specs[index] is not None:
            args["extra_args"] = {"intervention_spec": item_specs[index].to_wire()}
        prompt = TokensPrompt(prompt_token_ids=ids)
        if item_salts is not None:
            prompt["cache_salt"] = item_salts[index]
        prompts.append(prompt)
        sampling.append(SamplingParams(**args))
    self._generate_count += 1

    generate_kwargs: dict[str, Any] = {"use_tqdm": False}
    if self._backend._lora_request is not None:
        generate_kwargs["lora_request"] = self._backend._lora_request
    request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)

    results: list[ItemResult] = []
    for index, request_output in enumerate(request_outputs):
        candidates = [
            (
                list(candidate.token_ids),
                map_vllm_finish_reason(
                    candidate.finish_reason, getattr(candidate, "stop_reason", None),
                ),
            )
            for candidate in request_output.outputs
        ]
        results.append(self._pack_output(index, prompt_ids_per_item[index], candidates))
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

Each item's prompt and reference concatenate into one token-id prompt submitted with prompt_logprobs=0, and the reference positions' log-probabilities are read back.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	Must carry no extra keys; forward keyword arguments have no remote rendering.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len] on CPU.

Raises:

Type	Description
ValueError	If items carry differing reference lengths or params.extra is non-empty.

Source code in steerability/backends/vllm/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

    Each item's prompt and reference concatenate into one token-id prompt submitted with
    `prompt_logprobs=0`, and the reference positions' log-probabilities are read back.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: Must carry no `extra` keys; forward keyword arguments have no remote
            rendering.

    Returns:
        Log probabilities of shape `[num_items, ref_len]` on CPU.

    Raises:
        ValueError: If items carry differing reference lengths or `params.extra` is
            non-empty.
    """
    self._ensure_open()
    if params.extra:
        raise ValueError(
            f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
            "scoring accepts no forward keyword arguments."
        )
    if not items:
        return torch.zeros((0, 0), dtype=torch.float32)
    item_specs, _, item_salts = self._prepare_spec_submission(
        items, "vllm", allow_constraints=False,
    )
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
    ref_len = ref_lens.pop()
    if ref_len == 0:
        return torch.zeros((len(items), 0), dtype=torch.float32)

    from vllm import SamplingParams, TokensPrompt

    prompts = []
    sampling = []
    ref_ids_per_item: list[list[int]] = []
    for index, item in enumerate(items):
        prompt_ids = self._resolve_item_ids(item)
        ref_ids = item.ref_output_ids.reshape(-1).tolist()
        ref_ids_per_item.append(ref_ids)
        prompt = TokensPrompt(prompt_token_ids=[*prompt_ids, *ref_ids])
        args: dict[str, Any] = {"max_tokens": 1, "temperature": 0.0, "prompt_logprobs": 0}
        if item_specs[index] is not None:
            scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids))
            args["extra_args"] = {"intervention_spec": scoring_spec.to_wire()}
            if item_salts is not None:
                item_salts[index] = scoring_spec.salt()
        if item_salts is not None:
            prompt["cache_salt"] = item_salts[index]
        prompts.append(prompt)
        sampling.append(SamplingParams(**args))

    generate_kwargs: dict[str, Any] = {"use_tqdm": False}
    if self._backend._lora_request is not None:
        generate_kwargs["lora_request"] = self._backend._lora_request
    request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)
    rows = [
        extract_ref_logprobs(request_output.prompt_logprobs, ref_ids)
        for request_output, ref_ids in zip(request_outputs, ref_ids_per_item)
    ]
    return torch.tensor(rows, dtype=torch.float32)

VLLMServeBackend

Bases: Backend

The vLLM OpenAI-compatible server backend.

Targets a vLLM server rather than an arbitrary OpenAI-compatible endpoint: construction verifies the server's version surface (GET /version), fetches the plugin discovery payload (GET /v1/hook/capabilities) when the spec declares hook_plugin, and checks the served model id against the spec (or serves the pipeline's structural artifacts). Prompts submit as token ids on the completions endpoint with the token-id return option; the chat endpoint is not used. Requires no local vLLM installation.

Spec options: base_url (required, the server root), api_key, request_timeout, max_concurrency, max_retries, retry_backoff, tokenizer_name_or_path, trust_remote_code, hook_plugin.

Source code in steerability/backends/vllm/backend.py

class VLLMServeBackend(Backend):
    """The vLLM OpenAI-compatible server backend.

    Targets a vLLM server rather than an arbitrary OpenAI-compatible endpoint: construction
    verifies the server's version surface (`GET /version`), fetches the plugin discovery payload
    (`GET /v1/hook/capabilities`) when the spec declares `hook_plugin`, and checks the served
    model id against the spec (or serves the pipeline's structural artifacts). Prompts submit as
    token ids on the completions endpoint with the token-id return option; the chat endpoint is
    not used. Requires no local vLLM installation.

    Spec options: `base_url` (required, the server root), `api_key`, `request_timeout`,
    `max_concurrency`, `max_retries`, `retry_backoff`, `tokenizer_name_or_path`,
    `trust_remote_code`, `hook_plugin`.
    """

    def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> None:
        if spec.kind != "vllm-serve":
            raise ValueError(f"VLLMServeBackend requires a 'vllm-serve' spec; got kind {spec.kind!r}.")
        self.spec = spec
        base_url = spec.get_option("base_url")
        if not base_url:
            raise ValueError("VLLMServeBackend requires a 'base_url' option on the spec.")
        self._base_url = base_url.rstrip("/").removesuffix("/v1")
        self._api_key = spec.get_option("api_key")
        self._timeout = float(spec.get_option("request_timeout", default=_DEFAULT_REQUEST_TIMEOUT))
        self.max_concurrency = int(spec.get_option("max_concurrency", default=_DEFAULT_MAX_CONCURRENCY))
        self.max_attempts = int(spec.get_option("max_retries", default=_DEFAULT_MAX_ATTEMPTS))
        self.backoff_base = float(spec.get_option("retry_backoff", default=0.5))
        trust_remote_code = bool(spec.get_option("trust_remote_code", default=False))

        version = self._get_json("/version")
        if not isinstance(version, dict) or "version" not in version:
            raise ValueError(
                f"The endpoint at {self._base_url} does not expose the vLLM version surface; "
                "only vLLM servers are supported."
            )

        self._discovery: dict | None = None
        if spec.get_option("hook_plugin"):
            self._discovery = _DISCOVERY_CACHE.get(spec.spec_hash)
            if self._discovery is None:
                try:
                    self._discovery = self._get_json("/v1/hook/capabilities")
                except (TransportError, ValueError) as error:
                    raise ValueError(
                        f"The spec declares hook_plugin but {self._base_url} serves no "
                        f"/v1/hook/capabilities discovery surface: {error}"
                    ) from error
                _DISCOVERY_CACHE[spec.spec_hash] = self._discovery
                _reconcile_discovery(spec, self.capabilities_for_spec(spec), self._discovery)

        checkpoint, lora = _split_artifacts(artifacts)
        expected_model = checkpoint.path if checkpoint is not None else spec.model
        if lora is not None:
            self._served_model = self._load_lora_adapter(lora)
        else:
            served = self._served_model_ids()
            if expected_model is None:
                if len(served) != 1:
                    raise ValueError(
                        f"The spec names no model and the server serves {served}; set the "
                        "spec's model to disambiguate."
                    )
                self._served_model = served[0]
            elif expected_model in served:
                self._served_model = expected_model
            else:
                raise ValueError(
                    f"The server at {self._base_url} serves {served}, not the configured "
                    f"model {expected_model!r}."
                )

        tokenizer_source = (
            spec.get_option("tokenizer_name_or_path")
            or (checkpoint.path if checkpoint is not None else None)
            or (lora.base_model if lora is not None else None)
            or spec.model
        )
        if tokenizer_source is None:
            tokenizer_source = self._served_model
        self.tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code)
        self._layout = _config_layout(
            (checkpoint.path if checkpoint is not None else None) or spec.model or self._served_model,
            trust_remote_code,
        )
        self._plain_salt = uuid.uuid4().hex
        # spec artifacts write to the registry root the server reads; shared_fs visibility
        # is verified against the server after writing (see stage_artifacts)
        self._artifact_uploader = _ArtifactUploader(spec.get_option("artifact_dir"))
        if self._discovery is not None:
            self._verify_fingerprints()

    def _served_model_ids(self) -> list[str]:
        payload = self._get_json("/v1/models")
        return [entry.get("id") for entry in payload.get("data", []) if isinstance(entry, dict)]

    def stage_artifacts(self, payloads) -> None:
        """Make each content-addressed artifact available to the serving engine.

        With an `artifact_dir` option the payloads are written into that registry root, which
        must be the server's registry directory (its `VLLM_HOOK_REGISTRY_DIR`) on a shared
        filesystem; visibility is verified through the server's artifact route when the
        discovery payload advertises `artifact_registry_root`. Otherwise each payload is PUT
        to the plugin's artifact route (`/v1/hook/artifacts/{id}`, body safetensors bytes, id
        verified server-side); already-exists is success.
        """
        if not payloads:
            return
        if self.spec.get_option("artifact_dir"):
            self._artifact_uploader.upload_payloads(payloads)
            self._verify_shared_fs_visibility(payloads)
            return
        import safetensors.torch

        for artifact_id, tensors in payloads.items():
            if artifact_id in self._artifact_uploader._written:
                continue
            data = safetensors.torch.save({name: tensors[name] for name in sorted(tensors)})
            self._put_bytes(f"/v1/hook/artifacts/{artifact_id}", data)
            self._artifact_uploader._written.add(artifact_id)

    def _verify_shared_fs_visibility(self, payloads) -> None:
        """Probe that shared_fs artifacts are visible to the server's registry.

        Gated on the discovery payload advertising `artifact_registry_root` (servers that
        advertise it also serve `HEAD /v1/hook/artifacts/{id}`); older servers skip the
        probe and keep the write-and-trust behavior.
        """
        server_root = (self._discovery or {}).get("artifact_registry_root")
        if not server_root:
            return
        client_root = os.path.abspath(self.spec.get_option("artifact_dir"))
        for artifact_id in payloads:
            if self._head_ok(f"/v1/hook/artifacts/{artifact_id}"):
                continue
            raise ValueError(
                f"artifact {artifact_id} written under {client_root} is not visible to the "
                f"server's registry ({server_root}); the shared_fs transport requires "
                "artifact_dir and the server's VLLM_HOOK_REGISTRY_DIR to name the same "
                "directory. Set VLLM_HOOK_REGISTRY_DIR on the server, point artifact_dir at "
                "the server's registry root, or drop artifact_dir to use the HTTP artifact "
                "route."
            )

    def _head_ok(self, path: str) -> bool:
        """HEAD a server path; True on 200, False on 404, raise otherwise."""
        request = urllib.request.Request(f"{self._base_url}{path}", method="HEAD")
        if self._api_key:
            request.add_header("Authorization", f"Bearer {self._api_key}")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout):
                return True
        except urllib.error.HTTPError as error:
            if error.code == 404:
                return False
            body = error.read().decode("utf-8", errors="replace")
            raise ValueError(f"HTTP {error.code} from {self._base_url}{path}: {body}") from error

    def _put_bytes(self, path: str, data: bytes) -> None:
        """PUT raw bytes to the server, mapping a missing route to a configuration error."""
        import urllib.error
        import urllib.request

        url = f"{self._base_url}{path}"
        request = urllib.request.Request(url, data=data, method="PUT")
        request.add_header("Content-Type", "application/octet-stream")
        if self._api_key:
            request.add_header("Authorization", f"Bearer {self._api_key}")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout):
                return
        except urllib.error.HTTPError as error:
            body = error.read().decode("utf-8", errors="replace")
            if error.code in (404, 405):
                raise ValueError(
                    f"{self._base_url} serves no artifact route ({error.code}); update the "
                    "server's vllm_hook_plugins, or configure artifact_dir on a filesystem "
                    "shared with the server."
                ) from error
            raise_for_spec_rejection(body)
            raise ValueError(f"HTTP {error.code} from {url}: {body}") from error
        except (urllib.error.URLError, TimeoutError, OSError) as error:
            raise TransportError(f"artifact upload to {url} failed: {error}") from error

    def _load_lora_adapter(self, lora: LoRAArtifact) -> str:
        served = self._served_model_ids()
        base = lora.base_model or self.spec.model
        if base and base not in served:
            raise ValueError(
                f"The server at {self._base_url} serves {served}, not the adapter's base "
                f"model {base!r}."
            )
        # the adapter name keys on path plus provenance, so a retrained adapter at the same
        # path loads as a new server-side adapter rather than reusing stale weights
        identity = f"{lora.path}:{lora.provenance.model_fingerprint or ''}"
        adapter_name = f"steered-{hashlib.sha256(identity.encode('utf-8')).hexdigest()[:8]}"
        if adapter_name in served:
            return adapter_name
        try:
            self._post_json(
                "/v1/load_lora_adapter",
                {"lora_name": adapter_name, "lora_path": lora.path},
                expect_json=False,
            )
        except (TransportError, ValueError) as error:
            raise ValueError(
                f"Could not load the LoRA artifact at {lora.path!r} onto the server "
                f"(dynamic adapter loading requires VLLM_ALLOW_RUNTIME_LORA_UPDATING): {error}"
            ) from error
        return adapter_name

    def _verify_fingerprints(self) -> None:
        """Verify the client tokenizer against the discovery payload's fingerprint recipes.

        Uses the plugin's engine-free `core.fingerprints` when the `vllm_hook_plugins` package
        is installed; mismatches warn rather than raise. Without the package, verification is
        skipped with a warning. A served fingerprint equal to the absent-template digest means
        the server exposes no chat template, so the comparison is skipped since a mismatch
        against it would reflect exposure rather than divergence.
        """
        model_block = (self._discovery or {}).get("model", {})
        remote_chat = model_block.get("chat_template_fingerprint")
        if remote_chat is None:
            return
        if is_absent_chat_template_fingerprint(remote_chat):
            logger.debug(
                "The server at %s does not expose a chat template (fingerprint %s is the "
                "absent-template digest); skipping the chat template comparison.",
                self._base_url, remote_chat,
            )
            return
        try:
            from vllm_hook_plugins.core.fingerprints import chat_template_fingerprint
        except ImportError:
            logger.warning(
                "Install vllm-hook-plugins to verify the client tokenizer against the server's "
                "fingerprints; skipping verification."
            )
            return
        local_chat = chat_template_fingerprint(getattr(self.tokenizer, "chat_template", None))
        if local_chat != remote_chat:
            logger.warning(
                "Client chat template (fingerprint %s) differs from the served one (%s); "
                "templated prompts may diverge from server-side expectations.",
                local_chat, remote_chat,
            )

    def _request_json(self, path: str, payload: dict | None, expect_json: bool = True) -> dict:
        url = f"{self._base_url}{path}"
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        data = json.dumps(payload).encode("utf-8") if payload is not None else None
        request = urllib.request.Request(url, data=data, headers=headers)
        try:
            with urllib.request.urlopen(request, timeout=self._timeout) as response:
                body = response.read().decode("utf-8")
        except urllib.error.HTTPError as error:
            body = ""
            try:
                body = error.read().decode("utf-8", errors="replace")
            except Exception:
                pass
            # 5xx, timeouts, and rate limiting are transport-level and safe to retry
            if error.code >= 500 or error.code in (408, 429):
                raise TransportError(f"HTTP {error.code} from {url}: {body}") from error
            # admission rejections carry the plugin's E_* code and JSON path verbatim
            raise_for_spec_rejection(body)
            raise ValueError(f"HTTP {error.code} from {url}: {body}") from error
        except (urllib.error.URLError, TimeoutError, OSError) as error:
            raise TransportError(f"Request to {url} failed: {error}") from error
        if not expect_json:
            return {"text": body}
        try:
            return json.loads(body)
        except json.JSONDecodeError as error:
            raise ValueError(f"Non-JSON response from {url}: {error}") from error

    def _get_json(self, path: str) -> dict:
        return self._request_json(path, None)

    def _post_json(self, path: str, payload: dict, expect_json: bool = True) -> dict:
        return self._request_json(path, payload, expect_json=expect_json)

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The capability advertisement implied by `spec`."""
        return _vllm_capabilities(spec, offline=False)

    def open_session(self) -> "VLLMServeSession":
        """Open a request session over the shared connection."""
        return VLLMServeSession(self)

backoff_base = float(spec.get_option('retry_backoff', default=0.5)) instance-attribute

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

max_attempts = int(spec.get_option('max_retries', default=_DEFAULT_MAX_ATTEMPTS)) instance-attribute

max_concurrency = int(spec.get_option('max_concurrency', default=_DEFAULT_MAX_CONCURRENCY)) instance-attribute

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code) instance-attribute

capabilities_for_spec(spec) classmethod

The capability advertisement implied by spec.

Source code in steerability/backends/vllm/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The capability advertisement implied by `spec`."""
    return _vllm_capabilities(spec, offline=False)

open_session()

Open a request session over the shared connection.

Source code in steerability/backends/vllm/backend.py

def open_session(self) -> "VLLMServeSession":
    """Open a request session over the shared connection."""
    return VLLMServeSession(self)

release()

Free the resources this backend owns beyond what the caller owns.

The default is a no-op. Engine-owning backends override this to shut their engines down deterministically. Release is idempotent.

Source code in steerability/algorithms/core/execution/backend.py

def release(self) -> None:
    """Free the resources this backend owns beyond what the caller owns.

    The default is a no-op. Engine-owning backends override this to shut their engines down
    deterministically. Release is idempotent.
    """
    return None

stage_artifacts(payloads)

Make each content-addressed artifact available to the serving engine.

With an artifact_dir option the payloads are written into that registry root, which must be the server's registry directory (its VLLM_HOOK_REGISTRY_DIR) on a shared filesystem; visibility is verified through the server's artifact route when the discovery payload advertises artifact_registry_root. Otherwise each payload is PUT to the plugin's artifact route (/v1/hook/artifacts/{id}, body safetensors bytes, id verified server-side); already-exists is success.

Source code in steerability/backends/vllm/backend.py

def stage_artifacts(self, payloads) -> None:
    """Make each content-addressed artifact available to the serving engine.

    With an `artifact_dir` option the payloads are written into that registry root, which
    must be the server's registry directory (its `VLLM_HOOK_REGISTRY_DIR`) on a shared
    filesystem; visibility is verified through the server's artifact route when the
    discovery payload advertises `artifact_registry_root`. Otherwise each payload is PUT
    to the plugin's artifact route (`/v1/hook/artifacts/{id}`, body safetensors bytes, id
    verified server-side); already-exists is success.
    """
    if not payloads:
        return
    if self.spec.get_option("artifact_dir"):
        self._artifact_uploader.upload_payloads(payloads)
        self._verify_shared_fs_visibility(payloads)
        return
    import safetensors.torch

    for artifact_id, tensors in payloads.items():
        if artifact_id in self._artifact_uploader._written:
            continue
        data = safetensors.torch.save({name: tensors[name] for name in sorted(tensors)})
        self._put_bytes(f"/v1/hook/artifacts/{artifact_id}", data)
        self._artifact_uploader._written.add(artifact_id)

VLLMServeSession

Bases: _RequestSessionBase

Request session over a vLLM server's completions endpoint.

Items fan out concurrently under the backend's max_concurrency; transport failures retry with exponential backoff, and a batch whose items partially fail raises PartialBatchError carrying the successes and the re-issuable failures.

Source code in steerability/backends/vllm/session.py

class VLLMServeSession(_RequestSessionBase):
    """Request session over a vLLM server's completions endpoint.

    Items fan out concurrently under the backend's `max_concurrency`; transport failures retry
    with exponential backoff, and a batch whose items partially fail raises `PartialBatchError`
    carrying the successes and the re-issuable failures.
    """

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item through the completions endpoint.

        Args:
            items: The generation items; entries must be empty.
            params: Normalized generation parameters shared by all items; unmapped `extra` keys
                raise.

        Returns:
            One `ItemResult` per item, in item order.

        Raises:
            PartialBatchError: If some items failed after transport retries while others
                succeeded.
        """
        self._ensure_open()
        if not items:
            return []
        item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm-serve")
        base_args = render_vllm_sampling_args(params)
        backend = self._backend

        item_ids = [self._resolve_item_ids(item) for item in items]
        seeds = [self._item_seed(item, params, index) for index, item in enumerate(items)]
        self._generate_count += 1

        def make_task(index: int):
            def task() -> ItemResult:
                body: dict[str, Any] = {
                    "model": backend._served_model,
                    "prompt": item_ids[index],
                    "return_token_ids": True,
                    **base_args,
                }
                if seeds[index] is not None:
                    body["seed"] = seeds[index]
                if item_constraints[index] is not None:
                    field, value = render_guided_decoding_field(item_constraints[index])
                    body[f"guided_{field}"] = value
                if item_specs[index] is not None:
                    # vllm_xargs is scalar-only, so nested specs travel as JSON strings
                    body["vllm_xargs"] = {
                        "intervention_spec": item_specs[index].canonical(),
                    }
                if item_salts is not None:
                    body["cache_salt"] = item_salts[index]
                payload = with_transport_retries(
                    lambda: backend._post_json("/v1/completions", body),
                    max_attempts=backend.max_attempts,
                    backoff_base=backend.backoff_base,
                )
                choices = payload.get("choices", [])
                if not choices:
                    raise ValueError("The completions response carries no choices.")
                candidates = []
                for choice in choices:
                    token_ids = choice.get("token_ids")
                    if token_ids is None:
                        raise ValueError(
                            "The completions response carries no token_ids; the server does "
                            "not support the token-id return option (return_token_ids)."
                        )
                    candidates.append((
                        list(token_ids),
                        map_vllm_finish_reason(choice.get("finish_reason"), choice.get("stop_reason")),
                    ))
                return self._pack_output(index, item_ids[index], candidates)
            return task

        outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
        failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
        results = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
        if failures:
            raise PartialBatchError(results, failures)
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: Must carry no `extra` keys.

        Returns:
            Log probabilities of shape `[num_items, ref_len]` on CPU.

        Raises:
            ValueError: If items carry differing reference lengths or `params.extra` is
                non-empty.
            PartialBatchError: If some items failed after transport retries while others
                succeeded.
        """
        self._ensure_open()
        if params.extra:
            raise ValueError(
                f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
                "scoring accepts no forward keyword arguments."
            )
        if not items:
            return torch.zeros((0, 0), dtype=torch.float32)
        item_specs, _, item_salts = self._prepare_spec_submission(
            items, "vllm-serve", allow_constraints=False,
        )
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
        ref_len = ref_lens.pop()
        if ref_len == 0:
            return torch.zeros((len(items), 0), dtype=torch.float32)
        backend = self._backend

        prompt_ids = [self._resolve_item_ids(item) for item in items]
        ref_ids = [item.ref_output_ids.reshape(-1).tolist() for item in items]

        def make_task(index: int):
            def task() -> list[float]:
                body = {
                    "model": backend._served_model,
                    "prompt": [*prompt_ids[index], *ref_ids[index]],
                    "max_tokens": 1,
                    "temperature": 0.0,
                    "prompt_logprobs": 0,
                }
                if item_specs[index] is not None:
                    scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids[index]))
                    body["vllm_xargs"] = {"intervention_spec": scoring_spec.canonical()}
                    body["cache_salt"] = scoring_spec.salt()
                elif item_salts is not None:
                    body["cache_salt"] = item_salts[index]
                payload = with_transport_retries(
                    lambda: backend._post_json("/v1/completions", body),
                    max_attempts=backend.max_attempts,
                    backoff_base=backend.backoff_base,
                )
                choices = payload.get("choices", [])
                if not choices:
                    raise ValueError("The completions response carries no choices.")
                prompt_logprobs = choices[0].get("prompt_logprobs")
                return extract_ref_logprobs(prompt_logprobs, ref_ids[index])
            return task

        outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
        failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
        if failures:
            successes = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
            raise PartialBatchError(successes, failures)
        return torch.tensor(outcomes, dtype=torch.float32)

closed property

Whether the session has been closed.

layout property

Structural facts from the model config (client-side).

Raises:

Type	Description
RuntimeError	If the model config could not be resolved.

tokenizer property

The backend's client-side tokenizer.

capture(prompts, layers, mode, location='layer_output')

Hidden-state capture over the plugin is not implemented in this toolkit version.

Source code in steerability/backends/vllm/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Hidden-state capture over the plugin is not implemented in this toolkit version."""
    raise UnsupportedOperationError(
        "Hidden-state capture on vLLM backends is not implemented in this toolkit version."
    )

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/vllm/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

generate(items, params)

Generate one result per item through the completions endpoint.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items; entries must be empty.	required
params	GenerationParams	Normalized generation parameters shared by all items; unmapped extra keys raise.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order.

Raises:

Type	Description
PartialBatchError	If some items failed after transport retries while others succeeded.

Source code in steerability/backends/vllm/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item through the completions endpoint.

    Args:
        items: The generation items; entries must be empty.
        params: Normalized generation parameters shared by all items; unmapped `extra` keys
            raise.

    Returns:
        One `ItemResult` per item, in item order.

    Raises:
        PartialBatchError: If some items failed after transport retries while others
            succeeded.
    """
    self._ensure_open()
    if not items:
        return []
    item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm-serve")
    base_args = render_vllm_sampling_args(params)
    backend = self._backend

    item_ids = [self._resolve_item_ids(item) for item in items]
    seeds = [self._item_seed(item, params, index) for index, item in enumerate(items)]
    self._generate_count += 1

    def make_task(index: int):
        def task() -> ItemResult:
            body: dict[str, Any] = {
                "model": backend._served_model,
                "prompt": item_ids[index],
                "return_token_ids": True,
                **base_args,
            }
            if seeds[index] is not None:
                body["seed"] = seeds[index]
            if item_constraints[index] is not None:
                field, value = render_guided_decoding_field(item_constraints[index])
                body[f"guided_{field}"] = value
            if item_specs[index] is not None:
                # vllm_xargs is scalar-only, so nested specs travel as JSON strings
                body["vllm_xargs"] = {
                    "intervention_spec": item_specs[index].canonical(),
                }
            if item_salts is not None:
                body["cache_salt"] = item_salts[index]
            payload = with_transport_retries(
                lambda: backend._post_json("/v1/completions", body),
                max_attempts=backend.max_attempts,
                backoff_base=backend.backoff_base,
            )
            choices = payload.get("choices", [])
            if not choices:
                raise ValueError("The completions response carries no choices.")
            candidates = []
            for choice in choices:
                token_ids = choice.get("token_ids")
                if token_ids is None:
                    raise ValueError(
                        "The completions response carries no token_ids; the server does "
                        "not support the token-id return option (return_token_ids)."
                    )
                candidates.append((
                    list(token_ids),
                    map_vllm_finish_reason(choice.get("finish_reason"), choice.get("stop_reason")),
                ))
            return self._pack_output(index, item_ids[index], candidates)
        return task

    outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
    failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
    results = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
    if failures:
        raise PartialBatchError(results, failures)
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	Must carry no extra keys.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len] on CPU.

Raises:

Type	Description
ValueError	If items carry differing reference lengths or params.extra is non-empty.
PartialBatchError	If some items failed after transport retries while others succeeded.

Source code in steerability/backends/vllm/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: Must carry no `extra` keys.

    Returns:
        Log probabilities of shape `[num_items, ref_len]` on CPU.

    Raises:
        ValueError: If items carry differing reference lengths or `params.extra` is
            non-empty.
        PartialBatchError: If some items failed after transport retries while others
            succeeded.
    """
    self._ensure_open()
    if params.extra:
        raise ValueError(
            f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
            "scoring accepts no forward keyword arguments."
        )
    if not items:
        return torch.zeros((0, 0), dtype=torch.float32)
    item_specs, _, item_salts = self._prepare_spec_submission(
        items, "vllm-serve", allow_constraints=False,
    )
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
    ref_len = ref_lens.pop()
    if ref_len == 0:
        return torch.zeros((len(items), 0), dtype=torch.float32)
    backend = self._backend

    prompt_ids = [self._resolve_item_ids(item) for item in items]
    ref_ids = [item.ref_output_ids.reshape(-1).tolist() for item in items]

    def make_task(index: int):
        def task() -> list[float]:
            body = {
                "model": backend._served_model,
                "prompt": [*prompt_ids[index], *ref_ids[index]],
                "max_tokens": 1,
                "temperature": 0.0,
                "prompt_logprobs": 0,
            }
            if item_specs[index] is not None:
                scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids[index]))
                body["vllm_xargs"] = {"intervention_spec": scoring_spec.canonical()}
                body["cache_salt"] = scoring_spec.salt()
            elif item_salts is not None:
                body["cache_salt"] = item_salts[index]
            payload = with_transport_retries(
                lambda: backend._post_json("/v1/completions", body),
                max_attempts=backend.max_attempts,
                backoff_base=backend.backoff_base,
            )
            choices = payload.get("choices", [])
            if not choices:
                raise ValueError("The completions response carries no choices.")
            prompt_logprobs = choices[0].get("prompt_logprobs")
            return extract_ref_logprobs(prompt_logprobs, ref_ids[index])
        return task

    outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
    failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
    if failures:
        successes = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
        raise PartialBatchError(successes, failures)
    return torch.tensor(outcomes, dtype=torch.float32)

extract_ref_logprobs(prompt_logprobs, ref_ids)

Pull the reference tokens' log-probabilities from a prompt-logprobs structure.

Accepts both the offline shape (per-position mappings from token id to an object with a logprob attribute) and the serve JSON shape (string token-id keys mapping to dicts with a "logprob" entry). The reference occupies the last len(ref_ids) prompt positions.

Parameters:

Name	Type	Description	Default
prompt_logprobs	Sequence | None	The per-prompt-position logprob entries, aligned with the submitted prompt tokens (position 0 is None).	required
ref_ids	Sequence[int]	The reference token ids.	required

Returns:

Type	Description
list[float]	One log-probability per reference token.

Raises:

Type	Description
ValueError	If the structure is missing or a reference position lacks its token's entry.

Source code in steerability/backends/vllm/rendering.py

def extract_ref_logprobs(prompt_logprobs: Sequence | None, ref_ids: Sequence[int]) -> list[float]:
    """Pull the reference tokens' log-probabilities from a prompt-logprobs structure.

    Accepts both the offline shape (per-position mappings from token id to an object with a
    `logprob` attribute) and the serve JSON shape (string token-id keys mapping to dicts with a
    `"logprob"` entry). The reference occupies the last `len(ref_ids)` prompt positions.

    Args:
        prompt_logprobs: The per-prompt-position logprob entries, aligned with the submitted
            prompt tokens (position 0 is None).
        ref_ids: The reference token ids.

    Returns:
        One log-probability per reference token.

    Raises:
        ValueError: If the structure is missing or a reference position lacks its token's entry.
    """
    if prompt_logprobs is None:
        raise ValueError(
            "The response carries no prompt_logprobs; scoring requires prompt_logprobs=0 support."
        )
    if len(prompt_logprobs) < len(ref_ids):
        raise ValueError(
            f"prompt_logprobs has {len(prompt_logprobs)} positions for {len(ref_ids)} reference tokens."
        )
    values: list[float] = []
    offset = len(prompt_logprobs) - len(ref_ids)
    for position, token_id in enumerate(ref_ids):
        entry = prompt_logprobs[offset + position]
        if entry is None:
            raise ValueError(f"No logprob entry at reference position {position}.")
        record = entry.get(token_id, entry.get(str(token_id))) if hasattr(entry, "get") else None
        if record is None:
            raise ValueError(f"Token {token_id} missing from the logprob entry at position {position}.")
        if hasattr(record, "logprob"):
            values.append(float(record.logprob))
        elif isinstance(record, dict):
            values.append(float(record["logprob"]))
        else:
            values.append(float(record))
    return values

map_vllm_finish_reason(finish_reason, stop_reason)

Map a vLLM candidate's finish reason onto the toolkit vocabulary.

vLLM reports "stop" for EOS, stop strings, and stop token ids alike, with stop_reason None for EOS and the matched string or token id otherwise; "length" maps through unchanged, and anything else (e.g. "abort") maps to None.

Parameters:

Name	Type	Description	Default
finish_reason	str | None	The vLLM candidate's finish reason.	required
stop_reason	Any	The vLLM candidate's stop reason.	required

Returns:

Type	Description
str | None	One of "stop", "eos", "length", or None.

Source code in steerability/backends/vllm/rendering.py

def map_vllm_finish_reason(finish_reason: str | None, stop_reason: Any) -> str | None:
    """Map a vLLM candidate's finish reason onto the toolkit vocabulary.

    vLLM reports `"stop"` for EOS, stop strings, and stop token ids alike, with `stop_reason`
    None for EOS and the matched string or token id otherwise; `"length"` maps through
    unchanged, and anything else (e.g. `"abort"`) maps to None.

    Args:
        finish_reason: The vLLM candidate's finish reason.
        stop_reason: The vLLM candidate's stop reason.

    Returns:
        One of `"stop"`, `"eos"`, `"length"`, or None.
    """
    if finish_reason == "stop":
        return "eos" if stop_reason is None else "stop"
    if finish_reason == "length":
        return "length"
    return None

merge_intervention_specs(specs)

One spec carrying every op of specs, in order, with tensor payloads unioned.

Source code in steerability/backends/vllm/rendering.py

def merge_intervention_specs(specs: Sequence[InterventionSpec]) -> InterventionSpec:
    """One spec carrying every op of `specs`, in order, with tensor payloads unioned."""
    if len(specs) == 1:
        return specs[0]
    ops: list = []
    artifacts: dict = {}
    for spec in specs:
        ops.extend(spec.ops)
        artifacts.update(spec.artifacts)
    return InterventionSpec(ops=tuple(ops), artifacts=artifacts)

raise_for_spec_rejection(message)

Raise the toolkit error for a server-side spec rejection message carrying an E_* code.

Kind and constraint gaps (E_UNKNOWN_KIND, E_CONSTRAINT) are support facts a stale client missed and raise UnsupportedOperationError; every other E_* rejection is a malformed spec and raises ValueError. The code and JSON path are preserved verbatim. A message without an E_* code returns without raising.

Source code in steerability/backends/vllm/rendering.py

def raise_for_spec_rejection(message: str) -> None:
    """Raise the toolkit error for a server-side spec rejection message carrying an `E_*` code.

    Kind and constraint gaps (`E_UNKNOWN_KIND`, `E_CONSTRAINT`) are support facts a stale
    client missed and raise `UnsupportedOperationError`; every other `E_*` rejection is a
    malformed spec and raises `ValueError`. The code and JSON path are preserved verbatim.
    A message without an `E_*` code returns without raising.
    """
    if not _SPEC_ERROR_RE.search(message):
        return
    if any(code in message for code in _SUPPORT_FACT_CODES):
        raise UnsupportedOperationError(message)
    raise ValueError(message)

remap_spec_for_scoring(spec, prompt_len)

A scoring copy of spec with after_prompt scopes rewritten to from_position.

The teacher-forced reference is part of the server-side prompt, so the worker's "after the prompt" would select nothing; the rewrite anchors the scope at the original prompt length, the position of the first reference token in the submitted ids.

Source code in steerability/backends/vllm/rendering.py

def remap_spec_for_scoring(spec: InterventionSpec, prompt_len: int) -> InterventionSpec:
    """A scoring copy of `spec` with `after_prompt` scopes rewritten to `from_position`.

    The teacher-forced reference is part of the server-side prompt, so the worker's "after the
    prompt" would select nothing; the rewrite anchors the scope at the original prompt length,
    the position of the first reference token in the submitted ids.
    """
    ops = []
    changed = False
    for op in spec.to_wire()["ops"]:
        if op.get("scope", {}).get("kind") == "after_prompt":
            op = {**op, "scope": {"kind": "from_position", "position": int(prompt_len)}}
            changed = True
        ops.append(op)
    if not changed:
        return spec
    return InterventionSpec(ops=tuple(ops), artifacts=spec.artifacts)

render_constraint_sampling_args(field, value)

Constraint kwargs for SamplingParams, tolerant of the structured-outputs rename.

Newer vLLM removes GuidedDecodingParams in favor of StructuredOutputsParams passed as structured_outputs=; older versions serve guided_decoding=. The declarative field names (json, regex, grammar, choice) are shared by both surfaces.

Source code in steerability/backends/vllm/rendering.py

def render_constraint_sampling_args(field: str, value: Any) -> dict:
    """Constraint kwargs for `SamplingParams`, tolerant of the structured-outputs rename.

    Newer vLLM removes `GuidedDecodingParams` in favor of `StructuredOutputsParams` passed as
    `structured_outputs=`; older versions serve `guided_decoding=`. The declarative field names
    (`json`, `regex`, `grammar`, `choice`) are shared by both surfaces.
    """
    try:
        from vllm.sampling_params import StructuredOutputsParams
    except ImportError:
        # legacy api: compact whitespace is only enforced on the structured-outputs surface
        from vllm.sampling_params import GuidedDecodingParams
        return {"guided_decoding": GuidedDecodingParams(**{field: value})}
    return {"structured_outputs": StructuredOutputsParams(**{field: value})}

render_guided_decoding_field(source)

The vLLM structured-output parameter name and payload for a constraint source.

Source code in steerability/backends/vllm/rendering.py

def render_guided_decoding_field(source: ConstraintSource) -> tuple[str, Any]:
    """The vLLM structured-output parameter name and payload for a constraint source."""
    if source.kind == "json_schema":
        value = source.value if isinstance(source.value, str) else dict(source.value)
        return "json", value
    if source.kind == "regex":
        return "regex", source.value
    if source.kind == "grammar":
        return "grammar", source.value
    return "choice", list(source.value)

render_vllm_sampling_args(params)

Render normalized generation parameters onto vLLM sampling-parameter names.

The table is exhaustive on this arm. Every normalized field maps to its vLLM name (max_new_tokens to max_tokens, min_new_tokens to min_tokens, greedy=True to temperature=0.0, n to n, stop strings to stop with include_stop_str_in_output=True, extra stop ids to stop_token_ids), and any key left in extra raises rather than being dropped. seed is not rendered here; sessions derive and attach per-item seeds.

Parameters:

Name	Type	Description	Default
params	GenerationParams	The normalized parameters.	required

Returns:

Type	Description
dict[str, Any]	Keyword arguments for vllm.SamplingParams (also valid as vLLM completions-request
dict[str, Any]	fields).

Raises:

Type	Description
ValueError	If params.extra is non-empty; the message names the unmapped keys.
ValueError	If params.greedy is True while a non-zero temperature is also set.

Source code in steerability/backends/vllm/rendering.py

def render_vllm_sampling_args(params: GenerationParams) -> dict[str, Any]:
    """Render normalized generation parameters onto vLLM sampling-parameter names.

    The table is exhaustive on this arm. Every normalized field maps to its vLLM name
    (`max_new_tokens` to `max_tokens`, `min_new_tokens` to `min_tokens`, `greedy=True` to
    `temperature=0.0`, `n` to `n`, stop strings to `stop` with
    `include_stop_str_in_output=True`, extra stop ids to `stop_token_ids`), and any key left in
    `extra` raises rather than being dropped. `seed` is not rendered here; sessions derive and
    attach per-item seeds.

    Args:
        params: The normalized parameters.

    Returns:
        Keyword arguments for `vllm.SamplingParams` (also valid as vLLM completions-request
        fields).

    Raises:
        ValueError: If `params.extra` is non-empty; the message names the unmapped keys.
        ValueError: If `params.greedy` is True while a non-zero `temperature` is also set.
    """
    if params.extra:
        raise ValueError(
            f"Generation parameter(s) {sorted(params.extra)} have no vLLM rendering; the vLLM "
            "table is exhaustive and unmapped parameters are rejected rather than dropped."
        )
    args: dict[str, Any] = {}
    if params.max_new_tokens is not None:
        args["max_tokens"] = params.max_new_tokens
    if params.min_new_tokens is not None:
        args["min_tokens"] = params.min_new_tokens
    if params.temperature is not None:
        args["temperature"] = params.temperature
    if params.top_p is not None:
        args["top_p"] = params.top_p
    if params.top_k is not None:
        args["top_k"] = params.top_k
    if params.repetition_penalty is not None:
        args["repetition_penalty"] = params.repetition_penalty
    if params.n is not None:
        args["n"] = params.n
    if params.greedy is True:
        if params.temperature not in (None, 0.0):
            raise ValueError(
                "greedy decoding conflicts with a non-zero temperature; drop one of the two."
            )
        args["temperature"] = 0.0
    if params.stop_strings:
        args["stop"] = list(params.stop_strings)
        args["include_stop_str_in_output"] = True
    if params.stop_token_ids:
        args["stop_token_ids"] = list(params.stop_token_ids)
    return args

serve_environment(hook_plugin, base=None)

Return the environment for a vllm serve process under the engine boot policy.

Parameters:

Name	Type	Description	Default
hook_plugin	bool	Whether the server loads the vLLM-Hook unified worker.	required
base	Mapping[str, str] | None	Environment to start from; os.environ when omitted. Not mutated.	None

Returns:

Type	Description
dict[str, str]	A copy of base with the boot defaults filled where absent and the forced variables applied.

Source code in steerability/backends/vllm/environment.py

def serve_environment(hook_plugin: bool, base: Mapping[str, str] | None = None) -> dict[str, str]:
    """Return the environment for a `vllm serve` process under the engine boot policy.

    Args:
        hook_plugin: Whether the server loads the vLLM-Hook unified worker.
        base: Environment to start from; `os.environ` when omitted. Not mutated.

    Returns:
        A copy of `base` with the boot defaults filled where absent and the forced variables applied.
    """
    forced, defaults = engine_boot_environment(hook_plugin)
    env = dict(os.environ if base is None else base)
    for name, value in defaults.items():
        env.setdefault(name, value)
    env.update(forced)
    return env

backend

The vLLM backends: the offline engine ("vllm") and the OpenAI-compatible server ("vllm-serve").

This module holds the backend-side helpers (the artifact uploader, layout and tokenizer resolution, discovery reconciliation) and the two backend classes. It imports cleanly without vLLM installed: constructing VLLMBackend requires the vllm optional dependency (it boots an engine); VLLMServeBackend needs only a reachable vLLM server. Every from vllm ... and from vllm_hook_plugins ... import stays inside a function or method body.

logger = logging.getLogger(__name__) module-attribute

VLLMBackend

Bases: Backend

The offline vLLM engine backend.

Boots one engine per backend instance from the spec (engine_kwargs option forwarded to vllm.LLM); requires the vllm optional dependency. A CheckpointArtifact overrides the served model reference and a LoRAArtifact attaches as a LoRA request on every generation. When the spec declares hook_plugin, the unified worker is selected via VLLM_HOOK_WORKER=unified and the discovery payload is fetched once and cached by spec hash. Capability advertisement is available through capabilities_for_spec without constructing the backend.

Source code in steerability/backends/vllm/backend.py

class VLLMBackend(Backend):
    """The offline vLLM engine backend.

    Boots one engine per backend instance from the spec (`engine_kwargs` option forwarded to
    `vllm.LLM`); requires the `vllm` optional dependency. A `CheckpointArtifact` overrides the
    served model reference and a `LoRAArtifact` attaches as a LoRA request on every generation.
    When the spec declares `hook_plugin`, the unified worker is selected via
    `VLLM_HOOK_WORKER=unified` and the discovery payload is fetched once and cached by spec
    hash. Capability advertisement is available through `capabilities_for_spec` without
    constructing the backend.
    """

    def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> None:
        if spec.kind != "vllm":
            raise ValueError(f"VLLMBackend requires a 'vllm' spec; got kind {spec.kind!r}.")
        self.spec = spec
        self._released = False
        require("vllm")

        from vllm import LLM

        checkpoint, lora = _split_artifacts(artifacts)
        model_ref = checkpoint.path if checkpoint is not None else spec.model
        if model_ref is None:
            raise ValueError("VLLMBackend needs a model reference on the spec or a checkpoint artifact.")
        trust_remote_code = bool(spec.get_option("trust_remote_code", default=False))
        _reject_encoder_decoder(model_ref, trust_remote_code)

        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"})
        if lora is not None:
            engine_kwargs.setdefault("enable_lora", True)
        if trust_remote_code:
            engine_kwargs.setdefault("trust_remote_code", True)
        if spec.get_option("hook_plugin"):
            # worker hooks do not run under CUDA-graph replay; spec construction rejects an
            # explicit False, so this only fills the default
            engine_kwargs.setdefault("enforce_eager", True)

        # the boot environment is scoped to this engine's construction and restored afterwards,
        # so a later plugin-free engine in the same process is unaffected
        forced, defaults = engine_boot_environment(bool(spec.get_option("hook_plugin")))
        with engine_environment(forced, defaults):
            self._llm = LLM(model=model_ref, **engine_kwargs)

        # the pipeline records a backend only when the constructor returns, so a failure from
        # here on must release the engine before propagating
        try:
            self._lora_request = None
            if lora is not None:
                from vllm.lora.request import LoRARequest

                self._lora_request = LoRARequest("steered", 1, lora.path)

            tokenizer_source = (
                spec.get_option("tokenizer_name_or_path")
                or model_ref
            )
            self.tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code)
            self._layout = _config_layout(model_ref, trust_remote_code)
            self._plain_salt = uuid.uuid4().hex
            self._artifact_uploader = _ArtifactUploader(spec.get_option("artifact_dir"))
            self._discovery: dict | None = None
            if spec.get_option("hook_plugin"):
                self._discovery = self._fetch_discovery()
        except Exception:
            self.release()
            raise

    def stage_artifacts(self, payloads) -> None:
        """Write each content-addressed artifact into the plugin registry the engine reads.

        The offline engine shares the process's filesystem, so staging is a registry write
        (idempotent, verified against the content address).
        """
        self._artifact_uploader.upload_payloads(payloads)

    def _fetch_discovery(self) -> dict | None:
        cached = _DISCOVERY_CACHE.get(self.spec.spec_hash)
        if cached is not None:
            return cached
        payload = None
        for target in (self._llm, getattr(self._llm, "llm_engine", None)):
            rpc = getattr(target, "collective_rpc", None)
            if callable(rpc):
                try:
                    replies = rpc("hook_capabilities")
                except Exception as error:
                    logger.warning("vLLM-Hook discovery failed: %s", error)
                    return None
                payload = next((reply for reply in replies if reply), None)
                break
        if payload is None:
            logger.warning(
                "vLLM-Hook discovery returned no payload; is VLLM_HOOK_WORKER=unified active?"
            )
            return None
        _DISCOVERY_CACHE[self.spec.spec_hash] = payload
        _reconcile_discovery(self.spec, self.capabilities_for_spec(self.spec), payload)
        return payload

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The capability advertisement implied by `spec`."""
        return _vllm_capabilities(spec, offline=True)

    def open_session(self) -> "VLLMOfflineSession":
        """Open a request session over the shared engine.

        Raises:
            RuntimeError: If the backend has been released.
        """
        self._require_llm()
        return VLLMOfflineSession(self)

    def _require_llm(self):
        """The live engine, or a `RuntimeError` when the backend has been released."""
        if self._llm is None:
            raise RuntimeError(
                "This VLLMBackend was released; construct a new backend (or a new "
                "SteeringPipeline operation, which reconstructs backends automatically)."
            )
        return self._llm

    def release(self) -> None:
        """Shut the engine down explicitly and mark the backend unusable.

        Release is idempotent; after it, a new backend must be constructed. The distributed-state
        teardown is process-global, so release assumes no other live vLLM engine in the process.
        Ray-based executors are out of scope. Engine-touching calls on any still-open session raise
        after release.
        """
        if self._released:
            return
        self._released = True
        llm = self._llm
        self._llm = None
        self._lora_request = None

        for resolve in (
            lambda: getattr(llm, "shutdown", None),
            lambda: getattr(getattr(llm, "llm_engine", None), "shutdown", None),
            lambda: getattr(
                getattr(getattr(llm, "llm_engine", None), "engine_core", None), "shutdown", None
            ),
        ):
            shutdown = resolve()
            if callable(shutdown):
                try:
                    shutdown()
                except Exception:
                    logger.warning("vLLM engine shutdown hop failed; continuing.", exc_info=True)
                break

        del llm
        gc.collect()

        try:
            from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel

            destroy_model_parallel()
            destroy_distributed_environment()
        except Exception:
            logger.warning("vLLM distributed-state teardown failed; continuing.", exc_info=True)

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code) instance-attribute

capabilities_for_spec(spec) classmethod

The capability advertisement implied by spec.

Source code in steerability/backends/vllm/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The capability advertisement implied by `spec`."""
    return _vllm_capabilities(spec, offline=True)

open_session()

Open a request session over the shared engine.

Raises:

Type	Description
RuntimeError	If the backend has been released.

Source code in steerability/backends/vllm/backend.py

def open_session(self) -> "VLLMOfflineSession":
    """Open a request session over the shared engine.

    Raises:
        RuntimeError: If the backend has been released.
    """
    self._require_llm()
    return VLLMOfflineSession(self)

release()

Shut the engine down explicitly and mark the backend unusable.

Release is idempotent; after it, a new backend must be constructed. The distributed-state teardown is process-global, so release assumes no other live vLLM engine in the process. Ray-based executors are out of scope. Engine-touching calls on any still-open session raise after release.

Source code in steerability/backends/vllm/backend.py

def release(self) -> None:
    """Shut the engine down explicitly and mark the backend unusable.

    Release is idempotent; after it, a new backend must be constructed. The distributed-state
    teardown is process-global, so release assumes no other live vLLM engine in the process.
    Ray-based executors are out of scope. Engine-touching calls on any still-open session raise
    after release.
    """
    if self._released:
        return
    self._released = True
    llm = self._llm
    self._llm = None
    self._lora_request = None

    for resolve in (
        lambda: getattr(llm, "shutdown", None),
        lambda: getattr(getattr(llm, "llm_engine", None), "shutdown", None),
        lambda: getattr(
            getattr(getattr(llm, "llm_engine", None), "engine_core", None), "shutdown", None
        ),
    ):
        shutdown = resolve()
        if callable(shutdown):
            try:
                shutdown()
            except Exception:
                logger.warning("vLLM engine shutdown hop failed; continuing.", exc_info=True)
            break

    del llm
    gc.collect()

    try:
        from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel

        destroy_model_parallel()
        destroy_distributed_environment()
    except Exception:
        logger.warning("vLLM distributed-state teardown failed; continuing.", exc_info=True)

    if torch.cuda.is_available():
        torch.cuda.empty_cache()

stage_artifacts(payloads)

Write each content-addressed artifact into the plugin registry the engine reads.

The offline engine shares the process's filesystem, so staging is a registry write (idempotent, verified against the content address).

Source code in steerability/backends/vllm/backend.py

def stage_artifacts(self, payloads) -> None:
    """Write each content-addressed artifact into the plugin registry the engine reads.

    The offline engine shares the process's filesystem, so staging is a registry write
    (idempotent, verified against the content address).
    """
    self._artifact_uploader.upload_payloads(payloads)

VLLMServeBackend

Bases: Backend

The vLLM OpenAI-compatible server backend.

Targets a vLLM server rather than an arbitrary OpenAI-compatible endpoint: construction verifies the server's version surface (GET /version), fetches the plugin discovery payload (GET /v1/hook/capabilities) when the spec declares hook_plugin, and checks the served model id against the spec (or serves the pipeline's structural artifacts). Prompts submit as token ids on the completions endpoint with the token-id return option; the chat endpoint is not used. Requires no local vLLM installation.

Spec options: base_url (required, the server root), api_key, request_timeout, max_concurrency, max_retries, retry_backoff, tokenizer_name_or_path, trust_remote_code, hook_plugin.

Source code in steerability/backends/vllm/backend.py

class VLLMServeBackend(Backend):
    """The vLLM OpenAI-compatible server backend.

    Targets a vLLM server rather than an arbitrary OpenAI-compatible endpoint: construction
    verifies the server's version surface (`GET /version`), fetches the plugin discovery payload
    (`GET /v1/hook/capabilities`) when the spec declares `hook_plugin`, and checks the served
    model id against the spec (or serves the pipeline's structural artifacts). Prompts submit as
    token ids on the completions endpoint with the token-id return option; the chat endpoint is
    not used. Requires no local vLLM installation.

    Spec options: `base_url` (required, the server root), `api_key`, `request_timeout`,
    `max_concurrency`, `max_retries`, `retry_backoff`, `tokenizer_name_or_path`,
    `trust_remote_code`, `hook_plugin`.
    """

    def __init__(self, spec: BackendSpec, artifacts: Sequence[Artifact] = ()) -> None:
        if spec.kind != "vllm-serve":
            raise ValueError(f"VLLMServeBackend requires a 'vllm-serve' spec; got kind {spec.kind!r}.")
        self.spec = spec
        base_url = spec.get_option("base_url")
        if not base_url:
            raise ValueError("VLLMServeBackend requires a 'base_url' option on the spec.")
        self._base_url = base_url.rstrip("/").removesuffix("/v1")
        self._api_key = spec.get_option("api_key")
        self._timeout = float(spec.get_option("request_timeout", default=_DEFAULT_REQUEST_TIMEOUT))
        self.max_concurrency = int(spec.get_option("max_concurrency", default=_DEFAULT_MAX_CONCURRENCY))
        self.max_attempts = int(spec.get_option("max_retries", default=_DEFAULT_MAX_ATTEMPTS))
        self.backoff_base = float(spec.get_option("retry_backoff", default=0.5))
        trust_remote_code = bool(spec.get_option("trust_remote_code", default=False))

        version = self._get_json("/version")
        if not isinstance(version, dict) or "version" not in version:
            raise ValueError(
                f"The endpoint at {self._base_url} does not expose the vLLM version surface; "
                "only vLLM servers are supported."
            )

        self._discovery: dict | None = None
        if spec.get_option("hook_plugin"):
            self._discovery = _DISCOVERY_CACHE.get(spec.spec_hash)
            if self._discovery is None:
                try:
                    self._discovery = self._get_json("/v1/hook/capabilities")
                except (TransportError, ValueError) as error:
                    raise ValueError(
                        f"The spec declares hook_plugin but {self._base_url} serves no "
                        f"/v1/hook/capabilities discovery surface: {error}"
                    ) from error
                _DISCOVERY_CACHE[spec.spec_hash] = self._discovery
                _reconcile_discovery(spec, self.capabilities_for_spec(spec), self._discovery)

        checkpoint, lora = _split_artifacts(artifacts)
        expected_model = checkpoint.path if checkpoint is not None else spec.model
        if lora is not None:
            self._served_model = self._load_lora_adapter(lora)
        else:
            served = self._served_model_ids()
            if expected_model is None:
                if len(served) != 1:
                    raise ValueError(
                        f"The spec names no model and the server serves {served}; set the "
                        "spec's model to disambiguate."
                    )
                self._served_model = served[0]
            elif expected_model in served:
                self._served_model = expected_model
            else:
                raise ValueError(
                    f"The server at {self._base_url} serves {served}, not the configured "
                    f"model {expected_model!r}."
                )

        tokenizer_source = (
            spec.get_option("tokenizer_name_or_path")
            or (checkpoint.path if checkpoint is not None else None)
            or (lora.base_model if lora is not None else None)
            or spec.model
        )
        if tokenizer_source is None:
            tokenizer_source = self._served_model
        self.tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code)
        self._layout = _config_layout(
            (checkpoint.path if checkpoint is not None else None) or spec.model or self._served_model,
            trust_remote_code,
        )
        self._plain_salt = uuid.uuid4().hex
        # spec artifacts write to the registry root the server reads; shared_fs visibility
        # is verified against the server after writing (see stage_artifacts)
        self._artifact_uploader = _ArtifactUploader(spec.get_option("artifact_dir"))
        if self._discovery is not None:
            self._verify_fingerprints()

    def _served_model_ids(self) -> list[str]:
        payload = self._get_json("/v1/models")
        return [entry.get("id") for entry in payload.get("data", []) if isinstance(entry, dict)]

    def stage_artifacts(self, payloads) -> None:
        """Make each content-addressed artifact available to the serving engine.

        With an `artifact_dir` option the payloads are written into that registry root, which
        must be the server's registry directory (its `VLLM_HOOK_REGISTRY_DIR`) on a shared
        filesystem; visibility is verified through the server's artifact route when the
        discovery payload advertises `artifact_registry_root`. Otherwise each payload is PUT
        to the plugin's artifact route (`/v1/hook/artifacts/{id}`, body safetensors bytes, id
        verified server-side); already-exists is success.
        """
        if not payloads:
            return
        if self.spec.get_option("artifact_dir"):
            self._artifact_uploader.upload_payloads(payloads)
            self._verify_shared_fs_visibility(payloads)
            return
        import safetensors.torch

        for artifact_id, tensors in payloads.items():
            if artifact_id in self._artifact_uploader._written:
                continue
            data = safetensors.torch.save({name: tensors[name] for name in sorted(tensors)})
            self._put_bytes(f"/v1/hook/artifacts/{artifact_id}", data)
            self._artifact_uploader._written.add(artifact_id)

    def _verify_shared_fs_visibility(self, payloads) -> None:
        """Probe that shared_fs artifacts are visible to the server's registry.

        Gated on the discovery payload advertising `artifact_registry_root` (servers that
        advertise it also serve `HEAD /v1/hook/artifacts/{id}`); older servers skip the
        probe and keep the write-and-trust behavior.
        """
        server_root = (self._discovery or {}).get("artifact_registry_root")
        if not server_root:
            return
        client_root = os.path.abspath(self.spec.get_option("artifact_dir"))
        for artifact_id in payloads:
            if self._head_ok(f"/v1/hook/artifacts/{artifact_id}"):
                continue
            raise ValueError(
                f"artifact {artifact_id} written under {client_root} is not visible to the "
                f"server's registry ({server_root}); the shared_fs transport requires "
                "artifact_dir and the server's VLLM_HOOK_REGISTRY_DIR to name the same "
                "directory. Set VLLM_HOOK_REGISTRY_DIR on the server, point artifact_dir at "
                "the server's registry root, or drop artifact_dir to use the HTTP artifact "
                "route."
            )

    def _head_ok(self, path: str) -> bool:
        """HEAD a server path; True on 200, False on 404, raise otherwise."""
        request = urllib.request.Request(f"{self._base_url}{path}", method="HEAD")
        if self._api_key:
            request.add_header("Authorization", f"Bearer {self._api_key}")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout):
                return True
        except urllib.error.HTTPError as error:
            if error.code == 404:
                return False
            body = error.read().decode("utf-8", errors="replace")
            raise ValueError(f"HTTP {error.code} from {self._base_url}{path}: {body}") from error

    def _put_bytes(self, path: str, data: bytes) -> None:
        """PUT raw bytes to the server, mapping a missing route to a configuration error."""
        import urllib.error
        import urllib.request

        url = f"{self._base_url}{path}"
        request = urllib.request.Request(url, data=data, method="PUT")
        request.add_header("Content-Type", "application/octet-stream")
        if self._api_key:
            request.add_header("Authorization", f"Bearer {self._api_key}")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout):
                return
        except urllib.error.HTTPError as error:
            body = error.read().decode("utf-8", errors="replace")
            if error.code in (404, 405):
                raise ValueError(
                    f"{self._base_url} serves no artifact route ({error.code}); update the "
                    "server's vllm_hook_plugins, or configure artifact_dir on a filesystem "
                    "shared with the server."
                ) from error
            raise_for_spec_rejection(body)
            raise ValueError(f"HTTP {error.code} from {url}: {body}") from error
        except (urllib.error.URLError, TimeoutError, OSError) as error:
            raise TransportError(f"artifact upload to {url} failed: {error}") from error

    def _load_lora_adapter(self, lora: LoRAArtifact) -> str:
        served = self._served_model_ids()
        base = lora.base_model or self.spec.model
        if base and base not in served:
            raise ValueError(
                f"The server at {self._base_url} serves {served}, not the adapter's base "
                f"model {base!r}."
            )
        # the adapter name keys on path plus provenance, so a retrained adapter at the same
        # path loads as a new server-side adapter rather than reusing stale weights
        identity = f"{lora.path}:{lora.provenance.model_fingerprint or ''}"
        adapter_name = f"steered-{hashlib.sha256(identity.encode('utf-8')).hexdigest()[:8]}"
        if adapter_name in served:
            return adapter_name
        try:
            self._post_json(
                "/v1/load_lora_adapter",
                {"lora_name": adapter_name, "lora_path": lora.path},
                expect_json=False,
            )
        except (TransportError, ValueError) as error:
            raise ValueError(
                f"Could not load the LoRA artifact at {lora.path!r} onto the server "
                f"(dynamic adapter loading requires VLLM_ALLOW_RUNTIME_LORA_UPDATING): {error}"
            ) from error
        return adapter_name

    def _verify_fingerprints(self) -> None:
        """Verify the client tokenizer against the discovery payload's fingerprint recipes.

        Uses the plugin's engine-free `core.fingerprints` when the `vllm_hook_plugins` package
        is installed; mismatches warn rather than raise. Without the package, verification is
        skipped with a warning. A served fingerprint equal to the absent-template digest means
        the server exposes no chat template, so the comparison is skipped since a mismatch
        against it would reflect exposure rather than divergence.
        """
        model_block = (self._discovery or {}).get("model", {})
        remote_chat = model_block.get("chat_template_fingerprint")
        if remote_chat is None:
            return
        if is_absent_chat_template_fingerprint(remote_chat):
            logger.debug(
                "The server at %s does not expose a chat template (fingerprint %s is the "
                "absent-template digest); skipping the chat template comparison.",
                self._base_url, remote_chat,
            )
            return
        try:
            from vllm_hook_plugins.core.fingerprints import chat_template_fingerprint
        except ImportError:
            logger.warning(
                "Install vllm-hook-plugins to verify the client tokenizer against the server's "
                "fingerprints; skipping verification."
            )
            return
        local_chat = chat_template_fingerprint(getattr(self.tokenizer, "chat_template", None))
        if local_chat != remote_chat:
            logger.warning(
                "Client chat template (fingerprint %s) differs from the served one (%s); "
                "templated prompts may diverge from server-side expectations.",
                local_chat, remote_chat,
            )

    def _request_json(self, path: str, payload: dict | None, expect_json: bool = True) -> dict:
        url = f"{self._base_url}{path}"
        headers = {"Content-Type": "application/json"}
        if self._api_key:
            headers["Authorization"] = f"Bearer {self._api_key}"
        data = json.dumps(payload).encode("utf-8") if payload is not None else None
        request = urllib.request.Request(url, data=data, headers=headers)
        try:
            with urllib.request.urlopen(request, timeout=self._timeout) as response:
                body = response.read().decode("utf-8")
        except urllib.error.HTTPError as error:
            body = ""
            try:
                body = error.read().decode("utf-8", errors="replace")
            except Exception:
                pass
            # 5xx, timeouts, and rate limiting are transport-level and safe to retry
            if error.code >= 500 or error.code in (408, 429):
                raise TransportError(f"HTTP {error.code} from {url}: {body}") from error
            # admission rejections carry the plugin's E_* code and JSON path verbatim
            raise_for_spec_rejection(body)
            raise ValueError(f"HTTP {error.code} from {url}: {body}") from error
        except (urllib.error.URLError, TimeoutError, OSError) as error:
            raise TransportError(f"Request to {url} failed: {error}") from error
        if not expect_json:
            return {"text": body}
        try:
            return json.loads(body)
        except json.JSONDecodeError as error:
            raise ValueError(f"Non-JSON response from {url}: {error}") from error

    def _get_json(self, path: str) -> dict:
        return self._request_json(path, None)

    def _post_json(self, path: str, payload: dict, expect_json: bool = True) -> dict:
        return self._request_json(path, payload, expect_json=expect_json)

    @classmethod
    def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
        """The capability advertisement implied by `spec`."""
        return _vllm_capabilities(spec, offline=False)

    def open_session(self) -> "VLLMServeSession":
        """Open a request session over the shared connection."""
        return VLLMServeSession(self)

backoff_base = float(spec.get_option('retry_backoff', default=0.5)) instance-attribute

capabilities property

The advertised capability atoms.

capture_kinds property

The advertised capture kinds, when Capability.HIDDEN_CAPTURE is present.

intervention_kinds property

The advertised intervention kinds, when Capability.INTERVENTION_SPECS is present.

max_attempts = int(spec.get_option('max_retries', default=_DEFAULT_MAX_ATTEMPTS)) instance-attribute

max_concurrency = int(spec.get_option('max_concurrency', default=_DEFAULT_MAX_CONCURRENCY)) instance-attribute

processor_kinds property

The advertised processor kinds, when Capability.PER_STEP_LOGIT_SPECS is present.

spec = spec instance-attribute

tokenizer = _client_tokenizer(tokenizer_source, trust_remote_code) instance-attribute

capabilities_for_spec(spec) classmethod

The capability advertisement implied by spec.

Source code in steerability/backends/vllm/backend.py

@classmethod
def capabilities_for_spec(cls, spec: BackendSpec) -> BackendCapabilities:
    """The capability advertisement implied by `spec`."""
    return _vllm_capabilities(spec, offline=False)

open_session()

Open a request session over the shared connection.

Source code in steerability/backends/vllm/backend.py

def open_session(self) -> "VLLMServeSession":
    """Open a request session over the shared connection."""
    return VLLMServeSession(self)

release()

Free the resources this backend owns beyond what the caller owns.

The default is a no-op. Engine-owning backends override this to shut their engines down deterministically. Release is idempotent.

Source code in steerability/algorithms/core/execution/backend.py

def release(self) -> None:
    """Free the resources this backend owns beyond what the caller owns.

    The default is a no-op. Engine-owning backends override this to shut their engines down
    deterministically. Release is idempotent.
    """
    return None

stage_artifacts(payloads)

Make each content-addressed artifact available to the serving engine.

With an artifact_dir option the payloads are written into that registry root, which must be the server's registry directory (its VLLM_HOOK_REGISTRY_DIR) on a shared filesystem; visibility is verified through the server's artifact route when the discovery payload advertises artifact_registry_root. Otherwise each payload is PUT to the plugin's artifact route (/v1/hook/artifacts/{id}, body safetensors bytes, id verified server-side); already-exists is success.

Source code in steerability/backends/vllm/backend.py

def stage_artifacts(self, payloads) -> None:
    """Make each content-addressed artifact available to the serving engine.

    With an `artifact_dir` option the payloads are written into that registry root, which
    must be the server's registry directory (its `VLLM_HOOK_REGISTRY_DIR`) on a shared
    filesystem; visibility is verified through the server's artifact route when the
    discovery payload advertises `artifact_registry_root`. Otherwise each payload is PUT
    to the plugin's artifact route (`/v1/hook/artifacts/{id}`, body safetensors bytes, id
    verified server-side); already-exists is success.
    """
    if not payloads:
        return
    if self.spec.get_option("artifact_dir"):
        self._artifact_uploader.upload_payloads(payloads)
        self._verify_shared_fs_visibility(payloads)
        return
    import safetensors.torch

    for artifact_id, tensors in payloads.items():
        if artifact_id in self._artifact_uploader._written:
            continue
        data = safetensors.torch.save({name: tensors[name] for name in sorted(tensors)})
        self._put_bytes(f"/v1/hook/artifacts/{artifact_id}", data)
        self._artifact_uploader._written.add(artifact_id)

capabilities

The vLLM capability tables, discovery cache and negotiation, and capability refusals.

The kind tables and baseline capabilities are static data used by check(); the discovery cache and negotiation narrow them to what a live engine or server confirms. This module imports cleanly without vLLM installed.

VLLM_BASELINE_CAPABILITIES = BackendCapabilities(atoms=frozenset({Capability.SERVE_CHECKPOINT, Capability.SERVE_LORA, Capability.GUIDED_DECODING}), constraint_kinds=_VLLM_CONSTRAINT_KINDS) module-attribute

logger = logging.getLogger(__name__) module-attribute

environment

Engine-boot environment policy for the offline vLLM backend and a launched vllm serve.

vLLM reads some settings from process environment variables only, and its engine core is a spawned process that inherits the environment as it stands when LLM(...) builds its config. This module computes one boot policy with two consumers: the offline backend applies it to os.environ with save and restore semantics around LLM(...), so the settings are live while the engine is constructed and are cleared again immediately afterwards, and serve_environment returns the same policy as a fresh mapping for a vllm serve process.

Two variables are governed:

• VLLM_USE_FLASHINFER_SAMPLER defaults to "0" on every boot. vLLM selects the FlashInfer top-k/top-p sampler when it is available, and flashinfer-python ships that kernel JIT-only; the JIT compile runs at startup warmup and invokes nvcc at a path derived from CUDA_HOME, so a node whose CUDA toolkit does not match the installed torch build fails at boot. The native PyTorch sampler decodes the same greedy tokens. This is a default: an explicit caller setting wins.
• VLLM_HOOK_WORKER is forced to "unified" for a hook_plugin boot to select the plugin's unified worker, and is restored to its prior value after the boot.

This module does not set VLLM_USE_V2_MODEL_RUNNER; the vLLM-Hook plugin owns the model-runner constraint (it pins the legacy runner from inside the plugin, covering the vllm serve process the toolkit does not launch).

FLASHINFER_SAMPLER_VARIABLE = 'VLLM_USE_FLASHINFER_SAMPLER' module-attribute

HOOK_WORKER_VARIABLE = 'VLLM_HOOK_WORKER' module-attribute

engine_boot_environment(hook_plugin)

Return the forced and default environment mappings for one offline engine boot.

Parameters:

Name	Type	Description	Default
hook_plugin	bool	Whether the boot selects the vLLM-Hook unified worker.	required

Returns:

Type	Description
dict[str, str]	A (forced, defaults) pair. forced variables are applied regardless of any existing
dict[str, str]	value and restored afterwards; defaults are applied only when the variable is unset,
tuple[dict[str, str], dict[str, str]]	so an explicit caller value wins. defaults always carries
tuple[dict[str, str], dict[str, str]]	VLLM_USE_FLASHINFER_SAMPLER="0"; forced carries VLLM_HOOK_WORKER="unified" only
tuple[dict[str, str], dict[str, str]]	when hook_plugin is true.

Source code in steerability/backends/vllm/environment.py

def engine_boot_environment(hook_plugin: bool) -> tuple[dict[str, str], dict[str, str]]:
    """Return the forced and default environment mappings for one offline engine boot.

    Args:
        hook_plugin: Whether the boot selects the vLLM-Hook unified worker.

    Returns:
        A `(forced, defaults)` pair. `forced` variables are applied regardless of any existing
        value and restored afterwards; `defaults` are applied only when the variable is unset,
        so an explicit caller value wins. `defaults` always carries
        `VLLM_USE_FLASHINFER_SAMPLER="0"`; `forced` carries `VLLM_HOOK_WORKER="unified"` only
        when `hook_plugin` is true.
    """
    forced: dict[str, str] = {}
    defaults: dict[str, str] = {FLASHINFER_SAMPLER_VARIABLE: "0"}
    if hook_plugin:
        forced[HOOK_WORKER_VARIABLE] = "unified"
    return forced, defaults

engine_environment(forced, defaults)

Apply the boot environment for the span of the block and restore it on exit.

Forced variables are written regardless of any existing value; default variables are written only when absent from os.environ. Every variable written is restored on exit, including variables that were unset on entry (they are removed again). Restoration also runs when the block raises.

Parameters:

Name	Type	Description	Default
forced	Mapping[str, str]	Variables to set regardless of any existing value.	required
defaults	Mapping[str, str]	Variables to set only when they are absent from the environment.	required

Yields:

Type	Description
dict[str, str]	The variables actually written, as a name-to-value mapping.

Source code in steerability/backends/vllm/environment.py

@contextmanager
def engine_environment(forced: Mapping[str, str], defaults: Mapping[str, str]) -> Iterator[dict[str, str]]:
    """Apply the boot environment for the span of the block and restore it on exit.

    Forced variables are written regardless of any existing value; default variables are
    written only when absent from `os.environ`. Every variable written is restored on exit,
    including variables that were unset on entry (they are removed again). Restoration also
    runs when the block raises.

    Args:
        forced: Variables to set regardless of any existing value.
        defaults: Variables to set only when they are absent from the environment.

    Yields:
        The variables actually written, as a name-to-value mapping.
    """
    previous: dict[str, str | None] = {}
    applied: dict[str, str] = {name: value for name, value in defaults.items() if name not in os.environ}
    applied.update(forced)
    for name, value in applied.items():
        previous[name] = os.environ.get(name)
        os.environ[name] = value
    try:
        yield dict(applied)
    finally:
        for name, value in previous.items():
            if value is None:
                os.environ.pop(name, None)
            else:
                os.environ[name] = value

serve_environment(hook_plugin, base=None)

Return the environment for a vllm serve process under the engine boot policy.

Parameters:

Name	Type	Description	Default
hook_plugin	bool	Whether the server loads the vLLM-Hook unified worker.	required
base	Mapping[str, str] | None	Environment to start from; os.environ when omitted. Not mutated.	None

Returns:

Type	Description
dict[str, str]	A copy of base with the boot defaults filled where absent and the forced variables applied.

Source code in steerability/backends/vllm/environment.py

def serve_environment(hook_plugin: bool, base: Mapping[str, str] | None = None) -> dict[str, str]:
    """Return the environment for a `vllm serve` process under the engine boot policy.

    Args:
        hook_plugin: Whether the server loads the vLLM-Hook unified worker.
        base: Environment to start from; `os.environ` when omitted. Not mutated.

    Returns:
        A copy of `base` with the boot defaults filled where absent and the forced variables applied.
    """
    forced, defaults = engine_boot_environment(hook_plugin)
    env = dict(os.environ if base is None else base)
    for name, value in defaults.items():
        env.setdefault(name, value)
    env.update(forced)
    return env

rendering

Request and response rendering for the vLLM sessions.

Holds the strict sampling-argument table, the finish-reason and logprob mapping helpers, the per-item entry split and intervention-spec merge/remap, and the server-side spec-rejection parser. These are plain functions; this module imports cleanly without vLLM installed, and the vllm imports it needs stay function-local.

extract_ref_logprobs(prompt_logprobs, ref_ids)

Pull the reference tokens' log-probabilities from a prompt-logprobs structure.

Accepts both the offline shape (per-position mappings from token id to an object with a logprob attribute) and the serve JSON shape (string token-id keys mapping to dicts with a "logprob" entry). The reference occupies the last len(ref_ids) prompt positions.

Parameters:

Name	Type	Description	Default
prompt_logprobs	Sequence | None	The per-prompt-position logprob entries, aligned with the submitted prompt tokens (position 0 is None).	required
ref_ids	Sequence[int]	The reference token ids.	required

Returns:

Type	Description
list[float]	One log-probability per reference token.

Raises:

Type	Description
ValueError	If the structure is missing or a reference position lacks its token's entry.

Source code in steerability/backends/vllm/rendering.py

def extract_ref_logprobs(prompt_logprobs: Sequence | None, ref_ids: Sequence[int]) -> list[float]:
    """Pull the reference tokens' log-probabilities from a prompt-logprobs structure.

    Accepts both the offline shape (per-position mappings from token id to an object with a
    `logprob` attribute) and the serve JSON shape (string token-id keys mapping to dicts with a
    `"logprob"` entry). The reference occupies the last `len(ref_ids)` prompt positions.

    Args:
        prompt_logprobs: The per-prompt-position logprob entries, aligned with the submitted
            prompt tokens (position 0 is None).
        ref_ids: The reference token ids.

    Returns:
        One log-probability per reference token.

    Raises:
        ValueError: If the structure is missing or a reference position lacks its token's entry.
    """
    if prompt_logprobs is None:
        raise ValueError(
            "The response carries no prompt_logprobs; scoring requires prompt_logprobs=0 support."
        )
    if len(prompt_logprobs) < len(ref_ids):
        raise ValueError(
            f"prompt_logprobs has {len(prompt_logprobs)} positions for {len(ref_ids)} reference tokens."
        )
    values: list[float] = []
    offset = len(prompt_logprobs) - len(ref_ids)
    for position, token_id in enumerate(ref_ids):
        entry = prompt_logprobs[offset + position]
        if entry is None:
            raise ValueError(f"No logprob entry at reference position {position}.")
        record = entry.get(token_id, entry.get(str(token_id))) if hasattr(entry, "get") else None
        if record is None:
            raise ValueError(f"Token {token_id} missing from the logprob entry at position {position}.")
        if hasattr(record, "logprob"):
            values.append(float(record.logprob))
        elif isinstance(record, dict):
            values.append(float(record["logprob"]))
        else:
            values.append(float(record))
    return values

map_vllm_finish_reason(finish_reason, stop_reason)

Map a vLLM candidate's finish reason onto the toolkit vocabulary.

vLLM reports "stop" for EOS, stop strings, and stop token ids alike, with stop_reason None for EOS and the matched string or token id otherwise; "length" maps through unchanged, and anything else (e.g. "abort") maps to None.

Parameters:

Name	Type	Description	Default
finish_reason	str | None	The vLLM candidate's finish reason.	required
stop_reason	Any	The vLLM candidate's stop reason.	required

Returns:

Type	Description
str | None	One of "stop", "eos", "length", or None.

Source code in steerability/backends/vllm/rendering.py

def map_vllm_finish_reason(finish_reason: str | None, stop_reason: Any) -> str | None:
    """Map a vLLM candidate's finish reason onto the toolkit vocabulary.

    vLLM reports `"stop"` for EOS, stop strings, and stop token ids alike, with `stop_reason`
    None for EOS and the matched string or token id otherwise; `"length"` maps through
    unchanged, and anything else (e.g. `"abort"`) maps to None.

    Args:
        finish_reason: The vLLM candidate's finish reason.
        stop_reason: The vLLM candidate's stop reason.

    Returns:
        One of `"stop"`, `"eos"`, `"length"`, or None.
    """
    if finish_reason == "stop":
        return "eos" if stop_reason is None else "stop"
    if finish_reason == "length":
        return "length"
    return None

merge_intervention_specs(specs)

One spec carrying every op of specs, in order, with tensor payloads unioned.

Source code in steerability/backends/vllm/rendering.py

def merge_intervention_specs(specs: Sequence[InterventionSpec]) -> InterventionSpec:
    """One spec carrying every op of `specs`, in order, with tensor payloads unioned."""
    if len(specs) == 1:
        return specs[0]
    ops: list = []
    artifacts: dict = {}
    for spec in specs:
        ops.extend(spec.ops)
        artifacts.update(spec.artifacts)
    return InterventionSpec(ops=tuple(ops), artifacts=artifacts)

raise_for_spec_rejection(message)

Raise the toolkit error for a server-side spec rejection message carrying an E_* code.

Kind and constraint gaps (E_UNKNOWN_KIND, E_CONSTRAINT) are support facts a stale client missed and raise UnsupportedOperationError; every other E_* rejection is a malformed spec and raises ValueError. The code and JSON path are preserved verbatim. A message without an E_* code returns without raising.

Source code in steerability/backends/vllm/rendering.py

def raise_for_spec_rejection(message: str) -> None:
    """Raise the toolkit error for a server-side spec rejection message carrying an `E_*` code.

    Kind and constraint gaps (`E_UNKNOWN_KIND`, `E_CONSTRAINT`) are support facts a stale
    client missed and raise `UnsupportedOperationError`; every other `E_*` rejection is a
    malformed spec and raises `ValueError`. The code and JSON path are preserved verbatim.
    A message without an `E_*` code returns without raising.
    """
    if not _SPEC_ERROR_RE.search(message):
        return
    if any(code in message for code in _SUPPORT_FACT_CODES):
        raise UnsupportedOperationError(message)
    raise ValueError(message)

remap_spec_for_scoring(spec, prompt_len)

A scoring copy of spec with after_prompt scopes rewritten to from_position.

The teacher-forced reference is part of the server-side prompt, so the worker's "after the prompt" would select nothing; the rewrite anchors the scope at the original prompt length, the position of the first reference token in the submitted ids.

Source code in steerability/backends/vllm/rendering.py

def remap_spec_for_scoring(spec: InterventionSpec, prompt_len: int) -> InterventionSpec:
    """A scoring copy of `spec` with `after_prompt` scopes rewritten to `from_position`.

    The teacher-forced reference is part of the server-side prompt, so the worker's "after the
    prompt" would select nothing; the rewrite anchors the scope at the original prompt length,
    the position of the first reference token in the submitted ids.
    """
    ops = []
    changed = False
    for op in spec.to_wire()["ops"]:
        if op.get("scope", {}).get("kind") == "after_prompt":
            op = {**op, "scope": {"kind": "from_position", "position": int(prompt_len)}}
            changed = True
        ops.append(op)
    if not changed:
        return spec
    return InterventionSpec(ops=tuple(ops), artifacts=spec.artifacts)

render_constraint_sampling_args(field, value)

Constraint kwargs for SamplingParams, tolerant of the structured-outputs rename.

Newer vLLM removes GuidedDecodingParams in favor of StructuredOutputsParams passed as structured_outputs=; older versions serve guided_decoding=. The declarative field names (json, regex, grammar, choice) are shared by both surfaces.

Source code in steerability/backends/vllm/rendering.py

def render_constraint_sampling_args(field: str, value: Any) -> dict:
    """Constraint kwargs for `SamplingParams`, tolerant of the structured-outputs rename.

    Newer vLLM removes `GuidedDecodingParams` in favor of `StructuredOutputsParams` passed as
    `structured_outputs=`; older versions serve `guided_decoding=`. The declarative field names
    (`json`, `regex`, `grammar`, `choice`) are shared by both surfaces.
    """
    try:
        from vllm.sampling_params import StructuredOutputsParams
    except ImportError:
        # legacy api: compact whitespace is only enforced on the structured-outputs surface
        from vllm.sampling_params import GuidedDecodingParams
        return {"guided_decoding": GuidedDecodingParams(**{field: value})}
    return {"structured_outputs": StructuredOutputsParams(**{field: value})}

render_guided_decoding_field(source)

The vLLM structured-output parameter name and payload for a constraint source.

Source code in steerability/backends/vllm/rendering.py

def render_guided_decoding_field(source: ConstraintSource) -> tuple[str, Any]:
    """The vLLM structured-output parameter name and payload for a constraint source."""
    if source.kind == "json_schema":
        value = source.value if isinstance(source.value, str) else dict(source.value)
        return "json", value
    if source.kind == "regex":
        return "regex", source.value
    if source.kind == "grammar":
        return "grammar", source.value
    return "choice", list(source.value)

render_vllm_sampling_args(params)

Render normalized generation parameters onto vLLM sampling-parameter names.

The table is exhaustive on this arm. Every normalized field maps to its vLLM name (max_new_tokens to max_tokens, min_new_tokens to min_tokens, greedy=True to temperature=0.0, n to n, stop strings to stop with include_stop_str_in_output=True, extra stop ids to stop_token_ids), and any key left in extra raises rather than being dropped. seed is not rendered here; sessions derive and attach per-item seeds.

Parameters:

Name	Type	Description	Default
params	GenerationParams	The normalized parameters.	required

Returns:

Type	Description
dict[str, Any]	Keyword arguments for vllm.SamplingParams (also valid as vLLM completions-request
dict[str, Any]	fields).

Raises:

Type	Description
ValueError	If params.extra is non-empty; the message names the unmapped keys.
ValueError	If params.greedy is True while a non-zero temperature is also set.

Source code in steerability/backends/vllm/rendering.py

def render_vllm_sampling_args(params: GenerationParams) -> dict[str, Any]:
    """Render normalized generation parameters onto vLLM sampling-parameter names.

    The table is exhaustive on this arm. Every normalized field maps to its vLLM name
    (`max_new_tokens` to `max_tokens`, `min_new_tokens` to `min_tokens`, `greedy=True` to
    `temperature=0.0`, `n` to `n`, stop strings to `stop` with
    `include_stop_str_in_output=True`, extra stop ids to `stop_token_ids`), and any key left in
    `extra` raises rather than being dropped. `seed` is not rendered here; sessions derive and
    attach per-item seeds.

    Args:
        params: The normalized parameters.

    Returns:
        Keyword arguments for `vllm.SamplingParams` (also valid as vLLM completions-request
        fields).

    Raises:
        ValueError: If `params.extra` is non-empty; the message names the unmapped keys.
        ValueError: If `params.greedy` is True while a non-zero `temperature` is also set.
    """
    if params.extra:
        raise ValueError(
            f"Generation parameter(s) {sorted(params.extra)} have no vLLM rendering; the vLLM "
            "table is exhaustive and unmapped parameters are rejected rather than dropped."
        )
    args: dict[str, Any] = {}
    if params.max_new_tokens is not None:
        args["max_tokens"] = params.max_new_tokens
    if params.min_new_tokens is not None:
        args["min_tokens"] = params.min_new_tokens
    if params.temperature is not None:
        args["temperature"] = params.temperature
    if params.top_p is not None:
        args["top_p"] = params.top_p
    if params.top_k is not None:
        args["top_k"] = params.top_k
    if params.repetition_penalty is not None:
        args["repetition_penalty"] = params.repetition_penalty
    if params.n is not None:
        args["n"] = params.n
    if params.greedy is True:
        if params.temperature not in (None, 0.0):
            raise ValueError(
                "greedy decoding conflicts with a non-zero temperature; drop one of the two."
            )
        args["temperature"] = 0.0
    if params.stop_strings:
        args["stop"] = list(params.stop_strings)
        args["include_stop_str_in_output"] = True
    if params.stop_token_ids:
        args["stop_token_ids"] = list(params.stop_token_ids)
    return args

session

The vLLM request sessions: the offline-engine session and the server session.

_RequestSessionBase holds the lifecycle and layout shared by both; VLLMOfflineSession submits batched engine calls and serves plugin capture, and VLLMServeSession fans requests out over the completions endpoint. This module imports cleanly without vLLM installed; the vllm imports the sessions need stay method-local.

VLLMOfflineSession

Bases: _RequestSessionBase

Request session over the offline engine.

Token-id prompts submit as TokensPrompts in one engine call with per-item sampling parameters; the engine schedules the batch internally, so no client-side fan-out is needed.

Source code in steerability/backends/vllm/session.py

class VLLMOfflineSession(_RequestSessionBase):
    """Request session over the offline engine.

    Token-id prompts submit as `TokensPrompt`s in one engine call with per-item sampling
    parameters; the engine schedules the batch internally, so no client-side fan-out is needed.
    """

    def capture(
        self,
        prompts: list[PreparedPrompt],
        layers: list[int],
        mode: Literal["all_tokens", "last_token"],
        location: Literal["layer_output", "layer_input"] = "layer_output",
    ) -> CaptureResult:
        """Hidden-state capture over the plugin's capture surface.

        One request per prompt carries a `capture` spec and a fresh random `cache_salt`
        (a prefix-cache hit skips forward passes, so capture cannot tolerate reused salts) with
        `max_tokens=1`; the surplus decode position is truncated by the plugin. Per-layer
        tensors are stacked and right-padded to the batch's longest prompt.

        Args:
            prompts: The prompts to capture over.
            layers: 0-based decoder-layer indices to capture.
            mode: `"all_tokens"` for every prompt position, `"last_token"` for the final real
                position per row.
            location: The residual-stream boundary, `"layer_output"` or `"layer_input"`.

        Returns:
            The capture result: `[N, T, H]` per layer for `"all_tokens"` or `[N, H]` for
            `"last_token"`, on CPU in the engine's native dtype, with the derived `[N, T]`
            attention mask.

        Raises:
            UnsupportedOperationError: If the spec declares no `hook_plugin`, the negotiated
                capture kinds lack the requested mode or location, or the engine facts refuse
                capture (speculative decoding, non-eager execution).
            ValueError: If `prompts` is empty, a layer id is out of range, or the engine
                returned no capture payload.
        """
        self._ensure_open()
        backend = self._backend
        if not backend.spec.get_option("hook_plugin"):
            raise UnsupportedOperationError(
                "Hidden-state capture requires the vLLM-Hook plugin; declare hook_plugin=True "
                "on the vllm backend spec, or run capture on the huggingface backend."
            )
        capture_kinds = backend.capture_kinds
        required = CaptureKinds(
            kinds=frozenset({"residual"}),
            locations=frozenset({location}),
            modes=frozenset({mode}),
        )
        if capture_kinds is None or not capture_kinds.contains(required):
            raise UnsupportedOperationError(
                f"The serving backend does not advertise capture mode {mode!r} at location "
                f"{location!r}; update the server's vllm_hook_plugins or run capture on the "
                "huggingface backend."
            )
        _refuse_by_engine_facts(backend._discovery, "capture")
        if not prompts:
            raise ValueError("capture() requires at least one prompt.")
        num_layers = self.layout.num_layers
        missing = sorted(int(layer) for layer in layers if not 0 <= int(layer) < num_layers)
        if missing:
            raise ValueError(
                f"Requested layer ids {missing} are out of range; the model has {num_layers} layers."
            )

        from vllm import SamplingParams, TokensPrompt

        layer_ids = [int(layer) for layer in layers]
        # the client's validator and assembly expect full prompt coverage and pool the last real
        # position themselves, so every wire capture requests all_tokens
        wire_mode = "all_tokens" if mode == "last_token" else mode
        capture_spec = {"layers": layer_ids, "mode": wire_mode, "location": location}
        engine_prompts = []
        prompt_lens: list[int] = []
        for prompt in prompts:
            resolved = prompt.resolve_token_ids(self.tokenizer)
            ids = resolved.token_ids[0]
            if resolved.attention_mask is not None:
                ids = ids[resolved.attention_mask[0].bool()]
            ids = ids.tolist()
            prompt_lens.append(len(ids))
            engine_prompt = TokensPrompt(prompt_token_ids=ids)
            engine_prompt["cache_salt"] = uuid.uuid4().hex
            engine_prompts.append(engine_prompt)
        sampling = SamplingParams(max_tokens=1, temperature=0.0, extra_args={"capture": capture_spec})

        request_outputs = self._backend._require_llm().generate(engine_prompts, sampling, use_tqdm=False)

        rows_per_layer: dict[int, list[torch.Tensor]] = {layer: [] for layer in layer_ids}
        for index, request_output in enumerate(request_outputs):
            payload = getattr(request_output, "captures", None)
            if payload is None:
                raise ValueError(
                    "The engine returned no capture payload; is the vLLM-Hook unified worker "
                    "active on this engine?"
                )
            manifest_json, data = payload
            manifest = json.loads(manifest_json)
            tensors = _load_safetensors_bytes(data)
            for layer in layer_ids:
                stacked = tensors.get(f"layer_{layer}")
                if stacked is None or stacked.size(0) < prompt_lens[index]:
                    raise ValueError(
                        f"The capture payload covers layer {layer} at "
                        f"{0 if stacked is None else stacked.size(0)} of {prompt_lens[index]} "
                        f"prompt positions for prompt {index}; positions recorded: "
                        f"{manifest.get('positions', {}).get(str(layer))}."
                    )
                rows_per_layer[layer].append(stacked[: prompt_lens[index]])

        max_len = max(prompt_lens)
        attention_mask = torch.zeros(len(prompts), max_len, dtype=torch.long)
        for index, length in enumerate(prompt_lens):
            attention_mask[index, :length] = 1

        hidden: dict[int, torch.Tensor] = {}
        for layer, rows in rows_per_layer.items():
            if mode == "last_token":
                hidden[layer] = torch.stack([row[-1] for row in rows])
            else:
                padded = torch.zeros(len(rows), max_len, rows[0].size(-1), dtype=rows[0].dtype)
                for index, row in enumerate(rows):
                    padded[index, : row.size(0)] = row
                hidden[layer] = padded
        return CaptureResult(hidden=hidden, attention_mask=attention_mask, mode=mode, location=location)

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item through the engine.

        Args:
            items: The generation items; state entries lower as intervention specs on
                plugin-active backends, and no client-side hooks or live processors execute
                here.
            params: Normalized generation parameters shared by all items; unmapped `extra` keys
                raise.

        Returns:
            One `ItemResult` per item, in item order.
        """
        self._ensure_open()
        if not items:
            return []
        item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm")
        base_args = render_vllm_sampling_args(params)

        from vllm import SamplingParams, TokensPrompt

        prompts = []
        sampling = []
        prompt_ids_per_item: list[list[int]] = []
        for index, item in enumerate(items):
            ids = self._resolve_item_ids(item)
            prompt_ids_per_item.append(ids)
            args = dict(base_args)
            seed = self._item_seed(item, params, index)
            if seed is not None:
                args["seed"] = seed
            if item_constraints[index] is not None:
                field, value = render_guided_decoding_field(item_constraints[index])
                args.update(render_constraint_sampling_args(field, value))
            if item_specs[index] is not None:
                args["extra_args"] = {"intervention_spec": item_specs[index].to_wire()}
            prompt = TokensPrompt(prompt_token_ids=ids)
            if item_salts is not None:
                prompt["cache_salt"] = item_salts[index]
            prompts.append(prompt)
            sampling.append(SamplingParams(**args))
        self._generate_count += 1

        generate_kwargs: dict[str, Any] = {"use_tqdm": False}
        if self._backend._lora_request is not None:
            generate_kwargs["lora_request"] = self._backend._lora_request
        request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)

        results: list[ItemResult] = []
        for index, request_output in enumerate(request_outputs):
            candidates = [
                (
                    list(candidate.token_ids),
                    map_vllm_finish_reason(
                        candidate.finish_reason, getattr(candidate, "stop_reason", None),
                    ),
                )
                for candidate in request_output.outputs
            ]
            results.append(self._pack_output(index, prompt_ids_per_item[index], candidates))
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

        Each item's prompt and reference concatenate into one token-id prompt submitted with
        `prompt_logprobs=0`, and the reference positions' log-probabilities are read back.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: Must carry no `extra` keys; forward keyword arguments have no remote
                rendering.

        Returns:
            Log probabilities of shape `[num_items, ref_len]` on CPU.

        Raises:
            ValueError: If items carry differing reference lengths or `params.extra` is
                non-empty.
        """
        self._ensure_open()
        if params.extra:
            raise ValueError(
                f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
                "scoring accepts no forward keyword arguments."
            )
        if not items:
            return torch.zeros((0, 0), dtype=torch.float32)
        item_specs, _, item_salts = self._prepare_spec_submission(
            items, "vllm", allow_constraints=False,
        )
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
        ref_len = ref_lens.pop()
        if ref_len == 0:
            return torch.zeros((len(items), 0), dtype=torch.float32)

        from vllm import SamplingParams, TokensPrompt

        prompts = []
        sampling = []
        ref_ids_per_item: list[list[int]] = []
        for index, item in enumerate(items):
            prompt_ids = self._resolve_item_ids(item)
            ref_ids = item.ref_output_ids.reshape(-1).tolist()
            ref_ids_per_item.append(ref_ids)
            prompt = TokensPrompt(prompt_token_ids=[*prompt_ids, *ref_ids])
            args: dict[str, Any] = {"max_tokens": 1, "temperature": 0.0, "prompt_logprobs": 0}
            if item_specs[index] is not None:
                scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids))
                args["extra_args"] = {"intervention_spec": scoring_spec.to_wire()}
                if item_salts is not None:
                    item_salts[index] = scoring_spec.salt()
            if item_salts is not None:
                prompt["cache_salt"] = item_salts[index]
            prompts.append(prompt)
            sampling.append(SamplingParams(**args))

        generate_kwargs: dict[str, Any] = {"use_tqdm": False}
        if self._backend._lora_request is not None:
            generate_kwargs["lora_request"] = self._backend._lora_request
        request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)
        rows = [
            extract_ref_logprobs(request_output.prompt_logprobs, ref_ids)
            for request_output, ref_ids in zip(request_outputs, ref_ids_per_item)
        ]
        return torch.tensor(rows, dtype=torch.float32)

closed property

Whether the session has been closed.

layout property

Structural facts from the model config (client-side).

Raises:

Type	Description
RuntimeError	If the model config could not be resolved.

tokenizer property

The backend's client-side tokenizer.

capture(prompts, layers, mode, location='layer_output')

Hidden-state capture over the plugin's capture surface.

One request per prompt carries a capture spec and a fresh random cache_salt (a prefix-cache hit skips forward passes, so capture cannot tolerate reused salts) with max_tokens=1; the surplus decode position is truncated by the plugin. Per-layer tensors are stacked and right-padded to the batch's longest prompt.

Parameters:

Name	Type	Description	Default
prompts	list[PreparedPrompt]	The prompts to capture over.	required
layers	list[int]	0-based decoder-layer indices to capture.	required
mode	Literal['all_tokens', 'last_token']	"all_tokens" for every prompt position, "last_token" for the final real position per row.	required
location	Literal['layer_output', 'layer_input']	The residual-stream boundary, "layer_output" or "layer_input".	'layer_output'

Returns:

Type	Description
CaptureResult	The capture result: [N, T, H] per layer for "all_tokens" or [N, H] for
CaptureResult	"last_token", on CPU in the engine's native dtype, with the derived [N, T]
CaptureResult	attention mask.

Raises:

Type	Description
UnsupportedOperationError	If the spec declares no hook_plugin, the negotiated capture kinds lack the requested mode or location, or the engine facts refuse capture (speculative decoding, non-eager execution).
ValueError	If prompts is empty, a layer id is out of range, or the engine returned no capture payload.

Source code in steerability/backends/vllm/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Hidden-state capture over the plugin's capture surface.

    One request per prompt carries a `capture` spec and a fresh random `cache_salt`
    (a prefix-cache hit skips forward passes, so capture cannot tolerate reused salts) with
    `max_tokens=1`; the surplus decode position is truncated by the plugin. Per-layer
    tensors are stacked and right-padded to the batch's longest prompt.

    Args:
        prompts: The prompts to capture over.
        layers: 0-based decoder-layer indices to capture.
        mode: `"all_tokens"` for every prompt position, `"last_token"` for the final real
            position per row.
        location: The residual-stream boundary, `"layer_output"` or `"layer_input"`.

    Returns:
        The capture result: `[N, T, H]` per layer for `"all_tokens"` or `[N, H]` for
        `"last_token"`, on CPU in the engine's native dtype, with the derived `[N, T]`
        attention mask.

    Raises:
        UnsupportedOperationError: If the spec declares no `hook_plugin`, the negotiated
            capture kinds lack the requested mode or location, or the engine facts refuse
            capture (speculative decoding, non-eager execution).
        ValueError: If `prompts` is empty, a layer id is out of range, or the engine
            returned no capture payload.
    """
    self._ensure_open()
    backend = self._backend
    if not backend.spec.get_option("hook_plugin"):
        raise UnsupportedOperationError(
            "Hidden-state capture requires the vLLM-Hook plugin; declare hook_plugin=True "
            "on the vllm backend spec, or run capture on the huggingface backend."
        )
    capture_kinds = backend.capture_kinds
    required = CaptureKinds(
        kinds=frozenset({"residual"}),
        locations=frozenset({location}),
        modes=frozenset({mode}),
    )
    if capture_kinds is None or not capture_kinds.contains(required):
        raise UnsupportedOperationError(
            f"The serving backend does not advertise capture mode {mode!r} at location "
            f"{location!r}; update the server's vllm_hook_plugins or run capture on the "
            "huggingface backend."
        )
    _refuse_by_engine_facts(backend._discovery, "capture")
    if not prompts:
        raise ValueError("capture() requires at least one prompt.")
    num_layers = self.layout.num_layers
    missing = sorted(int(layer) for layer in layers if not 0 <= int(layer) < num_layers)
    if missing:
        raise ValueError(
            f"Requested layer ids {missing} are out of range; the model has {num_layers} layers."
        )

    from vllm import SamplingParams, TokensPrompt

    layer_ids = [int(layer) for layer in layers]
    # the client's validator and assembly expect full prompt coverage and pool the last real
    # position themselves, so every wire capture requests all_tokens
    wire_mode = "all_tokens" if mode == "last_token" else mode
    capture_spec = {"layers": layer_ids, "mode": wire_mode, "location": location}
    engine_prompts = []
    prompt_lens: list[int] = []
    for prompt in prompts:
        resolved = prompt.resolve_token_ids(self.tokenizer)
        ids = resolved.token_ids[0]
        if resolved.attention_mask is not None:
            ids = ids[resolved.attention_mask[0].bool()]
        ids = ids.tolist()
        prompt_lens.append(len(ids))
        engine_prompt = TokensPrompt(prompt_token_ids=ids)
        engine_prompt["cache_salt"] = uuid.uuid4().hex
        engine_prompts.append(engine_prompt)
    sampling = SamplingParams(max_tokens=1, temperature=0.0, extra_args={"capture": capture_spec})

    request_outputs = self._backend._require_llm().generate(engine_prompts, sampling, use_tqdm=False)

    rows_per_layer: dict[int, list[torch.Tensor]] = {layer: [] for layer in layer_ids}
    for index, request_output in enumerate(request_outputs):
        payload = getattr(request_output, "captures", None)
        if payload is None:
            raise ValueError(
                "The engine returned no capture payload; is the vLLM-Hook unified worker "
                "active on this engine?"
            )
        manifest_json, data = payload
        manifest = json.loads(manifest_json)
        tensors = _load_safetensors_bytes(data)
        for layer in layer_ids:
            stacked = tensors.get(f"layer_{layer}")
            if stacked is None or stacked.size(0) < prompt_lens[index]:
                raise ValueError(
                    f"The capture payload covers layer {layer} at "
                    f"{0 if stacked is None else stacked.size(0)} of {prompt_lens[index]} "
                    f"prompt positions for prompt {index}; positions recorded: "
                    f"{manifest.get('positions', {}).get(str(layer))}."
                )
            rows_per_layer[layer].append(stacked[: prompt_lens[index]])

    max_len = max(prompt_lens)
    attention_mask = torch.zeros(len(prompts), max_len, dtype=torch.long)
    for index, length in enumerate(prompt_lens):
        attention_mask[index, :length] = 1

    hidden: dict[int, torch.Tensor] = {}
    for layer, rows in rows_per_layer.items():
        if mode == "last_token":
            hidden[layer] = torch.stack([row[-1] for row in rows])
        else:
            padded = torch.zeros(len(rows), max_len, rows[0].size(-1), dtype=rows[0].dtype)
            for index, row in enumerate(rows):
                padded[index, : row.size(0)] = row
            hidden[layer] = padded
    return CaptureResult(hidden=hidden, attention_mask=attention_mask, mode=mode, location=location)

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/vllm/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

generate(items, params)

Generate one result per item through the engine.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items; state entries lower as intervention specs on plugin-active backends, and no client-side hooks or live processors execute here.	required
params	GenerationParams	Normalized generation parameters shared by all items; unmapped extra keys raise.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order.

Source code in steerability/backends/vllm/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item through the engine.

    Args:
        items: The generation items; state entries lower as intervention specs on
            plugin-active backends, and no client-side hooks or live processors execute
            here.
        params: Normalized generation parameters shared by all items; unmapped `extra` keys
            raise.

    Returns:
        One `ItemResult` per item, in item order.
    """
    self._ensure_open()
    if not items:
        return []
    item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm")
    base_args = render_vllm_sampling_args(params)

    from vllm import SamplingParams, TokensPrompt

    prompts = []
    sampling = []
    prompt_ids_per_item: list[list[int]] = []
    for index, item in enumerate(items):
        ids = self._resolve_item_ids(item)
        prompt_ids_per_item.append(ids)
        args = dict(base_args)
        seed = self._item_seed(item, params, index)
        if seed is not None:
            args["seed"] = seed
        if item_constraints[index] is not None:
            field, value = render_guided_decoding_field(item_constraints[index])
            args.update(render_constraint_sampling_args(field, value))
        if item_specs[index] is not None:
            args["extra_args"] = {"intervention_spec": item_specs[index].to_wire()}
        prompt = TokensPrompt(prompt_token_ids=ids)
        if item_salts is not None:
            prompt["cache_salt"] = item_salts[index]
        prompts.append(prompt)
        sampling.append(SamplingParams(**args))
    self._generate_count += 1

    generate_kwargs: dict[str, Any] = {"use_tqdm": False}
    if self._backend._lora_request is not None:
        generate_kwargs["lora_request"] = self._backend._lora_request
    request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)

    results: list[ItemResult] = []
    for index, request_output in enumerate(request_outputs):
        candidates = [
            (
                list(candidate.token_ids),
                map_vllm_finish_reason(
                    candidate.finish_reason, getattr(candidate, "stop_reason", None),
                ),
            )
            for candidate in request_output.outputs
        ]
        results.append(self._pack_output(index, prompt_ids_per_item[index], candidates))
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

Each item's prompt and reference concatenate into one token-id prompt submitted with prompt_logprobs=0, and the reference positions' log-probabilities are read back.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	Must carry no extra keys; forward keyword arguments have no remote rendering.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len] on CPU.

Raises:

Type	Description
ValueError	If items carry differing reference lengths or params.extra is non-empty.

Source code in steerability/backends/vllm/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

    Each item's prompt and reference concatenate into one token-id prompt submitted with
    `prompt_logprobs=0`, and the reference positions' log-probabilities are read back.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: Must carry no `extra` keys; forward keyword arguments have no remote
            rendering.

    Returns:
        Log probabilities of shape `[num_items, ref_len]` on CPU.

    Raises:
        ValueError: If items carry differing reference lengths or `params.extra` is
            non-empty.
    """
    self._ensure_open()
    if params.extra:
        raise ValueError(
            f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
            "scoring accepts no forward keyword arguments."
        )
    if not items:
        return torch.zeros((0, 0), dtype=torch.float32)
    item_specs, _, item_salts = self._prepare_spec_submission(
        items, "vllm", allow_constraints=False,
    )
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
    ref_len = ref_lens.pop()
    if ref_len == 0:
        return torch.zeros((len(items), 0), dtype=torch.float32)

    from vllm import SamplingParams, TokensPrompt

    prompts = []
    sampling = []
    ref_ids_per_item: list[list[int]] = []
    for index, item in enumerate(items):
        prompt_ids = self._resolve_item_ids(item)
        ref_ids = item.ref_output_ids.reshape(-1).tolist()
        ref_ids_per_item.append(ref_ids)
        prompt = TokensPrompt(prompt_token_ids=[*prompt_ids, *ref_ids])
        args: dict[str, Any] = {"max_tokens": 1, "temperature": 0.0, "prompt_logprobs": 0}
        if item_specs[index] is not None:
            scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids))
            args["extra_args"] = {"intervention_spec": scoring_spec.to_wire()}
            if item_salts is not None:
                item_salts[index] = scoring_spec.salt()
        if item_salts is not None:
            prompt["cache_salt"] = item_salts[index]
        prompts.append(prompt)
        sampling.append(SamplingParams(**args))

    generate_kwargs: dict[str, Any] = {"use_tqdm": False}
    if self._backend._lora_request is not None:
        generate_kwargs["lora_request"] = self._backend._lora_request
    request_outputs = self._backend._require_llm().generate(prompts, sampling, **generate_kwargs)
    rows = [
        extract_ref_logprobs(request_output.prompt_logprobs, ref_ids)
        for request_output, ref_ids in zip(request_outputs, ref_ids_per_item)
    ]
    return torch.tensor(rows, dtype=torch.float32)

VLLMServeSession

Bases: _RequestSessionBase

Request session over a vLLM server's completions endpoint.

Items fan out concurrently under the backend's max_concurrency; transport failures retry with exponential backoff, and a batch whose items partially fail raises PartialBatchError carrying the successes and the re-issuable failures.

Source code in steerability/backends/vllm/session.py

class VLLMServeSession(_RequestSessionBase):
    """Request session over a vLLM server's completions endpoint.

    Items fan out concurrently under the backend's `max_concurrency`; transport failures retry
    with exponential backoff, and a batch whose items partially fail raises `PartialBatchError`
    carrying the successes and the re-issuable failures.
    """

    def generate(
        self,
        items: Sequence[GenerationItem],
        params: GenerationParams,
    ) -> list[ItemResult]:
        """Generate one result per item through the completions endpoint.

        Args:
            items: The generation items; entries must be empty.
            params: Normalized generation parameters shared by all items; unmapped `extra` keys
                raise.

        Returns:
            One `ItemResult` per item, in item order.

        Raises:
            PartialBatchError: If some items failed after transport retries while others
                succeeded.
        """
        self._ensure_open()
        if not items:
            return []
        item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm-serve")
        base_args = render_vllm_sampling_args(params)
        backend = self._backend

        item_ids = [self._resolve_item_ids(item) for item in items]
        seeds = [self._item_seed(item, params, index) for index, item in enumerate(items)]
        self._generate_count += 1

        def make_task(index: int):
            def task() -> ItemResult:
                body: dict[str, Any] = {
                    "model": backend._served_model,
                    "prompt": item_ids[index],
                    "return_token_ids": True,
                    **base_args,
                }
                if seeds[index] is not None:
                    body["seed"] = seeds[index]
                if item_constraints[index] is not None:
                    field, value = render_guided_decoding_field(item_constraints[index])
                    body[f"guided_{field}"] = value
                if item_specs[index] is not None:
                    # vllm_xargs is scalar-only, so nested specs travel as JSON strings
                    body["vllm_xargs"] = {
                        "intervention_spec": item_specs[index].canonical(),
                    }
                if item_salts is not None:
                    body["cache_salt"] = item_salts[index]
                payload = with_transport_retries(
                    lambda: backend._post_json("/v1/completions", body),
                    max_attempts=backend.max_attempts,
                    backoff_base=backend.backoff_base,
                )
                choices = payload.get("choices", [])
                if not choices:
                    raise ValueError("The completions response carries no choices.")
                candidates = []
                for choice in choices:
                    token_ids = choice.get("token_ids")
                    if token_ids is None:
                        raise ValueError(
                            "The completions response carries no token_ids; the server does "
                            "not support the token-id return option (return_token_ids)."
                        )
                    candidates.append((
                        list(token_ids),
                        map_vllm_finish_reason(choice.get("finish_reason"), choice.get("stop_reason")),
                    ))
                return self._pack_output(index, item_ids[index], candidates)
            return task

        outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
        failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
        results = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
        if failures:
            raise PartialBatchError(results, failures)
        return results

    def score(
        self,
        items: Sequence[ScoringItem],
        params: GenerationParams,
    ) -> torch.Tensor:
        """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

        Args:
            items: The scoring items. Every item must carry the same reference length.
            params: Must carry no `extra` keys.

        Returns:
            Log probabilities of shape `[num_items, ref_len]` on CPU.

        Raises:
            ValueError: If items carry differing reference lengths or `params.extra` is
                non-empty.
            PartialBatchError: If some items failed after transport retries while others
                succeeded.
        """
        self._ensure_open()
        if params.extra:
            raise ValueError(
                f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
                "scoring accepts no forward keyword arguments."
            )
        if not items:
            return torch.zeros((0, 0), dtype=torch.float32)
        item_specs, _, item_salts = self._prepare_spec_submission(
            items, "vllm-serve", allow_constraints=False,
        )
        ref_lens = {item.ref_output_ids.shape[-1] for item in items}
        if len(ref_lens) > 1:
            raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
        ref_len = ref_lens.pop()
        if ref_len == 0:
            return torch.zeros((len(items), 0), dtype=torch.float32)
        backend = self._backend

        prompt_ids = [self._resolve_item_ids(item) for item in items]
        ref_ids = [item.ref_output_ids.reshape(-1).tolist() for item in items]

        def make_task(index: int):
            def task() -> list[float]:
                body = {
                    "model": backend._served_model,
                    "prompt": [*prompt_ids[index], *ref_ids[index]],
                    "max_tokens": 1,
                    "temperature": 0.0,
                    "prompt_logprobs": 0,
                }
                if item_specs[index] is not None:
                    scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids[index]))
                    body["vllm_xargs"] = {"intervention_spec": scoring_spec.canonical()}
                    body["cache_salt"] = scoring_spec.salt()
                elif item_salts is not None:
                    body["cache_salt"] = item_salts[index]
                payload = with_transport_retries(
                    lambda: backend._post_json("/v1/completions", body),
                    max_attempts=backend.max_attempts,
                    backoff_base=backend.backoff_base,
                )
                choices = payload.get("choices", [])
                if not choices:
                    raise ValueError("The completions response carries no choices.")
                prompt_logprobs = choices[0].get("prompt_logprobs")
                return extract_ref_logprobs(prompt_logprobs, ref_ids[index])
            return task

        outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
        failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
        if failures:
            successes = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
            raise PartialBatchError(successes, failures)
        return torch.tensor(outcomes, dtype=torch.float32)

closed property

Whether the session has been closed.

layout property

Structural facts from the model config (client-side).

Raises:

Type	Description
RuntimeError	If the model config could not be resolved.

tokenizer property

The backend's client-side tokenizer.

capture(prompts, layers, mode, location='layer_output')

Hidden-state capture over the plugin is not implemented in this toolkit version.

Source code in steerability/backends/vllm/session.py

def capture(
    self,
    prompts: list[PreparedPrompt],
    layers: list[int],
    mode: Literal["all_tokens", "last_token"],
    location: Literal["layer_output", "layer_input"] = "layer_output",
) -> CaptureResult:
    """Hidden-state capture over the plugin is not implemented in this toolkit version."""
    raise UnsupportedOperationError(
        "Hidden-state capture on vLLM backends is not implemented in this toolkit version."
    )

close()

Close the session; further use raises RuntimeError.

Source code in steerability/backends/vllm/session.py

def close(self) -> None:
    """Close the session; further use raises `RuntimeError`."""
    self._closed = True

generate(items, params)

Generate one result per item through the completions endpoint.

Parameters:

Name	Type	Description	Default
items	Sequence[GenerationItem]	The generation items; entries must be empty.	required
params	GenerationParams	Normalized generation parameters shared by all items; unmapped extra keys raise.	required

Returns:

Type	Description
list[ItemResult]	One ItemResult per item, in item order.

Raises:

Type	Description
PartialBatchError	If some items failed after transport retries while others succeeded.

Source code in steerability/backends/vllm/session.py

def generate(
    self,
    items: Sequence[GenerationItem],
    params: GenerationParams,
) -> list[ItemResult]:
    """Generate one result per item through the completions endpoint.

    Args:
        items: The generation items; entries must be empty.
        params: Normalized generation parameters shared by all items; unmapped `extra` keys
            raise.

    Returns:
        One `ItemResult` per item, in item order.

    Raises:
        PartialBatchError: If some items failed after transport retries while others
            succeeded.
    """
    self._ensure_open()
    if not items:
        return []
    item_specs, item_constraints, item_salts = self._prepare_spec_submission(items, "vllm-serve")
    base_args = render_vllm_sampling_args(params)
    backend = self._backend

    item_ids = [self._resolve_item_ids(item) for item in items]
    seeds = [self._item_seed(item, params, index) for index, item in enumerate(items)]
    self._generate_count += 1

    def make_task(index: int):
        def task() -> ItemResult:
            body: dict[str, Any] = {
                "model": backend._served_model,
                "prompt": item_ids[index],
                "return_token_ids": True,
                **base_args,
            }
            if seeds[index] is not None:
                body["seed"] = seeds[index]
            if item_constraints[index] is not None:
                field, value = render_guided_decoding_field(item_constraints[index])
                body[f"guided_{field}"] = value
            if item_specs[index] is not None:
                # vllm_xargs is scalar-only, so nested specs travel as JSON strings
                body["vllm_xargs"] = {
                    "intervention_spec": item_specs[index].canonical(),
                }
            if item_salts is not None:
                body["cache_salt"] = item_salts[index]
            payload = with_transport_retries(
                lambda: backend._post_json("/v1/completions", body),
                max_attempts=backend.max_attempts,
                backoff_base=backend.backoff_base,
            )
            choices = payload.get("choices", [])
            if not choices:
                raise ValueError("The completions response carries no choices.")
            candidates = []
            for choice in choices:
                token_ids = choice.get("token_ids")
                if token_ids is None:
                    raise ValueError(
                        "The completions response carries no token_ids; the server does "
                        "not support the token-id return option (return_token_ids)."
                    )
                candidates.append((
                    list(token_ids),
                    map_vllm_finish_reason(choice.get("finish_reason"), choice.get("stop_reason")),
                ))
            return self._pack_output(index, item_ids[index], candidates)
        return task

    outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
    failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
    results = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
    if failures:
        raise PartialBatchError(results, failures)
    return results

score(items, params)

Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

Parameters:

Name	Type	Description	Default
items	Sequence[ScoringItem]	The scoring items. Every item must carry the same reference length.	required
params	GenerationParams	Must carry no extra keys.	required

Returns:

Type	Description
Tensor	Log probabilities of shape [num_items, ref_len] on CPU.

Raises:

Type	Description
ValueError	If items carry differing reference lengths or params.extra is non-empty.
PartialBatchError	If some items failed after transport retries while others succeeded.

Source code in steerability/backends/vllm/session.py

def score(
    self,
    items: Sequence[ScoringItem],
    params: GenerationParams,
) -> torch.Tensor:
    """Teacher-forced log-probabilities of each item's reference tokens via prompt logprobs.

    Args:
        items: The scoring items. Every item must carry the same reference length.
        params: Must carry no `extra` keys.

    Returns:
        Log probabilities of shape `[num_items, ref_len]` on CPU.

    Raises:
        ValueError: If items carry differing reference lengths or `params.extra` is
            non-empty.
        PartialBatchError: If some items failed after transport retries while others
            succeeded.
    """
    self._ensure_open()
    if params.extra:
        raise ValueError(
            f"Scoring parameter(s) {sorted(params.extra)} have no vLLM rendering; remote "
            "scoring accepts no forward keyword arguments."
        )
    if not items:
        return torch.zeros((0, 0), dtype=torch.float32)
    item_specs, _, item_salts = self._prepare_spec_submission(
        items, "vllm-serve", allow_constraints=False,
    )
    ref_lens = {item.ref_output_ids.shape[-1] for item in items}
    if len(ref_lens) > 1:
        raise ValueError(f"All scoring items must share one reference length; got {sorted(ref_lens)}.")
    ref_len = ref_lens.pop()
    if ref_len == 0:
        return torch.zeros((len(items), 0), dtype=torch.float32)
    backend = self._backend

    prompt_ids = [self._resolve_item_ids(item) for item in items]
    ref_ids = [item.ref_output_ids.reshape(-1).tolist() for item in items]

    def make_task(index: int):
        def task() -> list[float]:
            body = {
                "model": backend._served_model,
                "prompt": [*prompt_ids[index], *ref_ids[index]],
                "max_tokens": 1,
                "temperature": 0.0,
                "prompt_logprobs": 0,
            }
            if item_specs[index] is not None:
                scoring_spec = remap_spec_for_scoring(item_specs[index], len(prompt_ids[index]))
                body["vllm_xargs"] = {"intervention_spec": scoring_spec.canonical()}
                body["cache_salt"] = scoring_spec.salt()
            elif item_salts is not None:
                body["cache_salt"] = item_salts[index]
            payload = with_transport_retries(
                lambda: backend._post_json("/v1/completions", body),
                max_attempts=backend.max_attempts,
                backoff_base=backend.backoff_base,
            )
            choices = payload.get("choices", [])
            if not choices:
                raise ValueError("The completions response carries no choices.")
            prompt_logprobs = choices[0].get("prompt_logprobs")
            return extract_ref_logprobs(prompt_logprobs, ref_ids[index])
        return task

    outcomes = run_bounded([make_task(i) for i in range(len(items))], backend.max_concurrency)
    failures = [(i, outcome) for i, outcome in enumerate(outcomes) if isinstance(outcome, Exception)]
    if failures:
        successes = [outcome for outcome in outcomes if not isinstance(outcome, Exception)]
        raise PartialBatchError(successes, failures)
    return torch.tensor(outcomes, dtype=torch.float32)