From 6e1b52eedf4b273374bbcda983291df1c8a3a513 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Tue, 25 Nov 2025 21:42:55 -0800 Subject: [PATCH] change allreduce workspace dtype to torch.int64 to avoid overflow Signed-off-by: Zhenhuan Chen --- tensorrt_llm/plugin/plugin.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/plugin/plugin.py b/tensorrt_llm/plugin/plugin.py index 60e12e982071..f3da8341d940 100644 --- a/tensorrt_llm/plugin/plugin.py +++ b/tensorrt_llm/plugin/plugin.py @@ -738,7 +738,7 @@ def allocate_allreduce_fusion_workspace( 3 * lamport_buffers_size, ) flag_buffer = torch.tensor([0, 0, 0, lamport_buffers_size, 0], - dtype=torch.int, + dtype=torch.int64, device="cuda") buffers = [ipc_buffers, ipc_barriers, lamport_buffers, flag_buffer]