Capture API (vllm.capture)

Applications normally call easysteer.hidden_states.capture(), which starts capture, generates requests, fetches labelled tensors, and stops the stream. The fork exports the selection, serialization, and worker-session types documented below for integrations.

from vllm.capture import (
    ATTENTION_HEADS,
    HIDDEN_STATES,
    ROUTER_LOGITS,
    CaptureMeta,
    SelectSpec,
    deserialize_captured,
    match_capture_request_id,
)

Components and selection

Constant	Value	Captured representation
HIDDEN_STATES	"hidden_states"	Decoder-layer output; models returning a separate residual are combined for capture.
ATTENTION_HEADS	"attention_heads"	Concatenated query-head outputs before the attention output projection.
ROUTER_LOGITS	"router_logits"	Accessible MoE gate scores before routing.

SelectSpec defines which prompt and generation rows to collect. It is the same class exported from vllm.steer_vectors; its field reference is rendered once on the specs page. See the capture guide for model support and graph/cache behavior.

Row labels and deserialization

The decoded layer tensors and their CaptureMeta labels have matching row counts. Request labels identify the owning request under continuous batching; sequence positions order rows within a sample.

vllm.capture.CaptureMeta

CaptureMeta(
    req_ids: list[str], positions: Tensor, token_ids: Tensor
)

Row labels for one captured layer.

Attributes:

Name	Type	Description
req_ids		engine-internal request id string per row (the client-visible id plus an 8-hex uniqueness suffix; match with :func:match_capture_request_id).
positions		absolute sequence position per row (int32; -1 for synthesized rows such as 'mean' reductions).
token_ids		input token id per row (int32; -1 for synthesized rows).

Source code in vllm-steer/vllm/model_hooks/capture/serialization.py

def __init__(
    self,
    req_ids: list[str],
    positions: torch.Tensor,
    token_ids: torch.Tensor,
):
    self.req_ids = req_ids
    self.positions = positions
    self.token_ids = token_ids

req_ids instance-attribute

req_ids = req_ids

positions instance-attribute

positions = positions

token_ids instance-attribute

token_ids = token_ids

__len__

__len__() -> int

Source code in vllm-steer/vllm/model_hooks/capture/serialization.py

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

vllm.capture.deserialize_captured

deserialize_captured(
    serialized_data: dict[int, dict[str, Any]],
) -> tuple[dict[int, torch.Tensor], dict[int, CaptureMeta]]

Rebuild per-layer tensors AND their row labels.

Tensors arrive as raw bytes of their stored dtype (bf16 rides as int16 bytes and is reinterpreted here) — see StreamStore.serialize.

Returns (tensors, meta) where meta[layer] labels each row of tensors[layer] with (request id, absolute position, token id). Every layer must carry one label per row. Incomplete or older unlabeled payloads raise instead of silently losing sample alignment.

Source code in vllm-steer/vllm/model_hooks/capture/serialization.py

