diff --git a/src/corridorkey_mlx/model/hiera.py b/src/corridorkey_mlx/model/hiera.py
index 1533270..46a5809 100644
--- a/src/corridorkey_mlx/model/hiera.py
+++ b/src/corridorkey_mlx/model/hiera.py
@@ -284,10 +284,26 @@ class MaskUnitAttention(nn.Module):
             q = q.reshape(batch_size, self.heads, num_windows, self.q_stride, -1, self.head_dim)
             q = mx.max(q, axis=3)
 
-        # Scaled dot-product attention (manual implementation)
-        attn = (q * self.scale) @ mx.transpose(k, axes=(0, 1, 2, 4, 3))
-        attn = mx.softmax(attn, axis=-1)
-        x = attn @ v
+        # Fold window dim into batch for SDPA (expects 4D).
+        # Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows
+        # next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D).
+        q_tokens = q.shape[3]
+        kv_tokens = k.shape[3]
+        q_4d = mx.transpose(q, axes=(0, 2, 1, 3, 4)).reshape(
+            batch_size * num_windows, self.heads, q_tokens, self.head_dim
+        )
+        k_4d = mx.transpose(k, axes=(0, 2, 1, 3, 4)).reshape(
+            batch_size * num_windows, self.heads, kv_tokens, self.head_dim
+        )
+        v_4d = mx.transpose(v, axes=(0, 2, 1, 3, 4)).reshape(
+            batch_size * num_windows, self.heads, kv_tokens, self.head_dim
+        )
+
+        x = mx.fast.scaled_dot_product_attention(q_4d, k_4d, v_4d, scale=self.scale)
+
+        # Unfold: (B*W, H, T, D) -> (B, W, H, T, D) -> (B, H, W, T, D)
+        x = x.reshape(batch_size, num_windows, self.heads, q_tokens, self.head_dim)
+        x = mx.transpose(x, axes=(0, 2, 1, 3, 4))
 
         # [B, heads, num_windows, tokens, head_dim] -> [B, tokens, num_windows, heads, head_dim]
         # -> [B, N', dim_out]  (matches PyTorch transpose(1, 3))