Apple Silicon (MPS): random all-NaN attention -> black videos with bf16 models (LTX-2.x); two suggested fixes
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Summary
On Apple Silicon (MPS, macOS ≥ 14.5), video generation with **bf16** checkpoints (all current LTX-2.x models) intermittently collapses mid-sampling: a single transformer forward returns all-NaN hidden states, after which every remaining sampler step and both VAE decodes stay NaN — the saved video is solid black (audio waveform included).
Observed failure rates on one machine (Apple M5 Max, macOS 26.6.2, torch 2.13.0, ComfyUI v0.33.3):
- LTX-2.5 t2v workflow: **19 of 26 runs** produced all-black videos
- LTX-2.5 i2v workflow: **8 of 8 runs** before mitigation
Same seed reproduces bit-identically; different seeds fail stochastically — the failure depends on memory contents, not on model math.
Instrumenting the LTXAV transformer with per-stage finite checks shows the first non-finite values appear as the **entire output of a single attention op while its inputs are still finite** (e.g. `audio_to_video_attn` returning `nan=14417920/14417920`), consistent with garbage being read from uninitialized memory inside attention.
## Root cause 1: sub-quadratic attention relies on `beta=0` semantics that MPS does not honor
`comfy/ldm/modules/sub_quadratic_attention.py` computes attention scores as:
```python
attn_weights = torch.baddbmm(
torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
query, key_t, alpha=scale, beta=0,
)
```
PyTorch documents that with `beta=0` the input tensor is ignored ("nan and infinity in it will not be propagated"). MPS violates this today:
```python
import torch
q = torch.randn(1, 64, 128, device="mps", dtype=torch.bfloat16)
k = torch.randn(1, 128, 3520, device="mps", dtype=torch.bfloat16)
nan_in = torch.full((1, 1, 1), float("nan"), device="mps", dtype=torch.bfloat16)
out = torch.baddbmm(nan_in, q, k, alpha=0.088, beta=0)
print(int(torch.isnan(out).sum()), "/", out.numel()) # MPS: 225280/225280 ; CPU: 0
```
Since `torch.empty(1, 1, 1)` may contain NaN bit patterns, any run can get its whole score matrix poisoned by one broadcast scalar.
**Suggested fix** (verified locally): initialize with `zeros` instead of `empty` — 4 occurrences (`_summarize_chunk` lines ~71/79 and `_get_attention_scores_no_kv_chunking` lines ~152/160). No measurable cost.
## Root cause 2: macOS fp32-attention workaround covers fp16 but not bf16
`comfy/model_management.py::force_upcast_attention_dtype()` already works around a known macOS black-image bug:
```python
if macos_version is not None and ((14, 5) <= macos_version):
upcast = True # black image bug on recent versions of macOS ...
if upcast:
return {torch.float16: torch.float32}
```
`bfloat16` is missing from the map, so fully-bf16 pipelines (all current LTX-2.x checkpoints) bypass the protection and keep running attention through the fragile path above.
**Suggested fix**: `{torch.float16: torch.float32, torch.bfloat16: torch.float32}`.
## Verification after applying both fixes locally
- 3 consecutive seeds that previously collapsed generated valid videos (~1 MB vs ~120 KB black files); step-level NaN probes report zero non-finite values across base sampling (8 steps), latent upsampler, refine stage (3 steps) and both VAE decodes.
- Workflow details: two-stage ManualSigmas + euler_ancestral + DualCFGGuider [1,1], joint audio-video latent; failure signature was identical for t2v and i2v, confirming the shared attention path as the culprit rather than anything i2v-specific.
Contributor guide
Research direction
Inspect comfy/ldm/modules/sub_quadratic_attention.py at the four baddbmm initializations in _summarize_chunk and _get_attention_scores_no_kv_chunking, then review force_upcast_attention_dtype() in comfy/model_management.py. Verify the changes on Apple Silicon with bf16 LTX-2.x workflows, confirming that sampling and both VAE decodes produce finite values and valid non-black videos.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- macos, python, pytorch
- Domain
- ai, machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 82/100