Skip to content

Add SM86 Ampere support and improve Ada CuTe attention - #3295

Merged
bghira merged 1 commit into
mainfrom
feature/kohaku-cute-ampere-ada
Oct 22, 2026
Merged

bghira merged 1 commit into
mainfrom
feature/kohaku-cute-ampere-ada

Conversation

@bghira

@bghira bghira commented Oct 22, 2026

Copy link
Copy Markdown
Owner

SM86 Ampere GPUs were rejected by SimpleTuner's attention controller and native entry point, while Ada retained the older wide-head layouts. Sync KohakuFA-CuTe c0786dd and enable the existing CuTe backend for SM86 in both strict and automatic training modes.

  • Port Ada D256/D512 tile selection, paired QK/dP score producers, packed DS/P stores, and wider SM86/SM89 D512 operand copies. Include the accompanying SM120 cache-policy and causal D128 query updates.
  • Extend native compatibility checks, automatic dispatch, centered autograd handling, and capability error messages. Test SM86 acceptance and SM80/SM87 rejection, plus training/evaluation policy transitions.
  • Generalize warm-cache and changed-input CUDA graph regressions across SM86/SM89/SM120. Expand compiled projection-layout and gradient-cancellation coverage, and test optimized selection against masked/unaligned layouts.
  • Update the vendored revision notice, WebUI attention tooltip, and existing option documentation in English, Chinese, Japanese, Portuguese, Spanish, and Hindi. No new backend option is required.

Validation:

  • .venv/bin/python -m unittest discover -v -f tests/: 7,091 tests run, 323 skipped, no failures.
  • RTX 3090 (SM86), RTX 4090 (SM89), and RTX 5090 (SM120): 31 targeted checks passed per device, with no skips, in isolated attention-component environments using PyTorch 2.11.0+cu130, CuTe DSL 4.8.1, TVM FFI 0.1.15.post0, and Diffusers 0.40.0.
  • Hardware coverage includes FP16/BF16 outputs and Q/K/V gradients, saturation and tiny-gradient preservation, masks/GQA/strides, eager and compiled backward, cache reuse, changed-input CUDA graph replay, key centering, strict/automatic training dispatch, and SimpleTuner/Diffusers integration.
  • Device VRAM returned to the pre-test baseline after each test process exited. All validation rentals were released.
  • Vendored kernels.py matches the pinned upstream file byte-for-byte. Native C++ matches after namespace adaptation, and imported regression methods match after import-path integration. Pre-commit formatting, import sorting, lint, and whitespace checks passed.

These are correctness and dispatch checks; this PR does not claim newly measured attention latency or full training throughput.

Sync KohakuFA-CuTe c0786dd, enable SM86 in native and automatic dispatch,
and extend architecture, graph replay, and gradient regressions. Update
the attention option documentation and translations.
@bghira
bghira merged commit 30184da into main Oct 22, 2026
4 checks passed
@bghira
bghira deleted the feature/kohaku-cute-ampere-ada branch October 22, 2026 16:46
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