From f567c20539030dac7b35d3321878a6234881fe6d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 30 May 2026 23:39:45 +0000 Subject: [PATCH] Fix: set active CUDA device per layer for BMM/FMHA --- single_shot_inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 07571807..1fb8ed80 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -366,9 +366,10 @@ def main(): target_gpu = li % NUM_GPUS target_device = f"cuda:{target_gpu}" - # Move activation to layer's GPU + # Move activation to layer's GPU and set as active device if x.device != torch.device(target_device): x = x.to(target_device) + torch.cuda.set_device(target_gpu) lw = layer_weights[li] rc, rs = rope_caches[target_gpu]