From 6c298be8423c2dd2af74273982e5fed9360ce07e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 18:47:39 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20use=20new=5Ftensor=20instead=20of=20torc?= =?UTF-8?q?h.tensor=20for=20cudagraph=20(no=20CPU=E2=86=92CUDA=20copy)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torch.tensor() creates on CPU then copies to CUDA, which is forbidden during cudagraph capture. new_tensor() creates directly on the source tensor's device. --- cutedsl/bridge.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cutedsl/bridge.py b/cutedsl/bridge.py index b96211b7..8a58edb3 100644 --- a/cutedsl/bridge.py +++ b/cutedsl/bridge.py @@ -85,7 +85,7 @@ def quantize_to_nvfp4(x_bf16, block_size=SF_VEC_SIZE): abs_scaled = x_scaled.abs().clamp(max=6.0) half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8) - step_to_idx = torch.tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], dtype=torch.int8, device=x_bf16.device) + step_to_idx = x_bf16.new_tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], dtype=torch.int8) indices = step_to_idx[half_steps.long()] nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8) @@ -141,7 +141,7 @@ def quantize_activation_nvfp4(x_bf16, global_scale, block_size=SF_VEC_SIZE): abs_scaled = x_scaled.abs().clamp(max=6.0) half_steps = (abs_scaled * 2.0).round().clamp(0, 12).to(torch.int8) - step_to_idx = torch.tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], dtype=torch.int8, device=x_bf16.device) + step_to_idx = x_bf16.new_tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], dtype=torch.int8) indices = step_to_idx[half_steps.long()] nibbles = torch.where(signs < 0, indices + 8, indices).to(torch.uint8)