def deserialize_captured(
    serialized_data: dict[int, dict[str, Any]],
) -> tuple[dict[int, torch.Tensor], dict[int, CaptureMeta]]:
    """Rebuild per-layer tensors AND their row labels.

    Tensors arrive as raw bytes of their stored dtype (bf16 rides as
    int16 bytes and is reinterpreted here) — see StreamStore.serialize.

    Returns ``(tensors, meta)`` where ``meta[layer]`` labels each row of
    ``tensors[layer]`` with (request id, absolute position, token id).
    Every layer must carry one label per row. Incomplete or older unlabeled
    payloads raise instead of silently losing sample alignment.
    """
    tensors: dict[int, torch.Tensor] = {}
    meta: dict[int, CaptureMeta] = {}
    for layer_id, info in serialized_data.items():
        encoding = info.get("encoding")
        if encoding != "raw":
            raise ValueError(
                f"unknown capture wire encoding {encoding!r} for layer "
                f"{layer_id}; engine and client versions do not match"
            )
        wire_dtype = info["dtype"]
        if wire_dtype not in _WIRE_DTYPES:
            raise ValueError(
                f"unknown capture wire dtype {wire_dtype!r} for layer {layer_id}"
            )
        shape = tuple(info["shape"])
        array = np.frombuffer(info["data"], dtype=_RAW_DTYPES[_WIRE_DTYPES[wire_dtype]])
        tensor = torch.from_numpy(array.reshape(shape).copy())
        if wire_dtype == "torch.bfloat16":
            tensor = tensor.view(torch.bfloat16)
        tensors[layer_id] = tensor

        m = info.get("meta")
        if m is None:
            raise ValueError(
                f"capture layer {layer_id} is missing mandatory row labels"
            )
        req_idx = np.frombuffer(m["req_idx"], dtype=np.int32)
        table = m["req_table"]
        positions = np.frombuffer(m["positions"], dtype=np.int32)
        token_ids = np.frombuffer(m["token_ids"], dtype=np.int32)
        rows = tensor.shape[0]
        if len(req_idx) != rows or len(positions) != rows or len(token_ids) != rows:
            raise ValueError(f"capture layer {layer_id} has inconsistent row labels")
        if (req_idx < 0).any() or (req_idx >= len(table)).any():
            raise ValueError(f"capture layer {layer_id} has invalid request indices")
        meta[layer_id] = CaptureMeta(
            req_ids=[table[i] for i in req_idx],
            positions=torch.from_numpy(positions.copy()),
            token_ids=torch.from_numpy(token_ids.copy()),
        )
    return tensors, meta

serialized_data is the per-layer dictionary returned by a worker's fetch_captured RPC. The function returns (tensors, meta), both keyed by true layer ID. Optional attention layout metadata stays in the original wire dictionary; the application helper attaches it to CaptureResult.layouts.

vllm.capture.match_capture_request_id

match_capture_request_id(
    label_req_id: str, request_id: str
) -> bool

Whether a captured row's label belongs to a client request.

Rows are labelled with the engine-internal request id, which is the client-visible id plus an 8-hex-char uniqueness suffix ({request_id}-{8 hex}, see InputProcessor) — or identical to it when suffixing is disabled.

Source code in vllm-steer/vllm/model_hooks/capture/serialization.py

def match_capture_request_id(label_req_id: str, request_id: str) -> bool:
    """Whether a captured row's label belongs to a client request.

    Rows are labelled with the engine-internal request id, which is the
    client-visible id plus an 8-hex-char uniqueness suffix
    (``{request_id}-{8 hex}``, see InputProcessor) — or identical to it
    when suffixing is disabled.
    """
    if label_req_id == request_id:
        return True
    return (
        label_req_id.startswith(request_id + "-")
        and len(label_req_id) == len(request_id) + 9
        and all(c in "0123456789abcdef" for c in label_req_id[-8:])
    )

Use this helper to compare an engine row label with a client request ID. It accepts an exact match or the engine's recognized uniqueness suffix and returns a boolean; arbitrary prefix matching can mix different requests.

Worker stream configuration

These classes are exported for engine integrations. Ordinary callers should use capture() instead of creating or attaching a second CaptureSession.

StreamConfig is constructed when a session enables a stream:

Parameter	Meaning
layers	True layer IDs; None selects all available layers.
dtype	Optional storage dtype, independent of model compute dtype.
reduce	"all" retains rows; "last" or "mean" reduces within a request.
select	A SelectSpec.to_wire() dictionary; requires reduce="all".
budget_rows	Optional nonnegative row limit per layer.

The constructor validates the selection and resolves dtype/layer storage. These low-level reduction and budget arguments are not arguments to the high-level hs.capture() helper, which returns selected rows without a reduction.

vllm.capture.StreamConfig

StreamConfig(
    layers: list[int] | None = None,
    dtype: str | None = None,
    reduce: str = "all",
    select: dict | None = None,
    budget_rows: int | None = None,
)

Per-stream capture configuration.

