kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-13 22:18:36 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[https://nvbugs/5772521][fix] Fix draft token tree chain crash (#10386)

Browse Source 
Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com>
This commit is contained in:
Mike Iovine 2026-01-05 11:18:44 -05:00 committed by GitHub
parent e98c27ee4f
commit bedfff4f00
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorrt_llm/_torch/pyexecutor/sampler.py
View File

@ -1298,14 +1298,13 @@ class TorchSampler(Sampler, AsyncWorkerMixin):
# Take the longest accepted path as the next new token.
num_accepted_draft_tokens = 0
for idx in eagle_paths[longest_match_path_idx][:longest_accepted_len]:
add_token(
request, new_tokens_list, beam_idx=DEFAULT_BEAM_IDX, step=cast(int, idx.item())
)
step = cast(int, idx.item())
add_token(request, new_tokens_list, beam_idx=DEFAULT_BEAM_IDX, step=step)
num_accepted_draft_tokens += 1
if self.finish_if_reason(
request,
finish_reasons,
step=num_accepted_draft_tokens,
step=step,
beam_idx=DEFAULT_BEAM_IDX,
):
break