From 5a322407f0e86690bf39873df0c5112ccfba1fb7 Mon Sep 17 00:00:00 2001 From: m3ultra Date: Sun, 19 Jul 2026 12:54:44 +1000 Subject: [PATCH] =?UTF-8?q?mlx-tune:=20T2=20whole-DiT=20mx.compile=20(same?= =?UTF-8?q?=20HY3D=5FMLX=5FCOMPILE=20flag;=20parity=207.3e-06,=201.01x=20m?= =?UTF-8?q?icro=20=E2=80=94=20GEMM-bound,=20kept=20as=20free=20insurance)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 --- .../hy3dshape/models/denoisers/hunyuandit_mlx.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py b/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py index 45658b6..73b0bae 100644 --- a/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py +++ b/hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py @@ -16,6 +16,8 @@ import math import re from typing import Optional +import os + import mlx.core as mx import mlx.nn as nn import numpy as np @@ -513,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: