Skip to content

sm_75 follow-ups: fp32 GEMM output in the workspace, attention reads bf16, FA2 causal alignment - #2

Merged
xycld merged 1 commit into
mainfrom
sm75-followups
Oct 31, 2026
Merged

xycld merged 1 commit into
mainfrom
sm75-followups

Conversation

@xycld

@xycld xycld commented Oct 31, 2026

Copy link
Copy Markdown
Collaborator

Follow-ups to #1 (Turing support), from its review.

  • fp16 linears (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 hidden cudaMalloc/cudaDeviceSynchronize, no scratch shared across streams, nothing outside the planner's budget.
  • Out-of-range activations: a GEMM input row above fp16's range gets a per-row power-of-two scale, undone exactly in the output pass; rows in range convert exactly as before.
  • Attention (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.
  • Forced Int8Tile refused on sm_75 instead of ignored; vae_weight_to_krsc checks 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.md updated.

Before merging

Compile-checked on Linux (sm_89 dev preset, MOSSINFER_SM75 ON and OFF, -Werror); not yet run on a Turing GPU. Needed on the RTX 2080:

  1. sm75_check all, sm75_check big, sm75_check int8 and sm75_check vae pass.
  2. Step times against e240308 (the sm_75 (Turing) support: fp16 GEMMs, INT8 GEMMs, exact attention, VAE #1 head) for reference, high and turbo at 256² and 1024². The GEMM now writes fp32 (more output traffic) and uses smaller row chunks, while attention drops its conversion passes. The docs/sm75.md performance table predates this and should be re-measured.

…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.
@xycld
xycld merged commit 57a2540 into main Oct 31, 2026
5 checks passed
@xycld xycld mentioned this pull request Oct 31, 2026
@xycld
xycld deleted the sm75-followups branch November 8, 2026 12:23
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