huggingface/transformers · error · RuntimeError

decompose_prefill_decode expected at least {num_new_tokens}

Error message

decompose_prefill_decode expected at least {num_new_tokens} calls to {type(model).__name__}.forward() during generate(max_new_tokens={num_new_tokens}), but captured {len(calls)}. This likely means generate() bypasses the top-level forward() (e.g. delegates to an inner model), so prefill/decode decomposition is not supported for this architecture.

What it means

decompose_prefill_decode instruments the model's top-level forward() and runs generate(max_new_tokens=2 or 3). It expects at least one call per generated token; if fewer forward() calls were captured, generate() is bypassing the top-level forward (delegating to an inner model), so prefill/decode splitting at this level is impossible and the architecture is unsupported.

Source

Thrown at src/transformers/exporters/utils.py:774

    # `generation_config` alongside generation kwargs is deprecated. Base it on the model's own config
    # when none is given (preserving its defaults), and deep-copy into a distinct `capture_config` so
    # the caller's `generation_config` is never mutated.
    num_new_tokens = 3 if multi_token_decode else 2
    capture_config = copy.deepcopy(generation_config if generation_config is not None else model.generation_config)
    capture_config.max_new_tokens = num_new_tokens
    capture_config.min_new_tokens = num_new_tokens
    try:
        with _capture_forward(model) as calls:
            model.generate(**copy.deepcopy(inputs), generation_config=capture_config)
    except Exception as e:
        raise RuntimeError(
            f"decompose_prefill_decode failed for {type(model).__name__}. "
            f"Inputs passed: {list(inputs.keys())}. "
            f"Make sure the inputs are compatible with model.generate()."
        ) from e

    if len(calls) < num_new_tokens:
        raise RuntimeError(
            f"decompose_prefill_decode expected at least {num_new_tokens} calls to "
            f"{type(model).__name__}.forward() during generate(max_new_tokens={num_new_tokens}), but "
            f"captured {len(calls)}. This likely means generate() bypasses the top-level forward() "
            "(e.g. delegates to an inner model), so prefill/decode decomposition is not supported "
            "for this architecture."
        )

    # Remove `logits_to_keep` from the captured calls — it's a generation-time hint for the model's
    # internal top-k pruning, not a forward input. The export graph should not depend on it.
    for call in calls:
        call.pop("logits_to_keep", None)

    # A single-token decode specializes its query-sequence axis to 1 (never dynamic). When
    # `multi_token_decode`, merge the two decode steps into one multi-token decode so that axis stays
    # symbolic (continuation-from-past, or a plain prefill when the cache is empty, and it still covers seq == 1).
    prefill_inputs = calls[0]
    decode_inputs = _merge_decode_calls(calls[1:num_new_tokens]) if multi_token_decode else calls[1]
    return {

View on GitHub (pinned to a597f97485)

Solutions

  1. Export the inner model directly (e.g. model.model or model.language_model) with decompose_prefill_decode instead of the wrapper
  2. Check for a transformers update where this architecture's export is supported
  3. Fall back to exporting a single non-decomposed generation graph if acceptable

Example fix

// before
decompose_for_generation(wrapper_model, inputs)

// after
# export the inner causal LM that generate() actually calls
decompose_for_generation(wrapper_model.language_model, inner_inputs)
Defensive patterns

Strategy: fallback

Validate before calling

from transformers.exporters.utils import _capture_forward

def forwards_are_top_level(model, inputs, n=2) -> bool:
    with _capture_forward(model) as calls:
        model(**inputs)
    return len(calls) >= 1

Try / catch

try:
    parts = decompose_for_generation(model, inputs)
except RuntimeError as e:
    if "bypasses the top-level forward" in str(e):
        parts = decompose_for_generation(model.model, inner_inputs)  # inner module
    else:
        raise

Prevention

When it happens

Trigger: Exporting an architecture whose generate() path calls an inner module (e.g. a wrapper model that calls self.model(...)/self.language_model(...)) instead of its own forward; models with custom generate loops or multi-token decode paths that skip the top-level forward.

Common situations: Newly added or custom wrapper architectures in an export pipeline; models refactored so forward() is a thin shim; version changes where generation dispatch moved to inner components.

Related errors

AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/945e53733269f43c. Report an issue: GitHub.