Skip to content

Added support for Apple Silicon (MPS) and dropped the flash-attn dependency - #5

Open
ap-401 wants to merge 1 commit into
inclusionAI:mainfrom
ap-401:main
Open

ap-401 wants to merge 1 commit into
inclusionAI:mainfrom
ap-401:main

Conversation

@ap-401

@ap-401 ap-401 commented Sep 9, 2026

Copy link
Copy Markdown

Backstory

I was attempting to test this model locally on an M5 Pro Mac and ran into issue with the flash-attn package requiring CUDA. flash-attn package was never actually used: every attention path in this codebase resolves to torch.nn.functional.scaled_dot_product_attention, either through diffusers' dispatch_attention_fn (default backend is native SDPA) or through the text encoder's ATTENTION_CLASSES, which maps "eager", "sdpa", and "flash_attention_2" to the same SDPA class. Since flash-attn ships no wheels, pip tried to compile it from source against nvcc and failed on any machine without the CUDA toolkit, making the package uninstallable on Apple Silicon and CPU-only setups.

Changes

  • requirements.txt: drop flash-attn, add accelerate (needed by
    device_map), relax torch pin to >=2.8.0
  • src/device_utils.py: new resolve_device()/resolve_dtype() helpers
    that prefer CUDA, then MPS, then CPU, with explicit overrides
    respected
  • pipeline_llada_image.py: from_pretrained() auto-resolves device and
    dtype instead of defaulting to CPU-only when device is unset
  • run_llada_image.sh: DEVICE defaults to "auto" instead of "cuda"
  • README: document Apple Silicon support, memory footprint (~49 GB
    resident in bf16), and the MoE eager fallback on non-CUDA devices

Testing

Verified: clean-venv install with no flash-attn succeeds; a scaled-down DiT with the real architecture runs correct forward passes on MPS, including the padded/masked ragged-batch attention path, with MPS bf16 output within bf16 rounding of CPU fp32 (max abs diff 0.0148, std 0.58). Full pipeline run against the real ~49 GB checkpoint exercised (on Mac).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant