diff --git a/cutedsl/native_swa_decode.py b/cutedsl/native_swa_decode.py
index 800628fc..ccab3133 100644
--- a/cutedsl/native_swa_decode.py
+++ b/cutedsl/native_swa_decode.py
@@ -94,6 +94,10 @@ def _fallback_batched_sdp(
     num_tokens, NH, HD = q.shape
     device = q.device
 
+    # swa_indices may be 3D (batch, tokens, window) — squeeze batch dim
+    if swa_indices.dim() == 3:
+        swa_indices = swa_indices.squeeze(0)  # (tokens, window)
+
     safe_indices = swa_indices[:num_tokens].clamp(min=0)
     block_indices = safe_indices // block_size
     offsets = safe_indices % block_size
@@ -116,7 +120,7 @@ def _fallback_batched_sdp(
     # Batched SDPA
     q_t = q.permute(1, 0, 2)
     q_batch = q_t.reshape(NH * num_tokens, 1, HD)
-    kv_expanded = kv_bf16.unsqueeze(0).expand(NH, -1, -1, -1)
+    kv_expanded = kv_bf16.unsqueeze(0).expand(NH, num_tokens, window_size, HD)
     k_batch = kv_expanded.reshape(NH * num_tokens, window_size, HD)
     v_batch = k_batch