Add NVIDIA Turing support to FlashMMA kernel - #1246
Merged
Merged
Conversation
Owner
|
Thanks! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The flash attention kernel in
cpp/neuralnet/cudaflashmma.cuhneeded sm_80+, because it usedmma.sync.aligned.m16n8k16andcp.async.On NVIDIA Turing GPUs (sm_75), transformer models fallback to a much slower attention path. cuDNN graph SDPA also needs sm_80, so it is not available either.
This PR added Turing support to
cudaflashmma.cuh, with a +70.3% speedup compared to original CUDA backend on transformer models.Environment
kata1-tf3-b11c768-s11001M-d5973M.bin.gzTesting
kata-raw-nnon the empty board (komi 7.5), mma on vs off: whiteWin 0.655850 vs 0.656328, whiteLead 0.876 vs 0.882.genmovesanity check passes../katago runtests: all pass../katago runnnlayertests: all pass.cpp/rungpuerrortest.sh: all pass.cpp/runcudaopttests.shwithEXPECT_FUSED_FFN=0 EXPECT_SDPA=0: all pass.Performance
./katago benchmark -config gtp.cfg -model <tf3 model> -v 5000 -n 5CUDA builds use
numNNServerThreadsPerModel = 2. TensorRT use1.This PR is also +15.5% faster than Origin TensorRT backend in nnEvals/s.
Full sweep (nnEvals/s):