Skip to content

Sync KohakuFA-CuTe forward tuning and launch-grid fixes - #3303

Merged
bghira merged 2 commits into
mainfrom
feature/kohaku-cute-latest-sync
Oct 28, 2026
Merged

bghira merged 2 commits into
mainfrom
feature/kohaku-cute-latest-sync

Conversation

@bghira

@bghira bghira commented Oct 28, 2026 •

Copy link
Copy Markdown
Owner

Sync the vendored CuTe implementation with merged upstream revision 1afb72698c74b3a1b5958d3b80948458c9bb6e84 from KohakuFA-CuTe.

This imports specialized Ada/Hopper causal GQA forward paths, split forward attention for short cross-attention queries, and flattened launch grids when batch × heads exceeds 65,535. Native extension build names now include the C++ source hash, preventing stale binaries when copied checkouts preserve modification times.

Port numerical, saved-statistics, cache/stride, graph replay, launch-boundary, and architecture-specific forward regressions. The comparison script reuses SimpleTuner's existing centered FP64 reference. Update the vendoring notice and all OPTIONS translations.

Validation:

  • .venv/bin/python -m unittest discover -v -f tests/ — 7,154 tests completed successfully, with 345 skips.
  • Local focused attention tests — 100 tests completed successfully, with 77 skips; native build-cache regression also passed.
  • H100 with PyTorch 2.13.0+cu130: 29 attention/forward tests completed successfully, with seven architecture-specific skips; all 11 dispatch configuration checks passed.
  • The previous implementation passes 65,535 batch/head groups and fails at 65,536 with CUDA error 9; the imported implementation passes both boundaries, including masked and grouped-query cases.
  • Kernel ASTs match upstream after namespace adaptation and two lint-only local variable renames.
  • Extreme saturated FP16 logits reproduce a pre-existing H100 precision limit: the baseline and imported implementations have bit-identical outputs/gradients, with approximately 0.96% dQ error against FP64 in the reported case. The split-forward regression compares gradients with the serial CuTe path there; ordinary-input coverage retains FP64 checks.
  • All four additional short-query split-forward tests passed on H100, covering FP16/BF16 saved statistics, graph replay, saturated logits against the serial control, and compiled dynamic query-length boundaries.
  • All pre-commit hooks passed.

The CPU CI job installs PyTorch-family wheels from the official CPU index and verifies Torch, TorchAudio, and TorchVision imports before test discovery, avoiding mixed CUDA runtime builds in a CPU test environment.

@bghira
bghira merged commit d4b6b7b into main Oct 28, 2026
4 checks passed
@bghira
bghira deleted the feature/kohaku-cute-latest-sync branch October 28, 2026 16:14
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