ConstrainedDecoding

steerability.algorithms.output_control.constrained_decoding

args

Constrained decoding argument validation.

control

Constrained decoding: declarative structured outputs rendered per execution arm.

ConstrainedDecoding

Bases: OutputControl

Constrained decoding from one declarative source, rendered per execution arm.

A declarative constraint (JSON schema, regex, EBNF grammar, or a choice set) renders two ways from one source: in process it compiles into a client-side automaton driving a ConstraintProcessor (masking every logit the grammar forbids), and on vLLM backends it renders onto the engine's native structured-output request parameters, consumed in place of the live processor. Both arms compile from the same source, so shared cases under greedy decode produce identically constrained outputs; the masking implementation (client automaton or engine grammar backend) is the documented difference between the arms.

A control constructed with a live automaton object has no declarative form and runs in process only. In process the constraint compiles to a client-side xgrammar automaton; a vLLM-only pipeline never compiles client-side and lowers the constraint to native structured outputs.

Structured outputs do not apply to prompt logprobs, so include_in_scoring=True requires the in-process backend at score; include_in_scoring=False opts out of scoring.

Source code in steerability/algorithms/output_control/constrained_decoding/control.py

class ConstrainedDecoding(OutputControl):
    """Constrained decoding from one declarative source, rendered per execution arm.

    A declarative constraint (JSON schema, regex, EBNF grammar, or a choice set) renders two
    ways from one source: in process it compiles into a client-side automaton driving a
    `ConstraintProcessor` (masking every logit the grammar forbids), and on vLLM backends it
    renders onto the engine's native structured-output request parameters, consumed in place of
    the live processor. Both arms compile from the same source, so shared cases under greedy
    decode produce identically constrained outputs; the masking implementation (client automaton
    or engine grammar backend) is the documented difference between the arms.

    A control constructed with a live `automaton` object has no declarative form and runs in
    process only. In process the constraint compiles to a client-side xgrammar automaton; a
    vLLM-only pipeline never compiles client-side and lowers the constraint to native
    structured outputs.

    Structured outputs do not apply to prompt logprobs, so `include_in_scoring=True` requires
    the in-process backend at score; `include_in_scoring=False` opts out of scoring.
    """

    Args = ConstrainedDecodingArgs
    supports_batching = False  # one automaton per processor; the allowed set applies batch-wide

    def _configure(self) -> None:
        self.tokenizer = None
        self._compiled_automaton = None

    def requirements(self) -> Requirements:
        """In-process compilation or engine-native structured outputs at generate."""
        score = needs(Capability.IN_PROCESS_TORCH) if self.include_in_scoring else ()
        if self.source is None:
            return Requirements(
                generate=needs(
                    Capability.IN_PROCESS_TORCH,
                    hint=(
                        "a live automaton object has no declarative form; construct the control "
                        "with a ConstraintSource (or json_schema/regex/grammar/choice) or run "
                        "this pipeline on the huggingface backend"
                    ),
                ),
                score=score,
            )
        return Requirements(
            generate=any_of(
                needs(Capability.IN_PROCESS_TORCH),
                needs(
                    Capability.GUIDED_DECODING,
                    kinds=ConstraintKinds(constraints=frozenset({self.source.kind})),
                ),
            ),
            score=score,
        )

    def export_constraint(self, runtime_kwargs: dict | None = None) -> ConstraintSource | None:
        """The declarative source, or None for automaton-object configurations."""
        return self.source

    def _automaton(self):
        if self.automaton is not None:
            return self.automaton
        if self._compiled_automaton is None:
            if self.tokenizer is None:
                raise RuntimeError(
                    "ConstrainedDecoding requires a tokenizer to compile its constraint; "
                    "steer() must run first."
                )
            from .utils.automaton import compile_constraint_automaton

            self._compiled_automaton = compile_constraint_automaton(self.source, self.tokenizer)
        return self._compiled_automaton

    def get_logits_processors(
        self,
        input_ids: torch.Tensor,
        runtime_kwargs: dict | None,
        **kwargs,
    ) -> list:
        """The in-process arm: one `ConstraintProcessor` over the (compiled) automaton."""
        return [ConstraintProcessor(self._automaton())]

