Assisted/speculative decoding generates past an EOS that is accepted mid-block

[pjordanandrsn]

Assisted/speculative decoding generates past an EOS that is accepted mid-block

System info

• transformers 5.15.0 (reproduced on the v5.15.0 tag)
• torch 2.13.0 — the minimal repro is CPU-only
• Originally observed on meta-models/Muse-Glimmer-30B + -assistant (DFlash), A100 80GB, bf16

Who can help?

@Cyrilvallez (generate). The DFlash generator came in with #47867.

Description

_assisted_decoding commits n_matches + 1 tokens per round. When an accepted
token in the middle of that block is an EOS, generation does not stop: the whole
block is appended (and streamed) first, and the stopping check afterwards only
looks at the last committed token. The result is output that continues past EOS
and does not match plain greedy.

It is position-dependent, which makes it look intermittent — an EOS that happens
to land on the target's bonus token stops correctly.

Minimal reproduction (CPU, tiny random model, no drafter needed)

The bug needs the EOS to be accepted, i.e. the target's own argmax must be EOS
at that position. Instead of hunting for such a prompt, run greedy first and then
declare one of the tokens the model actually produced to be the EOS.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.candidate_generator import CandidateGenerator
import transformers.generation.utils as gu

NAME, BLOCK, EOS_OFFSET = "hf-internal-testing/tiny-random-LlamaForCausalLM", 8, 3
model = AutoModelForCausalLM.from_pretrained(NAME).eval()
tok = AutoTokenizer.from_pretrained(NAME)
ids = tok("Hello world", return_tensors="pt").input_ids
plen = ids.shape[1]

cont = model.generate(ids, max_new_tokens=BLOCK * 2, do_sample=False,
                      eos_token_id=-1, pad_token_id=0)[0, plen:].tolist()
EOS = cont[EOS_OFFSET]                      # a token the model really emits

class BlockStub(CandidateGenerator):        # proposes the model's own continuation
    requires_model_outputs = False
    def get_candidates(self, input_ids, **kw):
        n = input_ids.shape[1] - plen
        nxt = cont[n:n + BLOCK]
        if not nxt:
            return input_ids, None
        return torch.cat([input_ids, torch.tensor([nxt])], dim=-1), None
    def update_candidate_strategy(self, input_ids, scores, num_matches):
        pass

base = model.generate(ids, max_new_tokens=BLOCK * 2, do_sample=False,
                      eos_token_id=EOS, pad_token_id=0)
gu.GenerationMixin._get_candidate_generator = lambda self, **kw: BlockStub()
spec = model.generate(ids, max_new_tokens=BLOCK * 2, do_sample=False,
                      eos_token_id=EOS, pad_token_id=0, prompt_lookup_num_tokens=BLOCK)

print(base[0, plen:].tolist())   # [30016, 7356, 9023, 25068]                   <- stops at EOS
print(spec[0, plen:].tolist())   # [30016, 7356, 9023, 25068, 22665, ... ]      <- 12 tokens past EOS

Observed: plain greedy returns 4 tokens, assisted returns 16 — 12 tokens after
the EOS.

On a real model

meta-models/Muse-Glimmer-30B + its DFlash drafter, greedy, max_new_tokens=128,
prompt "List the first eight prime numbers, one per line, with no commentary.":

	length	EOS index	ends with EOS
plain greedy	122	121	yes
speculation_type="dflash"	128	121	no

The first 122 tokens are identical; the speculative stream then emits
<|start|>assistant to=self<|message|>We. Instrumenting the candidate generator
to record (proposed, num_matches) per round and locating the EOS in its commit
round explains the intermittency:

• EOS at offset 6 of an 8-token commit (an accepted proposal) → ran 6 tokens past it.
• On another prompt, EOS at offset 1 of a 2-token commit (the bonus token) → stopped correctly.

Expected behavior

Greedy assisted decoding should return the same token stream as plain greedy, and
stop at the first EOS. Instead it continues past an EOS that was accepted inside a
committed block, so the two disagree on both content and length.

Cause

Two things compound, both in _assisted_decoding:

1. The committed block is appended and streamed before any stopping check:

   input_ids = torch.cat((input_ids, valid_tokens), dim=-1)   # utils.py:3808
   if streamer is not None:
       streamer.put(valid_tokens.cpu())

2. The check that follows uses the default new_token_length=1:

   unfinished_sequences = unfinished_sequences & ~stopping_criteria(input_ids, scores)  # utils.py:3859

   and EosTokenCriteria only inspects the last new_token_length tokens
   (stopping_criteria.py:580):

   is_done = torch.isin(input_ids[:, -new_token_length:], self.eos_token_id).any(dim=-1)

The existing guard at utils.py:3777 does not cover it: it fires only when the
entire candidate block was accepted (n_matches == candidate_length), and its
is_done_candidate (utils.py:3696) is itself computed at new_token_length=1.

Classic autoregressive drafters largely avoid this because the draft stops at EOS,
leaving it last. Block proposers (DFlash block_size 16, MTP) hit it routinely.

This failure mode is already known in-repo — PromptLookupCandidateGenerator
truncates its candidates at the first EOS for exactly this reason
(candidate_generator.py:1129):

# remove remaining candidate ids if an "eos" token is found, otherwise the target model may
# accept eos and the rest as valid, thus not stopping generation after "eos"

DFlashTokenCandidateGenerator has no equivalent.

Fix

Trim the committed block at the first EOS, mirroring what
PromptLookupCandidateGenerator already does to its drafts. That also leaves EOS
as the last token, so the existing criteria fire unchanged:

eos_tensor = generation_config._eos_token_tensor
if eos_tensor is not None and valid_tokens.shape[1] > 1:
    eos_hits = torch.nonzero(torch.isin(valid_tokens[0], eos_tensor.to(valid_tokens.device)))
    if eos_hits.numel() > 0:
        valid_tokens = valid_tokens[:, : eos_hits[0].item() + 1]
        n_matches = valid_tokens.shape[1] - 1

placed just after the existing length-budget clamp, recomputing n_matches the way
that clamp already does so number_of_tokens_to_crop stays consistent.

Note the obvious one-liner — passing new_token_length=n_matches + 1 to the
stopping criteria at utils.py:3859, as generation_diffusion_gemma.py:1088 does
for its canvas — is not sufficient on its own. Measured on the repro above it
takes the output from 16 tokens to 9, but 5 post-EOS tokens remain: it ends the
loop without undoing the commit.

min_new_tokens is unaffected: the logits processors suppress EOS until the
minimum is reached, so the target never accepts one for the trim to find (verified
at min_new_tokens=0 → 4 tokens and =10 → 16).

Verification

Branch is 61 added lines (11 source, 50 test) on top of main (dfb1af6).

pytest tests/generation/test_utils.py tests/generation/test_candidate_generator.py \
       tests/generation/test_stopping_criteria.py
# 115 passed, 21 skipped, 3 failed

The 3 failures (test_default_max_length_warning,
test_inputs_embeds_warn_without_ids_for_token_based_logit_processors,
test_validate_stopping_criteria) fail identically on clean main without this
change.

The new test is armed: reverting only generation/utils.py to main makes it fail
with generation continued past an accepted EOS: [750, 297, 297, 860, 860, 860, 860, 860], and it passes with the fix. ruff check / ruff format --check clean
on both touched files.

---

AI-assisted contribution, per CONTRIBUTING.md: the patch, the repro and the test
were drafted with Claude Code, and the results above are from actual local runs.
The changed lines and the behaviour have been reviewed by me before submitting.

---

Would you like a PR for this? I have the fix and the regression test above ready on
a branch off main (dfb1af6) — happy to open it, or to leave it if you would
rather handle the stopping logic differently.

[Jithendra2608]

@pjordanandrsn i would love to contribute to community through this

[pjordanandrsn]

Welcome! Fair warning that the fix and a regression test are already written and verified — they're the "Fix" section of the issue body above, so it'd be a shame for you to redo them from scratch.

I'm happy to open the PR and would genuinely welcome your review on it.

The one part worth knowing whoever writes it: the obvious one-liner — passing new_token_length=n_matches + 1 to the stopping criteria at utils.py:3859, which is what generation_diffusion_gemma.py does for its canvas — is not sufficient on its own. Measured on the CPU repro above it takes the output from 16 tokens down to 9, but 5 post-EOS tokens remain, because the block is appended (and streamed) before the check runs. It ends the loop without undoing the commit. What works is trimming the committed block at the first EOS, which also leaves EOS last so the existing criteria fire unchanged.

The repro in the issue is CPU-only and runs in seconds on a tiny random model, if you want to see the failure directly.

[pjordanandrsn]

Branch with the fix and the regression test, for anyone who wants to look before a PR exists: main...pjordanandrsn:transformers:fix-eos-main

[zucchini-nlp]

hey @pjordanandrsn , you can open a PR tho the solution should be inside Dflash.get_candidates making sure that any tokens/logits past EOS are either zeroed out or cropped. Take a look at how SinglePosition cand generator does it, most prob we can simply copy it

[pjordanandrsn]

PR is up: #47931, CI green.

@Jithendra2608 — the branch I linked earlier (fix-eos-main) is superseded, so don't spend time on it. @zucchini-nlp asked for the fix to sit in DFlashTokenCandidateGenerator.get_candidates, which made it a smaller and different change. Your review on the PR would be very welcome.

@zucchini-nlp — done as described: the draft is cropped at the first EOS, tokens and logits together. It gates on generation_config._eos_token_tensor, the same tensor EosTokenCriteria is built from, so a block can never be shortened for a token that wouldn't have stopped generation. On the real Muse-Glimmer-30B + -assistant pair (A100, bf16, greedy), 5 of 6 prompts ran past an accepted EOS before the change and 0 of 6 after.

Two caveats in the PR body: dflash isn't token-for-token identical to plain greedy here and wasn't before this change either (block vs single-token verification uses different bf16 kernel shapes), and MTPCandidateGenerator still has the same defect — so this shouldn't auto-close the issue unless you'd like me to extend the PR to cover it.

[pjordanandrsn]

@zucchini-nlp — before anything else, a correction to this issue. I wrote above that block proposers
"hit it routinely", and for MTP that was too strong. Measuring it, the defect fires only when the accepted
EOS is not the final drafted token, so num_mtp_layers == 1 cannot reach it at all — the single draft
position is also the last, and the is_done_candidate guard at utils.py:3777 already covers that.
DeepSeek-V3/V3.1 and GLM-4.5/4.5-Air/4.6 all ship exactly one layer, so they are unaffected.
thinkingmachines/Inkling and Inkling-Small ship 8 and are reachable today. DFlash is untouched by
this correction — block_size 16 reaches it easily, which is what #47931 measured on a real model.

So the MTP half is narrower than I implied: real, but latent on every checkpoint except Inkling, and I have
no MTP checkpoint I can run (Inkling-Small is 532 GB) to show it on real weights.

Given that, tell me which you'd prefer and I'll follow it:

1. I open the MTP PR — it's written and verified, 14 lines of candidate_generator.py plus two tests.
2. You'd rather not carry the extra test surface for a latent bug, and I open it as the crop alone.
3. Not worth it, and this issue closes with Crop DFlash drafts at the first EOS #47931 plus the caveat above recorded here.

Also worth settling either way: #47931's body says Fixes #47912, so it will auto-close this issue on
merge while the MTP half is unresolved. Happy to change that to Part of if you want the issue to stay open.

@Jithendra2608 — the measurement side of this is where the interesting part turned out to be, if you want a
look.

The change, and how it was measured

@@ class MTPCandidateGenerator(AssistedCandidateGenerator):
         self.do_sample = generation_config.do_sample
         self.logits_processor = logits_processor
 
+        # Same tensor the stopping criteria are built from, so a draft cropped here stops generation
+        self.eos_token_id = getattr(generation_config, "_eos_token_tensor", None)
+
         self.is_main_model_prefill = True

