fix: attention mask for MPS#252
Merged
Merged
Conversation
longniao
pushed a commit
to longniao/dia
that referenced
this pull request
Nov 19, 2025
This commit fixes several macOS compatibility issues and updates documentation to reflect existing MPS (Metal Performance Shaders) support. Changes: 1. Fix CUDA-specific operations to be properly guarded: - dia/model.py: Only call torch.compiler.cudagraph_mark_step_begin() on CUDA devices - cli.py: Move cudnn settings inside CUDA availability check - app.py: Move cudnn settings inside CUDA availability check 2. Update documentation: - README.md: Add comprehensive macOS/Apple Silicon section with usage examples - README.md: Update hardware support list to include MPS and CPU - README.md: Update TODO to reflect that macOS is already supported - README.md: Add performance notes and best practices for macOS users 3. Fix example code: - example/simple-mac.py: Use float32 instead of float16 for better MPS stability Key improvements for macOS users: - Prevents errors from CUDA-specific API calls on non-CUDA devices - Documents proper usage: use_torch_compile=False and compute_dtype="float32" - Clarifies that MPS backend provides GPU acceleration on Apple Silicon Fixes issues where the code would attempt CUDA operations on macOS, causing errors or warnings. Documentation now accurately reflects the current state of macOS support that was added in previous commits (nari-labs#197, nari-labs#252).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #251
I found the issue. It was because the attention mask changed and now padding tokens could attend to other padding tokens. On CUDA a softmax of
-infwill not result inNaNvalues but on MPS it will be NaN.I've reverted the mask back to the original mask. Another idea would be to check for MPS and create a different mask. Is there a specific reason to do
mask = p_mask_q & p_mask_kinstead of filtering out padding masks like before the change?