[Spec] Clamp multimodal pad sentinels in spec-v2 draft prefill embedding (#27512)

This commit is contained in:
Liangsheng Yin
2026-06-07 17:00:49 -07:00
committed by GitHub
parent 10d33bd77e
commit 70db73afce
2 changed files with 16 additions and 1 deletions
+6 -1
View File
@@ -201,7 +201,12 @@ class MiMoV2ModelNextN(nn.Module):
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
# Multimodal pad sentinels (MM_PAD_SHIFT_VALUE + hash) sit out of vocab;
# clamp to avoid an OOB gather. The draft gets visual semantics from target
# hidden_states, so the embedding at these positions is unused anyway.
hidden_states = self.embed_tokens(
input_ids.clamp(min=0, max=self.vocab_size - 1)
)
else:
hidden_states = input_embeds
if hidden_states.shape[0] > 0:
@@ -383,6 +383,16 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
target_hidden_states: Hidden states from the target model forward
next_token_ids: Next token ids generated from the target forward.
"""
# The draft embed clamps unconditionally (to tolerate multimodal pad
# sentinels), so probe next_token_ids here first -- otherwise a corrupted id
# would be clamped away instead of surfacing.
maybe_detect_oob(
next_token_ids,
0,
self.model_config.vocab_size,
"draft_extend_for_prefill: next_token_ids before draft embed",
)
# Construct spec_info
next_draft_input = EagleDraftInput(
hidden_states=target_hidden_states,