easysteer.hidden_states

Hidden-state, attention-head, and MoE router-logit capture from a running vLLM engine. The capture guide explains selection, graph execution, and prefix-cache behavior.

Capture (recommended)

easysteer.hidden_states.capture

capture(
    llm: Any,
    prompts: Any,
    max_tokens: int = 1,
    layers: list[int] | None = None,
    dtype: str | None = None,
    select: Any | None = None,
    per_prompt_selects: list[Any | None] | None = None,
    stream: str = "hidden_states",
    steering: Any | None = None,
    **generate_kwargs,
) -> CaptureResult

Capture intermediate state for a batch of prompts.

Parameters:

Name	Type	Description	Default
llm	Any	Single-worker vLLM LLM instance (compiled or eager, prefix caching on or off). Tensor-parallel capture is not supported.	required
prompts	Any	prompt list (text or multimodal dicts).	required
max_tokens	int	tokens to generate (1 = prompt-only forward).	1
layers	list[int] | None	layer-id subset (None = all hooked layers).	None
dtype	str | None	engine-side storage dtype (e.g. 'float16').	None
select	Any | None	global SelectSpec (or wire dict) row selection.	None
per_prompt_selects	list[Any | None] | None	one SelectSpec (or wire dict) per prompt, overriding the global selection for that prompt; None entries keep the global selection. The helper returns selected rows without engine-side reduction.	None
stream	str	'hidden_states', 'router_logits', or 'attention_heads'.	'hidden_states'
steering	Any | None	SteeringSpec (or per-prompt list), as in LLM.generate().	None
**generate_kwargs	Any	forwarded into SamplingParams.	{}

Returns:

Type	Description
CaptureResult	CaptureResult with exact per-sample views.

Source code in easysteer/hidden_states/capture_result.py

def capture(
    llm: Any,
    prompts: Any,
    max_tokens: int = 1,
    layers: list[int] | None = None,
    dtype: str | None = None,
    select: Any | None = None,
    per_prompt_selects: list[Any | None] | None = None,
    stream: str = "hidden_states",
    steering: Any | None = None,
    **generate_kwargs,
) -> CaptureResult:
    """Capture intermediate state for a batch of prompts.

    Args:
        llm: Single-worker vLLM LLM instance (compiled or eager,
            prefix caching on or off). Tensor-parallel capture is not supported.
        prompts: prompt list (text or multimodal dicts).
        max_tokens: tokens to generate (1 = prompt-only forward).
        layers: layer-id subset (None = all hooked layers).
        dtype: engine-side storage dtype (e.g. 'float16').
        select: global SelectSpec (or wire dict) row selection.
        per_prompt_selects: one SelectSpec (or wire dict) per prompt,
            overriding the global selection for that prompt; None
            entries keep the global selection. The helper returns selected rows
            without engine-side reduction.
        stream: 'hidden_states', 'router_logits', or 'attention_heads'.
        steering: SteeringSpec (or per-prompt list), as in LLM.generate().
        **generate_kwargs (Any): forwarded into SamplingParams.

    Returns:
        CaptureResult with exact per-sample views.
    """
    from vllm import SamplingParams
    from vllm.capture import deserialize_captured

    def to_wire(spec):
        if spec is None or isinstance(spec, dict):
            return spec
        return spec.to_wire()

    def rpc(method, *args, **kwargs):
        results = llm.llm_engine.collective_rpc(method, args=args, kwargs=kwargs)
        if len(results) != 1:
            raise RuntimeError(
                f"capture expects a single worker, got {len(results)} "
                "RPC results — tensor-parallel capture would return "
                "per-rank shards and is not supported"
            )
        return results

    enable_kwargs: dict[str, Any] = {}
    if layers is not None:
        enable_kwargs["layers"] = list(layers)
    if dtype is not None:
        enable_kwargs["dtype"] = dtype
    if select is not None:
        enable_kwargs["select"] = to_wire(select)

    capture_select = None
    if per_prompt_selects is not None:
        if len(per_prompt_selects) != len(prompts):
            raise ValueError(
                f"per_prompt_selects ({len(per_prompt_selects)}) must "
                f"match prompts ({len(prompts)})"
            )
        capture_select = [
            None if s is None else {stream: to_wire(s)} for s in per_prompt_selects
        ]

    sampling_params = SamplingParams(
        max_tokens=max_tokens,
        temperature=generate_kwargs.pop("temperature", 0.0),
        **generate_kwargs,
    )

    rpc("start_capture", stream, **enable_kwargs)
    try:
        outputs = llm.generate(
            prompts,
            sampling_params=sampling_params,
            capture_select=capture_select,
            steering=steering,
            use_tqdm=False,
        )
        raw = rpc("fetch_captured", stream, clear=True)[0]
    finally:
        rpc("stop_capture", stream)
    tensors, meta = deserialize_captured(raw)
    layouts = {lid: info["layout"] for lid, info in raw.items() if "layout" in info}
    return CaptureResult(tensors, meta, outputs, layouts=layouts)

