chunked_forward

Run Hugging Face transformers chunk by chunk, carrying their state from one call to the next.

A sequential chunked forward pass calls a transformer once per fixed-size chunk of a sequence. For the chunks to compose into the same computation as a single pass, each call must continue from the cache and positions the previous call left behind. This module attaches hooks that carry that state without changing how the transformer is called, and holds the adapters that describe what each transformer architecture needs in order to run this way.

Classes:

Name	Description
CausalDecoderChunkedForwardAdapter	Run estimators of causal decoder transformers, such as Llama, Mistral, and Qwen3, in a chunked forward pass.
ChunkedForwardAdapter	Describe how the estimators of one transformer architecture run in a chunked forward pass.
NemotronHChunkedForwardAdapter	Run NemotronH estimators of Mamba, attention, and MLP layers in a chunked forward pass.

Functions:

Name	Description
attach_chunked_estimator_hooks	Attach hooks that carry an estimator's state across the chunks of a chunked forward pass.
chunked_forward_adapter	Return the adapter registered for an estimator's transformer class.
estimator_supports_chunked_forward	Return whether a chunked forward pass supports an estimator's transformer.

Attributes:

Name	Type	Description
CHUNKED_FORWARD_ADAPTERS	Final[dict[type[PreTrainedModel], ChunkedForwardAdapter[Any]]]	The chunked forward adapter of each supported estimator transformer class. Each adapter handles the class it is registered under,

CHUNKED_FORWARD_ADAPTERS module-attribute

CHUNKED_FORWARD_ADAPTERS: Final[
    dict[type[PreTrainedModel], ChunkedForwardAdapter[Any]]
] = {
    transformers.LlamaModel: CausalDecoderChunkedForwardAdapter[
        transformers.LlamaModel
    ](),
    transformers.MistralModel: CausalDecoderChunkedForwardAdapter[
        transformers.MistralModel
    ](),
    transformers.Qwen3Model: CausalDecoderChunkedForwardAdapter[
        transformers.Qwen3Model
    ](),
    transformers.NemotronHModel: NemotronHChunkedForwardAdapter(),
}

The chunked forward adapter of each supported estimator transformer class. Each adapter handles the class it is registered under, which a type annotation cannot express per entry, so the values are typed as adapters of any class.

CausalDecoderChunkedForwardAdapter

Bases: ChunkedForwardAdapter[_TransformerT_contra]

Run estimators of causal decoder transformers, such as Llama, Mistral, and Qwen3, in a chunked forward pass.

Each chunk's call continues from the cache at the chunk's positions, and the transformer builds the chunk's attention masks from the cache itself, as generate() does for a chunked prefill: each chunk's queries attend to every cached key before them, or to those within the sliding window for sliding-window layers, and causally within the chunk. A chunk's padding follows its real positions, so the causal masks keep it from affecting them.

Dynamic RoPE keeps its rescaled frequencies on the rotary module rather than in the cache, and transformers restores the original ones only for a call shorter than max_position_embeddings. A chunked pass whose chunk_size is at least that long would keep a previous longer call's frequencies, so the same seeded prefix could be transformed differently. No supported estimator uses dynamic RoPE, and such a chunk size is unrealistic anyways, but the failure would be silent, so it is rejected outright.

Added in version v4.4.0.

Methods:

Name	Description
chunk_kwargs	Return the cache and positions of a causal decoder transformer's call for a chunk, and no attention mask.
unsupported_reason	Return why a causal decoder transformer is not supported, which is when its config makes its attention bidirectional or its

chunk_kwargs

chunk_kwargs(
    transformer: _TransformerT_contra,
    cache: Cache,
    position_ids: Tensor,
) -> dict[str, Any]

Return the cache and positions of a causal decoder transformer's call for a chunk, and no attention mask.

TransformerCloak rejects padded inputs before the pass, so the caller's attention mask attends everywhere. It also spans the whole sequence rather than the chunk, so it is replaced with None, and the transformer builds the chunk's causal masks from the cache.

