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

[argentumaurum-eth]

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.

Summary by CodeRabbit

Release Notes

• New Features

  • Added optional bfloat16 compute mode for the MLX DiT decoder on Apple Silicon, enabled via environment variable for improved performance and memory efficiency.
  • Decoder now automatically converts inputs to the selected compute precision during inference and returns outputs in the original precision for compatibility.

• Bug Fixes

  • Added safe fallback to float32 if bfloat16 conversion fails during MLX DiT initialization/runtime.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: b331c19f-88ed-4358-9004-490faed4581a

📥 Commits

Reviewing files that changed from the base of the PR and between a1b4382 and 47eac48.

📒 Files selected for processing (1)

• acestep/core/generation/handler/mlx_dit_init.py

🚧 Files skipped from review as they are similar to previous changes (1)

• acestep/core/generation/handler/mlx_dit_init.py

---

📝 Walkthrough

Walkthrough

Adds an optional bfloat16 compute path to the MLX DiT decoder on Apple Silicon. The MLXDiTDecoder gains a compute_dtype attribute and dtype-cast logic in its forward pass. A new _maybe_apply_mlx_dit_bf16 method in MlxDitInitMixin converts decoder parameters and activates the path when the ACESTEP_MLX_DIT_BF16 environment variable is set.

Changes

MLX DiT bfloat16 compute path

Layer / File(s)	Summary
MLXDiTDecoder compute_dtype attribute and forward dtype switching acestep/models/mlx/dit_model.py	Adds self.compute_dtype = mx.float32 to __init__. In __call__, records the caller's input dtype, casts all five decoder inputs to compute_dtype when they differ, casts RoPE cos/sin to the active hidden state dtype before attention, and casts the final output back to the caller's original dtype before returning.
Env var helper and MlxDitInitMixin bf16 wiring acestep/core/generation/handler/mlx_dit_init.py	Adds _mlx_dit_bf16_requested() reading ACESTEP_MLX_DIT_BF16. Extends _init_mlx_dit to call _maybe_apply_mlx_dit_bf16, store the result in self.mlx_dit_bf16, and log the chosen dtype; resets self.mlx_dit_bf16 to False on init failure. Adds _maybe_apply_mlx_dit_bf16 which casts floating parameter arrays to mx.bfloat16, sets mlx_decoder.compute_dtype, evaluates parameters, and returns False with a warning on any exception.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

🐇 Hoppity-hop through the number lane,
Where float32 once ruled the domain.
Now bfloat16 joins the dance,
Cast in, cast out—a precision prance!
A single env var sets the stage,
And the rabbit cheers from every page. 🎉

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: adding optional bf16 compute support for MLX DiT, which directly aligns with the PR's core objective of implementing environment-controlled bfloat16 precision for the Apple Silicon MLX decoder.
Docstring Coverage	✅ Passed	Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@acestep/core/generation/handler/mlx_dit_init.py`:
- Around line 84-93: The bf16 conversion in the _maybe_apply_mlx_dit_bf16
function is not atomic - the mlx_decoder parameters and compute_dtype can be
mutated by the tree_map and assignment operations before mx.eval is called, and
if mx.eval fails, the exception returns False while the decoder has already been
partially converted to bf16, creating a state mismatch. Fix this by moving the
mx.eval call before any mutations occur, so validation happens before any state
changes, ensuring that if any step fails the decoder remains completely
unmodified. Apply the same fix to the related code also referenced in the
comment (lines 37-46).

In `@acestep/models/mlx/dit_model.py`:
- Around line 454-460: The file acestep/models/mlx/dit_model.py exceeds the 200
line of code hard limit and this PR adds more logic to it. Per policy, add a
module-level comment or PR note that provides a brief justification for why the
current structure requires the module to exceed this limit and includes a
concrete, actionable follow-up plan describing how and when the module will be
split into smaller components (e.g., separating initialization logic, parameter
management, or forward pass computation into separate modules).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: b557d8cd-9d2c-4585-a91a-be29912bd295

📥 Commits

Reviewing files that changed from the base of the PR and between dce6214 and a1b4382.

📒 Files selected for processing (2)

• acestep/core/generation/handler/mlx_dit_init.py
• acestep/models/mlx/dit_model.py