Row selection uses the shared SelectSpec where-clause language (select, wire form of vllm.model_hooks.selection.spec.SelectSpec) — the same clause semantics steering resolves, so "which rows to capture" and "which tokens to steer" mean the same thing. reduce applies a within-sample reduction ("all" | "last" | "mean"); budget_rows caps stored rows per layer.

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def __init__(
    self,
    layers: list[int] | None = None,
    dtype: str | None = None,
    reduce: str = "all",
    select: dict | None = None,
    budget_rows: int | None = None,
):
    if reduce not in ("all", "last", "mean"):
        raise ValueError(f"reduce must be 'all', 'last' or 'mean', got {reduce!r}")
    if reduce != "all" and select is not None:
        raise ValueError(
            "row selection cannot combine with the 'last'/'mean' "
            "reductions; use reduce='all' with a select clause"
        )
    if select is not None:
        # Validate at enable time instead of failing mid-forward.
        select = SelectSpec.from_wire(select).to_wire()
    self.layers = set(layers) if layers is not None else None
    self.dtype = resolve_storage_dtype(dtype) if dtype is not None else None
    self.reduce = reduce
    self.select = select
    if budget_rows is not None and budget_rows < 0:
        raise ValueError("budget_rows must be nonnegative")
    self.budget_rows = budget_rows
    if reduce == "all" and select is None and budget_rows is None:
        logger.warning_once(
            "Capture enabled with reduce='all', no select clause and "
            "no budget_rows: every position of every request "
            "accumulates in CPU memory. Prefer a select clause, a "
            "reduction, or an explicit budget for long corpora."
        )

layers instance-attribute

layers = set(layers) if layers is not None else None

dtype instance-attribute

dtype = (
    resolve_storage_dtype(dtype)
    if dtype is not None
    else None
)

reduce instance-attribute

reduce = reduce

select instance-attribute

select = select

budget_rows instance-attribute

budget_rows = budget_rows

selects_rows property

selects_rows: bool

Whether a source-side row selection (not a reduction) is set.

Capture session

The model runner owns a CaptureSession, attaches the discovered components at model load, and calls its lifecycle methods for requests and batches. For stream management, the relevant operations are:

Method	Effect / result
enable_stream(stream, **config_kwargs)	Validate a StreamConfig and start a fresh store; hooks must already be attached.
disable_stream(stream)	Disable collection for that stream.
fetch_stream(stream, clear=True, layers=None, req_ids=None)	Return serialized rows, labels, and available layouts; optionally clear the selected data.
clear_stream(stream)	Clear retained rows and errors.
stream_status(stream)	Return enablement, hooked layers, layouts, storage counters, and capture-graph execution counters.

The supported RPCs (start_capture, stop_capture, fetch_captured, clear_captured, capture_status) delegate to these operations in the worker. The high-level helper also coordinates prefix-cache read policy before admission.

vllm.capture.CaptureSession

CaptureSession()

Owns capture hooks and streams for one worker's model.

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def __init__(self):
    self._streams: dict[str, StreamStore | None] = dict.fromkeys(COMPONENTS)
    # Engine-internal request id -> {stream: SelectSpec wire dict}.
    self._request_selects: dict[str, dict[str, dict]] = {}
    # Retain cache-hit metadata even with no stream enabled. A stream
    # started later must check rows that were already skipped.
    self._cache_hits: dict[str, tuple[list[int], int]] = {}
    self._unsupported_requests: set[str] = set()
    self._hook_handles: list = []
    self._hooked_layers: dict[str, set[int]] = {name: set() for name in COMPONENTS}
    self._layouts: dict[str, dict[int, dict[str, int]]] = {}
    self._attached = False
    self.graph_state: CaptureGraphState | None = None
    self._graph_replays = 0
    self._eager_capture_forwards = 0

attach

