monster/Hunyuan3D-2.2-mrp-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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)

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
m3ultra 2026-07-19 12:54:44 +10:00
parent 4b368eef55
commit 5a322407f0
1 changed files with 13 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
hy3dshape/hy3dshape/models/denoisers/hunyuandit_mlx.py
View File

@ -16,6 +16,8 @@ import math
import re import re
from typing import Optional from typing import Optional
import os
import mlx.core as mx import mlx.core as mx
import mlx.nn as nn import mlx.nn as nn
import numpy as np import numpy as np
@ -513,6 +515,17 @@ class HunYuanDiTPlain(nn.Module):
self.final_layer = FinalLayer(hidden_size, self.out_channels) self.final_layer = FinalLayer(hidden_size, self.out_channels)
def __call__(self, x: mx.array, t: mx.array, contexts: dict, **kwargs) -> mx.array: 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. """Forward pass.
Args: Args: