Skip to content

Algorithm extension API

An algorithm transforms selected rows of one model component. The controller handles hooks, row selection, request slots, and conflict resolution. Payload loading and validation happen before execution. The contribution guide shows the registration and capability declaration needed to add an algorithm.

from vllm.model_hooks.steering.algorithms.base import BaseSteerVectorAlgorithm
from vllm.model_hooks.steering.algorithms.registry import register_algorithm

Transformation and graph contract

Member Extension contract
_transform(hidden_state, params) Required. Return transformed selected rows with the same shape. params is one layer's prepared payload.
set_payload(payload, scale_factor=1.0) Shared preparation: scale tensor payloads or attach scale_factor to a dictionary without mutating cached source tensors.
_renormalize(original, transformed) Optional helper for algorithms declaring normalization support. Returns transformed rows rescaled to the original norms.
graph_family None uses eager/split execution. A supported family declares that the algorithm can lower to persistent in-graph tables.
graph_lower(payload, scale) Required for an in-graph family. Map the unscaled layer payload and request scale to that family's table keys. Its math must match _transform.
graph_payload_problem(request=None) Return a restriction/rejection reason, or None when admissible. With no request, describe conditions that cannot be guaranteed from the algorithm name.
wire_rank(wire) Low-rank families report payload rank so admission can enforce the engine's rank capacity.

The eager preparation and graph-lowering inputs differ: tensor payloads passed to _transform have already been scaled, while graph_lower receives the scale separately. Handle scaling once in each path. For dictionary payloads, the algorithm reads the prepared scale_factor as appropriate to its transformation.

normalize is a constructor option, not an automatic postprocessing step for every algorithm. The algorithm must implement it, and its capability declaration must permit it. For components whose decoder output is split into hidden and residual tensors, the shared controller presents the complete selected rows to the transformation and writes the resulting change back in the model's format.

BaseSteerVectorAlgorithm

BaseSteerVectorAlgorithm(*, normalize: bool = False)

Bases: ABC

Base class for steering algorithms (see module docstring).

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
def __init__(self, *, normalize: bool = False):
    # Payload of any type (Tensor, dict, ...); format is defined by
    # the canonical payload kind and consumed by _transform.
    self.payload: Any | None = None

    # Rescale transformed rows back to their original norm
    # (see _renormalize); honored by the dense-vector algorithms.
    self.normalize = normalize

graph_family class-attribute instance-attribute

graph_family: str | None = None

graph_payload_problem classmethod

graph_payload_problem(request=None) -> str | None

Payload restriction, or its condition when only a name is known.

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
@classmethod
def graph_payload_problem(cls, request=None) -> str | None:
    """Payload restriction, or its condition when only a name is known."""
    return None

graph_lower staticmethod

graph_lower(payload: Any, scale: float) -> dict[str, Any]

Lower one layer's (payload, scale) to family slot tensors.

Returns {table_key: tensor-or-None}; None entries keep the table's zero default (e.g. an absent bias).

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
@staticmethod
def graph_lower(payload: Any, scale: float) -> dict[str, Any]:
    """Lower one layer's (payload, scale) to family slot tensors.

    Returns {table_key: tensor-or-None}; None entries keep the
    table's zero default (e.g. an absent bias).
    """
    raise NotImplementedError

wire_rank staticmethod

wire_rank(wire) -> int | None

Rank of an inline wire payload (low-rank families only).

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
@staticmethod
def wire_rank(wire) -> int | None:
    """Rank of an inline wire payload (low-rank families only)."""
    return None

set_payload

set_payload(
    payload: Any, scale_factor: float = 1.0
) -> None

Set a private scaled payload without mutating cached source tensors.

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
def set_payload(self, payload: Any, scale_factor: float = 1.0) -> None:
    """Set a private scaled payload without mutating cached source tensors."""
    if payload is None:
        raise ValueError(f"{self.__class__.__name__} requires a payload")

    if isinstance(payload, torch.Tensor):
        payload = payload * scale_factor
    elif isinstance(payload, dict):
        payload = {**payload, "scale_factor": scale_factor}

    self.payload = payload

_transform abstractmethod

_transform(
    hidden_state: Tensor, params: Any
) -> torch.Tensor