attach(model: Module, components: ModelComponents) -> None

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def attach(self, model: nn.Module, components: ModelComponents) -> None:
    if self._attached:
        return
    for component in COMPONENTS.values():
        self._layouts[component.id] = {
            target.layer_id: {
                key: value
                for key in ("width", "num_heads", "head_size")
                if (value := getattr(target, key)) is not None
            }
            for target in components[component.id]
            if target.width is not None
        }
        self._hooked_layers[component.id] = self._attach_stream_hooks(
            components[component.id], component
        )

    def flush_hook(mod, args, output):
        if self.graph_state is not None and self.graph_state.recording:
            return
        if any(
            store is not None and store._pending for store in self._streams.values()
        ):
            self._eager_capture_forwards += 1
        for store in self._streams.values():
            if store is not None:
                store.flush()

    # One post-forward flush per step: per-layer hooks only stage
    # GPU rows; the D2H copies coalesce here (pinned, non-blocking,
    # single sync).
    self._hook_handles.append(model.register_forward_hook(flush_hook))
    self._attached = True

detach

detach() -> None

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def detach(self) -> None:
    for handle in self._hook_handles:
        handle.remove()
    self._hook_handles.clear()
    self._attached = False
    self.release_graph()
    self._streams = dict.fromkeys(COMPONENTS)
    self._request_selects.clear()
    self._cache_hits.clear()
    self._unsupported_requests.clear()
    self._layouts.clear()
    for layers in self._hooked_layers.values():
        layers.clear()

enable_stream

enable_stream(stream: str, **config_kwargs) -> None

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def enable_stream(self, stream: str, **config_kwargs) -> None:
    if stream not in self._streams:
        raise ValueError(f"Unknown capture stream: {stream}")
    if not self._attached:
        raise RuntimeError(
            "Capture hooks are not attached to this model; "
            "attachment happens once at model load."
        )
    config = StreamConfig(**config_kwargs)
    if stream == ATTENTION_HEADS:
        available = self._hooked_layers[stream]
        if not available:
            raise ValueError("No supported decoder attention head outputs found")
        if config.layers is not None and config.layers - available:
            raise ValueError(
                "Unsupported attention head capture layers: "
                f"{sorted(config.layers - available)}"
            )
    if not self._hooked_layers[stream]:
        logger.warning(
            "Enabling %s capture but no usable component hooks were found.",
            stream,
        )
    self._streams[stream] = StreamStore(config)

disable_stream

disable_stream(stream: str) -> None

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def disable_stream(self, stream: str) -> None:
    if stream in self._streams:
        self._streams[stream] = None

fetch_stream

fetch_stream(
    stream: str,
    clear: bool = True,
    layers: list[int] | None = None,
    req_ids: list[str] | None = None,
) -> dict[int, dict[str, Any]]

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def fetch_stream(
    self,
    stream: str,
    clear: bool = True,
    layers: list[int] | None = None,
    req_ids: list[str] | None = None,
) -> dict[int, dict[str, Any]]:
    store = self._streams.get(stream)
    if store is None:
        return {}
    store.flush()  # capture any rows staged since the last step
    if req_ids is not None:
        result = store.serialize(
            layers=layers, req_ids=req_ids, clear_selected=clear
        )
    else:
        result = store.serialize(layers=layers)
        if clear:
            if layers is None:
                store.clear()
            else:
                store.drop_layers(list(result))
    for layer, info in result.items():
        if layout := self._layouts.get(stream, {}).get(layer):
            info["layout"] = layout.copy()
    return result

clear_stream

clear_stream(stream: str) -> None

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def clear_stream(self, stream: str) -> None:
    store = self._streams.get(stream)
    if store is not None:
        store.clear()

stream_status

stream_status(stream: str) -> dict[str, Any]

Source code in vllm-steer/vllm/model_hooks/capture/session.py

def stream_status(self, stream: str) -> dict[str, Any]:
    store = self._streams.get(stream)
    hooked = len(self._hooked_layers.get(stream, ()))
    execution = {
        "layouts": self._layouts.get(stream, {}),
        "graph_ready": self.graph_state is not None and self.graph_state.ready,
        "graph_buffer_bytes": 0
        if self.graph_state is None
        else self.graph_state.buffer_bytes,
        "graph_allocation_bytes": 0
        if self.graph_state is None
        else self.graph_state.allocation_bytes,
        # Session-lifetime counters survive stop; a single inactive graph
        # variant stays cached for repeated capture() calls.
        "graph_replays": self._graph_replays,
        "eager_capture_forwards": self._eager_capture_forwards,
    }
    if store is None:
        return {"enabled": False, "hooked_layers": hooked, **execution}
    return {
        "enabled": True,
        "hooked_layers": hooked,
        "layers_captured": len(store.chunks),
        "tokens_stored": store.tokens_stored,
        "tokens_dropped": store.tokens_dropped,
        "reduce": store.config.reduce,
        "select": store.config.select,
        "meta_complete": True,
        "cache_elided_reqs": len(store.elided_reqs),
        **execution,
    }

