[Bug]: Interior control-token rejection can split a tool-call prelude and silently drop tool_calls

Relationship to existing issues and evidence status

• This issue is the acceptance-side mechanism analysis companion to #16377. #16377 is the user-visible Harmony manifestation (gpt-oss-120b + Eagle3 on DGX Spark) with a full end-to-end HTTP reproduction; this issue isolates the spec-decode acceptance invariant that manifestation depends on.
• Evidence status: this is a code review plus a deterministic numpy-only replay of the flashinfer chain_speculative_sampling decision rule. It is not yet an end-to-end tensorrt_llm reproduction (that needs HF weights and is tracked as a follow-up). Treat the claim below as an investigation until the end-to-end run confirms it.
• The acceptance kernel is hardware-independent, so the mechanism is not specific to any GPU SKI; reproduction should apply to consumer Blackwell (RTX 5090, sm_120) and datacenter GPUs alike.

Falsifiable claim (what would confirm or refute this)

For an interior rejection to be a real regression (as opposed to a property both decoding modes share), the engine must be decoded and emitted the accepted prefix (CH MSG...) as final output without any downstream guard or rollback. That specific invariant is testable independent of any specific model:

• Test 1 (prefix-emission): feed a synthetic input where the draft proposes MSG TOK END, the target disfavours the interior TOK slot, and assert whether the emitted prefix ends up as final new_tokens. In the current path this is always true because the accepted prefix is written directly as new_tokens with new_tokens_lens = num_accepted_tokens (no structural check).
• Test 2 (rollback-existence): assert that no downstream code path inspects accepted_tokens content or checks control-block balance before committing them to output. A clean negation — a guard that re-samples past the block opener — would refute the "no rollback" claim and make the fix a (b)-style hook rather than a core acceptance change.

The rollback-semantics claim is the part of this issue that is genuinely new relative to non-speculative decoding. In direct (non-speculative) decoding the model generates token-by-token and can equally emit an unclosed block, so "the identical well-formed draft yields a complete call without spec decoding" is not a rigorous proof of a spec-decode-only bug by itself; the distinguishing and testable fact is the engine's commit-accepted-prefix-without-rollback behavior shown above.

Summary

In one-model speculative decoding, the acceptance path returns only two values: accepted_tokens and num_accepted_tokens. It never inspects whether the accepted token stream is structurally complete. When a draft proposes a rare control preamble for a tool call (for example a Harmony-style CH MSG TOK END block) and the target disfavours one interior control token, the accepted prefix is left incomplete: the opened control block never closes, and the caller/parser reports finish_reason=stop with empty tool_calls and empty content. The model clearly intended a function call, yet the call is silently dropped. This is the acceptance-side mechanism behind issue 16377 for gpt-oss / Harmony.

This is a spec-decode acceptance issue, not a model/parser bug: the identical well-formed draft yields a complete call without speculative decoding and an empty tool_calls with it, differing only in where the draft/accept boundary falls.

System Info

• TensorRT-LLM branch: main @ 937bacc
• Backend: PyTorch (trtllm-serve --backend pytorch), one-model speculative decoding (Eagle one-model / MTP Eagle path)
• GPU: NVIDIA RTX 5090 (sm_120), 32 GB (local environment)
• Python: 3.12
• OS: Ubuntu 24.04 (WSL2). Note: this reproduction is a code-review plus deterministic numpy replay of the documented chain_speculative_sampling decision rule; a full end-to-end tensorrt_llm run requires HF weights and is tracked as a follow-up.

Steps to reproduce / Reproduction

Minimal numpy-only, dependency-free replay (fixed seed => identical output everywhere):

VOCAB mapping for the toy control preamble: 6=MSG, 7=TOK, 8=END, golden=CH. OPEN tokens open a control block (MSG=6); CLOSE tokens close it (END=8). A tool call is complete only if a block both opens and closes.

The replay models chain_speculative_sampling: golden token is sampled from target; draft token k is accepted with probability min(1, P_target/P_draft); the first rejection re-samples from target and terminates.

CASE A (baseline) - target fully agrees with draft:
draft : MSG TOK END
emit : CH MSG TOK END complete=True
result: {"tool_calls":[{"name":"call_fn"}], "content":"some", "finish_reason":"stop"}

CASE B (bug) - target disfavours the interior TOK slot (favours ordinary w3):
draft : MSG TOK END
emit : CH MSG w3 complete=False
result: {"tool_calls":[], "content":"", "finish_reason":"stop"}

A single interior rejection turns an intended tool call into a silently empty one, with no signal to the caller.

actual behavior

With speculative decoding enabled, a rejected interior control token leaves a structurally incomplete, unclosed control prefix while finish_reason=stop and empty tool_calls/content are returned. No warning is emitted.

Expected behavior

Acceptance must not be able to split a structurally required control sequence. Options that each restore a correctness invariant:
(a) on rejection inside a control block, roll back to before the block opener and let the target re-emit cleanly;
(b) expose a structured-parse validation hook so the executor rejects a partial block instead of returning finish_reason=stop with empty content;
(c) for tool-call channels, gate speculative acceptance on structural completeness.

Root cause (code)

1. The acceptance decision is delegated to flashinfer.chain_speculative_sampling; rejection_sampling_one_model trims the kernel result to exactly (accepted_tokens, output_emitted_draft_token_num + 1). No value describes structural balance.
   • tensorrt_llm/_torch/speculative/interface.py:53 rejection_sampling_one_model
   • tensorrt_llm/_torch/speculative/interface.py:1645 _sample_and_accept_draft_tokens_rejection
2. The accepted prefix is committed directly as output and the draft loop advances only by position using numeric num_accepted_draft_tokens:
   • tensorrt_llm/_torch/speculative/drafting_loops.py:102 LinearDraftingLoopWrapper
   • drafting_loops.py:171 last_tokens_idx = cumsum(seq_len) - seq_len + num_accepted_draft_tokens
     No token content is examined.
3. A rejection truncates exactly at the numeric prefix (new_tokens = accepted_tokens[:num_accepted]), which can leave the emitted prefix structurally unbalanced.

Start with the fixed-seed NumPy replay, then read rejection_sampling_one_model in tensorrt_llm/_torch/speculative/interface.py and _sample_and_accept_draft_tokens_rejection at line 1645. Trace how LinearDraftingLoopWrapper in drafting_loops.py commits accepted_tokens and advances last_tokens_idx. Done means confirming or refuting the prefix-emission and rollback claims and identifying the structural invariant that must be preserved.