huggingface/transformers · error · ValueError

Expected assistant_model to be a Gemma4AssistantForCausalLM

Error message

Expected assistant_model to be a Gemma4AssistantForCausalLM or Gemma4UnifiedAssistantForCausalLM. Got {} This candidate generator requires that the assistant model is able to work from a shared_kv_states dictionary. Currently, only the Gemma4AssistantForCausalLM and Gemma4UnifiedAssistantForCausalLM support this.

What it means

The Gemma4 shared-KV candidate generator (target-model-driven speculative decoding) drafts candidates from the main model's hidden states and shared KV cache, which only Gemma4AssistantForCausalLM / Gemma4UnifiedAssistantForCausalLM support. The constructor rejects any assistant model whose class name contains neither 'Gemma4Assistant' nor 'Gemma4UnifiedAssistant'.

Source

Thrown at src/transformers/generation/candidate_generator.py:1294

        "return_shared_kv_states": True,
    }

    def __init__(
        self,
        input_ids: torch.LongTensor,
        assistant_model: "PreTrainedModel",
        target_model_input_embeddings: nn.Embedding,
        generation_config: "GenerationConfig",
        model_kwargs: dict,
        inputs_tensor: torch.Tensor | None = None,
        logits_processor: Optional["LogitsProcessorList"] = None,
        eos_token_id: int | list[int] | torch.Tensor | None = None,
    ):
        if (
            "Gemma4Assistant" not in assistant_model.__class__.__name__
            and "Gemma4UnifiedAssistant" not in assistant_model.__class__.__name__
        ):
            raise ValueError(
                f"Expected assistant_model to be a Gemma4AssistantForCausalLM or Gemma4UnifiedAssistantForCausalLM. Got {assistant_model.__class__.__name__}"
                " This candidate generator requires that the assistant model is able to work from a shared_kv_states"
                " dictionary. Currently, only the Gemma4AssistantForCausalLM and Gemma4UnifiedAssistantForCausalLM support this."
            )

        super().__init__(input_ids, assistant_model, generation_config, model_kwargs, inputs_tensor, logits_processor)
        self.target_model_input_embeddings = target_model_input_embeddings

        if eos_token_id is None:
            eos_token_id: set = set()

            if isinstance(self.generation_config.eos_token_id, Iterable):
                eos_token_id.update(self.generation_config.eos_token_id)
            elif isinstance(self.generation_config.eos_token_id, int):
                eos_token_id.add(self.generation_config.eos_token_id)

            if isinstance(self.assistant_generation_config.eos_token_id, Iterable):
                eos_token_id.update(self.assistant_generation_config.eos_token_id)

View on GitHub (pinned to a597f97485)

Solutions

1. Use the standard AssistedCandidateGenerator path with any small draft model (pass it as assistant_model without enabling the shared-KV/Gemma4-specific config)
2. Use Gemma4AssistantForCausalLM or Gemma4UnifiedAssistantForCausalLM as the assistant for this generator
3. If you wrote a custom assistant that handles shared_kv_states, rename/ensure the class supports the contract instead of relying on the name check

Example fix

# before
assistant = LlamaForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
# then routed to Gemma4 shared-KV generator -> ValueError
# after: keep the generic assisted-decoding path
model.generate(inputs, assistant_model=assistant)  # without shared-KV generator selection

def supports_shared_kv_assistant(assistant_model) -> bool:
    name = assistant_model.__class__.__name__
    return "Gemma4Assistant" in name or "Gemma4UnifiedAssistant" in name

def is_gemma4_assistant(model) -> bool:
    name = model.__class__.__name__
    return ("Gemma4Assistant" in name) or ("Gemma4UnifiedAssistant" in name)

When it happens

Trigger: Wiring an ordinary draft model (e.g. a small LlamaForCausalLM or an older Gemma assistant) into the shared-KV candidate path — usually by passing a mismatched assistant_model when the generation config selects the shared-KV generator.

Common situations: Mixing model families in assisted decoding, upgrading transformers where Gemma4 introduced this path, or a custom assistant class that does not implement shared_kv_states handling.

Related errors

• `model_outputs` cannot be None, and they need to contain `hi
• Could not find `num_mtp_layers` in the model config. This mo
• `assistant_ensemble_weight` must be in the open interval `(0
• {} is an abstract class. Only classes inheriting this class
• {} is an abstract class. Only classes inheriting this class

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