easysteer.hidden_states.CaptureResult

CaptureResult(
    layers: dict[int, Tensor],
    meta: dict[int, Any],
    outputs: Any,
    layouts: dict[int, dict[str, int]] | None = None,
)

Result of one capture call.

Attributes:

Name	Type	Description
layers		{true_layer_id: Tensor(total_rows, dim)} in fetch order.
outputs		the vLLM RequestOutput list, prompt order.
layouts		component dimensions by layer. Attention-head outputs include width, query num_heads, and per-head value-output head_size.

Source code in easysteer/hidden_states/capture_result.py

def __init__(
    self,
    layers: dict[int, torch.Tensor],
    meta: dict[int, Any],
    outputs: Any,
    layouts: dict[int, dict[str, int]] | None = None,
):
    self.layers = layers
    self.outputs = outputs
    self.layouts = layouts or {}
    for lid, layout in self.layouts.items():
        if lid not in layers or layers[lid].shape[-1] != layout["width"]:
            raise ValueError(f"layer {lid}: capture layout does not match rows")
    if not isinstance(meta, dict) or set(meta) != set(layers):
        raise ValueError("capture requires row labels for every captured layer")
    self._meta = meta
    for lid, m in meta.items():
        if len(m) != layers[lid].shape[0]:
            raise RuntimeError(
                f"layer {lid}: {len(m)} row labels for "
                f"{layers[lid].shape[0]} rows — engine/client label desync"
            )
    self._sample_rows = self._index_samples()

layers instance-attribute

layers = layers

outputs instance-attribute

outputs = outputs

layouts instance-attribute

layouts = layouts or {}

layer_ids property

layer_ids: list[int]

labelled property

labelled: bool

rows

rows(layer: int) -> torch.Tensor

Source code in easysteer/hidden_states/capture_result.py

def rows(self, layer: int) -> torch.Tensor:
    return self.layers[layer]

meta

meta(layer: int)

Row labels (req_ids/positions/token_ids) for a layer.

Source code in easysteer/hidden_states/capture_result.py

def meta(self, layer: int):
    """Row labels (req_ids/positions/token_ids) for a layer."""
    return self._meta[layer]

sample

sample(i: int) -> dict[int, torch.Tensor]

One sample's rows for every layer: {layer_id: (rows, dim)}.

Source code in easysteer/hidden_states/capture_result.py

def sample(self, i: int) -> dict[int, torch.Tensor]:
    """One sample's rows for every layer: {layer_id: (rows, dim)}."""
    idx = torch.tensor(self._sample_rows[i], dtype=torch.long)
    return {lid: t[idx] for lid, t in self.layers.items()}

sample_positions

sample_positions(i: int) -> list[int]

Absolute sequence positions of sample i's rows (row order).

Source code in easysteer/hidden_states/capture_result.py

def sample_positions(self, i: int) -> list[int]:
    """Absolute sequence positions of sample i's rows (row order)."""
    if not self.layers:
        return []
    first = self.meta(self.layer_ids[0])
    return [int(first.positions[r]) for r in self._sample_rows[i]]

sample_token_ids

sample_token_ids(i: int) -> list[int]

Input token ids of sample i's rows (row order).

Source code in easysteer/hidden_states/capture_result.py

def sample_token_ids(self, i: int) -> list[int]:
    """Input token ids of sample i's rows (row order)."""
    if not self.layers:
        return []
    first = self.meta(self.layer_ids[0])
    return [int(first.token_ids[r]) for r in self._sample_rows[i]]

to_nested

to_nested() -> list[list[torch.Tensor]]

Return [sample][layer_pos] tensors with layers sorted by model id.

Source code in easysteer/hidden_states/capture_result.py

