[Bug]: VBWS returns empty outputs when generation ends on a widening step

System Info

• CPU architecture: x86_64
• GPU: NVIDIA RTX PRO 5000 72GB Blackwell
• NVIDIA driver: 580.126.09
• TensorRT-LLM: 1.3.0rc26
• Backend: PyTorch, single GPU
• PyTorch: 2.12.0+cu130
• Transformers: 5.5.4
• Python: 3.12.7
• OS: Linux 5.10.134, glibc 2.39
• Overlap scheduler: enabled (default)

Who can help?

@NVIDIA/trt-llm-devs

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (the problem is model- and dataset-independent after logits are produced)

Reproduction

Variable-Beam-Width Search returns the requested number of output slots, but many outputs have empty token_ids when generation terminates on a widening step.

A minimal LLM API setup is:

import os

from tensorrt_llm import LLM, SamplingParams

model = os.environ["MODEL_PATH"]  # Any locally available causal LM; observed with Qwen3.
llm = LLM(
    model=model,
    tokenizer=model,
    dtype="bfloat16",
    max_batch_size=1,
    max_input_len=128,
    max_seq_len=131,
    max_beam_width=100,
)

sampling_params = SamplingParams(
    n=100,
    best_of=100,
    beam_width_array=[16, 32, 100],
    use_beam_search=True,
    max_tokens=3,
    temperature=0.0,
    ignore_eos=True,
)

request_output = llm.generate(["Hello"], sampling_params)[0]
lengths = [len(output.token_ids) for output in request_output.outputs]
print(len(request_output.outputs))
print({length: lengths.count(length) for length in set(lengths)})

assert len(request_output.outputs) == 100
assert all(length == 3 for length in lengths)

With lightweight instrumentation around beam_search_sampling_batch_cba, the expected (beam_width_in, beam_width_out) transitions are:

(1, 16), (16, 32), (32, 100)

With the default overlap scheduler, the observed transitions are:

(1, 16), (16, 16), (16, 32)

The final step creates 32 beams, but _prepare_beam_history_cba() then collects only 16 active paths. In five repeated requests, the final output length histogram was consistently:

{3: 16, 0: 84}

The same request with a constant schedule [100, 100, 100] returns 100 non-empty three-token outputs.

Set TLLM_DEBUG_MODE=1 for additional logs.

Expected behavior

For beam_width_array=[16, 32, 100] and max_tokens=3:

1. Sampling should execute widths 1 -> 16, 16 -> 32, and 32 -> 100.
2. The final response should contain 100 outputs.
3. Every output should contain three generated tokens; no output should have an empty token_ids list.

actual behavior

The overlap path repeats the first configured width for the second sampling step, so the final step widens only from 16 to 32. Finalization then slices the current-step snapshot using the input width (16), dropping the additional 16 paths. The remaining output slots stay padding and are exposed as empty token_ids.

additional notes

This appears to be two issues in the Python PyExecutor VBWS path:

1. LlmRequest.get_beam_width_by_iter() indexes beam_width_array with self.decoding_iter. Under the overlap scheduler, the current Python sampling state is tracked by py_decoding_iter, while decoding_iter is synchronized later during response handling. This makes the width schedule lag by one step.

2. _prepare_beam_history_cba() computes:

   active_width = _get_beam_width_in(request)
   

   However, the snapshot being finalized already contains the current step's output beams. On a widening terminal step, it should collect the produced/output width, not the input width.

A local diagnostic patch that:

• indexes the Python width schedule with py_decoding_iter, and
• finalizes using request.get_beam_width_by_iter(for_next_iteration=True)

restores the expected transitions and produces 100 non-empty three-token outputs for all tested prompts. Applying only either change produces 32 valid outputs, while applying both produces all 100.

The relevant upstream implementation was introduced/completed in #16620. Current main still appears to use decoding_iter for the Python width lookup and _get_beam_width_in(request) for CBA finalization.

No private model paths, prompts, endpoints, or datasets are included in this report.

Before submitting a new issue...

• I searched existing and past issues, documentation, and examples for the same problem.

Start with the Python PyExecutor path, especially LlmRequest.get_beam_width_by_iter(), _prepare_beam_history_cba(), and beam_search_sampling_batch_cba(); compare decoding_iter with py_decoding_iter under the overlap scheduler. Run the minimal LLM API reproduction with TLLM_DEBUG_MODE=1 and verify the transitions reach 100 beams and all 100 outputs contain three tokens for the widening schedule.