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'.

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…
• `assistant_ensemble_weight` must be in the open interval…
• Could not find `num_mtp_layers` in the model config. This…
• `early_stopping` must be a boolean or 'never', but is
• Invalid `cache_implementation

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

Appendix: 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)