From a1b4382000df5ca0afd7160da21e6d6ef6b217b3 Mon Sep 17 00:00:00 2001 From: argentumaurum-eth Date: Mon, 22 Jun 2026 08:28:03 +0300 Subject: [PATCH 1/2] feat(mlx): opt-in bf16 compute for the native MLX DiT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add ACESTEP_MLX_DIT_BF16 (default off) to run the Apple Silicon MLX DiT decoder in bfloat16 instead of float32. bf16 keeps the full fp32 exponent range (no fp16 overflow) and matches the DiT's served precision. Measured on M4 Max (isolated DiT forward, xl-base shape): 1.07x @1024, 1.26x @2048, 1.33x @3072 latent frames — the win grows with sequence length, so long tracks benefit most. The default-off path stays float32 and byte-identical. - dit_model: add compute_dtype; cast inputs and RoPE cos/sin to it on forward entry, cast the velocity back to the caller dtype at the end so the diffusion sampler loop stays float32. - mlx_dit_init: when the flag is set, cast decoder params to bf16 via tree_map (mirrors the existing MLX VAE fp16 path); flag-off never imports mlx, so float32 behaviour is byte-identical. slop music deserves fast slop code. --- .../core/generation/handler/mlx_dit_init.py | 51 ++++++++++++++++++- acestep/models/mlx/dit_model.py | 34 ++++++++++++- 2 files changed, 83 insertions(+), 2 deletions(-) diff --git a/acestep/core/generation/handler/mlx_dit_init.py b/acestep/core/generation/handler/mlx_dit_init.py index 63cedde22..3d4afd61d 100644 --- a/acestep/core/generation/handler/mlx_dit_init.py +++ b/acestep/core/generation/handler/mlx_dit_init.py @@ -1,8 +1,15 @@ """MLX DiT initialization helpers for Apple Silicon acceleration.""" +import os + from loguru import logger +def _mlx_dit_bf16_requested() -> bool: + """Return True when ``ACESTEP_MLX_DIT_BF16`` opts into bf16 DiT compute.""" + return os.environ.get("ACESTEP_MLX_DIT_BF16", "0").lower() in ("1", "true", "yes") + + class MlxDitInitMixin: """Initialize native MLX DiT decoder state used by generation runtime.""" @@ -27,13 +34,15 @@ def _init_mlx_dit(self, compile_model: bool = False) -> bool: mlx_decoder = MLXDiTDecoder.from_config(self.config) convert_and_load(self.model, mlx_decoder) + bf16_applied = self._maybe_apply_mlx_dit_bf16(mlx_decoder) mlx_decoder.materialize_static_buffers() self.mlx_decoder = mlx_decoder self.use_mlx_dit = True self.mlx_dit_compiled = compile_model + self.mlx_dit_bf16 = bf16_applied logger.info( "[MLX-DiT] Native MLX DiT decoder initialized successfully " - f"(mx.compile={compile_model})." + f"(mx.compile={compile_model}, dtype={'bfloat16' if bf16_applied else 'float32'})." ) return True except Exception as exc: # noqa: BLE001 @@ -41,4 +50,44 @@ def _init_mlx_dit(self, compile_model: bool = False) -> bool: self.mlx_decoder = None self.use_mlx_dit = False self.mlx_dit_compiled = False + self.mlx_dit_bf16 = False + return False + + @staticmethod + def _maybe_apply_mlx_dit_bf16(mlx_decoder) -> bool: + """Optionally cast the MLX DiT to bfloat16 for faster Apple-Silicon compute. + + Controlled by the ``ACESTEP_MLX_DIT_BF16`` environment variable (off by + default). When disabled this returns immediately without importing + ``mlx`` so the float32 path — and unit tests that stub the MLX modules — + is completely unaffected. + + bf16 keeps the float32 exponent range (no overflow risk, unlike fp16), + matches the precision the DiT is trained/served at on CUDA, and roughly + halves both matmul time and weight bandwidth on the M-series GPU. + + Returns: + bool: ``True`` when the decoder was converted to bf16, else ``False``. + """ + if not _mlx_dit_bf16_requested(): + return False + try: + import mlx.core as mx + from mlx.utils import tree_map + + def _to_bf16(value): + """Cast floating MLX arrays to bfloat16, leaving others intact.""" + if isinstance(value, mx.array) and mx.issubdtype(value.dtype, mx.floating): + return value.astype(mx.bfloat16) + return value + + mlx_decoder.update(tree_map(_to_bf16, mlx_decoder.parameters())) + mlx_decoder.compute_dtype = mx.bfloat16 + mx.eval(mlx_decoder.parameters()) + logger.info("[MLX-DiT] Parameters converted to bfloat16 (ACESTEP_MLX_DIT_BF16=1).") + return True + except Exception as exc: # noqa: BLE001 + logger.warning( + f"[MLX-DiT] bfloat16 conversion failed ({exc}); staying on float32." + ) return False diff --git a/acestep/models/mlx/dit_model.py b/acestep/models/mlx/dit_model.py index 75c1ea05e..3cf06a46e 100644 --- a/acestep/models/mlx/dit_model.py +++ b/acestep/models/mlx/dit_model.py @@ -451,6 +451,13 @@ def __init__( super().__init__() self.hidden_size = hidden_size self.patch_size = patch_size + # Compute dtype for the diffusion forward pass. Defaults to float32 so + # behaviour is unchanged; ``mlx_dit_init`` flips this to bfloat16 (and + # casts the parameters to match) when ``ACESTEP_MLX_DIT_BF16`` is set. + # Not an ``mx.array`` and not underscore-free-by-accident: a plain dtype + # object is ignored by ``Module.parameters()`` so it never participates + # in weight loading or ``mx.eval``. + self.compute_dtype = mx.float32 inner_dim = hidden_size if layer_types is None: @@ -556,6 +563,21 @@ def __call__( Returns: (output_hidden_states, cache) """ + # Optional reduced-precision compute path. When ``compute_dtype`` is + # bfloat16 we cast the inputs here so every matmul inside the decoder + # runs in bf16, then cast the final velocity back to the caller's dtype + # at the end so the diffusion loop (CFG, ODE step, DCW, repaint) keeps + # running in float32 exactly as before. When ``compute_dtype`` is + # float32 (default) every branch below is a no-op. + cdt = getattr(self, "compute_dtype", mx.float32) + external_dtype = hidden_states.dtype + if cdt != external_dtype: + hidden_states = hidden_states.astype(cdt) + encoder_hidden_states = encoder_hidden_states.astype(cdt) + context_latents = context_latents.astype(cdt) + timestep = timestep.astype(cdt) + timestep_r = timestep_r.astype(cdt) + # Timestep embeddings temb_t, proj_t = self.time_embed(timestep) temb_r, proj_r = self.time_embed_r(timestep - timestep_r) @@ -587,8 +609,13 @@ def __call__( seq_len = hidden_states.shape[1] dtype = hidden_states.dtype - # Position embeddings (RoPE) + # Position embeddings (RoPE). The cached cos/sin tables live in + # float32; cast them to the compute dtype so rotary application and the + # subsequent SDPA see a single consistent dtype for q/k/v. cos, sin = self.rotary_emb(seq_len) + if cos.dtype != dtype: + cos = cos.astype(dtype) + sin = sin.astype(dtype) # Attention masks # Self-attention: full layers get None; sliding layers get windowed mask @@ -624,6 +651,11 @@ def __call__( # Crop back to original sequence length hidden_states = hidden_states[:, :original_seq_len, :] + # Hand the velocity back to the diffusion loop in its original dtype + # (float32) so the sampler math is unchanged. No-op on the fp32 path. + if hidden_states.dtype != external_dtype: + hidden_states = hidden_states.astype(external_dtype) + return hidden_states, cache @classmethod From 47eac488b85ca1cf6e53e256896d371d2f8297fe Mon Sep 17 00:00:00 2001 From: argentumaurum-eth Date: Mon, 22 Jun 2026 08:51:23 +0300 Subject: [PATCH 2/2] fix(mlx): make bf16 application atomic on the init failure path --- acestep/core/generation/handler/mlx_dit_init.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/acestep/core/generation/handler/mlx_dit_init.py b/acestep/core/generation/handler/mlx_dit_init.py index 3d4afd61d..5fbc1adef 100644 --- a/acestep/core/generation/handler/mlx_dit_init.py +++ b/acestep/core/generation/handler/mlx_dit_init.py @@ -81,12 +81,23 @@ def _to_bf16(value): return value.astype(mx.bfloat16) return value - mlx_decoder.update(tree_map(_to_bf16, mlx_decoder.parameters())) - mlx_decoder.compute_dtype = mx.bfloat16 + # Materialize the bf16 copy *before* mutating any decoder state, so a + # failure here leaves the decoder fully untouched (still float32). + # compute_dtype is flipped *last* — only after update + eval succeed — + # so the "bf16 applied" signal can never disagree with the actual + # parameter dtype on the exception path. + bf16_params = tree_map(_to_bf16, mlx_decoder.parameters()) + mx.eval(bf16_params) + mlx_decoder.update(bf16_params) mx.eval(mlx_decoder.parameters()) + mlx_decoder.compute_dtype = mx.bfloat16 logger.info("[MLX-DiT] Parameters converted to bfloat16 (ACESTEP_MLX_DIT_BF16=1).") return True except Exception as exc: # noqa: BLE001 + # Defensive: guarantee compute_dtype never reports bf16 after a failure, + # even if the cast partially applied before raising. + if "mx" in locals(): + mlx_decoder.compute_dtype = mx.float32 logger.warning( f"[MLX-DiT] bfloat16 conversion failed ({exc}); staying on float32." )