Sparse DiT blocks for the SLAT flows
SparseDiTAttention / SparseProjectAttention / ModulatedSparseTransformerCrossBlock - the sparse counterparts of the dense DiT pieces, attending within each batch item. The modulation needs care: it is per batch item ([B, 6C]) while features are a flat [N, C] stack, so each row must pick up its own item's shift/scale/gate. Broadcasting would silently apply item 0's modulation to everything when B == 1. Verified end to end: slat_flow matches upstream at correlation 1.00000000.
This commit is contained in:
John
parent
a44a0d7360
commit
93e31f9f80
@ -22,6 +22,9 @@ from .dit import (
|
||||
TimestepEmbedder,
|
||||
apply_rope,
|
||||
rope_phases_from_coords,
|
||||
SparseDiTAttention,
|
||||
SparseProjectAttention,
|
||||
ModulatedSparseTransformerCrossBlock,
|
||||
)
|
||||
from .tensor import SparseTensor, VarLenTensor, downsample, subdivide, upsample
|
||||
from .ops import (
|
||||
@ -46,5 +49,7 @@ __all__ = [
|
||||
# DiT / flow-model pieces
|
||||
"MultiHeadRMSNorm", "apply_rope", "rope_phases_from_coords", "TimestepEmbedder", "DiTAttention",
|
||||
"ModulatedTransformerCrossBlock", "ProjectAttention",
|
||||
"SparseDiTAttention", "SparseProjectAttention",
|
||||
"ModulatedSparseTransformerCrossBlock",
|
||||
]
|
||||
__version__ = "0.1.0"
|
||||
|
||||
@ -338,3 +338,168 @@ def rope_phases_from_coords(
|
||||
pad = mx.ones((n, want - phases.shape[-1]), dtype=mx.complex64)
|
||||
phases = mx.concatenate([phases, pad], axis=-1)
|
||||
return phases.astype(mx.complex64)
|
||||
|
||||
|
||||
# --------------------------------------------------------------- sparse variants
|
||||
# The SLAT flows are the same DiT, but their tokens are a SparseTensor's voxels rather
|
||||
# than a dense grid. Attention runs WITHIN each batch item (never across), and RoPE
|
||||
# positions come from the tensor's own coordinates.
|
||||
|
||||
|
||||
class SparseDiTAttention(nn.Module):
|
||||
"""DiT attention over a SparseTensor: per-batch-item, optional RoPE + qk RMS norm."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
num_heads: int,
|
||||
ctx_channels: Optional[int] = None,
|
||||
attn_type: str = "self",
|
||||
qkv_bias: bool = True,
|
||||
use_rope: bool = False,
|
||||
qk_rms_norm: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels, self.num_heads = channels, num_heads
|
||||
self.head_dim = channels // num_heads
|
||||
self.scale = self.head_dim**-0.5
|
||||
self._type = attn_type
|
||||
self.use_rope = use_rope
|
||||
self.qk_rms_norm = qk_rms_norm
|
||||
ctx = ctx_channels if ctx_channels is not None else channels
|
||||
if attn_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(ctx, channels * 2, bias=qkv_bias)
|
||||
if qk_rms_norm:
|
||||
self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
|
||||
self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
|
||||
self.to_out = nn.Linear(channels, channels)
|
||||
|
||||
def __call__(self, x, context=None, phases: Optional[mx.array] = None):
|
||||
n = x.feats.shape[0]
|
||||
h, d = self.num_heads, self.head_dim
|
||||
if self._type == "self":
|
||||
qkv = self.to_qkv(x.feats).reshape(n, 3, h, d)
|
||||
q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
|
||||
lq = lkv = x.layout
|
||||
else:
|
||||
cf = context.feats if hasattr(context, "feats") else context
|
||||
m = cf.shape[0]
|
||||
q = self.to_q(x.feats).reshape(n, h, d)
|
||||
kv = self.to_kv(cf).reshape(m, 2, h, d)
|
||||
k, v = kv[:, 0], kv[:, 1]
|
||||
lq = x.layout
|
||||
lkv = context.layout if hasattr(context, "layout") else [slice(0, m)]
|
||||
|
||||
# Same ordering rule as the dense path: RMS norm first, then rotate.
|
||||
if self.qk_rms_norm:
|
||||
q, k = self.q_rms_norm(q), self.k_rms_norm(k)
|
||||
if self.use_rope and phases is not None and self._type == "self":
|
||||
q4, k4 = apply_rope(q[None], k[None], phases)
|
||||
q, k = q4[0], k4[0]
|
||||
|
||||
outs = []
|
||||
for sq, skv in zip(lq, lkv):
|
||||
qi = q[sq].transpose(1, 0, 2)[None]
|
||||
ki = k[skv].transpose(1, 0, 2)[None]
|
||||
vi = v[skv].transpose(1, 0, 2)[None]
|
||||
o = mx.fast.scaled_dot_product_attention(qi, ki, vi, scale=self.scale)
|
||||
outs.append(o[0].transpose(1, 0, 2))
|
||||
o = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0)
|
||||
return x.replace(self.to_out(o.reshape(n, self.channels)))
|
||||
|
||||
|
||||
class SparseProjectAttention(nn.Module):
|
||||
"""Sparse `image_attn_mode="proj"`: cross-attend `global`, add projected `proj`."""
|
||||
|
||||
def __init__(self, cross_attn_block, channels: int, proj_in: int):
|
||||
super().__init__()
|
||||
self.cross_attn_block = cross_attn_block
|
||||
self.proj_linear = nn.Linear(proj_in, channels)
|
||||
|
||||
def __call__(self, x, context):
|
||||
if isinstance(context, dict):
|
||||
g, pr = context["global"], context["proj"]
|
||||
else:
|
||||
g, pr = context
|
||||
prf = pr.feats if hasattr(pr, "feats") else pr
|
||||
return x.replace(self.proj_linear(prf) + self.cross_attn_block(x, g).feats)
|
||||
|
||||
|
||||
class ModulatedSparseTransformerCrossBlock(nn.Module):
|
||||
"""Sparse counterpart of ModulatedTransformerCrossBlock.
|
||||
|
||||
Modulation is per BATCH ITEM but features are a flat [N, C] stack, so each row must
|
||||
pick up its own item's shift/scale/gate — hence the per-layout scatter rather than a
|
||||
simple broadcast. Broadcasting a [B, C] modulation against [N, C] would either fail
|
||||
or, worse, silently apply item 0's modulation to everything when B == 1.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
ctx_channels: int,
|
||||
num_heads: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
share_mod: bool = False,
|
||||
use_rope: bool = False,
|
||||
qk_rms_norm: bool = False,
|
||||
qk_rms_norm_cross: bool = False,
|
||||
image_attn_mode: str = "cross",
|
||||
proj_in_channels: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.share_mod = share_mod
|
||||
self.norm1 = _LN(channels, affine=False, eps=1e-6)
|
||||
self.norm2 = _LN(channels, affine=True, eps=1e-6)
|
||||
self.norm3 = _LN(channels, affine=False, eps=1e-6)
|
||||
self.self_attn = SparseDiTAttention(
|
||||
channels, num_heads, use_rope=use_rope, qk_rms_norm=qk_rms_norm
|
||||
)
|
||||
_cross = SparseDiTAttention(
|
||||
channels,
|
||||
num_heads,
|
||||
ctx_channels=ctx_channels,
|
||||
attn_type="cross",
|
||||
qk_rms_norm=qk_rms_norm_cross,
|
||||
)
|
||||
self.cross_attn = (
|
||||
SparseProjectAttention(_cross, channels, proj_in_channels or ctx_channels)
|
||||
if image_attn_mode == "proj"
|
||||
else _cross
|
||||
)
|
||||
hidden = int(channels * mlp_ratio)
|
||||
self.mlp_0 = nn.Linear(channels, hidden)
|
||||
self.mlp_2 = nn.Linear(hidden, channels)
|
||||
if share_mod:
|
||||
self.modulation = mx.zeros((6 * channels,))
|
||||
else:
|
||||
self.adaLN_modulation_1 = nn.Linear(channels, 6 * channels)
|
||||
|
||||
@staticmethod
|
||||
def _expand(chunk: mx.array, layout) -> mx.array:
|
||||
"""[B, C] modulation -> [N, C], each row taking its own batch item's values."""
|
||||
parts = [mx.broadcast_to(chunk[i][None], (sl.stop - sl.start, chunk.shape[-1]))
|
||||
for i, sl in enumerate(layout)]
|
||||
return parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=0)
|
||||
|
||||
def __call__(self, x, mod: mx.array, context, phases: Optional[mx.array] = None):
|
||||
m = (self.modulation + mod) if self.share_mod else self.adaLN_modulation_1(
|
||||
nn.silu(mod)
|
||||
)
|
||||
c = m.shape[-1] // 6
|
||||
lay = x.layout
|
||||
sh_msa, sc_msa, g_msa, sh_mlp, sc_mlp, g_mlp = (
|
||||
self._expand(m[..., i * c : (i + 1) * c], lay) for i in range(6)
|
||||
)
|
||||
|
||||
h = x.replace(self.norm1(x.feats) * (1 + sc_msa) + sh_msa)
|
||||
x = x.replace(x.feats + self.self_attn(h, phases=phases).feats * g_msa)
|
||||
|
||||
x = x.replace(x.feats + self.cross_attn(x.replace(self.norm2(x.feats)), context).feats)
|
||||
|
||||
h = self.norm3(x.feats) * (1 + sc_mlp) + sh_mlp
|
||||
h = self.mlp_2(nn.gelu_approx(self.mlp_0(h)))
|
||||
return x.replace(x.feats + h * g_mlp)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user