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

[#6425][fix] address CUDA stream sync issue in ModelRunnerCPP (#6426)

Browse Source 
Signed-off-by: yifei.w <yifei.w@bytedance.com>
This commit is contained in:
Yifei Wang 2025-12-11 21:33:22 -08:00 committed by GitHub
parent fded6c393d
commit 9d1f2a9925
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorrt_llm/runtime/model_runner.py
View File

@ -473,6 +473,7 @@ class ModelRunnerMixin:
prompt_table,
torch.Tensor), "Prompt table should be str or torch.Tensor"
prompt_table_data = prompt_table.to(dtype=self.dtype)
torch.cuda.current_stream().synchronize()
return prompt_table_data