@@ def get_candidates(...):
         # Once we arrive here the first time, it's no longer the case
         self.is_main_model_prefill = False
 
+        # Crop the draft after the first EOS, otherwise the target model may accept eos and the rest as valid,
+        # thus not stopping generation after "eos" -- the block is committed before the stopping criteria run,
+        # and they only look at the last committed token. Cropping leaves EOS last, so they fire unchanged.
+        if self.eos_token_id is not None:
+            drafted_tokens = candidate_ids[0]
+            eos_positions = (torch.isin(drafted_tokens, self.eos_token_id.to(drafted_tokens.device))).nonzero()
+            if eos_positions.numel() > 0:
+                num_drafted = eos_positions[0].item() + 1
+                candidate_ids = candidate_ids[:, :num_drafted]
+                candidate_logits = candidate_logits[:, :num_drafted]
+
         # cat everything back together (we need to return the full ids here)
         candidate_ids = torch.cat([input_ids, candidate_ids], dim=-1)
         return candidate_ids, candidate_logits

Same crop PromptLookupCandidateGenerator already applies (candidate_generator.py:1129) and the one
#47931 applies to DFlash. n_matches is the leading run where the target's argmax equals the drafted token
(utils.py:3774) and valid_tokens = selected_tokens[:, : n_matches + 1] (:3779), so every committed
position except the last is provably a drafted token — cropping at the first EOS makes an
EOS-before-the-end impossible. It is applied to the returned block rather than by exiting the drafting loop
early, so every mtp layer still records its position and the next call's mtp_cache.crop(-num_mtp_layers)
is unaffected.

When it fires

Tiny randomly-initialised DeepseekV3ForCausalLM, run on macOS/CPU/torch 2.11 and on Linux/CUDA (RTX
A2000)/torch 2.8 with identical results. A random mtp head never agrees with the main model — not one draft
accepted in 39 rounds — so these pin the shared vocab projection to make drafts accepted, which is the
precondition for the bug.

num_mtp_layers	EOS at draft index	committed after the first EOS
1	0 (also last)	0
2	0 / 1 (last)	38 / 0
3	0 / 1 / 2 (last)	38 / 37 / 0

With the fix, every row is 0. For num_mtp_layers=3 with the EOS at draft index 1 and
max_new_tokens=40: plain greedy emits 3 tokens and stops at EOS, use_mtp=True emitted 40 — 37 of them
after the EOS — and after the fix emits the same 3 as greedy. do_sample=True is 37 past on each of 3 seeds
before and 0 after. With a draft that is cropped and then rejected, generation runs 28 further rounds
behind a shortened block and stays token-for-token identical to plain greedy, which is what says the mtp
cache stays in sync.

Layer counts resolved with AutoConfig.from_pretrained(...).get_text_config().num_mtp_layers: DeepSeek-V3 /
V3.1 → 1; GLM-4.5 / 4.5-Air / 4.6 → 1; Inkling and Inkling-Small → 8. Nemotron-H ships no mtp weights at all,
and Solar-Open registers no mtp key patterns, so MtpModel.from_pretrained could not load any.

Two things noticed in passing

test_generate_with_mtp (test_utils.py:2476) gates on
any("mtp" in x or re.search(r"layers\.\d+", x) ...) over _keys_to_ignore_on_load_unexpected. That regex
does not match patterns written with escaped dots, so it skips for deepseek_v3 (model\.layers\.61.*),
glm4_moe and solar_open; it runs only for inkling. It asserts just that generation raises nothing, so
it would not have caught this either way.

MtpModel.forward reads next_token_scores, only assigned at modeling_layers.py:541 when
logits_processor is not None, so building MTPCandidateGenerator with the default logits_processor=None
raises UnboundLocalError. Not reachable through generate, which always passes a LogitsProcessorList.

Happy to file either separately.

[zucchini-nlp]

Oke, lets fix the EOS now for MTP and dflash by cropping it out. It is a known bug, we already fixed in the past for other cand generator. If you think there are other things in MTP model specifically, please open another issue and tag CyrilVallez for that (mtp maintainers)