Comfy-Org / Comfy-Org/ComfyUI

Silent output corruption on MPS (Apple Silicon) for large attention matrices

Open
#14,837 7 comments 0 reactions 0 assignees View on GitHub
Potential Bug
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Custom Node Testing

- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Expected Behavior

Attention computation on MPS produces correct output regardless of sequence length, resolution, or batch size, the same as it would on CUDA/CPU.

### Actual Behavior

When a single attention matrix (`batch * heads * seq_q * seq_k`) reaches roughly 2^31 elements, PyTorch's MPS backend silently corrupts the output instead of raising an error. There's no crash, no exception, no log line — the generation completes normally but produces wrong output (visually corrupted images/video frames). This is easy to misdiagnose as a model or workflow problem rather than a backend bug, since nothing in the logs indicates anything went wrong.

Affects every attention implementation that's reachable on MPS in `comfy/ldm/modules/attention.py` (`attention_basic`, `attention_split`, `attention_pytorch`, `attention_sub_quad`) — including `attention_sub_quad`, which is the default MPS path when no `--use-*-attention` flag is passed.

Root cause: MPS uses 32-bit indexing internally for many ops, and none of these implementations' existing chunking logic accounts for it — they're all free-memory-based, designed to avoid running out of memory, not to respect this indexing ceiling. On a machine with a lot of unified memory, "there's plenty of free memory" is true right up until the silent corruption, so the existing safeguards never trigger.

### Steps to Reproduce

Run any workflow on MPS where a single attention op's `batch * heads * seq_q * seq_k` exceeds ~2^31 (2,147,483,648) elements. Easiest to hit with:
- Long-sequence video diffusion models (self-attention over the full spatiotemporal latent scales quickly with frame count and resolution)
- High resolution + many attention heads
- Larger batch sizes

Concrete example: a 61-frame, 832x640 video-model generation (single self-attention call over the full latent) measures a ~7.68B-element attention matrix — 3.6x over the limit — and produces visibly corrupted output frames with no error at all.

### Debug Logs

```powershell
No error output of any kind is produced — that's the core of the bug. The generation completes and reports success; only the visual output is wrong.
```

### Other

_No response_

Contributor guide

Open the contributing guide

Research direction

Start in comfy/ldm/modules/attention.py and inspect attention_basic, attention_split, attention_pytorch, and attention_sub_quad, especially their existing free-memory-based chunking. Reproduce on MPS with an attention matrix above roughly 2^31 elements and compare the output with CUDA or CPU. Done means the affected path no longer silently produces corrupted output and has a verifiable failure or correct result at the indexing limit.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.