diff --git a/hy3dpaint/hunyuanpaintpbr_mlx/dino_mlx.py b/hy3dpaint/hunyuanpaintpbr_mlx/dino_mlx.py index dea8148..6238717 100644 --- a/hy3dpaint/hunyuanpaintpbr_mlx/dino_mlx.py +++ b/hy3dpaint/hunyuanpaintpbr_mlx/dino_mlx.py @@ -54,9 +54,9 @@ class ViTAttention(nn.Module): qkv = qkv.transpose(2, 0, 3, 1, 4) # (3, B, H, N, D) q, k, v = qkv[0], qkv[1], qkv[2] - attn = (q @ k.transpose(0, 1, 3, 2)) * self.scale - attn = mx.softmax(attn, axis=-1) - x = (attn @ v).transpose(0, 2, 1, 3).reshape(B, N, C) + # fused SDPA (head_dim 64 = fast path); replaces manual q@k.T softmax + x = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale) + x = x.transpose(0, 2, 1, 3).reshape(B, N, C) return self.proj(x) diff --git a/hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py b/hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py index af22508..d41ab77 100644 --- a/hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py +++ b/hy3dpaint/hunyuanpaintpbr_mlx/unet/unet_mlx.py @@ -6,6 +6,8 @@ modules attached externally via load_model._enhance_unet. from typing import Dict, Optional, Tuple +import os + import mlx.core as mx import mlx.nn as nn import numpy as np @@ -161,6 +163,24 @@ class UNet2DConditionModelMLX(nn.Module): timestep: mx.array, encoder_hidden_states: Optional[mx.array] = None, **kwargs, + ) -> mx.array: + # HY3D_MLX_COMPILE=1 routes through mx.compile (fixed-shape denoise + # loop -> one trace per chunk-shape/ctx-structure; corridorkey + # precedent 4c660df). Default off = numerics identical to upstream. + if os.environ.get("HY3D_MLX_COMPILE", "0") == "1": + fn = getattr(self, "_compiled_forward", None) + if fn is None: + fn = mx.compile(self._forward) + self._compiled_forward = fn + return fn(sample, timestep, encoder_hidden_states, **kwargs) + return self._forward(sample, timestep, encoder_hidden_states, **kwargs) + + def _forward( + self, + sample: mx.array, + timestep: mx.array, + encoder_hidden_states: Optional[mx.array] = None, + **kwargs, ) -> mx.array: """ Args: diff --git a/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py b/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py index f36d02c..45658b6 100644 --- a/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py +++ b/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py @@ -19,7 +19,8 @@ from typing import Optional import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_arsenal.moe import MoELayer +from mlx_arsenal.moe import MoELayer # noqa: F401 (kept for API compat) +from .moe_static import StaticMoELayer # --------------------------------------------------------------------------- # @@ -353,7 +354,7 @@ class HunYuanDiTBlock(nn.Module): return MLP(hidden_size) shared = MLP(hidden_size) - self.moe = MoELayer( + self.moe = StaticMoELayer( hidden_size=hidden_size, num_experts=num_experts, top_k=moe_top_k, diff --git a/hy3dshape/hy3dshape/models/denoisers/moe_static.py b/hy3dshape/hy3dshape/models/denoisers/moe_static.py new file mode 100644 index 0000000..cae5f7c --- /dev/null +++ b/hy3dshape/hy3dshape/models/denoisers/moe_static.py @@ -0,0 +1,37 @@ +"""Static MoE: mlx_arsenal.MoELayer minus the per-expert `.item()` early-skip. + +The skip costs a host sync per expert per block per step (~4800/gen) and +makes the graph dynamic (blocks mx.compile). With top-k routing over +thousands of tokens every expert almost always receives work, so the skip +almost never fires — running all experts with w=0 rows zeroed is +numerically identical (same accumulation order) and fully static. +Vendored subclass — mlx_arsenal itself is a SHARED dependency, unmodified. +""" +import mlx.core as mx + +from mlx_arsenal.moe import MoELayer + + +class StaticMoELayer(MoELayer): + def __call__(self, x: mx.array) -> mx.array: + orig_shape = x.shape + hidden = x.reshape(-1, orig_shape[-1]) + indices, weights = self.gate(hidden) + num_experts = len(self.experts) + num_tokens = hidden.shape[0] + expert_weights = mx.zeros((num_tokens, num_experts)) + for k in range(indices.shape[1]): + col_indices = indices[:, k : k + 1] + col_weights = weights[:, k : k + 1] + one_hot = mx.zeros((num_tokens, num_experts)) + rows = mx.arange(num_tokens).reshape(-1, 1) + one_hot[rows, col_indices] = col_weights + expert_weights = expert_weights + one_hot + output = mx.zeros_like(hidden) + for expert_idx, expert in enumerate(self.experts): + w = expert_weights[:, expert_idx : expert_idx + 1] + expert_output = expert(hidden) + output = output + w * expert_output + if self.shared_expert is not None: + output = output + self.shared_expert(hidden) + return output.reshape(orig_shape)