diff --git a/single_shot_inference.py b/single_shot_inference.py
index 0ae5dfc3..0661f8e2 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -432,6 +432,9 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
     ratio = compressor.ratio if compressor is not None else 0
     scale = 1.0 / math.sqrt(hd)
     pfx = f"model.layers.{li}.self_attn"
+    # Ensure positions is on the same device as rope caches
+    if positions.device != rope_cos.device:
+        positions = positions.to(rope_cos.device)
 
     # 1. Q projection: q_a → q_a_norm → q_b → q_b_norm
     q_a = do_nvfp4_linear(x_normed, w, pfx, 'q_a_proj')