Skip to content

Honor native attention backend selection in direct SDPA calls - #3217

Merged
bghira merged 1 commit into
mainfrom
investigate/3195-rdna2-training-crash
Sep 7, 2026
Merged

bghira merged 1 commit into
mainfrom
investigate/3195-rdna2-training-crash

Conversation

@bghira

@bghira bghira commented Sep 7, 2026

Copy link
Copy Markdown
Owner

Selecting native-math previously changed only Diffusers' dispatch registry. SDXL's AttnProcessor2_0 calls PyTorch SDPA directly, so it could still select flash attention. Scope PyTorch's SDPA selection alongside the existing Diffusers context for explicit native backends, restoring both when the backend changes.

The real CPU profiler regression previously observed aten::_scaled_dot_product_flash_attention_for_cpu and its backward despite selecting native-math. It now observes aten::_scaled_dot_product_attention_math with finite gradients. Additional tests cover backend switching and train/evaluation restoration for all explicit native SDPA choices.

Validation: .venv/bin/python -m unittest -v -f tests.test_attention_backend — 34 passed. Both the processor dispatch and backend-switch regressions failed before the fix and passed afterward.

Supplemental investigation: a tiny SDXL-shaped UNet with BF16 LoRA, math attention, gradient checkpointing, and SimpleTuner's AdamWBF16 completed three steps each on CPU and MPS. Convolution hooks observed recomputation during the first backward and every subsequent backward. This checks the control flow at reduced scale; it is not a full SDXL or RDNA2 hardware reproduction.

Related to #3195. This corrects the ineffective native-math diagnostic; resolution of the reported RDNA2 SIGSEGV requires testing on affected hardware.

@bghira
bghira merged commit 5356296 into main Sep 7, 2026
2 checks passed
@bghira
bghira deleted the investigate/3195-rdna2-training-crash branch September 7, 2026 01:41
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