#!/usr/bin/env python3
"""Repro: LogitsProcessor receives stale (non-ancestral) per-beam token histories
under beam-search parent switches on the PyTorch backend.
Mechanism
---------
A LogitsProcessor overwrites logits to force a deterministic beam trajectory
(beam width 2, 3 generation steps):
step 1 (context, 1 row): force top-2 = {A1, B1} -> beams [A1], [B1]
step 2 (2 rows): row(A1): force {A2, A2'};
row(B1): force nothing (uniform, far below A's)
-> BOTH survivors descend from parent A1:
beams become [A1, A2] and [A1, A2'] (parent switch:
the slot that held [B1] now continues A1's lineage)
step 3 (2 rows): record the `token_ids` the callback receives, then
force each beam to continue itself with a distinct
token (C1 after A2, C2 after A2') so both lineages
survive into the final output for verification.
Contract (tensorrt_llm/sampling_params.py, LogitsProcessor): token_ids is the
per-beam token history of the request. After the forced switch, the beam whose
last token is A2' has true history [..., A1, A2'].
Observed (bug): the callback sees [..., B1, A2'] -- the newest token appended
onto the PREVIOUS occupant of that beam slot. The final request outputs ARE
correct (histories are repaired at finalization), which is exactly why any
sequence-level constraint applied via LogitsProcessor silently mis-masks at
beam_width > 1: it acts on histories that are only fixed up later.
Run:
python3 repro_stale_beam_history.py [--model Qwen/Qwen3-0.6B]
"""
import argparse
import json
import os
from tensorrt_llm import LLM
import tensorrt_llm
from tensorrt_llm.llmapi import KvCacheConfig
from tensorrt_llm.sampling_params import SamplingParams
# Arbitrary distinct token ids used to force the trajectory.
A1, B1, A2, A2P, C1, C2 = 1000, 2000, 3000, 4000, 5000, 6000
PROMPT = [100, 200, 300]
class ForcingRecorder:
"""Forces the beam trajectory above and records what the callback sees.
The LLM API executes the callback in a worker process, so observations are
persisted to a JSONL file for the driver process to read back.
"""
def __init__(self, record_path):
self.record_path = record_path
self.calls = [] # worker-local
def __call__(self, req_id, logits, token_ids, stream_ptr, client_id):
# logits: [beam_width, 1, vocab]; token_ids: List[beam][tokens]
self.calls.append([list(beam) for beam in token_ids])
step = len(self.calls)
with open(self.record_path, "a") as f:
f.write(json.dumps({"call": step, "token_ids": self.calls[-1]}) + "\n")
rows = logits.view(logits.shape[0], -1)
if step == 1:
# Context step: single row. Force top-2 = {A1 (best), B1 (second)}.
rows[0].zero_()
rows[0, A1] = 20.0
rows[0, B1] = 18.0
elif step == 2:
for b in range(rows.shape[0]):
last = token_ids[b][-1]
rows[b].zero_()
if last == A1:
# Parent A offers the two best continuations overall.
rows[b, A2] = 20.0
rows[b, A2P] = 18.0
# Parent B: left uniform -> its best candidate scores ~ -log(V),
# far below A's two candidates -> both survivors come from A.
elif step == 3:
# Recorded above; now force each beam to continue itself with a
# distinct token so both step-2 lineages survive to the output.
for b in range(rows.shape[0]):
last = token_ids[b][-1]
rows[b].zero_()
if last == A2:
rows[b, C1] = 20.0
elif last == A2P:
rows[b, C2] = 20.0
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", default="Qwen/Qwen3-0.6B")
args = parser.parse_args()
print(f"tensorrt_llm version: {tensorrt_llm.__version__}")
record_path = os.path.abspath("stale_beam_calls.jsonl")
if os.path.exists(record_path):
os.remove(record_path)
proc = ForcingRecorder(record_path)
llm = LLM(
model=args.model,
max_beam_width=2,
max_batch_size=2,
disable_overlap_scheduler=True,
kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.3),
)
sampling = SamplingParams(
max_tokens=3,
n=2,
best_of=2,
use_beam_search=True,
logits_processor=proc,
ignore_eos=True,
detokenize=False,
)
out = llm.generate([{"prompt_token_ids": PROMPT}], sampling)[0]
llm.shutdown()
final = [list(seq.token_ids) for seq in out.outputs]
print(f"\nfinal (finalized/gathered) beams: {final}")
with open(record_path) as f:
calls = [json.loads(line)["token_ids"] for line in f]
print(f"callback calls: {len(calls)}")
for i, call in enumerate(calls, 1):
print(f" call {i} token_ids (generated part): "
f"{[hist[len(PROMPT):] for hist in call]}")
# 1) Confirm the forced parent switch actually happened: the finalized
# beams must be exactly the two forced lineages, both descending from A1.
assert sorted(final) == [[A1, A2, C1], [A1, A2P, C2]], (
f"forcing failed, got {final}; cannot evaluate the contract"
)
print("\nforced parent switch confirmed: both surviving beams descend from A1")
# 2) What did the step-3 callback see for the switched beam?
assert len(calls) == 3, f"expected 3 callback calls, got {len(calls)}"
step3 = calls[2]
switched = next((h for h in step3 if h[-1] == A2P), None)
assert switched is not None, (
f"no step-3 callback row ends in A2'={A2P}; rows seen: {step3} "
"(on builds where the callback receives only beam 0, this repro "
"cannot run -- it requires the per-beam callback of recent main)"
)
prev_token = switched[-2]
truth, stale = A1, B1
print(f"\nbeam ending in A2'={A2P}: previous token seen by callback = {prev_token}")
print(f" ancestral truth: {truth} (A1) stale slot buffer: {stale} (B1)")
if prev_token == stale:
print(
"\nCONTRACT VIOLATION: the callback received the new token appended onto "
"the previous slot occupant's lineage, not this beam's ancestral history. "
"Any sequence-level constraint (trie/grammar) driven by token_ids "
"silently mis-masks at beam_width > 1."
)
elif prev_token == truth:
print("\nOK: callback histories are ancestrally correct on this build.")
else:
print("\nUNEXPECTED: neither stale nor true value; investigate.")
if __name__ == "__main__":
main()
System Info
main@8e7b98c09b, built from source (build_wheel.py -a 80-real)Qwen/Qwen3-0.6B(any model reproduces; the repro forces the trajectory)Summary
On the PyTorch backend, the per-request
LogitsProcessor(
tensorrt_llm/sampling_params.py) is invoked with per-beamtoken_ids(
model_engine.py:request.get_tokens(beam_idx)for each beam). Under beamsearch, those histories are append-only per beam slot: when beam search
switches a slot's parent, the callback receives the newly committed token
appended onto the previous occupant's lineage, not the beam's true
ancestral history. The ancestral repair (
_gather_beam_path→set_generated_tokens) runs only at finalization, so the final outputs arecorrect — which makes the mid-generation staleness easy to miss.
Consequence: any sequence-level constraint driven by
token_ids(trie /grammar / DFA-style processors needing ≥2 tokens of history) silently masks
against the wrong history at
beam_width > 1. Last-token-only processors areunaffected. Nothing validates or rejects the combination, so it fails silently.
Reproduction
Deterministic, single-GPU, ~170-line script (full script below:
repro_stale_beam_history.py). ALogitsProcessorforces a beam-2 trajectorywith a guaranteed parent switch at step 2, records what the callback receives,
and cross-checks against the finalized outputs:
At step 2 both survivors were forced to descend from beam 0 (
A1=1000), asthe finalized outputs confirm (
[1000, 3000, ...]and[1000, 4000, ...]).Yet the step-3 callback sees beam slot 1 as
[2000, 4000]— the new token4000appended to the deadB1=2000lineage instead of the true[1000, 4000].Mechanism (code references @ 8e7b98c)
tensorrt_llm/_torch/pyexecutor/model_engine.py(_execute_logit_post_processors,~L5880-5925): passes
request.get_tokens(beam_idx)per beam and the[beam_width, 1, vocab]logits slice to the callback, pre-sampling.tensorrt_llm/_torch/pyexecutor/sampler/sampler.py: new tokens are appendedraw per (slot, beam) —
_update_original_tokens(~L2597-2611) andadd_token→request.add_new_token(new_token, beam_idx)(~L775-783) — withno per-step lineage reordering.
predecessor_beams/cache_indirectioninBeamSearchStore(~L2197-2227), updated each step inbeam_search_sampling_batch(sampler/ops/vanilla.py~L234-273). But it isapplied to token histories only at finalization:
_gather_beam_path(~L1141-1147) inside
_prepare_beam_history(~L3225-3370), whose docstringstates "a beam may change its previously sampled tokens. This function
corrects the stored tokens."
Related: the PyTorch
GuidedDecoder(
tensorrt_llm/_torch/pyexecutor/guided_decoder.py) has the same exposure —one grammar matcher per request advanced by
get_last_tokens(0)(beam 0 only)and bitmask rows sized without
× beam_width— and beam search + guideddecoding is likewise not rejected anywhere.
Expected behavior
Either:
through
predecessor_beams/cache_indirectionfor requests withregistered processors when
beam_width > 1— the machinery exists and isalready used at finalization), or
(assert/validation) instead of silently mis-masking; same for guided
decoding + beam search.
Questions / offer
beam search intended to be supported?
(b) an ancestry-aware per-beam constraint hook in the TorchSampler beam step
(mask applied to the
[batch, beam_in, vocab]logprobs before theflatten/topk, alongside the existing finished-beams masking) — we run a
production trie-constrained beam-search workload (beam width 128) and have a
working per-beam constraint implementation (currently downstream, on the
removed TRT-engine path) that we would like to bring to the maintained
backend.
Repro script
repro_stale_beam_history.pyFull run output (main @ 8e7b98c)