Stream store

StreamStore(config) holds one stream's captured rows and metadata. A session owns it; integrations can inspect the storage/serialization contract here. serialize() returns the per-layer wire dictionaries consumed by deserialize_captured(). Request and layer filters operate on retained data.

vllm.capture.StreamStore

StreamStore(config: StreamConfig)

Bounded, chunk-appending CPU store for one capture stream.

Every stored row carries labels — (req_idx, position, token_id) int32 columns, where req_idx indexes req_table (request id strings) — so clients never re-derive sample alignment. Missing or malformed labels are rejected before any rows enter the store.

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def __init__(self, config: StreamConfig):
    self.config = config
    self.chunks: dict[int, dict[int, _StoredChunk]] = {}
    self._next_chunk = 0
    self._request_chunks: dict[int, dict[int, set[int]]] = {}
    self.layer_names: dict[int, str] = {}
    self.req_table: dict[int, str] = {}
    self._req_index: dict[str, int] = {}
    self._request_aliases: dict[str, set[int]] = {}
    self._free_request_indices: list[int] = []
    self._next_request_index = 0
    self._unused_request_indices: set[int] = set()
    # Active requests retain capture history across draining and preemption.
    self._captured_requests: set[str] = set()
    # Cache hits that omitted selected prompt rows, including requests
    # already running when this stream started (fetch raises).
    self.elided_reqs: set[str] = set()
    self._failed_requests: dict[str, str] = {}
    # Rows stored per layer; the budget caps each layer at
    # budget_rows rows (i.e. sequence tokens, per layer).
    self._layer_rows: dict[int, int] = {}
    self.tokens_dropped = 0
    self._warned_budget = False
    self._pending: list[tuple[int, torch.Tensor, torch.Tensor, str]] = []
    # Geometry and request selections are fixed during one forward pass.
    # Layers share its row plans; a fresh geometry invalidates every plan.
    self._row_plan_geometry: BatchGeometry | None = None
    self._row_plan_request_selects: dict[str, dict[str, dict]] | None = None
    self._row_plans: dict = {}
    self.lock = threading.Lock()

tokens_stored property

tokens_stored: int

append

append(
    layer_id: int,
    tensor: Tensor,
    meta: Tensor,
    layer_name: str,
)

Stage one layer's selected rows for this step (GPU side).

meta is an int32 [rows, 3] tensor of (req_idx, position, token_id) labels. The device-to-host copy happens once per step in flush(): per-layer synchronous .cpu() calls stall the GPU stream once per hooked layer, which serializes capture-heavy runs.

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def append(
    self,
    layer_id: int,
    tensor: torch.Tensor,
    meta: torch.Tensor,
    layer_name: str,
):
    """Stage one layer's selected rows for this step (GPU side).

    ``meta`` is an int32 ``[rows, 3]`` tensor of
    (req_idx, position, token_id) labels. The device-to-host copy happens once
    per step in
    ``flush()``: per-layer synchronous ``.cpu()`` calls stall the
    GPU stream once per hooked layer, which serializes
    capture-heavy runs.
    """
    with self.lock:
        if (
            not isinstance(meta, torch.Tensor)
            or meta.ndim != 2
            or meta.shape[1] != 3
            or meta.dtype != torch.int32
        ):
            raise RuntimeError(
                f"capture layer {layer_id} requires int32 [rows, 3] row labels"
            )
        if meta.shape[0] != tensor.shape[0]:
            raise RuntimeError(
                f"capture meta rows ({meta.shape[0]}) != data rows "
                f"({tensor.shape[0]}) for layer {layer_id}"
            )
        keep = self._limit_rows(layer_id, tensor.shape[0])
        if keep == 0:
            return
        tensor = tensor[:keep]
        meta = meta[:keep]
        if self.config.dtype is not None:
            tensor = tensor.to(self.config.dtype)
        # `tensor` is owned by the hook (freshly materialized), so an
        # async copy in flush() cannot race buffer reuse.
        self._pending.append(
            (
                layer_id,
                tensor.detach(),
                meta.detach(),
                layer_name,
            )
        )