Parameters:

Name	Type	Description	Default
transformer	_TransformerT_contra	The estimator's transformer.	required
cache	Cache	The cache that carries the estimator's state from one chunk to the next.	required
position_ids	Tensor	The positions of the chunk.	required

Returns:

Type	Description
dict[str, Any]	The keyword arguments.

unsupported_reason

unsupported_reason(
    transformer: _TransformerT_contra,
) -> str | None

Return why a causal decoder transformer is not supported, which is when its config makes its attention bidirectional or its rotary embeddings use dynamic RoPE.

Parameters:

Name	Type	Description	Default
transformer	_TransformerT_contra	The estimator's transformer.	required

Returns:

Type	Description
str | None	Why the transformer is not supported, or None if it is.

ChunkedForwardAdapter

Bases: ABC, Generic[_TransformerT_contra]

Describe how the estimators of one transformer architecture run in a chunked forward pass.

The chunked forward pass itself does not depend on the architecture. An adapter supplies the parts that do: whether it supports a particular estimator, and the keyword arguments that carry an estimator's state into each chunk. Its type parameter is the transformer class it handles. Register an adapter in CHUNKED_FORWARD_ADAPTERS under that class.

Added in version v4.2.0.

Methods:

Name	Description
chunk_kwargs	Return the keyword arguments that carry an estimator's state into the transformer's call for one chunk of an unpadded input.
unsupported_reason	Return why a chunked forward pass does not support an estimator's transformer, or None if it does.

chunk_kwargs

chunk_kwargs(
    transformer: _TransformerT_contra,
    cache: Cache,
    position_ids: Tensor,
) -> dict[str, Any]

Return the keyword arguments that carry an estimator's state into the transformer's call for one chunk of an unpadded input.

They replace any caller arguments of the same names. By default they are the estimator's cache, with caching enabled, and the chunk's positions. Subclasses add what their architecture needs to the result of super().chunk_kwargs.

Parameters:

Name	Type	Description	Default
transformer	_TransformerT_contra	The estimator's transformer.	required
cache	Cache	The cache that carries the estimator's state from one chunk to the next.	required
position_ids	Tensor	The positions of the chunk, of shape (1, chunk_size).	required

Returns:

Type	Description
dict[str, Any]	The keyword arguments.

unsupported_reason abstractmethod

unsupported_reason(
    transformer: _TransformerT_contra,
) -> str | None

Return why a chunked forward pass does not support an estimator's transformer, or None if it does.

Parameters:

Name	Type	Description	Default
transformer	_TransformerT_contra	The estimator's transformer, an instance of the class this adapter is registered for.	required

Returns:

Type	Description
str | None	A reason that completes the sentence "Chunked forward passes are not supported for this estimator, because ...", or None.

Raises:

Type	Description
NotImplementedError	Always. Subclasses must implement it, and a super() call must not silently report support.

NemotronHChunkedForwardAdapter

Bases: CausalDecoderChunkedForwardAdapter[NemotronHModel]

Run NemotronH estimators of Mamba, attention, and MLP layers in a chunked forward pass.

Changed in version v4.4.0: Attention layers are now supported, including in hybrid estimators.

Added in version v4.2.0.

Methods:

Name	Description
chunk_kwargs	Return the cache and positions of a NemotronH transformer's call for a chunk, and its attention masks.
unsupported_reason	Return why a NemotronH transformer is not supported, which is when its attention is not causal or any of its layers uses an

chunk_kwargs

chunk_kwargs(
    transformer: NemotronHModel,
    cache: Cache,
    position_ids: Tensor,
) -> dict[str, Any]

Return the cache and positions of a NemotronH transformer's call for a chunk, and its attention masks.

