[aarushdubey]

What does this PR do?

Fixes #50598

Background & Problem

In universal speculative decoding with assistant LM head pruning enabled (assistant_prune_lm_head=True), AssistantToTargetTranslator prunes the assistant LM head using self.assistant_overlap_token_ids:

self.assistant_overlap_token_ids = torch.tensor(
    list(self.target_to_assistant_input_ids.values()),
    dtype=torch.long,
    device=self._assistant_model_device,
)

self.target_to_assistant_input_ids is constructed by iterating through assistant_vocab.items(). When an assistant tokenizer's dictionary order does not match ascending token-ID order (which frequently occurs with fast tokenizers or serialized tokenizers), self.assistant_overlap_token_ids follows dictionary key insertion order rather than sorted token ID order.

However, in get_target_logits:

assistant_indices_mask = self._assistant_to_target_input_ids != self.SUPPRESS_TOKEN_ID
target_logits_supported_indices = self._assistant_to_target_input_ids[assistant_indices_mask]
if self.assistant_prune_lm_head:
    target_logits[..., target_logits_supported_indices] = assistant_logits

target_logits_supported_indices assumes assistant_logits is ordered by ascending assistant token ID (torch.where(assistant_indices_mask)[0]).
When self.assistant_overlap_token_ids is ordered by dictionary insertion, the pruned LM head produces logits in dictionary order while get_target_logits maps them as if they were in token ID order, scrambling logits across target tokens. Consequently, candidate verification in speculative decoding evaluates candidate tokens using incorrect probabilities.

Fix

Initialize self.assistant_overlap_token_ids directly via torch.where:

self.assistant_overlap_token_ids = torch.where(
    self._assistant_to_target_input_ids != self.SUPPRESS_TOKEN_ID
)[0]

This guarantees self.assistant_overlap_token_ids strictly follows ascending assistant token IDs, aligning 1:1 with assistant_indices_mask and target_logits_supported_indices regardless of tokenizer vocabulary dictionary order.

Changes

• Updated self.assistant_overlap_token_ids initialization in AssistantToTargetTranslator (src/transformers/generation/candidate_generator.py).
• Added unit regression test test_pruned_lm_head_unordered_vocab in tests/generation/test_candidate_generator.py.
• Verified with pytest tests/generation/test_candidate_generator.py (14/14 passed), ruff check, and python utils/check_copies.py.

Before submitting

• This PR fixes a bug: Universal speculative decoding with a pruned assistant LM head can map logits to the wrong target tokens #50598
• Did you read the contributor guideline?
• Did you write any new necessary tests?