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

docs/sm75.md

@@ -27,8 +27,8 @@ presets stay Ampere-and-newer. Without the option the engine refuses compute cap
 
 | Component | sm_80+ build (unchanged) | sm_75 build |
 |---|---|---|
-| bf16 linears (DiT blocks and globals, text refiner, text encoder) | cuBLASLt bf16 tensor cores | fp16 tensor cores through cuBLASLt: weights are converted to fp16 once at load (or per streamed block), activations are converted in row chunks, accumulation stays fp32, the output is converted back to bf16 with the bias added in fp32 |
-| Exact attention | FlashAttention-2 (bf16, `cp.async`) | a dedicated `m16n8k8` fp16 kernel with an online softmax (head_dim 128 non-causal varlen GQA for the DiT, head_dim 256 causal for the text encoder) |
+| bf16 linears (DiT blocks and globals, text refiner, text encoder) | cuBLASLt bf16 tensor cores | fp16 tensor cores through cuBLASLt: weights are converted to fp16 once at load (or per streamed block), activations are converted in row chunks inside the caller's GEMM workspace (a row above fp16's range is scaled by a power of two, undone exactly), the GEMM writes fp32 and one pass adds the bias in fp32 and rounds once to the output dtype |
+| Exact attention | FlashAttention-2 (bf16, `cp.async`) | a dedicated `m16n8k8` fp16 kernel with an online softmax (head_dim 128 non-causal varlen GQA for the DiT, head_dim 256 causal for the text encoder); it reads the bf16 q/k/v directly, converting them as it stages them, and writes bf16 |
 | INT8 ConvRot GEMMs (`high` tier) | CUTLASS `m16n8k32` | CUTLASS `m8n8k16` (Turing's INT8 MMA), same epilogues; bit-exact against a host int32 reference |
 | INT4 ConvRot GEMMs (`turbo` tier) | CUTLASS `m16n8k64` | the engine's existing INT4 → INT8 unpack fallback (bit-exact, the same path Hopper and Blackwell use) |
 | 8-bit attention (`q8_fp16`, `q8_fp8`) | SageAttention kernels | not built: an unset `attention` resolves to `exact` with the reason saying so (the tier is otherwise unchanged) |
@@ -41,8 +41,8 @@ its CUDA graphs are not used. Every block path produces the same image, so this
 
 ### Tiers on sm_75
 
-`reference` (bf16 bundle), `high` (INT8 ConvRot DiT) and `turbo` (INT4 ConvRot DiT) all run. As on any device without
-an 8-bit attention kernel, an unset `attention` resolves to `exact` and the report's `quality.reason` says
+`reference` (bf16 bundle), `high` (INT8 ConvRot DiT) and `turbo` (INT4 ConvRot DiT) all run. On sm_75, which has no
+8-bit attention kernel, an unset `attention` resolves to `exact` and the report's `quality.reason` says
 "(no 8-bit attention kernel on this device)"; a caller may still pass `attention=exact` explicitly. NVFP4 DiTs are
 refused as on every non-Blackwell GPU.
 
@@ -64,23 +64,28 @@ roughly 3–5 s per prompt when streamed (its ~16 GB do not fit 8 GiB of VRAM);
 
 ## Numerics
 
-- bf16 → fp16 is exact for values in fp16's normal range, and fp16 → bf16 recovers the original bf16 value, so the
-  converted operands are the same numbers the sm_80 kernels see. Accumulation stays fp32 where the sm_80 kernel
-  accumulates in fp32.
+- bf16 → fp16 is exact for values in fp16's normal range, so the converted operands are the same numbers the sm_80
+  kernels see. Accumulation stays fp32 where the sm_80 kernel accumulates in fp32, and every output is rounded once
+  from fp32 (the GEMM's bias is added in fp32 first, as cuBLASLt's bias epilogue does on sm_80).
 - Attention, GEMMs and the VAE therefore differ from an sm_80 run only in accumulation order and rounding, not in the
-  values fed in. INT8 GEMM results are bit-identical (exact int32 accumulation).
+  values fed in. Attention's probabilities enter P·V as fp16 (bf16 on sm_80), and the VAE computes in fp32 (TF32 on
+  sm_80). INT8 GEMM results are bit-identical (exact int32 accumulation).
 - The reference vectors under `tests/data` are recorded on sm_89; as on sm_120, tests that compare against them run
   with noise-class gates off that architecture.
