Repository navigation
sm_75 follow-ups: fp32 GEMM output in the workspace, attention reads bf16, FA2 causal alignment - #2
Merged
Merged
Conversation
…16 directly, FA2's causal alignment - linear_f16: the GEMM writes fp32 (fp16 overflowed to inf above 65504) and one pass adds the bias in fp32 and rounds once, as the sm_80 BIAS epilogue. The fp16 rows, their scales and the fp32 result live in the caller's workspace after cuBLASLt's share (1/8), in equal 128-row-tile chunks: no hidden cudaMalloc/cudaDeviceSynchronize, no scratch shared across streams, nothing outside the planner's budget. - GEMM input rows above fp16's range get a per-row power-of-two scale, undone exactly in the output pass (rows in range convert as before). - sm_75 attention reads the bf16 q/k/v as it stages them and writes bf16 from fp32: the caller's tensors are no longer converted in place and back (seven passes per call, lossy outside fp16's range). The causal mask is aligned bottom-right like FlashAttention-2's, masked scores give p = 0 before a row has a real max, and a row that sees no key gets a zero output and an +inf log-sum-exp, as FA2's. The in-place 3-D conversion API is gone with its last use. - The in-place flat conversions drop __restrict__ (src == dst). - A forced Int8Tile is refused on sm_75 (one tile) instead of ignored; vae_weight_to_krsc refuses a destination dtype other than f32/bf16. - sm75_check: f32-out, out-of-range-row and multi-chunk GEMMs; causal attention with key lengths unlike the query lengths; attention inputs left untouched. test_sm75: the per-row scale. docs/sm75.md follows.
Merged
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.
Follow-ups to #1 (Turing support), from its review.
kernels/basic/gemm.cpp): the GEMM writes fp32 instead of fp16 (outputs above 65504 became inf), and one pass adds the bias in fp32 and rounds once, as the sm_80 BIAS epilogue. The fp16 input rows, their scales and the fp32 result now live in the caller's workspace (Gemm::kWorkspaceBytes, cuBLASLt keeps 1/8), in equal chunks of whole 128-row tiles: no hiddencudaMalloc/cudaDeviceSynchronize, no scratch shared across streams, nothing outside the planner's budget.kernels/sm75/attention_sm75.cu): reads the bf16 q/k/v as it stages them and writes bf16 from fp32. The caller's tensors are no longer converted in place and back (seven extra passes per call, lossy outside fp16's range). The causal mask is aligned bottom-right like FlashAttention-2's; a row that sees no key gets a zero output and an +inf LSE, as FA2's.Int8Tilerefused on sm_75 instead of ignored;vae_weight_to_krscchecks its destination dtype; the in-place flat conversions drop__restrict__(src == dst).tools/sm75_check.cu: f32-out, out-of-range-row and multi-chunk GEMMs, causal attention with unequal key/query lengths, attention inputs left untouched.tests/unit/test_sm75.cpp: the per-row scale.docs/sm75.mdupdated.Before merging
Compile-checked on Linux (sm_89 dev preset,
MOSSINFER_SM75ON and OFF,-Werror); not yet run on a Turing GPU. Needed on the RTX 2080:sm75_check all,sm75_check big,sm75_check int8andsm75_check vaepass.e240308(the sm_75 (Turing) support: fp16 GEMMs, INT8 GEMMs, exact attention, VAE #1 head) forreference,highandturboat 256² and 1024². The GEMM now writes fp32 (more output traffic) and uses smaller row chunks, while attention drops its conversion passes. Thedocs/sm75.mdperformance table predates this and should be re-measured.