monster/corridorkey-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: replace manual attention with mx.fast.scaled_dot_product_attention

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
cmoyates 2026-03-09 17:04:51 -02:30
parent 6787a7fd21
commit f80bec92b6
No known key found for this signature in database
GPG Key ID: F65E08480FDC22D2
1 changed files with 20 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
src/corridorkey_mlx/model/hiera.py
View File

@ -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))