mlx-tune: T2 whole-DiT mx.compile (same HY3D_MLX_COMPILE flag; parity 7.3e-06, 1.01x micro — GEMM-bound, kept as free insurance)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

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
+        # fused SDPA (head_dim 64 = fast path); replaces manual q@k.T softmax
-        attn = mx.softmax(attn, axis=-1)
+        x = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale)
-        x = (attn @ v).transpose(0, 2, 1, 3).reshape(B, N, C)
+        x = x.transpose(0, 2, 1, 3).reshape(B, N, C)
         return self.proj(x)
 
 

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:

hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py

@@ -16,10 +16,13 @@ import math
 import re
 from typing import Optional
 
+import os
+
 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 +356,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,
@@ -512,6 +515,17 @@ class HunYuanDiTPlain(nn.Module):
         self.final_layer = FinalLayer(hidden_size, self.out_channels)
 
     def __call__(self, x: mx.array, t: mx.array, contexts: dict, **kwargs) -> mx.array:
+        # HY3D_MLX_COMPILE=1: whole-DiT mx.compile — legal now that
+        # StaticMoELayer removed the data-dependent .item() skips (T3).
+        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(x, t, contexts, **kwargs)
+        return self._forward(x, t, contexts, **kwargs)
+
+    def _forward(self, x: mx.array, t: mx.array, contexts: dict, **kwargs) -> mx.array:
         """Forward pass.
 
         Args:

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)