Summary
The LogitsPostProcessor callback's per-beam token histories (BeamTokens, cpp/include/tensorrt_llm/executor/types.h:54,62-63) go stale under beam-search parent switches. The documented contract — "the tokens produced by the request so far" (docs/source/legacy/advanced/executor.md:67), delivered per beam — is violated mid-generation whenever a surviving hypothesis descends from a different parent beam: the callback sees the newest token correctly, stapled onto the previous occupant's lineage in that beam slot.
Consequence: any logits post-processor whose decision depends on ≥2 tokens of history is silently wrong at beamWidth > 1. Last-token-only processors are safe. The final request outputs are correct (gatherTree at finalize repairs lineages), which makes the mid-flight staleness easy to miss — the constraint just quietly leaks.
Root cause (code level)
Per-step token updates are append-only per-(slot, beam) buffers (updateRequests → addNewToken); beam parent relationships live separately as decoder parent pointers and are only applied lazily by postProcessRequest → decoder finalize(seqSlot, ...) → gathered ids — normally only at request completion. Mid-generation, nothing reorders the slot buffers after a parent switch, but the callback (and C++ llmReq->getTokens() / getLastTokens()) read exactly those buffers.
Evidence
Two controlled, weights-independent experiments (executor bindings, enc-dec engine):
- Singleton-mask test (exactly one allowed token per row per step → pure lineages, no parent switches): callback
token_ids are correct and aligned every step. Bookkeeping is fine when no switch occurs.
- Forced-parent-switch test (scores manipulated via the mask so both survivors descend from one parent): the surviving slot's
token_ids show the new token on top of the stale old lineage ([2, B1, A2'] instead of the true [2, A1, A2']), while the request's final output is correctly parent-gathered.
Measured end-to-end impact, per-beam trie constraint implemented in a LogitsPostProcessor (non-streaming, 20 queries, trained enc-dec engine):
| beamWidth |
valid constrained beams |
| 8 |
123/160 (77%) |
| 16 |
225/320 (70%) |
| 32 |
396/640 (62%) |
| 128 |
1206/2560 (47%) |
Validity collapses monotonically with beam width — exactly the signature of leakage on parent switches (all valid beams remain unique, so this is not duplicate collapse).
Confirming the diagnosis: with streaming=True and beamWidth > 1, trtGptModelInflightBatching already runs the non-destructive per-step gather (postProcessRequest on every step), which lands before the next step's callback — and the same processor then produces 100% valid beams (verified at beam 2/4/8/16). So the repair machinery exists; it just doesn't run for non-streaming requests.
Environment
TensorRT-LLM 1.3.0rc21 (d567f924b7), C++ executor, A100, fp16 encoder-decoder engine, beam search. The staleness mechanism is model-independent.
Questions for maintainers
- Fix the contract or document the limitation? A contract fix looks small: run the existing non-destructive per-step gather when
beamWidth > 1 and a logits post-processor is registered for the request (the streaming path proves it works and lands in time). Cost is per-step gatherTree + copies for affected requests only. Alternatively, document that BeamTokens beyond the last token are unreliable at beamWidth > 1 for non-streaming requests.
- Is guided decoding + beam search intended to be supported at all? Today the C++
GuidedDecoder (xgrammar) keeps one matcher per request and masks only beam 0's logits row, so grammar-constrained beam search silently constrains a single beam. If per-beam guided decoding is in scope, I have design notes and a working downstream implementation (per-beam state advanced through the decoder's parent pointers) I'd be glad to discuss.
I can provide a minimal repro script (forced-parent-switch with a tiny engine) on request, and I'm happy to contribute the gather-on-callback fix if that's the preferred direction.
Summary
The
LogitsPostProcessorcallback's per-beam token histories (BeamTokens,cpp/include/tensorrt_llm/executor/types.h:54,62-63) go stale under beam-search parent switches. The documented contract — "the tokens produced by the request so far" (docs/source/legacy/advanced/executor.md:67), delivered per beam — is violated mid-generation whenever a surviving hypothesis descends from a different parent beam: the callback sees the newest token correctly, stapled onto the previous occupant's lineage in that beam slot.Consequence: any logits post-processor whose decision depends on ≥2 tokens of history is silently wrong at
beamWidth > 1. Last-token-only processors are safe. The final request outputs are correct (gatherTree at finalize repairs lineages), which makes the mid-flight staleness easy to miss — the constraint just quietly leaks.Root cause (code level)
Per-step token updates are append-only per-(slot, beam) buffers (
updateRequests→addNewToken); beam parent relationships live separately as decoder parent pointers and are only applied lazily bypostProcessRequest→ decoderfinalize(seqSlot, ...)→ gathered ids — normally only at request completion. Mid-generation, nothing reorders the slot buffers after a parent switch, but the callback (and C++llmReq->getTokens()/getLastTokens()) read exactly those buffers.Evidence
Two controlled, weights-independent experiments (executor bindings, enc-dec engine):
token_idsare correct and aligned every step. Bookkeeping is fine when no switch occurs.token_idsshow the new token on top of the stale old lineage ([2, B1, A2']instead of the true[2, A1, A2']), while the request's final output is correctly parent-gathered.Measured end-to-end impact, per-beam trie constraint implemented in a
LogitsPostProcessor(non-streaming, 20 queries, trained enc-dec engine):Validity collapses monotonically with beam width — exactly the signature of leakage on parent switches (all valid beams remain unique, so this is not duplicate collapse).
Confirming the diagnosis: with
streaming=TrueandbeamWidth > 1,trtGptModelInflightBatchingalready runs the non-destructive per-step gather (postProcessRequeston every step), which lands before the next step's callback — and the same processor then produces 100% valid beams (verified at beam 2/4/8/16). So the repair machinery exists; it just doesn't run for non-streaming requests.Environment
TensorRT-LLM 1.3.0rc21 (
d567f924b7), C++ executor, A100, fp16 encoder-decoder engine, beam search. The staleness mechanism is model-independent.Questions for maintainers
beamWidth > 1and a logits post-processor is registered for the request (the streaming path proves it works and lands in time). Cost is per-step gatherTree + copies for affected requests only. Alternatively, document thatBeamTokensbeyond the last token are unreliable atbeamWidth > 1for non-streaming requests.GuidedDecoder(xgrammar) keeps one matcher per request and masks only beam 0's logits row, so grammar-constrained beam search silently constrains a single beam. If per-beam guided decoding is in scope, I have design notes and a working downstream implementation (per-beam state advanced through the decoder's parent pointers) I'd be glad to discuss.I can provide a minimal repro script (forced-parent-switch with a tiny engine) on request, and I'm happy to contribute the gather-on-callback fix if that's the preferred direction.