RUNTIME_KWARGS_SCHEMA = [] class-attribute instance-attribute

args = self.Args.validate(*args, **kwargs) instance-attribute

enabled = True class-attribute instance-attribute

include_in_scoring = True class-attribute instance-attribute

same_model_forwards = False class-attribute instance-attribute

supports_batching = False class-attribute instance-attribute

cleanup()

Release resources allocated during steer().

Override in subclasses that allocate GPU memory or other resources during steering. Default no-op.

Source code in steerability/algorithms/core/base_control.py

def cleanup(self) -> None:
    """Release resources allocated during `steer()`.

    Override in subclasses that allocate GPU memory or other resources during steering. Default
    no-op.
    """
    pass

clone_for_call(seed=None)

A configuration-preserving shallow clone for one generation call.

The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers) with the original but has its own attribute namespace, so per-call attribute mutation on the clone never races another call using the original. When seed is given and the control defines reseed(seed), the clone's client-side RNG is re-seeded.

Parameters:

Name	Type	Description	Default
seed	int | None	Optional seed forwarded to the clone's reseed().	None

Returns:

Type	Description
BaseControl	The clone.

Source code in steerability/algorithms/core/base_control.py

def clone_for_call(self, seed: int | None = None) -> "BaseControl":
    """A configuration-preserving shallow clone for one generation call.

    The clone shares steer-time artifacts (memories, steering vectors, attached tokenizers)
    with the original but has its own attribute namespace, so per-call attribute mutation on
    the clone never races another call using the original. When `seed` is given and the
    control defines `reseed(seed)`, the clone's client-side RNG is re-seeded.

    Args:
        seed: Optional seed forwarded to the clone's `reseed()`.

    Returns:
        The clone.
    """
    clone = copy.copy(self)
    if seed is not None:
        reseed = getattr(clone, "reseed", None)
        if callable(reseed):
            reseed(seed)
    return clone

export_constraint(runtime_kwargs=None)

The declarative source, or None for automaton-object configurations.

Source code in steerability/algorithms/output_control/constrained_decoding/control.py

def export_constraint(self, runtime_kwargs: dict | None = None) -> ConstraintSource | None:
    """The declarative source, or None for automaton-object configurations."""
    return self.source

export_generation_params(runtime_kwargs=None)

The control's sampling-expressible contribution, or None.

A control whose behavior is expressible as normalized generation parameters returns a mapping over a subset of stop_strings, stop_token_ids, max_new_tokens, and min_new_tokens; the pipeline merges it into the call's GenerationParams (stop rules union with the caller's; token bounds only tighten) and does not additionally collect the control's live processors and criteria for that call, so the control executes on every backend through the session's composed stop rules. The default returns None, which keeps the control on the live processor/criteria mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
Mapping[str, Any] | None	The parameter contribution, or None.

Source code in steerability/algorithms/output_control/base.py

def export_generation_params(self, runtime_kwargs: dict | None = None) -> Mapping[str, Any] | None:
    """The control's sampling-expressible contribution, or None.

    A control whose behavior is expressible as normalized generation parameters returns a
    mapping over a subset of `stop_strings`, `stop_token_ids`, `max_new_tokens`, and
    `min_new_tokens`; the pipeline merges it into the call's `GenerationParams` (stop rules
    union with the caller's; token bounds only tighten) and does not additionally collect
    the control's live processors and criteria for that call, so the control executes on
    every backend through the session's composed stop rules. The default returns None, which
    keeps the control on the live processor/criteria mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The parameter contribution, or None.
    """
    return None

export_processor_spec(runtime_kwargs=None)

The control's engine-hosted processor form, or None.

A control whose per-step logit math is expressible in an engine's served processor vocabulary returns a ProcessorSpec; on a backend advertising Capability.PER_STEP_LOGIT_SPECS with the spec's kind, the pipeline submits it as a ProcessorSpecEntry in place of the control's live processor. The default returns None, which keeps the control on the live processor mechanism.

