146 lines
4.9 KiB
Python
146 lines
4.9 KiB
Python
"""
|
|
Multi-head attention for MLX backend.
|
|
Fused scaled dot-product attention via Metal kernel, with QK-RMSNorm and RoPE.
|
|
|
|
Supports two modes:
|
|
- Sparse (unbatched): x is (N, C) — used for sparse flow models
|
|
- Dense (batched): x is (B, N, C) — used for dense structure flow with batched CFG
|
|
"""
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
from .norm import SparseMultiHeadRMSNorm
|
|
from .rope import build_rope_freqs, compute_rope_phases, apply_rope
|
|
|
|
|
|
class MlxMultiHeadAttention(nn.Module):
|
|
"""
|
|
Multi-head attention matching the PyTorch SparseMultiHeadAttention.
|
|
|
|
For self-attention: fused QKV projection, optional QK-RMSNorm, optional RoPE.
|
|
For cross-attention: separate Q and KV projections.
|
|
|
|
Supports both (N, C) unbatched and (B, N, C) batched input.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
num_heads: int,
|
|
ctx_channels: int = None,
|
|
type: str = "self",
|
|
qkv_bias: bool = True,
|
|
use_rope: bool = False,
|
|
rope_freq: tuple = (1.0, 10000.0),
|
|
qk_rms_norm: bool = False,
|
|
):
|
|
super().__init__()
|
|
assert channels % num_heads == 0
|
|
self.channels = channels
|
|
self.num_heads = num_heads
|
|
self.head_dim = channels // num_heads
|
|
self.ctx_channels = ctx_channels or channels
|
|
self._type = type
|
|
self.use_rope = use_rope
|
|
self.qk_rms_norm = qk_rms_norm
|
|
|
|
if type == "self":
|
|
self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
|
|
else:
|
|
self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
|
|
self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
|
|
|
|
if qk_rms_norm:
|
|
self.q_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads)
|
|
self.k_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads)
|
|
|
|
self.to_out = nn.Linear(channels, channels)
|
|
|
|
if use_rope:
|
|
self._rope_freqs = build_rope_freqs(self.head_dim, dim=3, rope_freq=rope_freq)
|
|
|
|
def __call__(
|
|
self,
|
|
x: mx.array,
|
|
context: mx.array = None,
|
|
rope_cache: tuple = None,
|
|
) -> mx.array:
|
|
"""
|
|
Args:
|
|
x: (N, C) or (B, N, C) input features
|
|
context: (M, ctx_C) or (B, M, ctx_C) cross-attention context
|
|
rope_cache: (cos, sin) precomputed RoPE phases
|
|
|
|
Returns:
|
|
Same shape as x
|
|
"""
|
|
H = self.num_heads
|
|
D = self.head_dim
|
|
batched = x.ndim == 3
|
|
|
|
if self._type == "self":
|
|
qkv = self.to_qkv(x) # (..., 3*C)
|
|
if batched:
|
|
B, N, _ = qkv.shape
|
|
qkv = qkv.reshape(B, N, 3, H, D)
|
|
q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] # (B, N, H, D)
|
|
else:
|
|
qkv = qkv.reshape(-1, 3, H, D)
|
|
q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] # (N, H, D)
|
|
|
|
if self.qk_rms_norm:
|
|
q = self.q_rms_norm(q)
|
|
k = self.k_rms_norm(k)
|
|
|
|
if self.use_rope and rope_cache is not None:
|
|
cos, sin = rope_cache
|
|
q = apply_rope(q, cos, sin)
|
|
k = apply_rope(k, cos, sin)
|
|
|
|
out = self._sdpa(q, k, v, batched)
|
|
else:
|
|
q = self.to_q(x)
|
|
kv = self.to_kv(context)
|
|
if batched:
|
|
B, N, _ = q.shape
|
|
q = q.reshape(B, N, H, D)
|
|
M = kv.shape[1]
|
|
kv = kv.reshape(B, M, 2, H, D)
|
|
k, v = kv[:, :, 0], kv[:, :, 1] # (B, M, H, D)
|
|
else:
|
|
q = q.reshape(-1, H, D)
|
|
kv = kv.reshape(-1, 2, H, D)
|
|
k, v = kv[:, 0], kv[:, 1]
|
|
|
|
if self.qk_rms_norm:
|
|
q = self.q_rms_norm(q)
|
|
k = self.k_rms_norm(k)
|
|
|
|
out = self._sdpa(q, k, v, batched)
|
|
|
|
if batched:
|
|
out = out.reshape(B, -1, self.channels)
|
|
else:
|
|
out = out.reshape(-1, self.channels)
|
|
return self.to_out(out)
|
|
|
|
def _sdpa(self, q: mx.array, k: mx.array, v: mx.array, batched: bool = False) -> mx.array:
|
|
"""
|
|
Fused scaled dot-product attention via Metal kernel.
|
|
Inputs: (N, H, D) unbatched or (B, N, H, D) batched.
|
|
"""
|
|
scale = self.head_dim ** -0.5
|
|
if batched:
|
|
# (B, N, H, D) -> (B, H, N, D)
|
|
q = q.transpose(0, 2, 1, 3)
|
|
k = k.transpose(0, 2, 1, 3)
|
|
v = v.transpose(0, 2, 1, 3)
|
|
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
|
|
return out.transpose(0, 2, 1, 3) # (B, N, H, D)
|
|
else:
|
|
# (N, H, D) -> (1, H, N, D)
|
|
q = q.transpose(1, 0, 2)[None]
|
|
k = k.transpose(1, 0, 2)[None]
|
|
v = v.transpose(1, 0, 2)[None]
|
|
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
|
|
return out[0].transpose(1, 0, 2) # (N, H, D)
|