diff --git a/trellis_sparse_mlx/__init__.py b/trellis_sparse_mlx/__init__.py index 89fb01b..3506f5a 100644 --- a/trellis_sparse_mlx/__init__.py +++ b/trellis_sparse_mlx/__init__.py @@ -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" diff --git a/trellis_sparse_mlx/dit.py b/trellis_sparse_mlx/dit.py index 40cfcec..c76ab2b 100644 --- a/trellis_sparse_mlx/dit.py +++ b/trellis_sparse_mlx/dit.py @@ -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)