-- Values outside fp16's range (|x| > 65504 or |x| < ~6e-5) do not survive the bf16 → fp16 conversion; the model's
-  activations and weights are inside that range in practice.
+- A GEMM input row whose largest magnitude exceeds fp16's 65504 is multiplied by the power of two that brings it in
+  range and the output row by its inverse, both exact; only the row's values that the scale pushes below ~6e-5 lose
+  precision. Weights and attention operands are converted unscaled: above 65504 they would become inf, below ~6e-5
+  they lose precision (fp16 subnormals); the model's weights and q/k/v are inside that range in practice.
 
 ## Verification
 
 - `tools/sm75_check.cu` (not installed, built with `MOSSINFER_SM75=ON`) runs the Turing kernels against host
-  references: dtype conversion, exact attention at both head dimensions and engine-shaped varlen/GQA cases, the fp16
-  GEMM, the INT8 GEMM (bit-exact), and the VAE convolution and attention.
-- `tests/unit/test_sm75.cpp` covers the build's presence and the in-place conversions on a device.
-- The two refusals are exercised on the RTX 2080: `moss-infer generate --init-image …` and `--lora …` end with
+  references: dtype conversion, exact attention at both head dimensions and engine-shaped varlen/GQA cases (also
+  causal with key lengths unlike the query lengths, and that q/k/v are left untouched), the fp16 GEMM (bf16 and f32
+  out, a row above fp16's range, several row chunks), the INT8 GEMM (bit-exact), and the VAE convolution and attention.
+- `tests/unit/test_sm75.cpp` covers the build's presence and, on a Turing device, the weight conversion round trip and
+  the GEMM rows' per-row scale.
+- The refusals are exercised on the RTX 2080: `moss-infer generate --init-image …` and `--lora …` end with
   `not_supported` (CLI exit status 1, nothing written), and `moss-server` reports the same code in the `done` event.
 - End to end on an RTX 2080: `moss-infer generate` and `moss-server` (protocol 2) produce the same PNG bytes for the
   `reference`, `high` and `turbo` tiers; the report's `quality.gpu_arch` is `sm_75`.
@@ -98,6 +103,9 @@ roughly 3–5 s per prompt when streamed (its ~16 GB do not fit 8 GiB of VRAM);
   `mask`), `references` or `hires` is refused up front with the code `not_supported` and a message saying why, so
   nothing is loaded or planned for it.
 - Runtime LoRAs and control adapters: their GEMM kernels are built for sm_80+ only. A request that carries `loras` or a
-  control adapter is refused up front with `not_supported`, again before any weight is read or uploaded.
+  control adapter is refused up front with `not_supported`, again before any weight is read or uploaded; a component
+  session refuses them the same way, when it opens and for every run.
+- Forced INT8 GEMM tiles (`Int8Tile`, a test and benchmark option): sm_75 runs one tile (64x128x64), so a forced tile
+  is refused rather than ignored.
 
 The VAE decode is fp32 SIMT on sm_75, so a 1024² decode costs about 10–12 s; a future fp16 path would cut that.

src/kernels/basic/gemm.cpp