Parameters:

Name	Type	Description	Default
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	None

Returns:

Type	Description
ProcessorSpec | None	The processor spec, or None.

Source code in steerability/algorithms/output_control/base.py

def export_processor_spec(self, runtime_kwargs: dict | None = None) -> ProcessorSpec | None:
    """The control's engine-hosted processor form, or None.

    A control whose per-step logit math is expressible in an engine's served processor
    vocabulary returns a `ProcessorSpec`; on a backend advertising
    `Capability.PER_STEP_LOGIT_SPECS` with the spec's kind, the pipeline submits it as a
    `ProcessorSpecEntry` in place of the control's live processor. The default returns
    None, which keeps the control on the live processor mechanism.

    Args:
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        The processor spec, or None.
    """
    return None

export_state()

Steer-time state to persist when freezing, keyed by logical name.

Values are typed artifacts (SteeringVector, Probe, ProbeSet, a Memory implementation, CheckpointArtifact, LoRAArtifact, torch.Tensor, or a Path to an on-disk product). Must be called after steer(). The default returns an empty mapping (nothing to persist).

Returns:

Type	Description
dict[str, Any]	Mapping from logical name to artifact value.

Source code in steerability/algorithms/core/base_control.py

def export_state(self) -> dict[str, Any]:
    """Steer-time state to persist when freezing, keyed by logical name.

    Values are typed artifacts (`SteeringVector`, `Probe`, `ProbeSet`, a `Memory`
    implementation, `CheckpointArtifact`, `LoRAArtifact`, `torch.Tensor`, or a `Path` to an
    on-disk product). Must be called after `steer()`. The default returns an empty mapping
    (nothing to persist).

    Returns:
        Mapping from logical name to artifact value.
    """
    return {}

fit_identity()

The object whose canonical form defines this control's fit digest, or None.

The digest detects recipe edits that invalidate frozen artifacts, so the returned object should cover exactly the inputs a re-fit would consume (training data, fit specs) and exclude inert application parameters (strengths, multipliers). The default returns None (no fits).

Returns:

Type	Description
Any | None	The fit-identity object, or None.

Source code in steerability/algorithms/core/base_control.py

def fit_identity(self) -> Any | None:
    """The object whose canonical form defines this control's fit digest, or None.

    The digest detects recipe edits that invalidate frozen artifacts, so the returned
    object should cover exactly the inputs a re-fit would consume (training data, fit
    specs) and exclude inert application parameters (strengths, multipliers). The default
    returns None (no fits).

    Returns:
        The fit-identity object, or None.
    """
    return None

frozen_form(state)

The (registry method key, constructor kwargs) of this control's frozen form.

state is the output of export_state(); artifact values inside the returned kwargs are replaced by store references at encode time. When steer_fits() is empty and state is empty, the recipe is the frozen form and the control's own key and recipe args are returned unchanged. A control whose freeze expands into several entries may return a list of (method key, kwargs) pairs instead.

Parameters:

Name	Type	Description	Default
state	dict[str, Any]	The output of export_state().	required

Returns:

Type	Description
tuple[str, dict[str, Any]]	The frozen method key and its constructor kwargs.

Raises:

Type	Description
NotFreezableError	If the control produces fits or state but declares no frozen form.

Source code in steerability/algorithms/core/base_control.py

def frozen_form(self, state: dict[str, Any]) -> tuple[str, dict[str, Any]]:
    """The `(registry method key, constructor kwargs)` of this control's frozen form.

    `state` is the output of `export_state()`; artifact values inside the returned kwargs
    are replaced by store references at encode time. When `steer_fits()` is empty and
    `state` is empty, the recipe is the frozen form and the control's own key and recipe
    args are returned unchanged. A control whose freeze expands into several entries may
    return a list of `(method key, kwargs)` pairs instead.

    Args:
        state: The output of `export_state()`.

    Returns:
        The frozen method key and its constructor kwargs.

    Raises:
        NotFreezableError: If the control produces fits or state but declares no frozen
            form.
    """
    if not state and not self.steer_fits():
        from steerability.algorithms.core.registry import method_key_for

        args = getattr(self, "args", None)
        kwargs = {field.name: getattr(args, field.name) for field in fields(args) if field.init} \
            if args is not None else {}
        return method_key_for(type(self)), kwargs
    raise NotFreezableError(
        f"{type(self).__name__} produces steer-time state but declares no frozen form; "
        "override frozen_form() (and export_state()) to support freezing."
    )

