Skip to content

Disable cuDNN SDPA to prevent nonfinite attention gradients - #3296

Merged
bghira merged 1 commit into
mainfrom
fix/disable-cudnn-sdpa
Oct 26, 2026
Merged

bghira merged 1 commit into
mainfrom
fix/disable-cudnn-sdpa

Conversation

@bghira

@bghira bghira commented Oct 26, 2026 •

Copy link
Copy Markdown
Owner

cuDNN SDPA can return NaN gradients from finite inputs, including with cuDNN 9.27 (NVIDIA #1072, NVIDIA #160). Disable cuDNN SDPA during attention initialization, backend selection and context cleanup so native attention and automatic Kohaku delegation cannot select it during training or validation.

Reject explicit cuDNN aliases, remove the CLI/WebUI choice and MiniMax-H3 context-parallel allowance, and remove the obsolete cuDNN backward-layout override. Switch the Krea2 example and quickstart to Kohaku-FA and update OPTIONS, Krea2 and MiniMax-H3 documentation in all five translations.

Validation:

  • .venv/bin/python -m unittest -v -f: 7,095 tests passed, 323 skipped.
  • Black, staged diff checks and GitHub CI passed.
  • L40S Krea2 Raw LoRA: 12 optimizer updates per cell, batch 1, rank 64, INT8 SDNQ weights, BF16 compute, gradient checkpointing and regional reduce-overhead compilation. All gradients were finite. Kohaku profiles confirmed CuTe kernels with cuDNN SDPA disabled. The cuDNN baseline uses the preceding main-branch policy in an isolated process.
Resolution Backend Peak Torch allocation Observed GPU memory peak Warm step time
512px Kohaku-FA 16.2 GiB 17.9 GiB 0.70 s
512px cuDNN baseline 16.3 GiB 17.6 GiB 0.71 s
1024px Kohaku-FA 17.8 GiB 19.9 GiB 2.20 s
1024px cuDNN baseline 17.9 GiB 19.9 GiB 2.13 s
1536px Kohaku-FA 20.7 GiB 23.6 GiB 5.88 s
1536px cuDNN baseline 20.8 GiB 23.9 GiB 5.32 s
2048px Kohaku-FA 24.7 GiB 27.8 GiB 13.28 s
2048px cuDNN baseline 24.8 GiB 28.4 GiB 11.37 s

2048px uses vae_enable_tiling=true for both backends: untiled VAE cache creation runs out of memory on the L40S before denoiser training begins. The smaller resolutions use untiled VAE preprocessing. Peak training memory includes first-use compilation; warm timing excludes the first three updates and the profiled update.

@bghira
bghira merged commit 0ab5386 into main Oct 26, 2026
4 checks passed
@bghira
bghira deleted the fix/disable-cudnn-sdpa branch October 28, 2026 02:26
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