Skip to content

Sync CuTe attention with KohakuFA-CuTe 9280c20 - #3293

Merged
bghira merged 3 commits into
mainfrom
feature/kohaku-cute-sync
Oct 21, 2026
Merged

bghira merged 3 commits into
mainfrom
feature/kohaku-cute-sync

Conversation

@bghira

@bghira bghira commented Oct 20, 2026 •

Copy link
Copy Markdown
Owner

SimpleTuner's vendored CuTe implementation predates upstream's centered RTX Blackwell backward and native datacenter Blackwell pipelines. Synchronize it through 9280c20, including upstream PRs #15–18.

  • Port RTX 5090 tuning, native B200/B300 producer/consumer pipelines, fused backward consumers, packed FP32 arithmetic, adjacent D512 gradient slices, and overlapped operand/statistic transfers.
  • Add the B300 register-resident FP32 auxiliary and 256-wide backward accumulators for BF16 D512 multi-head attention at 512 tokens and above, including vectorized gradient staging and causal register allocation, preserving the original dispatch for other shapes and layouts.
  • Preserve SimpleTuner's eager/compiled attention integration. Version custom operators for the expanded saved-output contract and include sequence configuration thresholds in native program cache keys.
  • Enable centered CuTe selection under kohaku-fa-auto on RTX Blackwell. Explicit kohaku-fa continues to use Gluon on B200/B300; automatic training keeps native SDPA there. The vendored CuTe API also supports their native kernels.
  • Import upstream cache-reuse, CUDA graph replay, cancellation and zero-covariance regressions. Add architecture/transfer-width checks, package the Apache license, and update the options reference, MOSS quickstart and UI tooltip in all six languages.

Validation:

  • Full .venv/bin/python -m unittest discover -s tests -v -f at the preceding dee9d84 sync: 7,083 tests run, 315 skipped, no failures.
  • Focused attention suite after the 9280c20 follow-up: 78 tests run, 58 skipped, no failures. Its B300 cases now cover 511/512/513-token cache dispatch, graph replay at 512/2048 tokens, and zero covariance at 513/2049 tokens.
  • Five CuTe configuration tests passed with the optional DSL installed, including B200/B300 transfer widths and the B300 BF16 D512 selector. Compilation was mocked for these configuration checks.
  • All three runtime modules changed since the previous upstream snapshot match the pinned upstream AST after namespace integration and lint-compatible variable normalization.
  • Earlier H100 validation covered FP16/BF16 outputs and gradients, compiled masks/tails/causal attention/GQA, CUDA graph mask replay, saturated tiny gradients, head dimensions from 1 to 512, and compiled checkpointed LoRA updates. It predates the final Blackwell-only follow-up.
  • Formatting, import sorting, lint and whitespace checks passed.

RTX 5090/B200/B300 runtime kernels were not rerun on physical Blackwell hardware during this port; their dedicated regression tests are included.

@bghira bghira changed the title Sync CuTe attention with KohakuFA-CuTe 3a398da Sync CuTe attention with KohakuFA-CuTe dee9d84 Oct 21, 2026
@bghira bghira changed the title Sync CuTe attention with KohakuFA-CuTe dee9d84 Sync CuTe attention with KohakuFA-CuTe 9280c20 Oct 21, 2026
@bghira
bghira merged commit fbbf4c0 into main Oct 21, 2026
4 checks passed
@bghira
bghira deleted the feature/kohaku-cute-sync branch October 21, 2026 15:01
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