get_logits_processors(input_ids, runtime_kwargs, **kwargs)

The in-process arm: one ConstraintProcessor over the (compiled) automaton.

Source code in steerability/algorithms/output_control/constrained_decoding/control.py

def get_logits_processors(
    self,
    input_ids: torch.Tensor,
    runtime_kwargs: dict | None,
    **kwargs,
) -> list:
    """The in-process arm: one `ConstraintProcessor` over the (compiled) automaton."""
    return [ConstraintProcessor(self._automaton())]

get_stopping_criteria(input_ids, runtime_kwargs, **kwargs)

The control's stopping criteria.

Not applied during scoring (there is no loop to stop). Same call convention as get_logits_processors.

Parameters:

Name	Type	Description	Default
input_ids	Tensor	The steered prompt token ids [batch, seq_len].	required
runtime_kwargs	dict | None	Per-call parameters supplied to generate().	required

Returns:

Type	Description
list	A list of HF StoppingCriteria-style objects.

Source code in steerability/algorithms/output_control/base.py

def get_stopping_criteria(self, input_ids: torch.Tensor, runtime_kwargs: dict | None, **kwargs) -> list:
    """The control's stopping criteria.

    Not applied during scoring (there is no loop to stop). Same call convention as
    `get_logits_processors`.

    Args:
        input_ids: The steered prompt token ids `[batch, seq_len]`.
        runtime_kwargs: Per-call parameters supplied to `generate()`.

    Returns:
        A list of HF `StoppingCriteria`-style objects.
    """
    return []

requirements()

In-process compilation or engine-native structured outputs at generate.

Source code in steerability/algorithms/output_control/constrained_decoding/control.py

def requirements(self) -> Requirements:
    """In-process compilation or engine-native structured outputs at generate."""
    score = needs(Capability.IN_PROCESS_TORCH) if self.include_in_scoring else ()
    if self.source is None:
        return Requirements(
            generate=needs(
                Capability.IN_PROCESS_TORCH,
                hint=(
                    "a live automaton object has no declarative form; construct the control "
                    "with a ConstraintSource (or json_schema/regex/grammar/choice) or run "
                    "this pipeline on the huggingface backend"
                ),
            ),
            score=score,
        )
    return Requirements(
        generate=any_of(
            needs(Capability.IN_PROCESS_TORCH),
            needs(
                Capability.GUIDED_DECODING,
                kinds=ConstraintKinds(constraints=frozenset({self.source.kind})),
            ),
        ),
        score=score,
    )

steer(model, tokenizer=None, session=None, **kwargs)

Optional one-time preparation (e.g., load a reward model, fit a probe).

session is a SteeringSession on the steering backend, provided by the pipeline.

Source code in steerability/algorithms/output_control/base.py

def steer(self, model: PreTrainedModel, tokenizer=None, session=None, **kwargs) -> None:
    """Optional one-time preparation (e.g., load a reward model, fit a probe).

    `session` is a `SteeringSession` on the steering backend, provided by the pipeline.
    """
    pass

steer_access()

The model access this instance's steer step requires, on the ModelAccess ladder.

The default is ModelAccess.FACTS (layout and tokenizer only). A control whose steer step generates or scores through the session declares ROLLOUTS, one that captures hidden states declares CAPTURE, and one that touches the model as a live torch.nn.Module declares MODULE. Configuration determines the result. The pipeline hands steer() a session scoped to the declared rung, and the live model only at MODULE. A control may retain the pipeline model beyond steer() only if its generate phase requires Capability.IN_PROCESS_TORCH.

Returns:

Type	Description
ModelAccess	The declared access rung.

Source code in steerability/algorithms/core/base_control.py

def steer_access(self) -> ModelAccess:
    """The model access this instance's steer step requires, on the `ModelAccess` ladder.

    The default is `ModelAccess.FACTS` (layout and tokenizer only). A control whose steer
    step generates or scores through the session declares `ROLLOUTS`, one that captures
    hidden states declares `CAPTURE`, and one that touches the model as a live
    `torch.nn.Module` declares `MODULE`. Configuration determines the result. The pipeline
    hands `steer()` a session scoped to the declared rung, and the live model only at
    `MODULE`. A control may retain the pipeline model beyond `steer()` only if its
    generate phase requires `Capability.IN_PROCESS_TORCH`.

    Returns:
        The declared access rung.
    """
    return ModelAccess.FACTS

steer_fits()

The fit artifacts this instance's steer step will produce, for the steer plan.

Each entry is (artifact, artifact_class), where artifact is the fit source or recipe class name and artifact_class is "direction" or "calibrated". The default is an empty tuple (no fits).

Returns:

Type	Description
tuple[tuple[str, str], ...]	The declared fit artifacts, in declaration order.

Source code in steerability/algorithms/core/base_control.py

def steer_fits(self) -> tuple[tuple[str, str], ...]:
    """The fit artifacts this instance's steer step will produce, for the steer plan.

    Each entry is `(artifact, artifact_class)`, where `artifact` is the fit source or
    recipe class name and `artifact_class` is `"direction"` or `"calibrated"`. The default
    is an empty tuple (no fits).

    Returns:
        The declared fit artifacts, in declaration order.
    """
    return ()

utils

automaton

Client-side automaton compilation for declarative constraints, over xgrammar.

Json-schema constraints compile compact (any_whitespace=False) so the grammar matches the whitespace policy applied on every venue.

XGrammarAutomaton

A ConstraintAutomaton over an xgrammar GrammarMatcher.

reset starts a fresh matcher at the prompt boundary (prompt tokens are not part of the constrained output); allowed feeds any newly generated tokens since the last call and returns the token ids the grammar permits next. Once the grammar terminates, only the tokenizer's stop tokens are permitted.

Parameters:

Name	Type	Description	Default
compiled	Any	The compiled xgrammar grammar.	required
vocab_size	int	The tokenizer's full vocabulary size.	required
stop_token_ids	list[int]	Token ids permitted after grammar termination.	required

Source code in steerability/algorithms/output_control/constrained_decoding/utils/automaton.py

class XGrammarAutomaton:
    """A `ConstraintAutomaton` over an xgrammar `GrammarMatcher`.

    `reset` starts a fresh matcher at the prompt boundary (prompt tokens are not part of the
    constrained output); `allowed` feeds any newly generated tokens since the last call and
    returns the token ids the grammar permits next. Once the grammar terminates, only the
    tokenizer's stop tokens are permitted.

    Args:
        compiled: The compiled xgrammar grammar.
        vocab_size: The tokenizer's full vocabulary size.
        stop_token_ids: Token ids permitted after grammar termination.
    """

    def __init__(self, compiled: Any, vocab_size: int, stop_token_ids: list[int]):
        self._compiled = compiled
        self._vocab_size = vocab_size
        self._stop_token_ids = list(stop_token_ids)
        self._bitmask = xgrammar.allocate_token_bitmask(1, vocab_size)
        self._matcher = None
        self._consumed = 0

    def reset(self, prefix_ids: torch.Tensor) -> None:
        """Start a fresh matcher; `prefix_ids` is the prompt the constraint begins after."""
        self._matcher = xgrammar.GrammarMatcher(self._compiled)
        self._consumed = prefix_ids.size(-1)

    def allowed(self, prefix_ids: torch.Tensor) -> torch.Tensor:
        """Token ids the grammar permits at the current step."""
        row = prefix_ids[0] if prefix_ids.ndim == 2 else prefix_ids
        for token_id in row[self._consumed:].tolist():
            self._matcher.accept_token(int(token_id))
        self._consumed = row.size(-1)
        if self._matcher.is_terminated():
            return torch.tensor(self._stop_token_ids, dtype=torch.long)
        xgrammar.reset_token_bitmask(self._bitmask)
        self._matcher.fill_next_token_bitmask(self._bitmask)
        mask_row = self._bitmask[0]
        bits = ((mask_row.unsqueeze(1) >> torch.arange(32)) & 1).to(torch.bool)
        return torch.nonzero(bits.reshape(-1)[: self._vocab_size], as_tuple=True)[0]