def to_nested(self) -> list[list[torch.Tensor]]:
    """Return `[sample][layer_pos]` tensors with layers sorted by model id."""
    layer_ids = self.layer_ids
    samples = []
    for i in range(len(self)):
        sample = self.sample(i)
        samples.append([sample[lid] for lid in layer_ids])
    return samples

__len__

__len__() -> int

Source code in easysteer/hidden_states/capture_result.py

def __len__(self) -> int:
    return len(self.outputs)

capture() returns a CaptureResult. Its methods select rows by true layer ID or by input sample; len(result) is the number of generated request outputs. An empty selection still preserves the outputs and empty per-sample views. Row-label types, component constants, serialization, and worker stream/session classes are documented in the fork capture API.

Compatibility wrappers

Nested-list wrappers over capture(); prefer capture() for new code.

easysteer.hidden_states.get_all_hidden_states_generate

get_all_hidden_states_generate(
    llm: Any,
    prompts: list[str] | list[dict[str, Any]],
    max_tokens: int = 1,
    split_by_samples: bool = True,
    token_ids: list[int] | None = None,
    positions: list[int] | None = None,
    layers: list[int] | None = None,
    dtype: str | None = None,
    select: dict | Any | None = None,
    **generate_kwargs,
) -> (
    tuple[list[list[torch.Tensor]], Any]
    | tuple[list[torch.Tensor], Any]
)

Capture every layer's hidden states while running generate.

Works for any generate-capable model, including multimodal models (Qwen-VL, LLaVA, ...) that do not support the embed task. With the default max_tokens=1 only the prompt forward is captured, matching what an embed task would produce.

Parameters:

Name	Type	Description	Default
llm	Any	Single-worker vLLM LLM instance (compiled or eager, prefix caching on or off).	required
prompts	list[str] | list[dict[str, Any]]	text prompts, or multimodal dicts with prompt and multi_modal_data keys.	required
max_tokens	int	tokens to generate (1 = prompt-only forward).	1
split_by_samples	bool	if True return [sample][layer] tensors; if False return per-layer tensors concatenated over samples.	True
token_ids	list[int] | None	only capture rows whose input token id is in this list (source-side filter; unions with positions).	None
positions	list[int] | None	only capture these prompt positions (negatives resolve from the prompt end; unions with token_ids).	None
layers	list[int] | None	layer-id subset (None = all hooked layers).	None
dtype	str | None	engine-side storage dtype (e.g. 'float16').	None
select	dict | Any | None	SelectSpec (or wire dict) — the full where-clause selection language; cannot combine with the shortcuts.	None
**generate_kwargs	Any	forwarded into SamplingParams.	{}

Returns:

Type	Description
tuple[list[list[Tensor]], Any] | tuple[list[Tensor], Any]	(hidden_states, outputs) where hidden_states is
tuple[list[list[Tensor]], Any] | tuple[list[Tensor], Any]	[sample][layer] (split) or [layer] (concatenated),
tuple[list[list[Tensor]], Any] | tuple[list[Tensor], Any]	layers ordered by layer id.

Source code in easysteer/hidden_states/capture_generate.py

def get_all_hidden_states_generate(
    llm: Any,
    prompts: list[str] | list[dict[str, Any]],
    max_tokens: int = 1,
    split_by_samples: bool = True,
    token_ids: list[int] | None = None,
    positions: list[int] | None = None,
    layers: list[int] | None = None,
    dtype: str | None = None,
    select: dict | Any | None = None,
    **generate_kwargs,
) -> tuple[list[list[torch.Tensor]], Any] | tuple[list[torch.Tensor], Any]:
    """Capture every layer's hidden states while running generate.

    Works for any generate-capable model, including multimodal models
    (Qwen-VL, LLaVA, ...) that do not support the embed task. With the
    default ``max_tokens=1`` only the prompt forward is captured,
    matching what an embed task would produce.

    Args:
        llm: Single-worker vLLM LLM instance (compiled or eager,
            prefix caching on or off).
        prompts: text prompts, or multimodal dicts with ``prompt`` and
            ``multi_modal_data`` keys.
        max_tokens: tokens to generate (1 = prompt-only forward).
        split_by_samples: if True return ``[sample][layer]`` tensors;
            if False return per-layer tensors concatenated over samples.
        token_ids: only capture rows whose input token id is in this
            list (source-side filter; unions with ``positions``).
        positions: only capture these prompt positions (negatives
            resolve from the prompt end; unions with ``token_ids``).
        layers: layer-id subset (None = all hooked layers).
        dtype: engine-side storage dtype (e.g. ``'float16'``).
        select: SelectSpec (or wire dict) — the full where-clause
            selection language; cannot combine with the shortcuts.
        **generate_kwargs (Any): forwarded into SamplingParams.

    Returns:
        ``(hidden_states, outputs)`` where hidden_states is
        ``[sample][layer]`` (split) or ``[layer]`` (concatenated),
        layers ordered by layer id.
    """
    result = capture(
        llm,
        prompts,
        max_tokens=max_tokens,
        layers=layers,
        dtype=dtype,
        select=_sugar_select(token_ids, positions, select),
        **generate_kwargs,
    )
    if split_by_samples:
        return result.to_nested(), result.outputs
    return [result.layers[lid] for lid in result.layer_ids], result.outputs

