From f80bec92b6c715e231fac42c85ce491a2acb7f16 Mon Sep 17 00:00:00 2001 From: cmoyates Date: Mon, 9 Mar 2026 17:04:51 -0230 Subject: [PATCH] feat: replace manual attention with mx.fast.scaled_dot_product_attention Fold window dim into batch (transpose before reshape), call SDPA, unfold back. All parity tests pass with 0 regression. Co-Authored-By: Claude Opus 4.6 --- src/corridorkey_mlx/model/hiera.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) 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))