From 7c3bb8534d230ecdd183344309fa81f0df2e13d1 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Fri, 28 Nov 2025 16:45:23 +0800 Subject: [PATCH] =?UTF-8?q?[None][chore]=20Revert=20"[None][fix]=20change?= =?UTF-8?q?=20allreduce=20workspace=20dtype=20to=20torch.int64=20t?= =?UTF-8?q?=E2=80=A6=20(#9538)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 f3da8341d9..60e12e9820 100644 --- a/tensorrt_llm/plugin/plugin.py +++ b/tensorrt_llm/plugin/plugin.py @@ -738,7 +738,7 @@ class CustomAllReduceHelper: 3 * lamport_buffers_size, ) flag_buffer = torch.tensor([0, 0, 0, lamport_buffers_size, 0], - dtype=torch.int64, + dtype=torch.int, device="cuda") buffers = [ipc_buffers, ipc_barriers, lamport_buffers, flag_buffer]