flush

flush()

Move this step's staged rows to CPU in one coalesced pass.

Copies go through pinned staging buffers with non_blocking=True (one stream sync at the end), amortizing D2H latency across all hooked layers instead of paying it per layer.

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def flush(self):
    """Move this step's staged rows to CPU in one coalesced pass.

    Copies go through pinned staging buffers with non_blocking=True
    (one stream sync at the end), amortizing D2H latency across all
    hooked layers instead of paying it per layer.
    """
    with self.lock:
        if not self._pending:
            return
        pinned = []
        any_cuda = False
        for layer_id, tensor, meta, layer_name in self._pending:
            host_pair = []
            for t in (tensor, meta):
                if t.is_cuda:
                    host = torch.empty_like(t, device="cpu", pin_memory=True)
                    host.copy_(t, non_blocking=True)
                    any_cuda = True
                else:
                    host = t
                host_pair.append(host)
            pinned.append((layer_id, host_pair[0], host_pair[1], layer_name))
        if any_cuda:
            torch.cuda.current_stream().synchronize()
        self._pending.clear()
        for layer_id, host, meta_host, layer_name in pinned:
            requests, counts = torch.unique_consecutive(
                meta_host[:, 0], return_counts=True
            )
            start = 0
            for request_index, count in zip(requests.tolist(), counts.tolist()):
                tensor = host[start : start + count]
                meta = meta_host[start : start + count]
                # Each request owns its storage. Draining one must release
                # its values without copying or pinning another request's rows.
                if len(counts) > 1:
                    tensor, meta = tensor.clone(), meta.clone()
                chunk_id = self._next_chunk
                self._next_chunk += 1
                self.chunks.setdefault(layer_id, {})[chunk_id] = _StoredChunk(
                    tensor, meta, request_index
                )
                self._request_chunks.setdefault(request_index, {}).setdefault(
                    layer_id, set()
                ).add(chunk_id)
                self._unused_request_indices.discard(request_index)
                start += count
            self.layer_names[layer_id] = layer_name
            self._layer_rows[layer_id] = (
                self._layer_rows.get(layer_id, 0) + host.shape[0]
            )
        self._reclaim_requests()

serialize

serialize(
    layers: list[int] | None = None,
    req_ids: list[str] | None = None,
    clear_selected: bool = False,
) -> dict[int, dict[str, Any]]

Concatenate selected chunks and pack for RPC transmission.

Tensors ship as raw bytes of their stored dtype (bf16 rides as int16 bytes and is reinterpreted client-side) — no float32 upcast, so the wire volume equals the stored volume. Passing layers serializes a subset, letting clients fetch layer by layer instead of one monolithic message.

Each layer entry carries a meta sub-dict labelling its rows (req_table request-id strings + int32 req_idx / positions / token_ids columns).

