Fix: set active CUDA device per layer for BMM/FMHA

2026-05-30 23:39:45 +00:00
parent 7a95983e0f
commit f567c20539
1 changed files with 2 additions and 1 deletions
--- 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]