feat(mlx): opt-in bf16 compute for the native MLX DiT

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,18 +34,71 @@ 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
             logger.warning(f"[MLX-DiT] Failed to initialize MLX decoder (non-fatal): {exc}")
             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
+
+            # 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."
+            )
             return False

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