allowed(prefix_ids)

Token ids the grammar permits at the current step.

Source code in steerability/algorithms/output_control/constrained_decoding/utils/automaton.py

def allowed(self, prefix_ids: torch.Tensor) -> torch.Tensor:
    """Token ids the grammar permits at the current step."""
    row = prefix_ids[0] if prefix_ids.ndim == 2 else prefix_ids
    for token_id in row[self._consumed:].tolist():
        self._matcher.accept_token(int(token_id))
    self._consumed = row.size(-1)
    if self._matcher.is_terminated():
        return torch.tensor(self._stop_token_ids, dtype=torch.long)
    xgrammar.reset_token_bitmask(self._bitmask)
    self._matcher.fill_next_token_bitmask(self._bitmask)
    mask_row = self._bitmask[0]
    bits = ((mask_row.unsqueeze(1) >> torch.arange(32)) & 1).to(torch.bool)
    return torch.nonzero(bits.reshape(-1)[: self._vocab_size], as_tuple=True)[0]

reset(prefix_ids)

Start a fresh matcher; prefix_ids is the prompt the constraint begins after.

Source code in steerability/algorithms/output_control/constrained_decoding/utils/automaton.py

def reset(self, prefix_ids: torch.Tensor) -> None:
    """Start a fresh matcher; `prefix_ids` is the prompt the constraint begins after."""
    self._matcher = xgrammar.GrammarMatcher(self._compiled)
    self._consumed = prefix_ids.size(-1)

compile_constraint_automaton(source, tokenizer)

Compile a declarative constraint into a client-side automaton.

Parameters:

Name	Type	Description	Default
source	ConstraintSource	The declarative constraint.	required
tokenizer	PreTrainedTokenizerBase	The tokenizer the automaton masks against.	required

Returns:

Type	Description
XGrammarAutomaton	The compiled automaton.

Source code in steerability/algorithms/output_control/constrained_decoding/utils/automaton.py

def compile_constraint_automaton(source: ConstraintSource, tokenizer: PreTrainedTokenizerBase) -> XGrammarAutomaton:
    """Compile a declarative constraint into a client-side automaton.

    Args:
        source: The declarative constraint.
        tokenizer: The tokenizer the automaton masks against.

    Returns:
        The compiled automaton.
    """
    vocab_size = max(len(tokenizer), getattr(tokenizer, "vocab_size", 0) or 0)
    tokenizer_info = xgrammar.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
    compiler = xgrammar.GrammarCompiler(tokenizer_info)
    if source.kind == "json_schema":
        schema = source.value if isinstance(source.value, str) else json.dumps(dict(source.value))
        # compile compact on every venue rather than inheriting each backend's whitespace default
        compiled = compiler.compile_json_schema(schema, any_whitespace=False)
    elif source.kind == "regex":
        compiled = compiler.compile_regex(source.value)
    elif source.kind == "grammar":
        compiled = compiler.compile_grammar(source.value)
    else:
        pattern = "(" + "|".join(re.escape(candidate) for candidate in source.value) + ")"
        compiled = compiler.compile_regex(pattern)
    stop_token_ids = list(tokenizer_info.stop_token_ids)
    if not stop_token_ids and tokenizer.eos_token_id is not None:
        stop_token_ids = [tokenizer.eos_token_id]
    return XGrammarAutomaton(compiled, vocab_size, stop_token_ids)