Transform the selected token rows (the algorithm's core math).

Parameters:

Name Type Description Default
hidden_state Tensor

[num_positions, hidden_dim] selected rows

required
params Any

this intervention's scaled payload

required
Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
@abstractmethod
def _transform(self, hidden_state: torch.Tensor, params: Any) -> torch.Tensor:
    """Transform the selected token rows (the algorithm's core math).

    Args:
        hidden_state: [num_positions, hidden_dim] selected rows
        params: this intervention's scaled payload
    """
    pass

_renormalize

_renormalize(
    original: Tensor, transformed: Tensor
) -> torch.Tensor

Rescale transformed rows to the norms of original rows.

Computed in float32: hidden-state norms can reach ~1e4, so the intermediate product overflows float16 (max 65504).

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/base.py
def _renormalize(
    self, original: torch.Tensor, transformed: torch.Tensor
) -> torch.Tensor:
    """Rescale `transformed` rows to the norms of `original` rows.

    Computed in float32: hidden-state norms can reach ~1e4, so the
    intermediate product overflows float16 (max 65504).
    """
    norm_pre = torch.norm(original, dim=-1, keepdim=True).float()
    norm_post = torch.norm(transformed, dim=-1, keepdim=True).float()
    scaled = transformed.float() * norm_pre / (norm_post + 1e-8)
    return scaled.to(original.dtype)

Registration and construction

The registered name is used by VectorSpec.algorithm and the engine's steer_algorithms declaration. Import the implementation in the algorithms package so every engine process registers it.

model_hooks.steering.algorithms.registry.register_algorithm

register_algorithm(name: str)

Class decorator registering an algorithm under a unique name.

Parameters:

Name Type Description Default
name str

Unique name of the algorithm (e.g., "direct", "loreft").

required
Source code in vllm-steer/vllm/model_hooks/steering/algorithms/registry.py
def register_algorithm(name: str):
    """Class decorator registering an algorithm under a unique name.

    Args:
        name: Unique name of the algorithm (e.g., "direct", "loreft").
    """

    def decorator(cls: type["BaseSteerVectorAlgorithm"]):
        if name in ALGORITHM_REGISTRY:
            raise ValueError(f"Algorithm '{name}' is already registered.")
        ALGORITHM_REGISTRY[name] = cls
        return cls

    return decorator

model_hooks.steering.algorithms.registry.get_algorithm

get_algorithm(name: str) -> type[BaseSteerVectorAlgorithm]

Look up a registered algorithm class by name.

Raises:

Type Description
ValueError

If the algorithm name is not registered.

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/registry.py
def get_algorithm(name: str) -> type["BaseSteerVectorAlgorithm"]:
    """Look up a registered algorithm class by name.

    Raises:
        ValueError: If the algorithm name is not registered.
    """
    if name not in ALGORITHM_REGISTRY:
        raise ValueError(
            f"Unknown algorithm: '{name}'. Available algorithms: "
            f"{list(ALGORITHM_REGISTRY.keys())}"
        )
    return ALGORITHM_REGISTRY[name]

model_hooks.steering.algorithms.registry.create_algorithm

create_algorithm(
    name: str, *, normalize: bool = False
) -> BaseSteerVectorAlgorithm

Create an algorithm instance by registered name.

Source code in vllm-steer/vllm/model_hooks/steering/algorithms/registry.py
def create_algorithm(
    name: str, *, normalize: bool = False
) -> "BaseSteerVectorAlgorithm":
    """Create an algorithm instance by registered name."""
    return get_algorithm(name)(normalize=normalize)

register_algorithm(name) returns a class decorator and rejects duplicate names. get_algorithm(name) returns the class; create_algorithm(name, normalize=False) constructs it. Neither function loads weights.

Authoring capabilities

The entry in ALGORITHM_CAPABILITIES defines accepted authoring inputs and the component an algorithm edits. The runtime class above owns graph behavior. After changing a capability, regenerate the packaged client contracts with python tools/export_client_contracts.py.

AlgorithmCapabilities dataclass

Source code in vllm-steer/vllm/model_hooks/steering/capabilities.py
@dataclass(frozen=True)
class AlgorithmCapabilities:
    payload_kind: str | None
    source: str
    normalize: bool = False
    params: tuple[str, ...] = ()
    target_component: str = "hidden_states"

payload_kind instance-attribute

payload_kind: str | None

source instance-attribute

source: str

normalize class-attribute instance-attribute

normalize: bool = False

params class-attribute instance-attribute

params: tuple[str, ...] = ()

target_component class-attribute instance-attribute

target_component: str = 'hidden_states'

payload_kind identifies a canonical payload; source declares native-source support; normalize and params constrain request fields; target_component selects hidden_states, attention_heads, or router_logits. See the algorithm table for current entries.