Skip to content
This repository was archived by the owner on Mar 21, 2026. It is now read-only.

fix: remove memory_efficient_threshold incompatible with MLX 0.22+#56

Closed
PatrickHeneise wants to merge 2 commits intoargmaxinc:mainfrom
PatrickHeneise:fix/mlx-sdpa-memory-efficient-threshold
Closed

fix: remove memory_efficient_threshold incompatible with MLX 0.22+#56
PatrickHeneise wants to merge 2 commits intoargmaxinc:mainfrom
PatrickHeneise:fix/mlx-sdpa-memory-efficient-threshold

Conversation

@PatrickHeneise
Copy link
Copy Markdown

Summary

  • Removes memory_efficient_threshold from the two scaled_dot_product_attention call sites in mmdit.py — this parameter was removed in MLX 0.22+ and causes a TypeError at the start of denoising
  • Guards DiffusionKitInferenceContext.os_spec() in __init__.py with a try/except — it raises IndexError on some macOS configurations where the system profiler output doesn't match the expected tab-delimited layout, crashing generation after the image is already decoded (the image was never saved)

Repro

TypeError: scaled_dot_product_attention(): incompatible function arguments
...
Invoked with types: kwargs = { q: mlx.core.array, k: mlx.core.array, v: mlx.core.array, scale: float64, memory_efficient_threshold: int }

and after patching mmdit.py:

IndexError: list index out of range
  File ".../argmaxtools/test_utils.py", line 596, in os_spec
    line.rsplit("\t\t")[1]

Test plan

  • Run diffusionkit-cli with argmaxinc/mlx-stable-diffusion-3.5-large on MLX 0.31.1 — image generates and saves successfully
  • Verify no regression on older MLX versions (flash attention path no longer explicitly gated, but MLX handles memory management internally)

🤖 Generated with Claude Code

The `memory_efficient_threshold` parameter was removed from
`scaled_dot_product_attention` in MLX 0.22+. Remove it from the
two SDPA call sites in mmdit.py to restore compatibility.

Also guard the `DiffusionKitInferenceContext.os_spec()` call in
`__init__.py` — it raises an IndexError on macOS when the system
profiler output format doesn't match the expected tab layout,
crashing generation after the image is already decoded.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@PatrickHeneise PatrickHeneise marked this pull request as ready for review March 21, 2026 06:22
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@PatrickHeneise PatrickHeneise closed this by deleting the head repository Mar 30, 2026
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant