diff --git a/docs/sm75.md b/docs/sm75.md index 0a174f0..5239d83 100644 --- a/docs/sm75.md +++ b/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. diff --git a/src/kernels/basic/gemm.cpp b/src/kernels/basic/gemm.cpp index 5abff99..5ea5912 100644 --- a/src/kernels/basic/gemm.cpp +++ b/src/kernels/basic/gemm.cpp @@ -92,22 +92,8 @@ struct Gemm::Impl { cublasLtHandle_t handle = nullptr; mutable std::mutex mutex; std::unordered_map, 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::max(n, k)) * 2; - std::int64_t chunk = static_cast(kScratchBytes / std::max(1, per_row)); - chunk = std::max(1, std::min(chunk, m)); - ensure_scratch(x16, x16_bytes, static_cast(chunk) * static_cast(k) * 2); - if (out_f32) { - ensure_scratch(acc16, acc16_bytes, static_cast(chunk) * static_cast(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(k) * 2 + static_cast(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((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(workspace) + lt_bytes; + void* const x16 = base; + auto* const acc = reinterpret_cast(base + align_up(static_cast(chunk) * k * 2)); + auto* const inv_scale = acc + align_up(static_cast(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(x.data) + r0 * ldx * 2, rows, k, ldx, k, - stream); - void* dptr = out_f32 ? acc16 : (static_cast(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(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(x.data) + r0 * ldx * 2, + rows, k, ldx, k, stream); + const Key key{rows, n, k, k, ldw, n, static_cast(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(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(out.data) + r0 * ldo * static_cast(out_elem), + out_f32, acc, inv_scale, bias_ptr, rows, n, ldo, n, stream); } } }; @@ -223,8 +226,6 @@ Gemm::Gemm() : impl_(std::make_unique()) { 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); } diff --git a/src/kernels/basic/gemm.h b/src/kernels/basic/gemm.h index 2f0d74f..a52bb64 100644 --- a/src/kernels/basic/gemm.h +++ b/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); diff --git a/src/kernels/quant/int8_gemm.cu b/src/kernels/quant/int8_gemm.cu index 30a8394..68a7078 100644 --- a/src/kernels/quant/int8_gemm.cu +++ b/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; } diff --git a/src/kernels/quant/int8_gemm_impl.cuh b/src/kernels/quant/int8_gemm_impl.cuh index e82a257..cec18a9 100644 --- a/src/kernels/quant/int8_gemm_impl.cuh +++ b/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); diff --git a/src/kernels/quant/int8_gemm_residual.cu b/src/kernels/quant/int8_gemm_residual.cu index 7892872..6eb2b01 100644 --- a/src/kernels/quant/int8_gemm_residual.cu +++ b/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; } diff --git a/src/kernels/quant/int8_gemm_situ.cu b/src/kernels/quant/int8_gemm_situ.cu index 424e443..6562b6c 100644 --- a/src/kernels/quant/int8_gemm_situ.cu +++ b/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; } diff --git a/src/kernels/sm75/attention_sm75.cu b/src/kernels/sm75/attention_sm75.cu index 0ee6269..d7ef6bb 100644 --- a/src/kernels/sm75/attention_sm75.cu +++ b/src/kernels/sm75/attention_sm75.cu @@ -8,11 +8,12 @@ // scores in the m16n8k8 accumulator layout, row max/sum reduced across the four lanes of a row group, and the P·V // product reusing the normalized scores directly as the A operand (the m16n8k8 accumulator layout of Q·Kᵀ is exactly // the A-operand layout of P·V for the next MMA). GQA maps query head h to KV head h / (heads / kv_heads); sequences are -// varlen through cu_seqlens. +// varlen through cu_seqlens. The causal mask is aligned bottom-right (key <= query + len_k - len_q), as +// FlashAttention-2's; a row that sees no key gets a zero output and an +inf log-sum-exp, as FA2's. // -// Precision: the engine's tensors are bf16 and sm_75's MMA is fp16. bf16 -> fp16 is exact for values in fp16's normal -// range and fp16 -> bf16 recovers the original value, so q/k/v are converted in place before the launch and restored -// after it; `out` is written as fp16 by the kernel and converted to bf16 in place. +// Precision: the engine's tensors are bf16 and sm_75's MMA is fp16. q/k/v are read as bf16 and converted to fp16 as +// they are staged into shared memory (exact for values in fp16's normal range); the caller's tensors are never +// written. `out` is rounded once, fp32 -> bf16, as FA2's. #include #include #include @@ -52,6 +53,17 @@ __device__ __forceinline__ std::uint32_t pack_half2(float lo, float hi) { return c.u; } +__device__ __forceinline__ std::uint32_t pack_bf162(float lo, float hi) { + union { + __nv_bfloat162 h; + std::uint32_t u; + } c; + c.h = __floats2bfloat162_rn(lo, hi); + return c.u; +} + +__device__ __forceinline__ __half to_half(__nv_bfloat16 v) { return __float2half_rn(__bfloat162float(v)); } + __device__ __forceinline__ std::uint32_t load_u32(const void* p) { return *static_cast(p); } @@ -61,10 +73,10 @@ __device__ __forceinline__ void store_u32(void* p, std::uint32_t v) { } struct Params { - const __half* q = nullptr; - const __half* k = nullptr; - const __half* v = nullptr; - __half* o = nullptr; + const __nv_bfloat16* q = nullptr; + const __nv_bfloat16* k = nullptr; + const __nv_bfloat16* v = nullptr; + __nv_bfloat16* o = nullptr; float* lse = nullptr; std::int64_t q_row = 0, q_head = 0; std::int64_t k_row = 0, k_head = 0; @@ -101,6 +113,7 @@ __global__ void attention_kernel_sm75(const Params p) { if (row_base >= len_q) return; const int rows = min(kBM, len_q - row_base); const int kv_head = head / (p.heads / p.kv_heads); + const int causal_shift = len_k - len_q; // bottom-right: query row i sees keys 0 .. i + causal_shift const int tid = threadIdx.x; const int lane = tid % 32; const int warp = tid / 32; @@ -110,7 +123,7 @@ __global__ void attention_kernel_sm75(const Params p) { // Q tile into shared memory (rows past this sequence's end are zero). for (int idx = tid; idx < kBM * D; idx += kThreads) { const int r = idx / D, d = idx % D; - qs[r * kQStride + d] = (r < rows) ? p.q[(q_start + row_base + r) * p.q_row + head * p.q_head + d] + qs[r * kQStride + d] = (r < rows) ? to_half(p.q[(q_start + row_base + r) * p.q_row + head * p.q_head + d]) : __float2half_rn(0.f); } __syncthreads(); @@ -133,7 +146,7 @@ __global__ void attention_kernel_sm75(const Params p) { float m0 = -1e30f, m1 = -1e30f; // running row max for r0 / r1 float l0 = 0.f, l1 = 0.f; // running row sum - const int key_tiles = p.causal ? min((len_k + BN - 1) / BN, (row_base + rows + BN - 1) / BN) + const int key_tiles = p.causal ? min((len_k + BN - 1) / BN, max(0, (row_base + rows + causal_shift + BN - 1) / BN)) : (len_k + BN - 1) / BN; for (int tile = 0; tile < key_tiles; ++tile) { const int j0 = tile * BN; @@ -142,8 +155,10 @@ __global__ void attention_kernel_sm75(const Params p) { const int j = idx / D, d = idx % D; const int token = j0 + j; const bool valid = token < len_k; - ks[j * kQStride + d] = valid ? p.k[(k_start + token) * p.k_row + kv_head * p.k_head + d] : __float2half_rn(0.f); - vt[d * kVStride + j] = valid ? p.v[(k_start + token) * p.v_row + kv_head * p.v_head + d] : __float2half_rn(0.f); + ks[j * kQStride + d] = + valid ? to_half(p.k[(k_start + token) * p.k_row + kv_head * p.k_head + d]) : __float2half_rn(0.f); + vt[d * kVStride + j] = + valid ? to_half(p.v[(k_start + token) * p.v_row + kv_head * p.v_head + d]) : __float2half_rn(0.f); } __syncthreads(); @@ -170,7 +185,7 @@ __global__ void attention_kernel_sm75(const Params p) { for (int e = 0; e < 4; ++e) { const int key = key0 + (e & 1); const int row_q = (e < 2) ? row_q0 : row_q1; - const bool valid = key < len_k && (!p.causal || key <= row_q); + const bool valid = key < len_k && (!p.causal || key <= row_q + causal_shift); s[t][e] = valid ? s[t][e] * p.scale : -1e30f; } tile_max0 = fmaxf(tile_max0, fmaxf(s[t][0], s[t][1])); @@ -186,14 +201,16 @@ __global__ void attention_kernel_sm75(const Params p) { const float alpha0 = (m0 <= -1e29f) ? 0.f : exp2f((m0 - m_new0) * kLog2e); const float alpha1 = (m1 <= -1e29f) ? 0.f : exp2f((m1 - m_new1) * kLog2e); - // P = exp(S - m_new) as fp16, and the running sum. + // P = exp(S - m_new) as fp16, and the running sum. A masked score is 0 even while its row has seen no key yet + // (m_new still the -1e30 sentinel, where exp(S - m_new) would be 1). float tile_sum0 = 0.f, tile_sum1 = 0.f; + const auto prob = [](float sv, float mv) { return sv <= -1e29f ? 0.f : exp2f((sv - mv) * kLog2e); }; #pragma unroll for (int t = 0; t < kSubtiles; ++t) { - const float p0 = exp2f((s[t][0] - m_new0) * kLog2e); - const float p1 = exp2f((s[t][1] - m_new0) * kLog2e); - const float p2 = exp2f((s[t][2] - m_new1) * kLog2e); - const float p3 = exp2f((s[t][3] - m_new1) * kLog2e); + const float p0 = prob(s[t][0], m_new0); + const float p1 = prob(s[t][1], m_new0); + const float p2 = prob(s[t][2], m_new1); + const float p3 = prob(s[t][3], m_new1); tile_sum0 += p0 + p1; tile_sum1 += p2 + p3; s[t][0] = p0; @@ -231,26 +248,27 @@ __global__ void attention_kernel_sm75(const Params p) { __syncthreads(); } - // Normalize, then out (fp16) and the log-sum-exp workspace (fp32 [heads, total_q], unpadded). + // Normalize, then out (bf16) and the log-sum-exp workspace (fp32 [heads, total_q], unpadded; +inf for a row that saw + // no key, as FA2's). const float inv0 = l0 > 0.f ? 1.f / l0 : 0.f; const float inv1 = l1 > 0.f ? 1.f / l1 : 0.f; if (r0 < rows) { #pragma unroll for (int i = 0; i < kDimChunks; ++i) { store_u32(&p.o[(q_start + row_base + r0) * p.o_row + head * p.o_head + i * 8 + tig * 2], - pack_half2(o[i][0] * inv0, o[i][1] * inv0)); + pack_bf162(o[i][0] * inv0, o[i][1] * inv0)); } if (tig == 0 && p.lse != nullptr) p.lse[static_cast(head) * p.total_q + q_start + row_base + r0] = - m0 + logf(l0); + l0 > 0.f ? m0 + logf(l0) : INFINITY; } if (r1 < rows) { #pragma unroll for (int i = 0; i < kDimChunks; ++i) { store_u32(&p.o[(q_start + row_base + r1) * p.o_row + head * p.o_head + i * 8 + tig * 2], - pack_half2(o[i][2] * inv1, o[i][3] * inv1)); + pack_bf162(o[i][2] * inv1, o[i][3] * inv1)); } if (tig == 0 && p.lse != nullptr) p.lse[static_cast(head) * p.total_q + q_start + row_base + r1] = - m1 + logf(l1); + l1 > 0.f ? m1 + logf(l1) : INFINITY; } } @@ -292,15 +310,11 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& throw std::invalid_argument("sm75::flash_fwd: workspace too small for the log-sum-exp"); } - convert_tensor_bf16_to_f16_inplace(q, total_q, heads, head_dim, stream); - convert_tensor_bf16_to_f16_inplace(k, k.dim(0), kv_heads, head_dim, stream); - convert_tensor_bf16_to_f16_inplace(v, v.dim(0), kv_heads, head_dim, stream); - Params p; - p.q = static_cast(q.data); - p.k = static_cast(k.data); - p.v = static_cast(v.data); - p.o = static_cast<__half*>(out.data); + p.q = static_cast(q.data); + p.k = static_cast(k.data); + p.v = static_cast(v.data); + p.o = static_cast<__nv_bfloat16*>(out.data); p.lse = static_cast(workspace); p.q_row = q.stride[0]; p.q_head = q.stride[1]; @@ -324,11 +338,6 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& } else { launch<256, 16>(p, max_seqlen, batch, stream); } - - convert_tensor_f16_to_bf16_inplace(out, out.dim(0), out.dim(1), out.dim(2), stream); - convert_tensor_f16_to_bf16_inplace(q, total_q, heads, head_dim, stream); - convert_tensor_f16_to_bf16_inplace(k, k.dim(0), kv_heads, head_dim, stream); - convert_tensor_f16_to_bf16_inplace(v, v.dim(0), kv_heads, head_dim, stream); } } // namespace mi::kernels::sm75 diff --git a/src/kernels/sm75/convert.cu b/src/kernels/sm75/convert.cu index af7df05..90db435 100644 --- a/src/kernels/sm75/convert.cu +++ b/src/kernels/sm75/convert.cu @@ -1,13 +1,14 @@ // sm_75 dtype conversion kernels: every tensor-core operand in the engine is bf16, while Turing's tensor cores take -// fp16. These kernels convert in place (the same 2-byte slot) or into a separate buffer; the attention wrapper relies -// on the round trip being exact for in-range values (fp16 has more mantissa bits than bf16, so bf16 -> fp16 is lossless -// and fp16 -> bf16 recovers the original value as long as it stayed in fp16's normal range). +// fp16. bf16 -> fp16 is lossless for values in fp16's normal range (fp16 has more mantissa bits than bf16), and +// fp16 -> bf16 recovers the original value. The weight conversions run in place (the same 2-byte slot); the GEMM's +// activation rows are converted into its scratch with a per-row power-of-two scale that keeps them in range. #include #include #include #include #include +#include #include "core/cuda_check.h" #include "kernels/sm75/sm75.h" @@ -17,67 +18,28 @@ namespace { constexpr int kThreads = 256; constexpr int kMaxBlocks = 65535; +constexpr float kF16Max = 65504.f; -__global__ void bf16_to_f16_flat(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, std::int64_t n) { +// In place when dst == src (no __restrict__): each element is read and written by the same thread. +__global__ void bf16_to_f16_flat(const __nv_bfloat16* src, __half* dst, std::int64_t n) { const std::int64_t stride = static_cast(gridDim.x) * blockDim.x; for (std::int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += stride) { dst[i] = __float2half_rn(__bfloat162float(src[i])); } } -__global__ void f16_to_bf16_flat(const __half* __restrict__ src, __nv_bfloat16* __restrict__ dst, std::int64_t n) { +__global__ void f16_to_bf16_flat(const __half* src, __nv_bfloat16* dst, std::int64_t n) { const std::int64_t stride = static_cast(gridDim.x) * blockDim.x; for (std::int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += stride) { dst[i] = __float2bfloat16_rn(__half2float(src[i])); } } -// One block per token row, `dim` contiguous elements per (token, head) run. -__global__ void bf16_to_f16_row(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, std::int64_t stride, - std::int64_t dim) { - const __nv_bfloat16* s = src + static_cast(blockIdx.x) * stride; - __half* d = dst + static_cast(blockIdx.x) * stride; - for (std::int64_t i = threadIdx.x; i < dim; i += blockDim.x) { - d[i] = __float2half_rn(__bfloat162float(s[i])); - } -} - -__global__ void f16_to_bf16_row(const __half* __restrict__ src, __nv_bfloat16* __restrict__ dst, std::int64_t stride, - std::int64_t dim) { - const __half* s = src + static_cast(blockIdx.x) * stride; - __nv_bfloat16* d = dst + static_cast(blockIdx.x) * stride; - for (std::int64_t i = threadIdx.x; i < dim; i += blockDim.x) { - d[i] = __float2bfloat16_rn(__half2float(s[i])); - } -} - int blocks_for(std::int64_t n) { const std::int64_t blocks = (n + kThreads - 1) / kThreads; return static_cast(blocks < kMaxBlocks ? (blocks > 0 ? blocks : 1) : kMaxBlocks); } -// The attention tensors are [tokens, heads, dim] with a contiguous last dim, so every (token, head) pair is one run of -// `dim` elements at token * stride[0] + head * stride[1]. q/k/v may be views of a fused QKV output (non-contiguous -// heads), which the per-head loop handles. -void convert_3d_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, bool to_f16, - cudaStream_t stream) { - if (rows == 0 || heads == 0 || dim == 0) return; - // One block per row on grid.x, whose limit is 2^31 - 1 (kMaxBlocks is the grid-stride cap of the flat kernels). - if (rows > 0x7fffffff) throw std::length_error("sm75::convert_3d_inplace: more rows than one grid dimension holds"); - auto* base = static_cast(t.data); - for (std::int64_t h = 0; h < heads; ++h) { - char* p = base + h * static_cast(t.stride[1]) * 2; - if (to_f16) { - bf16_to_f16_row<<(rows), kThreads, 0, stream>>>( - reinterpret_cast(p), reinterpret_cast<__half*>(p), t.stride[0], dim); - } else { - f16_to_bf16_row<<(rows), kThreads, 0, stream>>>( - reinterpret_cast(p), reinterpret_cast<__nv_bfloat16*>(p), t.stride[0], dim); - } - } - MI_CUDA_CHECK_LAUNCH(); -} - } // namespace bool available() { return true; } @@ -136,70 +98,82 @@ void convert_f16_to_bf16(void* dst, const void* src, std::int64_t n, cudaStream_ MI_CUDA_CHECK_LAUNCH(); } -void convert_tensor_bf16_to_f16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, - cudaStream_t stream) { - convert_3d_inplace(t, rows, heads, dim, /*to_f16=*/true, stream); -} - -void convert_tensor_f16_to_bf16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, - cudaStream_t stream) { - convert_3d_inplace(t, rows, heads, dim, /*to_f16=*/false, stream); -} - namespace { -__global__ void bf16_to_f16_rows(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, std::int64_t cols, - std::int64_t src_stride, std::int64_t dst_stride) { +// One block per row: the row's largest magnitude, then the row converted with scale 2^-e, the smallest power of two +// that brings it within fp16's range (1 for every row already in range, so those convert exactly as plain bf16 -> +// fp16). Scaling by a power of two is exact, so a scaled row loses nothing but the low values it pushes below fp16's +// normal range; inv_scale[row] = 2^e undoes it in the GEMM's output pass. A row holding inf is left unscaled (a NaN +// is skipped by the max and carried through either way). +__global__ void bf16_to_f16_rows_scaled(const __nv_bfloat16* __restrict__ src, __half* __restrict__ dst, + float* __restrict__ inv_scale, std::int64_t cols, std::int64_t src_stride, + std::int64_t dst_stride) { + __shared__ float warp_max[kThreads / 32]; + __shared__ float row_scale; const __nv_bfloat16* s = src + static_cast(blockIdx.x) * src_stride; __half* d = dst + static_cast(blockIdx.x) * dst_stride; - for (std::int64_t i = threadIdx.x; i < cols; i += blockDim.x) { - d[i] = __float2half_rn(__bfloat162float(s[i])); - } -} - -// out is bf16: `acc` may alias it (the GEMM writes fp16 into the output buffer, this pass converts in place). -__global__ void gemm_out_f16_bf16(__half* __restrict__ acc, __nv_bfloat16* __restrict__ out, - const __nv_bfloat16* __restrict__ bias, std::int64_t cols, std::int64_t out_stride, - std::int64_t acc_stride) { - const __half* a = acc + static_cast(blockIdx.x) * acc_stride; - __nv_bfloat16* o = out + static_cast(blockIdx.x) * out_stride; - for (std::int64_t c = threadIdx.x; c < cols; c += blockDim.x) { - const float v = __half2float(a[c]) + (bias != nullptr ? __bfloat162float(bias[c]) : 0.f); - o[c] = __float2bfloat16_rn(v); + float amax = 0.f; + for (std::int64_t i = threadIdx.x; i < cols; i += blockDim.x) amax = fmaxf(amax, fabsf(__bfloat162float(s[i]))); + for (int off = 16; off > 0; off >>= 1) amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, off)); + if (threadIdx.x % 32 == 0) warp_max[threadIdx.x / 32] = amax; + __syncthreads(); + if (threadIdx.x == 0) { + float m = 0.f; + for (int w = 0; w < kThreads / 32; ++w) m = fmaxf(m, warp_max[w]); + int e = 0; + if (m > kF16Max && isfinite(m)) e = static_cast(ceilf(log2f(m / kF16Max))); + while (e > 0 && ldexpf(m, -e) > kF16Max) ++e; // log2f's rounding at an exact power of two + row_scale = ldexpf(1.f, -e); + inv_scale[blockIdx.x] = ldexpf(1.f, e); } -} - -__global__ void gemm_out_f16_f32(const __half* __restrict__ acc, float* __restrict__ out, const float* __restrict__ bias, - std::int64_t cols, std::int64_t out_stride, std::int64_t acc_stride) { - const __half* a = acc + static_cast(blockIdx.x) * acc_stride; - float* o = out + static_cast(blockIdx.x) * out_stride; + __syncthreads(); + const float scale = row_scale; + for (std::int64_t i = threadIdx.x; i < cols; i += blockDim.x) d[i] = __float2half_rn(__bfloat162float(s[i]) * scale); +} + +// out[r][c] = out_type(acc[r][c] * inv_scale[r] + bias[c]): the fp32 add before the single final rounding. +template +__global__ void gemm_out_f32(const float* __restrict__ acc, const float* __restrict__ inv_scale, Out* __restrict__ out, + const Bias* __restrict__ bias, std::int64_t cols, std::int64_t out_stride, + std::int64_t acc_stride) { + const float* a = acc + static_cast(blockIdx.x) * acc_stride; + Out* o = out + static_cast(blockIdx.x) * out_stride; + const float sc = inv_scale[blockIdx.x]; for (std::int64_t c = threadIdx.x; c < cols; c += blockDim.x) { - o[c] = __half2float(a[c]) + (bias != nullptr ? bias[c] : 0.f); + float v = a[c] * sc; + if (bias != nullptr) v += static_cast(bias[c]); + if constexpr (std::is_same_v) { + o[c] = v; + } else { + o[c] = __float2bfloat16_rn(v); + } } } } // namespace -void convert_rows_bf16_to_f16(void* dst, const void* src, std::int64_t rows, std::int64_t cols, - std::int64_t src_stride, std::int64_t dst_stride, cudaStream_t stream) { +void convert_rows_bf16_to_f16_scaled(void* dst, float* inv_scale, const void* src, std::int64_t rows, + std::int64_t cols, std::int64_t src_stride, std::int64_t dst_stride, + cudaStream_t stream) { if (rows <= 0 || cols <= 0) return; - bf16_to_f16_rows<<(rows), kThreads, 0, stream>>>( - static_cast(src), static_cast<__half*>(dst), cols, src_stride, dst_stride); + bf16_to_f16_rows_scaled<<(rows), kThreads, 0, stream>>>( + static_cast(src), static_cast<__half*>(dst), inv_scale, cols, src_stride, dst_stride); MI_CUDA_CHECK_LAUNCH(); } -void convert_gemm_out_f16(void* out, const void* acc, const void* bias, bool bias_f32, bool out_f32, +void convert_gemm_out_f32(void* out, bool out_f32, const float* acc, const float* inv_scale, const void* bias, std::int64_t rows, std::int64_t cols, std::int64_t out_stride, std::int64_t acc_stride, cudaStream_t stream) { if (rows <= 0 || cols <= 0) return; + const auto grid = static_cast(rows); if (out_f32) { - gemm_out_f16_f32<<(rows), kThreads, 0, stream>>>( - static_cast(acc), static_cast(out), static_cast(bias_f32 ? bias : nullptr), - cols, out_stride, acc_stride); + gemm_out_f32<<>>(acc, inv_scale, static_cast(out), + static_cast(bias), cols, out_stride, + acc_stride); } else { - gemm_out_f16_bf16<<(rows), kThreads, 0, stream>>>( - static_cast<__half*>(const_cast(acc)), static_cast<__nv_bfloat16*>(out), - static_cast(bias), cols, out_stride, acc_stride); + gemm_out_f32<__nv_bfloat16, __nv_bfloat16><<>>( + acc, inv_scale, static_cast<__nv_bfloat16*>(out), static_cast(bias), cols, out_stride, + acc_stride); } MI_CUDA_CHECK_LAUNCH(); } diff --git a/src/kernels/sm75/sm75.h b/src/kernels/sm75/sm75.h index f32c031..29b6ac1 100644 --- a/src/kernels/sm75/sm75.h +++ b/src/kernels/sm75/sm75.h @@ -7,8 +7,10 @@ // that dispatches into it lives in the sm_80 translation units and is arch-independent. // // Numerics: sm_75 has fp16 tensor cores but no bf16 ones, so every tensor-core operand is converted bf16 -> fp16 at -// the boundary (lossless for |x| in fp16's normal range; values outside it saturate/round, which the engine reports -// as the sm_75 numerics). Accumulation stays fp32 where the sm_80 kernel had fp32 accumulation. +// the boundary (lossless for |x| in fp16's normal range). GEMM activations get a per-row power-of-two scale that keeps +// rows above 65504 in range; weights and attention operands above it would become inf, values below about 6e-5 lose +// precision (fp16 subnormals). Accumulation stays fp32 where the sm_80 kernel had fp32 accumulation, and outputs are +// rounded once, from fp32. #pragma once #include @@ -46,22 +48,17 @@ void convert_weight_storage_inplace(Tensor& w, cudaStream_t stream); void convert_bf16_to_f16(void* dst, const void* src, std::int64_t n, cudaStream_t stream); void convert_f16_to_bf16(void* dst, const void* src, std::int64_t n, cudaStream_t stream); -// bf16 -> fp16 over a 3-D tensor [rows, heads, dim] with free row/head strides, in place: every element is replaced by -// its fp16 bit pattern in the same 2-byte slot. Used by the sm_75 attention, which needs fp16 operands and restores -// bf16 after the launch so callers that reuse q/k/v are unaffected. -void convert_tensor_bf16_to_f16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, - cudaStream_t stream); -void convert_tensor_f16_to_bf16_inplace(const Tensor& t, std::int64_t rows, std::int64_t heads, std::int64_t dim, - cudaStream_t stream); - -// bf16 [rows, cols] (free row strides) -> fp16, into `dst` (dst_stride elements) or in place when dst == src. -void convert_rows_bf16_to_f16(void* dst, const void* src, std::int64_t rows, std::int64_t cols, - std::int64_t src_stride, std::int64_t dst_stride, cudaStream_t stream); - -// GEMM output conversion: `acc` is fp16 [rows, cols] (acc_stride elements; may alias `out` when `out` is bf16), -// `bias` is nullptr or fp32/bf16 [cols] (bias_f32 selects), `out` is bf16 or f32 [rows, cols] (out_stride). -// out[r][c] = out_type(float(acc[r][c]) + (bias ? bias[c] : 0)), the fp32 add before the final rounding. -void convert_gemm_out_f16(void* out, const void* acc, const void* bias, bool bias_f32, bool out_f32, +// GEMM input rows: bf16 [rows, cols] (src_stride elements) -> fp16 `dst` (dst_stride), each row multiplied by the +// power of two that brings its largest magnitude within fp16's range (1 for in-range rows); inv_scale[r] (fp32, +// device) receives the inverse. +void convert_rows_bf16_to_f16_scaled(void* dst, float* inv_scale, const void* src, std::int64_t rows, + std::int64_t cols, std::int64_t src_stride, std::int64_t dst_stride, + cudaStream_t stream); + +// GEMM output pass: `acc` is fp32 [rows, cols] (acc_stride), `out` bf16 or f32 (out_f32) [rows, cols] (out_stride), +// `bias` nullptr or [cols] of the output dtype. out[r][c] = out_type(acc[r][c] * inv_scale[r] + bias[c]), the fp32 add +// before the single final rounding (as the sm_80 path's cuBLASLt BIAS epilogue). +void convert_gemm_out_f32(void* out, bool out_f32, const float* acc, const float* inv_scale, const void* bias, std::int64_t rows, std::int64_t cols, std::int64_t out_stride, std::int64_t acc_stride, cudaStream_t stream); @@ -78,10 +75,9 @@ void vae_attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const void vae_weight_to_krsc(const Tensor& w_kcrs, const Tensor& w_krsc, cudaStream_t stream); // Exact (non-approximate) attention, the sm_75 counterpart of attention::flash_fwd for the same arguments and -// supported configurations (head_dim 128 non-causal for the DiT, head_dim 256 causal for the text encoder). q, k, v -// and out are bf16; they are converted to fp16 in place around the launch and restored to bf16 (the values are -// recoverable exactly for in-range inputs: fp16 has more mantissa bits than bf16). The kernel itself is the vendored -// FlashAttention-2 with its own sm_75 traits (fp16 m16n8k8 MMA, no cp.async). +// supported configurations (head_dim 128 non-causal for the DiT, head_dim 256 causal for the text encoder, whose +// causal mask is aligned bottom-right as FlashAttention-2's). q, k, v and out are bf16 and never modified but `out`: +// the kernel converts q/k/v to fp16 as it stages them and writes `out` from fp32 (attention_sm75.cu, fp16 m16n8k8 MMA). void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& out, const Tensor& cu_seqlens_q, const Tensor& cu_seqlens_k, int max_seqlen, bool causal, float scale, void* workspace, std::size_t workspace_bytes, cudaStream_t stream); diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index 24356da..14364ef 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -7,6 +7,7 @@ #include #include +#include #include "core/cuda_check.h" #include "kernels/sm75/sm75.h" @@ -171,6 +172,8 @@ void vae_weight_to_krsc(const Tensor& w_kcrs, const Tensor& w_krsc, cudaStream_t const int blocks = static_cast(std::min((total + 255) / 256, 4096)); if (w_krsc.dtype == DType::kF32) { kcrs_to_krsc_sm75<<>>(w_kcrs.ptr(), w_krsc.ptr(), k, c, rs); + } else if (w_krsc.dtype != DType::kBF16) { + throw std::invalid_argument("sm75::vae_weight_to_krsc: destination must be f32 or bf16"); } else { kcrs_to_krsc_sm75<__nv_bfloat16><<>>(w_kcrs.ptr(), w_krsc.ptr<__nv_bfloat16>(), k, c, rs); diff --git a/tests/unit/test_attention.cu b/tests/unit/test_attention.cu index 4eda288..02ac2e3 100644 --- a/tests/unit/test_attention.cu +++ b/tests/unit/test_attention.cu @@ -383,7 +383,7 @@ int max_smem_per_block() { } // Compute capability major of the current device: the four forced-tile tests below name sm_80 instantiations, which an -// sm_75 build does not carry (its dispatcher runs the vendored kernel's own Turing traits instead). +// sm_75 device does not run (its dispatcher runs src/kernels/sm75/attention_sm75.cu instead). int compute_major() { int device = 0, v = 0; MI_CUDA_CHECK(cudaGetDevice(&device)); diff --git a/tests/unit/test_sm75.cpp b/tests/unit/test_sm75.cpp index b7a0846..884668b 100644 --- a/tests/unit/test_sm75.cpp +++ b/tests/unit/test_sm75.cpp @@ -100,48 +100,47 @@ TEST(GpuSm75Convert, RoundTripIsExactInFp16Range) { #endif } -// The attention tensors are [tokens, heads, dim] views with free row/head strides: the in-place conversion must touch -// exactly the view's elements and no neighbour. -TEST(GpuSm75Convert, StridedInPlaceConversionTouchesOnlyTheView) { +// GEMM input rows: an in-range row converts exactly with scale 1; a row above fp16's range is scaled by the power of +// two that brings it in range, and inv_scale undoes it exactly (bf16 values times a power of two stay representable). +TEST(GpuSm75Convert, GemmRowsScaleOnlyOutOfRangeRows) { #ifndef MOSSINFER_HAVE_SM75 GTEST_SKIP() << "built without MOSSINFER_SM75"; #else if (!mi::kernels::sm75::turing()) GTEST_SKIP() << "the sm_75 kernels are built 75-real: needs a Turing device"; - constexpr int kRows = 3, kHeads = 2, kDim = 8, kHeadStride = 16; // heads interleaved: [token][head][16] - std::vector storage(static_cast(kRows) * kHeads * kHeadStride, -12345.f); - DeviceBf16 buf(storage); - Tensor view = Tensor::view(buf.data, DType::kBF16, Device::kCUDA, {kRows, kHeads, kDim}); - view.stride[0] = kHeads * kHeadStride; - view.stride[1] = kHeadStride; - view.stride[2] = 1; - const std::vector before = buf.host(); - mi::kernels::sm75::convert_tensor_bf16_to_f16_inplace(view, kRows, kHeads, kDim, nullptr); - MI_CUDA_CHECK(cudaDeviceSynchronize()); - // Converted: the view holds fp16 bits (a different pattern from bf16 for -12345), the padding is untouched. - const std::vector mid = buf.host(); - const std::uint16_t f16 = float_to_f16_bits(bf16_to_float(before[0])); - ASSERT_NE(f16, before[0]); + constexpr int kRows = 3, kCols = 8, kSrcStride = 12; + const std::vector rows[kRows] = {{1.f, -2.5f, 3.f, 0.f, 1e-3f, -7e-2f, 123.f, 32768.f}, + {1e5f, -2.5f, 3.f, 0.f, 1e-2f, -7e-2f, 123.f, 65536.f}, + {3e38f, 1.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}}; + std::vector src(static_cast(kRows) * kSrcStride, 0.f); for (int r = 0; r < kRows; ++r) { - for (int h = 0; h < kHeads; ++h) { - for (int d = 0; d < kHeadStride; ++d) { - const std::size_t i = static_cast(r) * kHeads * kHeadStride + h * kHeadStride + d; - EXPECT_EQ(mid[i], d < kDim ? f16 : before[i]) << "converted r " << r << " h " << h << " d " << d; - } + for (int c = 0; c < kCols; ++c) { + src[static_cast(r) * kSrcStride + c] = rows[r][static_cast(c)]; } } - mi::kernels::sm75::convert_tensor_f16_to_bf16_inplace(view, kRows, kHeads, kDim, nullptr); + DeviceBf16 in(src); + void* dst = nullptr; + float* inv = nullptr; + MI_CUDA_CHECK(cudaMalloc(&dst, kRows * kCols * 2)); + MI_CUDA_CHECK(cudaMalloc(&inv, kRows * sizeof(float))); + mi::kernels::sm75::convert_rows_bf16_to_f16_scaled(dst, inv, in.data, kRows, kCols, kSrcStride, kCols, nullptr); MI_CUDA_CHECK(cudaDeviceSynchronize()); - const std::vector after = buf.host(); + std::vector got(kRows * kCols); + std::vector inv_host(kRows); + MI_CUDA_CHECK(cudaMemcpy(got.data(), dst, got.size() * 2, cudaMemcpyDeviceToHost)); + MI_CUDA_CHECK(cudaMemcpy(inv_host.data(), inv, inv_host.size() * sizeof(float), cudaMemcpyDeviceToHost)); + cudaFree(dst); + cudaFree(inv); + EXPECT_EQ(inv_host[0], 1.f); + EXPECT_EQ(inv_host[1], 2.f); // max 1e5 (bf16 99840): 99840 / 2 fits, 99840 does not + EXPECT_GT(inv_host[2], 1e33f); for (int r = 0; r < kRows; ++r) { - for (int h = 0; h < kHeads; ++h) { - for (int d = 0; d < kDim; ++d) { - const std::size_t i = static_cast(r) * kHeads * kHeadStride + h * kHeadStride + d; - EXPECT_EQ(after[i], before[i]) << "r " << r << " h " << h << " d " << d; - } - for (int d = kDim; d < kHeadStride; ++d) { - const std::size_t i = static_cast(r) * kHeads * kHeadStride + h * kHeadStride + d; - EXPECT_EQ(after[i], before[i]) << "padding r " << r << " h " << h << " d " << d; - } + for (int c = 0; c < kCols; ++c) { + const float bf = bf16_to_float(float_to_bf16_bits(rows[r][static_cast(c)])); + __half_raw raw; + raw.x = got[static_cast(r) * kCols + c]; + const float back = __half2float(__half(raw)) * inv_host[r]; + // Row 2's 1.0 times 2^-112 is below fp16's range: the documented loss of a scaled row's low values. + EXPECT_EQ(back, r == 2 && c == 1 ? 0.f : bf) << "row " << r << " col " << c; } } #endif diff --git a/tools/sm75_check.cu b/tools/sm75_check.cu index a0daa54..6fc6f02 100644 --- a/tools/sm75_check.cu +++ b/tools/sm75_check.cu @@ -3,7 +3,9 @@ // // Usage: sm75_check [attn128|attn256|convert|all] // attn128 varlen non-causal attention, head_dim 128 (the DiT's joint attention shape), GQA 2:1 -// attn256 varlen causal attention, head_dim 256 (the text encoder's full-attention shape), GQA 4:1 +// attn256 varlen causal attention, head_dim 256 (the text encoder's full-attention shape), GQA 4:1, and causal +// with key lengths unlike the query lengths (bottom-right alignment) +// gemm the fp16 linear: bf16 and f32 out, a row above fp16's range, several row chunks // convert bf16 <-> fp16 round trip on device // The engine's own GPU tests compile their reference kernels for the build's CUDA architectures (sm_80+), so the // unit-test binary cannot exercise sm_75 kernels on a Turing device; this tool builds its own reference in host code. @@ -90,10 +92,12 @@ class DevBuf { void* p_ = nullptr; }; -// Host reference: one (row, head) at a time, fp32 scores/softmax/PV, bf16 inputs as the device sees them. +// Host reference: one (row, head) at a time, fp32 scores/softmax/PV, bf16 inputs as the device sees them. Causal is +// aligned bottom-right (FlashAttention-2): query i of a sequence sees keys 0 .. i + len_k - len_q; none: a zero row. std::vector reference_attention(const std::vector& q, const std::vector& k, const std::vector& v, const std::vector& cu_q, - int heads, int kv_heads, int dim, bool causal, float scale) { + const std::vector& cu_k, int heads, int kv_heads, int dim, bool causal, + float scale) { const int batch = static_cast(cu_q.size()) - 1; const int total_q = cu_q.back(); std::vector out(static_cast(total_q) * heads * dim, 0.f); @@ -101,14 +105,16 @@ std::vector reference_attention(const std::vector& q, cons std::vector scores; for (int b = 0; b < batch; ++b) { const int q0 = cu_q[b], q1 = cu_q[b + 1]; + const int k0 = cu_k[b], len_k = cu_k[b + 1] - k0; for (int h = 0; h < heads; ++h) { const int kh = h / group; for (int i = q0; i < q1; ++i) { - const int n = causal ? (i - q0 + 1) : (q1 - q0); + const int n = causal ? std::clamp(i - q0 + 1 + len_k - (q1 - q0), 0, len_k) : len_k; + if (n == 0) continue; scores.assign(static_cast(n), 0.f); float max_s = -INFINITY; for (int j = 0; j < n; ++j) { - const int row = q0 + j; + const int row = k0 + j; float s = 0.f; for (int d = 0; d < dim; ++d) { s += bf16_to_float(q[(static_cast(i) * heads + h) * dim + d]) * @@ -126,7 +132,7 @@ std::vector reference_attention(const std::vector& q, cons std::size_t base = (static_cast(i) * heads + h) * dim; for (int j = 0; j < n; ++j) { const float p = scores[static_cast(j)] / sum; - const int row = q0 + j; + const int row = k0 + j; for (int d = 0; d < dim; ++d) { out[base + d] += p * bf16_to_float(v[(static_cast(row) * kv_heads + kh) * dim + d]); } @@ -137,28 +143,36 @@ std::vector reference_attention(const std::vector& q, cons return out; } +// k_lengths: the keys' per-sequence lengths (empty: the queries' own). int run_attention(const char* name, int heads, int kv_heads, int dim, bool causal, const std::vector& lengths, - std::uint64_t seed, bool dump = false) { + std::uint64_t seed, bool dump = false, const std::vector& k_lengths = {}) { const int batch = static_cast(lengths.size()); - std::vector cu(static_cast(batch) + 1, 0); - for (int b = 0; b < batch; ++b) cu[static_cast(b) + 1] = cu[static_cast(b)] + lengths[static_cast(b)]; + const std::vector& kl = k_lengths.empty() ? lengths : k_lengths; + std::vector cu(static_cast(batch) + 1, 0), cu_k(cu); + for (int b = 0; b < batch; ++b) { + cu[static_cast(b) + 1] = cu[static_cast(b)] + lengths[static_cast(b)]; + cu_k[static_cast(b) + 1] = cu_k[static_cast(b)] + kl[static_cast(b)]; + } const int total_q = cu.back(); + const int total_k = cu_k.back(); const int max_len = *std::max_element(lengths.begin(), lengths.end()); const float scale = 1.f / std::sqrt(static_cast(dim)); const auto q = normal(static_cast(total_q) * heads * dim, seed, 1.f); - const auto k = normal(static_cast(total_q) * kv_heads * dim, seed + 1, 1.f); - const auto v = normal(static_cast(total_q) * kv_heads * dim, seed + 2, 1.f); + const auto k = normal(static_cast(total_k) * kv_heads * dim, seed + 1, 1.f); + const auto v = normal(static_cast(total_k) * kv_heads * dim, seed + 2, 1.f); std::vector qb(q.size()), kb(k.size()), vb(v.size()); for (std::size_t i = 0; i < q.size(); ++i) qb[i] = float_to_bf16_bits(q[i]); for (std::size_t i = 0; i < k.size(); ++i) kb[i] = float_to_bf16_bits(k[i]); for (std::size_t i = 0; i < v.size(); ++i) vb[i] = float_to_bf16_bits(v[i]); - DevBuf dq(qb.size() * 2), dk(kb.size() * 2), dv(vb.size() * 2), dout(qb.size() * 2), dcu(cu.size() * sizeof(int)); + DevBuf dq(qb.size() * 2), dk(kb.size() * 2), dv(vb.size() * 2), dout(qb.size() * 2), dcu(cu.size() * sizeof(int)), + dcu_k(cu_k.size() * sizeof(int)); dq.upload(qb); dk.upload(kb); dv.upload(vb); dcu.upload(cu); + dcu_k.upload(cu_k); MI_CUDA_CHECK(cudaMemset(dout.get(), 0, qb.size() * 2)); std::vector lse(static_cast(total_q) * heads, 0.f); DevBuf dls(lse.size() * sizeof(float)); @@ -168,22 +182,26 @@ int run_attention(const char* name, int heads, int kv_heads, int dim, bool causa return Tensor::view(p, dt, Device::kCUDA, shape); }; const Tensor tq = t(dq.get(), DType::kBF16, {total_q, heads, dim}); - const Tensor tk = t(dk.get(), DType::kBF16, {total_q, kv_heads, dim}); - const Tensor tv = t(dv.get(), DType::kBF16, {total_q, kv_heads, dim}); + const Tensor tk = t(dk.get(), DType::kBF16, {total_k, kv_heads, dim}); + const Tensor tv = t(dv.get(), DType::kBF16, {total_k, kv_heads, dim}); const Tensor to = t(dout.get(), DType::kBF16, {total_q, heads, dim}); const Tensor tcu = t(dcu.get(), DType::kI32, {batch + 1}); + const Tensor tcu_k = t(dcu_k.get(), DType::kI32, {batch + 1}); const std::size_t ws_bytes = mi::attention::flash_fwd_workspace_bytes(total_q, heads); MI_CUDA_CHECK(cudaGetLastError()); std::printf(" launch: tokens %d heads %d kv %d dim %d causal %d maxlen %d ws %zu\n", total_q, heads, kv_heads, dim, causal ? 1 : 0, max_len, ws_bytes); std::fflush(stdout); - mi::attention::flash_fwd(tq, tk, tv, to, tcu, tcu, max_len, causal, scale, dls.get(), ws_bytes, nullptr); + mi::attention::flash_fwd(tq, tk, tv, to, tcu, tcu_k, max_len, causal, scale, dls.get(), ws_bytes, nullptr); MI_CUDA_CHECK(cudaDeviceSynchronize()); + // The inputs are read only: q/k/v hold their bf16 bits after the call. + const bool inputs_kept = dq.download(qb.size()) == qb && dk.download(kb.size()) == kb && + dv.download(vb.size()) == vb; std::printf(" kernel done\n"); std::fflush(stdout); const auto got = dout.download(qb.size()); - const auto want = reference_attention(qb, kb, vb, cu, heads, kv_heads, dim, causal, scale); + const auto want = reference_attention(qb, kb, vb, cu, cu_k, heads, kv_heads, dim, causal, scale); double dot = 0, ng = 0, nw = 0; float max_abs = 0.f; for (std::size_t i = 0; i < want.size(); ++i) { @@ -194,8 +212,9 @@ int run_attention(const char* name, int heads, int kv_heads, int dim, bool causa max_abs = std::max(max_abs, std::fabs(g - want[i])); } const double cosine = dot / std::sqrt(ng * nw); - std::printf("%-8s tokens %5d heads %2d kv %2d dim %3d causal %d: cosine %.9f max|diff| %.3g\n", name, total_q, heads, - kv_heads, dim, causal ? 1 : 0, cosine, max_abs); + std::printf("%-8s tokens %5d keys %5d heads %2d kv %2d dim %3d causal %d: cosine %.9f max|diff| %.3g inputs %s\n", + name, total_q, total_k, heads, kv_heads, dim, causal ? 1 : 0, cosine, max_abs, + inputs_kept ? "kept" : "MODIFIED"); if (dump) { for (std::size_t i = 0; i < std::min(want.size(), 12); ++i) { std::printf(" [%zu] got %.6f want %.6f (q %.6f)\n", i, bf16_to_float(got[i]), want[i], q[i]); @@ -203,7 +222,7 @@ int run_attention(const char* name, int heads, int kv_heads, int dim, bool causa // Diagnostic: is the mismatch a different softmax temperature? const float log2e = 1.4426950408889634f; for (float f : {1.f, log2e, 1.f / log2e, 2.f, 0.5f, 1.f / std::sqrt(static_cast(dim))}) { - const auto alt = reference_attention(qb, kb, vb, cu, heads, kv_heads, dim, causal, scale * f); + const auto alt = reference_attention(qb, kb, vb, cu, cu_k, heads, kv_heads, dim, causal, scale * f); double d2 = 0, g2 = 0, w2 = 0; for (std::size_t i = 0; i < alt.size(); ++i) { const float g = bf16_to_float(got[i]); @@ -215,7 +234,7 @@ int run_attention(const char* name, int heads, int kv_heads, int dim, bool causa } } MI_CUDA_CHECK(cudaGetLastError()); - return cosine >= 0.9999 ? 0 : 1; + return cosine >= 0.9999 && inputs_kept ? 0 : 1; } int run_convert() { @@ -234,11 +253,17 @@ int run_convert() { return bad == 0 ? 0 : 1; } -// The sm_75 GEMM (cuBLASLt fp16 in / fp32 accumulate, bf16 activation input): bf16 x, fp16 weights, bf16 out with a -// bias added in fp32. Cross-checked against a host accumulation of the stored operand values. -int run_gemm(int M, int N, int K) { +// The sm_75 GEMM (cuBLASLt fp16 in / fp32 out, bf16 activation input): bf16 x, fp16 weights, bf16 or f32 out with a +// bias of the output dtype added in fp32. Cross-checked against a host accumulation of the stored operand values. +// big_row >= 0 scales that row of x far above fp16's range (the per-row power-of-two scale); a small ws_bytes makes +// the GEMM run several row chunks (the workspace holds its fp16 rows and fp32 result). +int run_gemm(int M, int N, int K, bool out_f32 = false, int big_row = -1, + std::size_t ws_bytes = std::size_t{32} << 20) { mi::backend::load_cublaslt(); - const auto xv = normal(static_cast(M) * K, 51, 1.f); + auto xv = normal(static_cast(M) * K, 51, 1.f); + if (big_row >= 0) { + for (int d = 0; d < K; ++d) xv[static_cast(big_row) * K + d] *= 3e5f; + } const auto wv = normal(static_cast(N) * K, 52, 0.05f); const auto bv = normal(static_cast(N), 53, 0.01f); std::vector xb(xv.size()), wb(wv.size()), bb(bv.size()); @@ -248,25 +273,41 @@ int run_gemm(int M, int N, int K) { std::memcpy(&wb[i], &h, sizeof h); } for (std::size_t i = 0; i < bv.size(); ++i) bb[i] = float_to_bf16_bits(bv[i]); + std::vector bf(bv.size()); + for (std::size_t i = 0; i < bv.size(); ++i) bf[i] = bf16_to_float(bb[i]); // the same bias values for an f32 out - DevBuf dx(xb.size() * 2), dw(wb.size() * 2), db(bb.size() * 2), dout(static_cast(M) * N * 2); - DevBuf dws(std::size_t{32} << 20); + const std::size_t out_elem = out_f32 ? 4 : 2; + DevBuf dx(xb.size() * 2), dw(wb.size() * 2), db(bb.size() * out_elem); + DevBuf dout(static_cast(M) * N * out_elem); + DevBuf dws(ws_bytes); dx.upload(xb); dw.upload(wb); - db.upload(bb); - MI_CUDA_CHECK(cudaMemset(dout.get(), 0, static_cast(M) * N * 2)); + if (out_f32) { + db.upload(bf); + } else { + db.upload(bb); + } + MI_CUDA_CHECK(cudaMemset(dout.get(), 0, static_cast(M) * N * out_elem)); + const DType out_type = out_f32 ? DType::kF32 : DType::kBF16; const Tensor tx = Tensor::view(dx.get(), DType::kBF16, Device::kCUDA, {M, K}); const Tensor tw = Tensor::view(dw.get(), DType::kF16, Device::kCUDA, {N, K}); - const Tensor tb = Tensor::view(db.get(), DType::kBF16, Device::kCUDA, {N}); - const Tensor to = Tensor::view(dout.get(), DType::kBF16, Device::kCUDA, {M, N}); + const Tensor tb = Tensor::view(db.get(), out_type, Device::kCUDA, {N}); + const Tensor to = Tensor::view(dout.get(), out_type, Device::kCUDA, {M, N}); mi::kernels::Gemm gemm; - gemm.linear(to, tx, tw, &tb, dws.get(), std::size_t{32} << 20, nullptr); + gemm.linear(to, tx, tw, &tb, dws.get(), ws_bytes, nullptr); MI_CUDA_CHECK(cudaDeviceSynchronize()); - const auto got = dout.download(static_cast(M) * N); + std::vector got(static_cast(M) * N); + if (out_f32) { + got = dout.download(got.size()); + } else { + const auto bits = dout.download(got.size()); + for (std::size_t i = 0; i < got.size(); ++i) got[i] = bf16_to_float(bits[i]); + } double dot = 0, ng = 0, nw = 0; float max_abs = 0.f, max_rel = 0.f; + bool finite = true; for (int r = 0; r < M; ++r) { for (int c = 0; c < N; ++c) { double acc = 0; @@ -278,8 +319,9 @@ int run_gemm(int M, int N, int K) { }(); acc += static_cast(bf16_to_float(xb[static_cast(r) * K + d])) * w; } - const float want = static_cast(acc) + bf16_to_float(bb[static_cast(c)]); - const float g = bf16_to_float(got[static_cast(r) * N + c]); + const float want = static_cast(acc) + bf[static_cast(c)]; + const float g = got[static_cast(r) * N + c]; + finite = finite && std::isfinite(g); dot += static_cast(g) * want; ng += static_cast(g) * g; nw += static_cast(want) * want; @@ -288,8 +330,9 @@ int run_gemm(int M, int N, int K) { } } const double cosine = dot / std::sqrt(ng * nw); - std::printf("gemm M %d N %d K %d: cosine %.9f max|diff| %.3g max|rel| %.3g\n", M, N, K, cosine, max_abs, max_rel); - return cosine >= 0.9999 && max_rel < 0.01 ? 0 : 1; + std::printf("gemm M %d N %d K %d out %s big row %d ws %zu: cosine %.9f max|diff| %.3g max|rel| %.3g%s\n", M, N, K, + out_f32 ? "f32" : "bf16", big_row, ws_bytes, cosine, max_abs, max_rel, finite ? "" : " NON-FINITE"); + return finite && cosine >= 0.9999 && max_rel < 0.01 ? 0 : 1; } // VAE sm_75 conv2d: NHWC fp32 with KRSC filters, bias + residual, against a host reference. @@ -485,6 +528,11 @@ int main(int argc, char** argv) { if (what == "gemm" || what == "all") { failures += run_gemm(64, 128, 256); failures += run_gemm(3, 8, 32); // a tail case: m not a multiple of anything + failures += run_gemm(64, 128, 256, /*out_f32=*/true); + failures += run_gemm(64, 128, 256, false, /*big_row=*/5); // a row above fp16's range (scaled, not inf) + // Several row chunks (1 MiB workspace: 357 rows of N 512 / K 256 each), f32 and bf16 out. + failures += run_gemm(1000, 512, 256, /*out_f32=*/true, 999, std::size_t{1} << 20); + failures += run_gemm(1000, 512, 256, false, 7, std::size_t{1} << 20); } if (what == "vae") { failures += run_vae_conv_case(1, 6, 6, 8, 8, 3, 61); @@ -512,6 +560,10 @@ int main(int argc, char** argv) { if (what == "attn256" || what == "all") { failures += run_attention("attn256", 4, 1, 256, true, {32, 100}, 13); failures += run_attention("attn256", 4, 1, 256, true, {256}, 14); + // Causal with more keys than queries (bottom-right: the last query sees every key) and fewer (the first queries + // see none: zero rows). + failures += run_attention("attn256", 4, 1, 256, true, {40, 64}, 15, false, {100, 64}); + failures += run_attention("attn256", 4, 1, 256, true, {100}, 16, false, {30}); } if (what == "big") { // Engine-shaped cases: 24 heads at head_dim 128 (the DiT's joint attention) and the text encoder's 16/4 GQA at