easysteer.hidden_states.get_moe_router_logits_generate

get_moe_router_logits_generate(
    llm: Any,
    prompts: list[str] | list[dict[str, Any]],
    max_tokens: int = 1,
    split_by_samples: bool = False,
    **generate_kwargs,
) -> (
    tuple[dict[int, torch.Tensor], Any]
    | tuple[list[dict[int, torch.Tensor]], Any]
)

Capture MoE router logits while running generate.

Works for any generate-capable MoE model, including multimodal ones (e.g. Qwen3-VL). With the default max_tokens=1 only the prompt forward is captured. When router-logits steering is active, the captured logits are the post-steering ones.

Parameters:

Name	Type	Description	Default
llm	Any	Single-worker vLLM LLM instance (compiled or eager, prefix caching on or off).	required
prompts	list[str] | list[dict[str, Any]]	text prompts, or multimodal dicts with prompt and multi_modal_data keys.	required
max_tokens	int	tokens to generate (1 = prompt-only forward).	1
split_by_samples	bool	if True return one {layer_id: tensor} dict per sample; if False return a single dict with all samples' rows concatenated per layer.	False
**generate_kwargs	Any	forwarded into SamplingParams.	{}

Returns:

Type	Description
tuple[dict[int, Tensor], Any] | tuple[list[dict[int, Tensor]], Any]	(router_logits, outputs) where router_logits is
tuple[dict[int, Tensor], Any] | tuple[list[dict[int, Tensor]], Any]	{layer_id: (rows, n_experts)} (concatenated) or
tuple[dict[int, Tensor], Any] | tuple[list[dict[int, Tensor]], Any]	[sample_idx]{layer_id: (rows, n_experts)} (split).

Source code in easysteer/hidden_states/moe_capture_generate.py

def get_moe_router_logits_generate(
    llm: Any,
    prompts: list[str] | list[dict[str, Any]],
    max_tokens: int = 1,
    split_by_samples: bool = False,
    **generate_kwargs,
) -> tuple[dict[int, torch.Tensor], Any] | tuple[list[dict[int, torch.Tensor]], Any]:
    """Capture MoE router logits while running generate.

    Works for any generate-capable MoE model, including multimodal
    ones (e.g. Qwen3-VL). With the default ``max_tokens=1`` only the
    prompt forward is captured. When router-logits steering is active,
    the captured logits are the post-steering ones.

    Args:
        llm: Single-worker vLLM LLM instance (compiled or eager,
            prefix caching on or off).
        prompts: text prompts, or multimodal dicts with ``prompt`` and
            ``multi_modal_data`` keys.
        max_tokens: tokens to generate (1 = prompt-only forward).
        split_by_samples: if True return one ``{layer_id: tensor}``
            dict per sample; if False return a single dict with all
            samples' rows concatenated per layer.
        **generate_kwargs (Any): forwarded into SamplingParams.

    Returns:
        ``(router_logits, outputs)`` where router_logits is
        ``{layer_id: (rows, n_experts)}`` (concatenated) or
        ``[sample_idx]{layer_id: (rows, n_experts)}`` (split).
    """
    result = capture(
        llm,
        prompts,
        max_tokens=max_tokens,
        stream="router_logits",
        **generate_kwargs,
    )
    if split_by_samples:
        return [result.sample(i) for i in range(len(result))], result.outputs
    return dict(result.layers), result.outputs