@@ -92,22 +92,8 @@ struct Gemm::Impl {
   cublasLtHandle_t handle = nullptr;
   mutable std::mutex mutex;
   std::unordered_map<Key, std::unique_ptr<Plan>, KeyHash> plans;
-  // sm_75: fp16 tensor cores with bf16 activations. x is converted to fp16 in a chunked scratch buffer, the weights
-  // are fp16 (converted once by the loader) and the GEMM writes fp16 (into the bf16 output buffer, or into `acc16`
-  // for an f32 output) which a second pass converts to the output dtype and adds the bias.
+  // sm_75: fp16 tensor cores with bf16 activations (linear_f16). The weights are fp16 (converted once by the loader).
   bool fp16 = false;
-  void* x16 = nullptr;
-  std::size_t x16_bytes = 0;
-  void* acc16 = nullptr;
-  std::size_t acc16_bytes = 0;
-
-  static void ensure_scratch(void*& p, std::size_t& cap, std::size_t need) {
-    if (cap >= need) return;
-    MI_CUDA_CHECK(cudaDeviceSynchronize());  // an earlier launch on another stream may still read it
-    if (p != nullptr) MI_CUDA_CHECK(cudaFree(p));
-    MI_CUDA_CHECK(cudaMalloc(&p, need));
-    cap = need;
-  }
 
   Plan& plan_for(const Key& key, const void* bias_ptr) {
     auto it = plans.find(key);
@@ -167,8 +153,11 @@ struct Gemm::Impl {
     return *plans.emplace(key, std::move(plan)).first->second;
   }
 
-  // sm_75 path: x is bf16, the GEMM runs fp16 in / fp32 out on Turing's tensor cores, and the result is converted to
-  // the output dtype with the bias added in fp32. Row-chunked so the fp16 copies stay in a small scratch buffer.
+  // sm_75 path: x is bf16, the GEMM runs fp16 in / fp32 out on Turing's tensor cores, and a second pass converts the
+  // result to the output dtype with the bias added in fp32 (one rounding, as the sm_80 path's BIAS epilogue). x's rows
+  // are converted to fp16 with a per-row power-of-two scale (1 for rows in fp16's range), undone in that pass. The
+  // fp16 rows, the scales and the fp32 result live in the caller's workspace, after the share cuBLASLt gets: row
+  // chunks of what fits, so the GEMM needs no memory of its own and each stream's workspace keeps its calls apart.
   void linear_f16(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, void* workspace,
                   std::size_t workspace_bytes, cudaStream_t stream) {
     const std::int64_t m = detail::rows_of(x);
@@ -179,31 +168,45 @@ struct Gemm::Impl {
     const std::int64_t ldw = row_stride(w);
     if (m == 0 || n == 0) return;
     const bool out_f32 = out.dtype == DType::kF32;
-    constexpr std::size_t kScratchBytes = std::size_t{16} << 20;
-    const std::size_t per_row = static_cast<std::size_t>(std::max(n, k)) * 2;
-    std::int64_t chunk = static_cast<std::int64_t>(kScratchBytes / std::max<std::size_t>(1, per_row));
-    chunk = std::max<std::int64_t>(1, std::min(chunk, m));
-    ensure_scratch(x16, x16_bytes, static_cast<std::size_t>(chunk) * static_cast<std::size_t>(k) * 2);
-    if (out_f32) {
-      ensure_scratch(acc16, acc16_bytes, static_cast<std::size_t>(chunk) * static_cast<std::size_t>(n) * 2);
+    constexpr std::size_t kAlign = 256;
+    const auto align_up = [](std::size_t v) { return (v + kAlign - 1) / kAlign * kAlign; };
+    const std::size_t lt_bytes = workspace_bytes / 8 / kAlign * kAlign;
+    const std::size_t per_row = static_cast<std::size_t>(k) * 2 + static_cast<std::size_t>(n) * 4 + sizeof(float);
+    const std::size_t scratch = workspace_bytes - lt_bytes;
+    const std::size_t slack = 3 * kAlign;
+    if (scratch < slack + per_row) {
+      detail::fail("linear", "the sm_75 path needs a workspace (Gemm::kWorkspaceBytes) for its fp16/fp32 rows, got " +
+                                 std::to_string(workspace_bytes) + " bytes");
+    }
+    // Equal chunks in whole 128-row tiles (no short tail that rereads the weight for a few rows) when the scratch
+    // holds at least one such tile.
+    std::int64_t chunk = static_cast<std::int64_t>((scratch - slack) / per_row);
+    constexpr std::int64_t kTileRows = 128;
+    if (chunk < m) {
+      if (chunk >= kTileRows) chunk = chunk / kTileRows * kTileRows;
+      const std::int64_t pieces = (m + chunk - 1) / chunk;
+      const std::int64_t even = (m + pieces - 1) / pieces;
+      chunk = std::min(chunk, chunk >= kTileRows ? (even + kTileRows - 1) / kTileRows * kTileRows : even);
     }
+    chunk = std::min(chunk, m);
+    char* const base = static_cast<char*>(workspace) + lt_bytes;
+    void* const x16 = base;
+    auto* const acc = reinterpret_cast<float*>(base + align_up(static_cast<std::size_t>(chunk) * k * 2));
+    auto* const inv_scale = acc + align_up(static_cast<std::size_t>(chunk) * n * 4) / sizeof(float);
     const void* bias_ptr = bias != nullptr ? bias->data : nullptr;
-    const bool bias_f32 = bias != nullptr && bias->dtype == DType::kF32;
+    const std::size_t out_elem = out_f32 ? 4 : 2;
     for (std::int64_t r0 = 0; r0 < m; r0 += chunk) {
       const std::int64_t rows = std::min(chunk, m - r0);
-      kernels::sm75::convert_rows_bf16_to_f16(x16, static_cast<const char*>(x.data) + r0 * ldx * 2, rows, k, ldx, k,
-                                              stream);
-      void* dptr = out_f32 ? acc16 : (static_cast<char*>(out.data) + r0 * ldo * 2);
-      const std::int64_t acc_ld = out_f32 ? n : ldo;
-      const Key key{rows, n, k, k, ldw, acc_ld, static_cast<int>(DType::kF16), -1, 1, workspace_bytes,
-                    alignment(x16), alignment(w.data), alignment(dptr), 0u};
+      kernels::sm75::convert_rows_bf16_to_f16_scaled(x16, inv_scale, static_cast<const char*>(x.data) + r0 * ldx * 2,
+                                                     rows, k, ldx, k, stream);
+      const Key key{rows, n, k, k, ldw, n, static_cast<int>(DType::kF32), -1, 1, lt_bytes,
+                    alignment(x16), alignment(w.data), alignment(acc), 0u};
       Plan& plan = plan_for(key, nullptr);
       const float alpha = 1.0f, beta = 0.0f;
-      MI_LT_CHECK(cublasLtMatmul(handle, plan.desc, &alpha, w.data, plan.a, x16, plan.b, &beta, dptr, plan.d, dptr,
+      MI_LT_CHECK(cublasLtMatmul(handle, plan.desc, &alpha, w.data, plan.a, x16, plan.b, &beta, acc, plan.d, acc,
                                  plan.d, &plan.algo, workspace, plan.workspace, stream));
-      // The output row offset in bytes: f32 rows are 4 bytes per element (dptr above is acc16 for them).
-      kernels::sm75::convert_gemm_out_f16(static_cast<char*>(out.data) + r0 * ldo * (out_f32 ? 4 : 2), dptr, bias_ptr,
-                                          bias_f32, out_f32, rows, n, ldo, acc_ld, stream);
+      kernels::sm75::convert_gemm_out_f32(static_cast<char*>(out.data) + r0 * ldo * static_cast<std::int64_t>(out_elem),
+                                          out_f32, acc, inv_scale, bias_ptr, rows, n, ldo, n, stream);
     }
   }
 };
@@ -223,8 +226,6 @@ Gemm::Gemm() : impl_(std::make_unique<Impl>()) {
 
 Gemm::~Gemm() {
   impl_->plans.clear();
-  if (impl_->x16 != nullptr) cudaFree(impl_->x16);
-  if (impl_->acc16 != nullptr) cudaFree(impl_->acc16);
   if (impl_->handle != nullptr) cublasLtDestroy(impl_->handle);
 }
 

src/kernels/basic/gemm.h

@@ -31,7 +31,9 @@ class Gemm {
   // x: bf16, rows of K with a contiguous last dim whose leading dims flatten to one row stride (M rows).
   // w: bf16 [N, K], contiguous rows. bias: nullptr or [N] of the output dtype (cuBLASLt has no bf16 bias
   // with f32 output; widen it with kernels::cast).
-  // out: bf16 or f32 [..., N], M rows with one row stride. workspace: 256-byte aligned, or null with 0 bytes.
+  // out: bf16 or f32 [..., N], M rows with one row stride. workspace: 256-byte aligned, or null with 0 bytes. On sm_75
+  // (w fp16, kernels/sm75) the workspace also holds the fp16 input rows and the fp32 result, so it is required there
+  // (Gemm::kWorkspaceBytes) and must not be shared by calls in flight on other streams.
   // Thread-safe; launches on `stream` only.
   void linear(const Tensor& out, const Tensor& x, const Tensor& w, const Tensor* bias, void* workspace,
               std::size_t workspace_bytes, cudaStream_t stream);

src/kernels/quant/int8_gemm.cu

@@ -92,6 +92,12 @@ Int8Tile resolve_tile(Int8Tile tile, const Problem& p) {
   return tile;
 }
 
+void check_sm75_tile(const char* op, Int8Tile tile) {
+  if (tile != Int8Tile::kAuto) {
+    fail(op, std::string("a forced tile (") + int8_tile_name(tile) + ") names an sm_80 kernel; sm_75 runs one tile");
+  }
+}
+
 void check_status(cutlass::Status status, const char* op, const char* stage) {
   if (status != cutlass::Status::kSuccess) {
     throw CudaError(std::string("mi::kernels::") + op + ": CUTLASS " + stage + " failed: " +
@@ -151,6 +157,7 @@ void int8_gemm(const Int8Operands& op, const __nv_bfloat16* bias, __nv_bfloat16*
   detail::check_vector(name, "bias", bias, /*optional=*/true);
   if (p.m == 0) return;
   if (kernels::sm75::turing()) {  // Turing's 8x8x16 INT8 MMA (src/kernels/sm75/int8_gemm_sm75.cu)
+    detail::check_sm75_tile(name, tile);
     kernels::sm75::int8_gemm(op, bias, out, ldo, stream);
     return;
   }

src/kernels/quant/int8_gemm_impl.cuh

@@ -344,6 +344,9 @@ void check_output(const char* op, const char* name, const void* p, std::int64_t
 void check_vector(const char* op, const char* name, const void* p, bool optional);
 
 Int8Tile resolve_tile(Int8Tile tile, const Problem& p);
+// sm_75 runs one tile (src/kernels/sm75/int8_gemm_sm75.cu, 64x128x64): a forced tile names an sm_80 instantiation and
+// is refused rather than ignored.
+void check_sm75_tile(const char* op, Int8Tile tile);
 
 void check_status(cutlass::Status status, const char* op, const char* stage);
 

src/kernels/quant/int8_gemm_residual.cu

@@ -39,6 +39,7 @@ void int8_gemm_gated_residual(const Int8Operands& op, const __nv_bfloat16* bias,
   detail::check_vector(name, "gate", gate, /*optional=*/false);
   if (p.m == 0) return;
   if (kernels::sm75::turing()) {
+    detail::check_sm75_tile(name, tile);
     kernels::sm75::int8_gemm_gated_residual(op, bias, gate, residual, ldr, stream);
     return;
   }

src/kernels/quant/int8_gemm_situ.cu

@@ -105,6 +105,7 @@ void int8_gemm_situ_glu(const Int8Operands& op, __nv_bfloat16* out, std::int64_t
   }
   if (p.m == 0) return;
   if (kernels::sm75::turing()) {
+    detail::check_sm75_tile(name, tile);
     kernels::sm75::int8_gemm_situ_glu(op, out, ldo, stream);
     return;
   }