req_ids (client-visible ids) restricts the payload to rows of those requests; with clear_selected the emitted rows are also removed from the store, so clients can drain request by request with bounded peak message size.

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def serialize(
    self,
    layers: list[int] | None = None,
    req_ids: list[str] | None = None,
    clear_selected: bool = False,
) -> dict[int, dict[str, Any]]:
    """Concatenate selected chunks and pack for RPC transmission.

    Tensors ship as raw bytes of their stored dtype (bf16 rides as
    int16 bytes and is reinterpreted client-side) — no float32
    upcast, so the wire volume equals the stored volume. Passing
    ``layers`` serializes a subset, letting clients fetch layer by
    layer instead of one monolithic message.

    Each layer entry carries a ``meta`` sub-dict labelling its rows
    (``req_table`` request-id strings + int32 ``req_idx`` /
    ``positions`` / ``token_ids`` columns).

    ``req_ids`` (client-visible ids) restricts the payload to rows
    of those requests; with ``clear_selected`` the emitted rows are
    also removed from the store, so clients can drain request by
    request with bounded peak message size.
    """
    with self.lock:
        failed = {
            rid: reason
            for rid, reason in self._failed_requests.items()
            if req_ids is None
            or any(match_capture_request_id(rid, ext) for ext in req_ids)
        }
        if failed:
            raise RuntimeError(
                f"capture contains failed request(s) {list(failed.items())[:5]}; "
                "their rows do not represent valid capture. Fetch with "
                "req_ids excluding them, or clear the stream."
            )
        elided = self.elided_reqs
        if req_ids is not None and elided:
            elided = {
                rid
                for rid in elided
                if any(match_capture_request_id(rid, ext) for ext in req_ids)
            }
        if elided:
            raise RuntimeError(
                "prefix-cache hits skipped recomputation of the prompt "
                f"head for request(s) {sorted(elided)[:5]}; their early "
                "rows were never computed, so this capture is "
                "incomplete. Submit capture requests with "
                "SamplingParams(skip_reading_prefix_cache=True), "
                "fetch with req_ids excluding them, or "
                "clear the stream."
            )
        table_match: set[int] | None = None
        if req_ids is not None:
            table_match = {
                i
                for req_id in req_ids
                for i in self._request_aliases.get(req_id, ())
            }
        result: dict[int, dict[str, Any]] = {}
        wanted = (
            sorted(self.chunks)
            if layers is None
            else [lid for lid in layers if lid in self.chunks]
        )
        for layer_id in wanted:
            layer_name = self.layer_names[layer_id]
            chunks = self.chunks[layer_id]
            if table_match is not None:
                chunk_ids = sorted(
                    chunk
                    for index in table_match
                    for chunk in self._request_chunks.get(index, {}).get(
                        layer_id, ()
                    )
                )
            else:
                chunk_ids = list(chunks)
            if not chunk_ids:
                continue
            selected = [chunks[index] for index in chunk_ids]
            tensor = torch.cat([chunk.tensor for chunk in selected], dim=0)
            meta = torch.cat([chunk.meta for chunk in selected], dim=0)
            result[layer_id] = serialize_capture_layer(
                tensor, meta, self.req_table, layer_name
            )
            if clear_selected:
                for index in chunk_ids:
                    self._remove_chunk(layer_id, index)
        if clear_selected:
            self._reclaim_requests()
        return result

clear

clear() -> None

Reset captured rows and counters, retaining active request history.

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def clear(self) -> None:
    """Reset captured rows and counters, retaining active request history."""
    with self.lock:
        self.chunks.clear()
        self._request_chunks.clear()
        self._next_chunk = 0
        self.layer_names.clear()
        self.req_table.clear()
        self._req_index.clear()
        self._request_aliases.clear()
        self._free_request_indices.clear()
        self._next_request_index = 0
        self._unused_request_indices.clear()
        self.elided_reqs.clear()
        self._failed_requests.clear()
        self._layer_rows.clear()
        self.tokens_dropped = 0
        self._warned_budget = False
        self._pending.clear()
        self._row_plans.clear()
        self._row_plan_geometry = None
        self._row_plan_request_selects = None

drop_layers

drop_layers(layers: list[int]) -> None

Source code in vllm-steer/vllm/model_hooks/capture/store.py

def drop_layers(self, layers: list[int]) -> None:
    with self.lock:
        for layer_id in layers:
            for chunk_id in list(self.chunks.get(layer_id, ())):
                self._remove_chunk(layer_id, chunk_id)
        dropped = set(layers)
        self._pending = [
            entry for entry in self._pending if entry[0] not in dropped
        ]
        self._reclaim_requests()