With attention layers, the transformer builds the chunk's masks from the cache, as for any causal decoder. Without them, the masks are passed precomputed, as generate() does, which stops NemotronH from building them from the cache: that would query the sequence length that a cache of only Mamba layers cannot report. Passing position_ids likewise stops it from deriving the positions from that sequence length. The layer types are then taken from the built model, and since TransformerCloak rejects padded inputs before the pass, every mask is None.

Parameters:

Name	Type	Description	Default
transformer	NemotronHModel	The estimator's NemotronH transformer.	required
cache	Cache	The cache that carries the estimator's state from one chunk to the next.	required
position_ids	Tensor	The positions of the chunk.	required

Returns:

Type	Description
dict[str, Any]	The keyword arguments.

unsupported_reason

unsupported_reason(
    transformer: NemotronHModel,
) -> str | None

Return why a NemotronH transformer is not supported, which is when its attention is not causal or any of its layers uses an unsupported mixer.

Parameters:

Name	Type	Description	Default
transformer	NemotronHModel	The estimator's NemotronH transformer.	required

Returns:

Type	Description
str | None	Why the transformer is not supported, or None if it is.

attach_chunked_estimator_hooks

attach_chunked_estimator_hooks(
    transformer: _TransformerT,
    adapter: ChunkedForwardAdapter[_TransformerT],
    chunk_size: int,
) -> contextlib.ExitStack

Attach hooks that carry an estimator's state across the chunks of a chunked forward pass.

This function's job is to make a Hugging Face transformer stateful for the duration of a with block, without changing how it's called. The estimator code calls transformer(inputs_embeds=...) exactly as in a single pass, and the hooks fill in the cache, positions and masks behind it. Each transformer gets its own cache and offset, so the mean and standard deviation estimators carry independent state.

A forward pre-hook replaces the caller's arguments of the same names with the adapter's chunk_kwargs: the estimator's own cache, the chunk's position_ids, and whatever else the architecture needs. A forward hook then advances the position offset by chunk_size, so each call continues from where the previous chunk ended.

The hooks track the positions themselves instead of leaving the transformer to derive them. Without position_ids, a transformer computes a call's positions from the sequence length its cache reports, and a cache of only Mamba layers cannot report one: NemotronH then fails before running any layer. Mamba and MLP layers do not use the positions, but attention layers do, through their rotary embeddings, so the positions must continue across chunks rather than restart at 0 for each chunk.

Parameters:

Name	Type	Description	Default
transformer	_TransformerT	The estimator's transformer.	required
adapter	ChunkedForwardAdapter[_TransformerT]	The adapter for the transformer's architecture.	required
chunk_size	int	The number of positions per chunk.	required

Returns:

Type	Description
contextlib.ExitStack	A context manager that, upon exiting, detaches the hooks.

Raises:

Type	Description
ValueError	If chunk_size is not positive.
NotImplementedError	If adapter does not support transformer.

Added in version v4.2.0.

chunked_forward_adapter

chunked_forward_adapter(
    transformer: _TransformerT,
) -> ChunkedForwardAdapter[_TransformerT] | None

Return the adapter registered for an estimator's transformer class.

Only the exact class is looked up, not its base classes. A subclass can change what the adapter relies on, such as how its transformer builds masks or uses its cache, so it is not supported until an adapter is registered for it.

Parameters:

Name	Type	Description	Default
transformer	_TransformerT	The estimator's transformer.	required

Returns:

Type	Description
ChunkedForwardAdapter[_TransformerT] | None	The adapter, or None if no adapter is registered for the transformer's class.

Added in version v4.2.0.

estimator_supports_chunked_forward

estimator_supports_chunked_forward(
    transformer: PreTrainedModel,
) -> bool

Return whether a chunked forward pass supports an estimator's transformer.

Parameters:

Name	Type	Description	Default
transformer	PreTrainedModel	The estimator's transformer.	required

Returns:

Type	Description
bool	Whether an adapter is registered for the transformer and supports it.

Added in version v4.2.0.