From 239db09ac05e6eb53a6681740b10e12cf52e26e9 Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 11:52:26 +0800 Subject: [PATCH 01/12] sm75: run the VAE encoder and runtime LoRAs on Turing - encoder_conv2d and conv2d_down dispatch to fp32 SIMT kernels on cc<8, with a new stride-2 3x3 downsample kernel in kernels/sm75/vae_sm75.cu. img2img, inpainting, reference images and the hires fix now work on RTX 20xx instead of being refused with not_supported. - Runtime LoRAs are merged into the weights at load on sm7x: the fused side-branch GEMMs are sm_80+ CUTLASS, but the merge kernels are plain SIMT, so the stack is forced to Placement::kMerge before the bf16-to-fp16 conversion. A strength change drops the DiT and reloads it (the resident weights are fp16, so an in-place re-merge would mix dtypes) instead of re-merging. - lora_merge.cu joins the portable sm75 sources; docs/sm75.md updated (only control adapters remain refused). Verified on the RTX 2080: --init-image, --reference and --hires produce images; --lora reports its 6 adapters as merged and changes the image; the reference-tier apple renders correctly with and without a LoRA. --- cmake/modules/sm75.cmake | 1 + docs/sm75.md | 21 +++++----- src/kernels/sm75/sm75.h | 4 ++ src/kernels/sm75/vae_sm75.cu | 70 +++++++++++++++++++++++++++++++ src/kernels/vae/encoder_conv2d.cu | 13 ++++++ src/pipeline/pipeline.cpp | 43 +++++++++---------- 6 files changed, 119 insertions(+), 33 deletions(-) diff --git a/cmake/modules/sm75.cmake b/cmake/modules/sm75.cmake index fbda063..868e53a 100644 --- a/cmake/modules/sm75.cmake +++ b/cmake/modules/sm75.cmake @@ -66,6 +66,7 @@ if((MOSSINFER_SM75 OR MOSSINFER_SM70) AND _mi_sm75_present) src/kernels/gdn/chunk_rule.cu src/kernels/gdn/conv_gates.cu src/kernels/lora/lora_kernels.cu + src/kernels/lora/lora_merge.cu src/kernels/quant/rotate_quantize.cu src/kernels/quant/w4a8_decode.cu src/kernels/quant4/convrot_quantize.cu diff --git a/docs/sm75.md b/docs/sm75.md index a7fdf2f..7f61d0a 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -85,27 +85,28 @@ tiers, with the bf16 text encoder streamed. 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. Both pass on an RTX 2080 Ti. -- 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. +- The VAE encoder and the merged-LoRA path are exercised on the RTX 2080: `--init-image`, `--reference` and `--hires` + produce images (the encoder runs the same fp32 SIMT convolutions as the decoder), and `--lora` reports its adapters as + `merged` in the report and changes the image at the expected strength. - 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`. On an RTX 2080 Ti, separate `moss-server` processes produce the same PNG bytes at all three tiers, at 256² and (`high`, `turbo`) 1024². - A build with `MOSSINFER_SM75=OFF` compiles and links unchanged and keeps the same CPU test baseline. -## Not supported on sm_75 yet +## Limits on sm_75 - The 8-bit attention kernels (`q8_fp16`, `q8_fp8`): the tier default resolves to `exact`; passing an 8-bit mode explicitly is refused. - Native INT4 MMA and NVFP4: INT4 files run through the INT8 unpack fallback (correct and bit-exact; the planner prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. - The fused DiT blocks and their CUDA graphs (the blocks run op by op). -- img2img, inpainting, reference images and the hires fix: the VAE **encoder** is still an sm_80 kernel, so only - text-to-image runs (the hires fix's second pass encodes the upscaled image). A request with `init_image` (or a - `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; a component - session refuses them the same way, when it opens and for every run. +- img2img, inpainting, reference images and the hires fix run: the VAE **encoder** uses the same fp32 SIMT + convolutions as the decoder on Turing (a stride-2 downsample kernel included). It is not fast (about 5 s for a + 256x256 encode, 16 s for a 384x384 hires second pass), but it works. +- Runtime LoRAs are **merged into the weights at load** instead of the sm_80+ fused side branches: the merge kernels + are plain SIMT, run before the bf16-to-fp16 conversion, and the report says `merged` for their adapters. A strength + change drops the DiT and reloads it (the resident weights are fp16, so an in-place re-merge would mix dtypes), which + costs a load. Control adapters still need the fused LoRA path and are refused. - 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. diff --git a/src/kernels/sm75/sm75.h b/src/kernels/sm75/sm75.h index 29b6ac1..3cfef53 100644 --- a/src/kernels/sm75/sm75.h +++ b/src/kernels/sm75/sm75.h @@ -69,6 +69,10 @@ void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tens cudaStream_t stream); // Mid-block attention: o = softmax(q k^T * scale) v, fp32, one head, [B, L, D] with D a multiple of 32 and any row // stride (>= D) for q/k/v; o is contiguous. +// Stride-2 3x3 downsample conv (VAE encoder's Downsample2D): y [N,P,Q,K] = bias + conv(x, w) over the window at +// (2i, 2j), x zero outside [0,H) x [0,W); R must be 3. +void vae_conv2d_down(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor& y, cudaStream_t stream); + void vae_attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& o, float scale, cudaStream_t stream); // torch conv weight [K,C,R,S] fp32 -> [K,R,S,C] in dst's dtype (kF32 or kBF16). diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index 14364ef..e5ce87d 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -138,6 +138,76 @@ void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tens MI_CUDA_CHECK_LAUNCH(); } + +// ---------------------------------------------------------------- stride-2 downsample conv (VAE encoder) +// Same layout as conv2d_f32_sm75 (one block: an 8x8 output tile of 32 output channels, input channels in groups of 8), +// with the encoder's Downsample2D semantics: out[i,j] accumulates the 3x3 window at (2i, 2j) of the input, treating the +// input as zero outside [0,H) x [0,W) (the reference pads one row/column at the bottom/right, nothing at the top/left). +constexpr int kDownTW = 8, kDownTH = 8, kDownTC = 8; +constexpr int kDownPX = kDownTW * 2 + 2; // 18 +constexpr int kDownPY = kDownTH * 2 + 2; + +__global__ void conv2d_down_f32_sm75(const float* __restrict__ x, const float* __restrict__ w, + const float* __restrict__ bias, float* __restrict__ y, int H, int W, int C, int K, + int P, int Q, int kb_count) { + __shared__ float xs[kDownPY * kDownPX * kDownTC]; + __shared__ float ws[kTK * 9 * kDownTC]; + const int n = blockIdx.z / kb_count; + const int kb = blockIdx.z % kb_count; + const int ox0 = blockIdx.x * kDownTW, oy0 = blockIdx.y * kDownTH; + const int tid = threadIdx.x; + const int pixel = tid >> 1; + const int cg = (tid & 1) * 16; + const int py = pixel / kDownTW, px = pixel % kDownTW; + float acc[16] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + + for (int c0 = 0; c0 < C; c0 += kDownTC) { + for (int i = tid; i < kDownPY * kDownPX * kDownTC; i += kConvThreads) { + const int cc = i % kDownTC, p = i / kDownTC; + const int yy = p / kDownPX, xx = p % kDownPX; + const int gy = oy0 * 2 + yy, gx = ox0 * 2 + xx, gc = c0 + cc; + xs[i] = (gy < H && gx < W && gc < C) ? x[((static_cast(n) * H + gy) * W + gx) * C + gc] : 0.f; + } + for (int i = tid; i < kTK * 9 * kDownTC; i += kConvThreads) { + const int cc = i % kDownTC, t = i / kDownTC, rs = t % 9, kk = t / 9; + const int kch = kb * kTK + kk, gc = c0 + cc; + ws[i] = (kch < K && gc < C) ? w[(static_cast(kch) * 3 + rs / 3) * 3 * C + (rs % 3) * C + gc] : 0.f; + } + __syncthreads(); + for (int cc = 0; cc < kDownTC; ++cc) { + if (c0 + cc >= C) break; + for (int r = 0; r < 3; ++r) { + for (int s = 0; s < 3; ++s) { + const float xv = xs[((py * 2 + r) * kDownPX + (px * 2 + s)) * kDownTC + cc]; +#pragma unroll + for (int ch = 0; ch < 16; ++ch) acc[ch] += xv * ws[((cg + ch) * 9 + r * 3 + s) * kDownTC + cc]; + } + } + } + __syncthreads(); + } + + const int ox = ox0 + px, oy = oy0 + py; + if (ox >= Q || oy >= P) return; +#pragma unroll + for (int ch = 0; ch < 16; ++ch) { + const int kch = kb * kTK + cg + ch; + if (kch >= K) break; + const std::int64_t idx = ((static_cast(n) * P + oy) * Q + ox) * K + kch; + y[idx] = acc[ch] + bias[kch]; + } +} +void vae_conv2d_down(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor& y, cudaStream_t stream) { + const int H = static_cast(x.dim(1)), W = static_cast(x.dim(2)), C = static_cast(x.dim(3)); + const int K = static_cast(w.dim(0)); + const int P = static_cast(y.dim(1)), Q = static_cast(y.dim(2)); + const int kb_count = (K + kTK - 1) / kTK; + const dim3 grid(static_cast((Q + kDownTW - 1) / kDownTW), static_cast((P + kDownTH - 1) / kDownTH), + static_cast(x.dim(0) * kb_count)); + conv2d_down_f32_sm75<<>>(x.ptr(), w.ptr(), bias.ptr(), + y.ptr(), H, W, C, K, P, Q, kb_count); + MI_CUDA_CHECK_LAUNCH(); +} void vae_attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& o, float scale, cudaStream_t stream) { const int L = static_cast(q.dim(1)), D = static_cast(q.dim(2)); diff --git a/src/kernels/vae/encoder_conv2d.cu b/src/kernels/vae/encoder_conv2d.cu index 5198044..596e91f 100644 --- a/src/kernels/vae/encoder_conv2d.cu +++ b/src/kernels/vae/encoder_conv2d.cu @@ -11,6 +11,7 @@ #include #include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" #include "kernels/vae/encoder_kernels.h" #include "cutlass/conv/conv2d_problem_size.h" @@ -163,6 +164,7 @@ void run(const char* name, const Shape& s, const Tensor& x, const Tensor& w, con } // namespace std::size_t encoder_conv2d_workspace_bytes(const Tensor& x, const Tensor& w, ConvMath math) { + if (kernels::sm75::turing()) return 0; // the fp32 SIMT path keeps no workspace if (math == ConvMath::kTf32) return conv2d_workspace_bytes(x, w); const Tensor y = Tensor::view(nullptr, DType::kF32, Device::kCUDA, {x.dim(0), x.dim(1), x.dim(2), w.dim(0)}); return workspace_for(check("encoder_conv2d", x, w, y, 1), x, w, y); @@ -170,6 +172,11 @@ std::size_t encoder_conv2d_workspace_bytes(const Tensor& x, const Tensor& w, Con void encoder_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, void* workspace, std::size_t workspace_bytes, ConvMath math, cudaStream_t stream) { + if (kernels::sm75::turing()) { // fp32 SIMT conv: Turing has no tensor cores for this path + check("encoder_conv2d", x, w, y, 1); + kernels::sm75::vae_conv2d(x, w, bias, residual, y, stream); + return; + } if (math == ConvMath::kTf32) { conv2d(x, w, bias, residual, y, workspace, workspace_bytes, stream); return; @@ -179,12 +186,18 @@ void encoder_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const } std::size_t conv2d_down_workspace_bytes(const Tensor& x, const Tensor& w, const Tensor& y, ConvMath math) { + if (kernels::sm75::turing()) return 0; const Shape s = check("conv2d_down", x, w, y, 2); return math == ConvMath::kTf32 ? workspace_for(s, x, w, y) : workspace_for(s, x, w, y); } void conv2d_down(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor& y, void* workspace, std::size_t workspace_bytes, ConvMath math, cudaStream_t stream) { + if (kernels::sm75::turing()) { + check("conv2d_down", x, w, y, 2); + kernels::sm75::vae_conv2d_down(x, w, bias, y, stream); + return; + } const Shape s = check("conv2d_down", x, w, y, 2); if (math == ConvMath::kTf32) run("conv2d_down", s, x, w, bias, nullptr, y, workspace, workspace_bytes, stream); diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index ff21b15..6fa4ea4 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -1063,7 +1063,12 @@ Pipeline::Impl::CallModels Pipeline::Impl::setup_models(const Request& req, Resu // Merged adapters: re-merged from the pristine weights kept at load, in place (no reload). // A model merged without pristine copies (an edit adapter's blocks) is reloaded instead (its key cleared). if (m.dit && m.dit_merge_key == files && !m.dit->can_remerge()) m.dit_merge_key.clear(); - if (m.dit && m.dit_merge_key == files && !m.lora_stack->merge_modules().empty()) { + if (kernels::sm75::turing() && !m.lora_stack->merge_modules().empty()) { + // The merge is in bf16 at load, and the weights are converted to fp16 afterwards: re-merging in place would + // mix the two, so drop the DiT and let the next call reload (its load re-merges, then re-converts). + m.dit_merge_key.clear(); + m.dit.reset(); + } else if (m.dit && m.dit_merge_key == files && !m.lora_stack->merge_modules().empty()) { const auto t0m = Clock::now(); try { m.dit->remerge_lora(*m.lora_stack); @@ -1500,36 +1505,28 @@ void Pipeline::Impl::explain_control(const Request::Control& control, Result& r) "): " + (inputs.empty() ? std::string("no images (zero control latents)") : inputs) + "."; } -// Two features still need sm_80+ kernels and are refused on Turing before anything is planned, read or uploaded (by -// Pipeline::run, and by a Session for its options and every run's spec): the VAE encoder's convolutions (img2img, -// inpainting, reference images, and the hires fix, whose second pass encodes the upscaled image) and the LoRA GEMMs -// (runtime LoRAs, and control adapters, which carry a LoRA). The message names the reason and what does work, instead -// of failing at a missing kernel image mid-request. +// Features that still need sm_80+ kernels are refused on Turing before anything is planned, read or uploaded (by +// Pipeline::run, and by a Session for its options and every run's spec): control adapters (a LoRA with a replacement +// embedding, fused into the block GEMMs). The VAE encoder runs on Turing (fp32 SIMT), and runtime LoRAs are merged at +// load (Pipeline::run sets the placement). void detail::refuse_unsupported_on_sm7x(const Request& req) { if (!kernels::sm75::turing()) return; - if (req.init_image != nullptr) { - throw NotSupportedError("img2img and inpainting encode the input with the VAE encoder, which has no sm_7x kernel " - "in this build; text-to-image works"); - } - if (req.hires.enabled()) { - throw NotSupportedError("the hires fix encodes the upscaled image with the VAE encoder, which has no sm_7x kernel " - "in this build; text-to-image without it works"); - } - if (req.references.enabled()) { - throw NotSupportedError("reference images are encoded with the VAE encoder, which has no sm_7x kernel in this " - "build; text-to-image without them works"); - } - if (!req.loras.empty()) { - throw NotSupportedError("runtime LoRAs need sm_80+ GEMM kernels; this sm_7x build cannot apply them"); - } + // The VAE encoder runs on Turing (fp32 SIMT convs, the same kernels as the decoder), so img2img, inpainting, + // reference images and the hires fix all work. Runtime LoRAs are merged into the weights instead of the sm_80+ fused + // side branches (Pipeline::run sets the placement); control adapters still need the fused LoRA path. if (req.control.enabled()) { throw NotSupportedError("control adapters are LoRAs and need sm_80+ GEMM kernels; this sm_7x build cannot run them"); } } - // generate(), or with `preload` the same up to the plan and the arena, then the loads of the models the plan keeps. -Result Pipeline::run(const Request& req, const GenerateControl& control, bool preload) { +Result Pipeline::run(const Request& request, const GenerateControl& control, bool preload) { Impl& m = *impl_; + Request req = request; + // Turing has no fused LoRA side-branch kernels (sm_80+ CUTLASS), so a runtime LoRA is merged into the weights at + // load instead (its load-time merge then feeds the fp16 conversion the same way as the bf16 bundle). + if (kernels::sm75::turing() && !req.loras.empty() && req.lora_placement != lora::Placement::kMerge) { + req.lora_placement = lora::Placement::kMerge; + } check_user_paths(m.options, req); validate(req); refuse_unsupported_on_sm7x(req); From d466475d807811396997c22ec30c77d9d40c6ef5 Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 12:30:57 +0800 Subject: [PATCH 02/12] sm75: vectorize the VAE conv and the GEMM output and weight conversions Profiling on the RTX 2080 (Nsight Systems) showed the VAE convolutions as the largest single kernel (52.6 % of the img2img GPU time, 10.7 ms per launch) and the sm75 GEMM output/bias conversion at 176 ms per four-step generation. - conv2d_f32_sm75: float4 shared-memory loads and float4 accumulators (four input channels per FMA), vectorized float4 output stores with the bias and residual added in the same order. VAE decode 0.37 -> 0.20 s at 256^2 (10-12 -> 7.8 s at 1024^2), VAE encode 0.39 -> 0.19 s. - gemm_out_f32: float4 accumulator loads and packed bf16 stores, 176 -> 93 ms per four-step run. - bf16_to_f16_rows_scaled: pair-wise bf16 loads and packed half2 stores in both passes. --- docs/sm75.md | 17 +++++++- src/kernels/sm75/convert.cu | 46 +++++++++++++++++---- src/kernels/sm75/vae_sm75.cu | 80 ++++++++++++++++++++++++++++-------- 3 files changed, 116 insertions(+), 27 deletions(-) diff --git a/docs/sm75.md b/docs/sm75.md index 7f61d0a..d71351e 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -59,9 +59,24 @@ bf16 text encoder from the bundle: | INT8 ConvRot DiT, `high` | 1024² | **2.58 s** | 7.5 GiB | | INT4 ConvRot DiT, `turbo` | 1024² | 2.59 s | 5.3 GiB | -The VAE decodes in fp32 on sm_75: about 0.5 s at 256² and 11 s at 1024². An 8 GiB card (RTX 2080) also runs all three +The VAE decodes in fp32 on sm_75: about 0.2 s at 256² and 7.8 s at 1024². An 8 GiB card (RTX 2080) also runs all three tiers, with the bf16 text encoder streamed. +### Profiling on Turing + +Nsight Systems profiles the whole pipeline (`nsys profile --trace=cuda,cublas,nvtx --stats=true moss-infer generate ...`); +Nsight Compute's counters need `ERR_NVGPUCTRPERM` permission that an unprivileged account does not have. + +What the timeline shows on an RTX 2080 8 GiB, `reference` tier, 256^2, four steps: about 3.0 s of host-to-device copies +(36 GB of streamed weights at ~12 GB/s) against ~3.1 s of kernels, so streaming is half the work. The largest kernels +are cuBLASLt's fp16 GEMMs (~35 %), the VAE convolutions (the single biggest kernel in img2img, 52.6 % of its kernel +time before the vectorisation below) and the sm75 conversion kernels. At 1024^2 the decode's fp32 convolutions cost +10-12 s before the optimisations below. + +The kernels this backend owns were optimised from that profile: `conv2d_f32_sm75` now uses float4 shared-memory loads +and float4 accumulators (four input channels per FMA) with vectorised float4 output stores (VAE decode 0.37 -> 0.20 s at +256^2 and 10-12 -> 7.8 s at 1024^2; encode 0.39 -> 0.19 s), `gemm_out_f32` loads float4 and stores packed bf16 +(176 -> 93 ms per four-step run), and `bf16_to_f16_rows_scaled` uses packed pair loads/stores. ## Numerics - bf16 → fp16 is exact for values in fp16's normal range, so the converted operands are the same numbers the sm_80 diff --git a/src/kernels/sm75/convert.cu b/src/kernels/sm75/convert.cu index 90db435..be293aa 100644 --- a/src/kernels/sm75/convert.cu +++ b/src/kernels/sm75/convert.cu @@ -113,7 +113,12 @@ __global__ void bf16_to_f16_rows_scaled(const __nv_bfloat16* __restrict__ src, _ const __nv_bfloat16* s = src + static_cast(blockIdx.x) * src_stride; __half* d = dst + static_cast(blockIdx.x) * dst_stride; float amax = 0.f; - for (std::int64_t i = threadIdx.x; i < cols; i += blockDim.x) amax = fmaxf(amax, fabsf(__bfloat162float(s[i]))); + std::int64_t i = 2 * static_cast(threadIdx.x); + for (; i + 1 < cols; i += 2 * static_cast(blockDim.x)) { + const float2 f = __bfloat1622float2(*reinterpret_cast(s + i)); + amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); + } + for (; 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(); @@ -128,9 +133,13 @@ __global__ void bf16_to_f16_rows_scaled(const __nv_bfloat16* __restrict__ src, _ } __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); + std::int64_t j = 2 * static_cast(threadIdx.x); + for (; j + 1 < cols; j += 2 * static_cast(blockDim.x)) { + const float2 f = __bfloat1622float2(*reinterpret_cast(s + j)); + *reinterpret_cast<__half2*>(d + j) = __floats2half2_rn(f.x * scale, f.y * scale); + } + for (; j < cols; j += blockDim.x) d[j] = __float2half_rn(__bfloat162float(s[j]) * 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, @@ -139,17 +148,36 @@ __global__ void gemm_out_f32(const float* __restrict__ acc, const float* __restr 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) { - float v = a[c] * sc; - if (bias != nullptr) v += static_cast(bias[c]); - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + for (std::int64_t c = threadIdx.x; c < cols; c += blockDim.x) { + float v = a[c] * sc; + if (bias != nullptr) v += static_cast(bias[c]); o[c] = v; - } else { + } + } else { + // Four columns per iteration: one float4 accumulator load and two packed bf16 stores. + std::int64_t c = 4 * static_cast(threadIdx.x); + for (; c + 3 < cols; c += 4 * static_cast(blockDim.x)) { + const float4 a4 = *reinterpret_cast(a + c); + float v0 = a4.x * sc, v1 = a4.y * sc, v2 = a4.z * sc, v3 = a4.w * sc; + if (bias != nullptr) { + const float2 b01 = __bfloat1622float2(*reinterpret_cast(bias + c)); + const float2 b23 = __bfloat1622float2(*reinterpret_cast(bias + c + 2)); + v0 += b01.x; + v1 += b01.y; + v2 += b23.x; + v3 += b23.y; + } + *reinterpret_cast<__nv_bfloat162*>(o + c) = __floats2bfloat162_rn(v0, v1); + *reinterpret_cast<__nv_bfloat162*>(o + c + 2) = __floats2bfloat162_rn(v2, v3); + } + for (; c < cols; c += blockDim.x) { + float v = a[c] * sc; + if (bias != nullptr) v += static_cast(bias[c]); o[c] = __float2bfloat16_rn(v); } } } - } // namespace void convert_rows_bf16_to_f16_scaled(void* dst, float* inv_scale, const void* src, std::int64_t rows, diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index e5ce87d..f010986 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -25,8 +25,8 @@ constexpr int kConvThreads = 128; __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias, const float* __restrict__ residual, float* __restrict__ y, int H, int W, int C, int K, int R, int kb_count) { - __shared__ float xs[kPY * kPX * kTC]; - __shared__ float ws[kTK * 9 * kTC]; + __shared__ __align__(16) float xs[kPY * kPX * kTC]; + __shared__ __align__(16) float ws[kTK * 9 * kTC]; const int n = blockIdx.z / kb_count; const int kb = blockIdx.z % kb_count; const int ox0 = blockIdx.x * kTW, oy0 = blockIdx.y * kTH; @@ -35,7 +35,9 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest const int pixel = tid >> 1; // 0..63 const int cg = (tid & 1) * 16; // first / second 16 output channels const int py = pixel / kTW, px = pixel % kTW; - float acc[16] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float4 acc[16]; +#pragma unroll + for (int ch = 0; ch < 16; ++ch) acc[ch] = make_float4(0.f, 0.f, 0.f, 0.f); for (int c0 = 0; c0 < C; c0 += kTC) { for (int i = tid; i < kPY * kPX * kTC; i += kConvThreads) { @@ -55,32 +57,76 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest : 0.f; } __syncthreads(); - for (int cc = 0; cc < kTC; ++cc) { - if (c0 + cc >= C) break; - for (int r = 0; r < R; ++r) { - for (int s = 0; s < R; ++s) { - const float xv = xs[((py + r) * kPX + (px + s)) * kTC + cc]; + // Four input channels at a time: one float4 shared load per (channel, tap) feeds four FMAs into the float4 + // accumulator, which is four times fewer shared load and arithmetic instructions than the scalar loop. + const int taps = R * R; + for (int cc4 = 0; cc4 < kTC; cc4 += 4) { + if (c0 + cc4 + 3 < C) { + for (int r = 0; r < R; ++r) { + for (int s = 0; s < R; ++s) { + const float4 xv = *reinterpret_cast(&xs[((py + r) * kPX + (px + s)) * kTC + cc4]); +#pragma unroll + for (int ch = 0; ch < 16; ++ch) { + const float4 wv = *reinterpret_cast(&ws[((cg + ch) * 9 + r * R + s) * kTC + cc4]); + acc[ch].x += xv.x * wv.x; + acc[ch].y += xv.y * wv.y; + acc[ch].z += xv.z * wv.z; + acc[ch].w += xv.w * wv.w; + } + } + } + } else { + for (int t = cc4; t < kTC && c0 + t < C; ++t) { + for (int r = 0; r < R; ++r) { + for (int s = 0; s < R; ++s) { + const float xv = xs[((py + r) * kPX + (px + s)) * kTC + t]; + const float xw = xv; #pragma unroll - for (int ch = 0; ch < 16; ++ch) acc[ch] += xv * ws[((cg + ch) * 9 + r * R + s) * kTC + cc]; + for (int ch = 0; ch < 16; ++ch) { + const float v = ws[((cg + ch) * 9 + r * R + s) * kTC + t]; + acc[ch].x += xw * v; + } + } + } } } } + (void)taps; __syncthreads(); } + float h[16]; +#pragma unroll + for (int ch = 0; ch < 16; ++ch) h[ch] = (acc[ch].x + acc[ch].y) + (acc[ch].z + acc[ch].w); const int ox = ox0 + px, oy = oy0 + py; if (ox >= W || oy >= H) return; + const std::int64_t out_base = ((static_cast(n) * H + oy) * W + ox) * K; + if ((K & 3) == 0 && kb * kTK + cg + 16 <= K) { + const int kch = kb * kTK + cg; #pragma unroll - for (int ch = 0; ch < 16; ++ch) { - const int kch = kb * kTK + cg + ch; - if (kch >= K) break; - const std::int64_t idx = ((static_cast(n) * H + oy) * W + ox) * K + kch; - float v = acc[ch] + bias[kch]; - if (residual != nullptr) v += residual[idx]; - y[idx] = v; + for (int ch = 0; ch < 16; ch += 4) { + const float4 b = *reinterpret_cast(&bias[kch + ch]); + float4 v = make_float4(h[ch] + b.x, h[ch + 1] + b.y, h[ch + 2] + b.z, h[ch + 3] + b.w); + if (residual != nullptr) { + const float4 r4 = *reinterpret_cast(&residual[out_base + kch + ch]); + v.x += r4.x; + v.y += r4.y; + v.z += r4.z; + v.w += r4.w; + } + *reinterpret_cast(&y[out_base + kch + ch]) = v; + } + } else { +#pragma unroll + for (int ch = 0; ch < 16; ++ch) { + const int kch = kb * kTK + cg + ch; + if (kch >= K) break; + float v = h[ch] + bias[kch]; + if (residual != nullptr) v += residual[out_base + kch]; + y[out_base + kch] = v; + } } } - // ---------------------------------------------------------------- mid-block attention (single head, fp32) // One warp per query row: the lane holds D/32 elements of q, v and the output; the score is a warp reduction and the // softmax runs online, so the L x L matrix never exists. D is a multiple of 32 (384 or 512). From 1e3f8a557936862de6d21136b0bbfce6e00003a6 Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 12:53:22 +0800 Subject: [PATCH 03/12] sm75: keep attention's Q in registers, three blocks per SM, eight warps for head_dim 256 The exact attention kernel staged the 64x64 Q tile in shared memory, which left one block per SM and four warps running. The Q fragments are now loaded straight from global memory (bf16 -> fp16 per lane), the block's shared memory is only the K and V tiles (about 19 KiB), and three blocks fit an SM. head_dim 256 (the text encoder's causal attention) uses eight warps per block. Measured with tools/sm75_check bench (RTX 2080): the DiT's 1024^2 attention call 32.8 -> 14.3 ms (3.2 -> 7.3 TFLOPS), the 256^2 call 0.23 -> 0.13 ms, the text encoder's 4096-token causal attention 29.4 -> 17.3 ms (7.9 TFLOPS). End to end at 1024^2 the denoise step drops from 6.4 to 5.2 s (two steps: 28.2 -> 23.6 s); 256^2 is unchanged within noise. tools/sm75_check gains a bench command for these shapes. --- docs/sm75.md | 2 +- src/kernels/sm75/attention_sm75.cu | 47 +++++++++++++------------- tools/sm75_check.cu | 53 +++++++++++++++++++++++++++++- 3 files changed, 78 insertions(+), 24 deletions(-) diff --git a/docs/sm75.md b/docs/sm75.md index d71351e..38f5e1f 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -77,7 +77,7 @@ The kernels this backend owns were optimised from that profile: `conv2d_f32_sm75 and float4 accumulators (four input channels per FMA) with vectorised float4 output stores (VAE decode 0.37 -> 0.20 s at 256^2 and 10-12 -> 7.8 s at 1024^2; encode 0.39 -> 0.19 s), `gemm_out_f32` loads float4 and stores packed bf16 (176 -> 93 ms per four-step run), and `bf16_to_f16_rows_scaled` uses packed pair loads/stores. -## Numerics +## Numerics The exact attention kernel is next: it kept its 64-query tile but now loads the Q fragments straight from global memory instead of staging them in shared memory, which lets three blocks share an SM instead of one, and it uses eight warps for the text encoder's head_dim 256 (four for the DiT's 128): the 1024^2 DiT attention call went from 32.8 to 14.3 ms (3.2 to 7.3 TFLOPS) and the text encoder's from 29.4 to 17.3 ms (7.9 TFLOPS). - 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 diff --git a/src/kernels/sm75/attention_sm75.cu b/src/kernels/sm75/attention_sm75.cu index d7ef6bb..5ab407b 100644 --- a/src/kernels/sm75/attention_sm75.cu +++ b/src/kernels/sm75/attention_sm75.cu @@ -30,9 +30,7 @@ namespace mi::kernels::sm75 { namespace { -constexpr int kThreads = 128; // 4 warps -constexpr int kWarps = 4; -constexpr int kBM = 64; // query rows per CTA (16 per warp) + constexpr float kLog2e = 1.4426950408889634f; // m16n8k8 fp16 -> fp32 MMA. a = {row group, row group + 8} 2-half registers, b = 2-half register, d = 4 fp32. @@ -89,16 +87,17 @@ struct Params { bool causal = false; }; -template +template __global__ void attention_kernel_sm75(const Params p) { + constexpr int kThreads = WARPS * 32; + constexpr int kBM = WARPS * 16; // 16 query rows per warp constexpr int kQStride = D + 8; constexpr int kVStride = BN + 8; constexpr int kSubtiles = BN / 8; constexpr int kDimChunks = D / 8; extern __shared__ __align__(16) char smem_raw[]; - __half* qs = reinterpret_cast<__half*>(smem_raw); // [kBM][D] - __half* ks = qs + kBM * kQStride; // [BN][D] + __half* ks = reinterpret_cast<__half*>(smem_raw); // [BN][D] __half* vt = ks + BN * kQStride; // [D][BN] const int batch = blockIdx.y; @@ -120,22 +119,24 @@ __global__ void attention_kernel_sm75(const Params p) { const int group = lane >> 2; // 0..7: the row of the 8x8 fragment const int tig = lane & 3; // 0..3: two consecutive columns per lane - // 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) ? to_half(p.q[(q_start + row_base + r) * p.q_row + head * p.q_head + d]) - : __float2half_rn(0.f); - } - __syncthreads(); - - // This warp's Q fragment: rows {warp*16 + group, +8}, dims {i*8 + 2*tig, +1}. + // This warp's Q fragment: rows {warp*16 + group, +8}, dims {i*8 + 2*tig, +1}, loaded straight from global memory + // (bf16 -> fp16). Keeping Q out of shared memory frees 17 KiB per block, so three blocks fit an SM instead of one. std::uint32_t qf[kDimChunks][2]; const int r0 = warp * 16 + group; const int r1 = r0 + 8; + const std::int64_t q_head_off = static_cast(head) * p.q_head; + const auto load_q_pair = [&](int row, int chunk) -> std::uint32_t { + if (row >= rows) return 0u; + const __nv_bfloat16* s = + p.q + static_cast(q_start + row_base + row) * p.q_row + q_head_off + chunk * 8 + tig * 2; + const float2 f = __bfloat1622float2(*reinterpret_cast(s)); + const __half2 h = __floats2half2_rn(f.x, f.y); + return *reinterpret_cast(&h); + }; #pragma unroll for (int i = 0; i < kDimChunks; ++i) { - qf[i][0] = load_u32(&qs[r0 * kQStride + i * 8 + tig * 2]); - qf[i][1] = load_u32(&qs[r1 * kQStride + i * 8 + tig * 2]); + qf[i][0] = load_q_pair(r0, i); + qf[i][1] = load_q_pair(r1, i); } float o[kDimChunks][4]; // P·V accumulators, 4 floats per 8-dim chunk @@ -272,13 +273,15 @@ __global__ void attention_kernel_sm75(const Params p) { } } -template +template void launch(const Params& p, int max_seqlen, int batch, cudaStream_t stream) { + constexpr int kThreads = WARPS * 32; + constexpr int kBM = WARPS * 16; constexpr int kQStride = D + 8; constexpr int kVStride = BN + 8; - constexpr int kSmem = (kBM * kQStride + BN * kQStride + D * kVStride) * static_cast(sizeof(__half)); + constexpr int kSmem = (BN * kQStride + D * kVStride) * static_cast(sizeof(__half)); static_assert(kSmem <= 64 * 1024, "sm_75 allows 64 KiB of opt-in shared memory per block"); - auto kernel = &attention_kernel_sm75; + auto kernel = &attention_kernel_sm75; if constexpr (kSmem > 48 * 1024) { MI_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmem)); } @@ -334,9 +337,9 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& const int batch = static_cast(cu_seqlens_q.dim(0) - 1); if (head_dim == 128) { - launch<128, 32>(p, max_seqlen, batch, stream); + launch<128, 32, 4>(p, max_seqlen, batch, stream); } else { - launch<256, 16>(p, max_seqlen, batch, stream); + launch<256, 16, 8>(p, max_seqlen, batch, stream); } } diff --git a/tools/sm75_check.cu b/tools/sm75_check.cu index 262e1c9..c3f756f 100644 --- a/tools/sm75_check.cu +++ b/tools/sm75_check.cu @@ -526,6 +526,53 @@ int run_int8_gemm(int M, int N, int K) { } // namespace +// Times the sm_75 exact attention at engine shapes: one sequence, `len` tokens, `iters` launches, CUDA events. +int run_attn_bench(const char* name, int heads, int kv_heads, int dim, bool causal, int len, int iters) { + const auto q = normal(static_cast(len) * heads * dim, 71, 1.f); + const auto k = normal(static_cast(len) * kv_heads * dim, 72, 1.f); + const auto v = normal(static_cast(len) * kv_heads * dim, 73, 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(2 * sizeof(int)); + dq.upload(qb); + dk.upload(kb); + dv.upload(vb); + const std::vector cu{0, len}; + dcu.upload(cu); + MI_CUDA_CHECK(cudaMemset(dout.get(), 0, qb.size() * 2)); + auto t = [](void* ptr, DType dt, std::initializer_list shape) { + return Tensor::view(ptr, dt, Device::kCUDA, shape); + }; + const Tensor tq = t(dq.get(), DType::kBF16, {len, heads, dim}); + const Tensor tk = t(dk.get(), DType::kBF16, {len, kv_heads, dim}); + const Tensor tv = t(dv.get(), DType::kBF16, {len, kv_heads, dim}); + const Tensor to = t(dout.get(), DType::kBF16, {len, heads, dim}); + const Tensor tcu = t(dcu.get(), DType::kI32, {2}); + const float scale = 1.f / std::sqrt(static_cast(dim)); + const std::size_t ws_bytes = static_cast(len) * heads * sizeof(float); + DevBuf dls(ws_bytes); + MI_CUDA_CHECK(cudaMemset(dls.get(), 0, ws_bytes)); + mi::kernels::sm75::flash_fwd(tq, tk, tv, to, tcu, tcu, len, causal, scale, dls.get(), ws_bytes, nullptr); + MI_CUDA_CHECK(cudaDeviceSynchronize()); + cudaEvent_t start, stop; + MI_CUDA_CHECK(cudaEventCreate(&start)); + MI_CUDA_CHECK(cudaEventCreate(&stop)); + MI_CUDA_CHECK(cudaEventRecord(start)); + for (int i = 0; i < iters; ++i) mi::kernels::sm75::flash_fwd(tq, tk, tv, to, tcu, tcu, len, causal, scale, dls.get(), ws_bytes, nullptr); + MI_CUDA_CHECK(cudaEventRecord(stop)); + MI_CUDA_CHECK(cudaEventSynchronize(stop)); + float ms = 0.f; + MI_CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop)); + const double seconds = ms / 1000.0 / iters; + const double flops = (causal ? 0.5 : 1.0) * 2.0 * len * len * dim * heads; + std::printf("attn_bench %-8s L %5d heads %2d kv %2d dim %3d causal %d: %8.3f ms %7.1f TFLOPS\n", name, len, heads, + kv_heads, dim, causal ? 1 : 0, seconds * 1e3, flops / seconds / 1e12); + MI_CUDA_CHECK(cudaEventDestroy(start)); + MI_CUDA_CHECK(cudaEventDestroy(stop)); + return 0; +} int main(int argc, char** argv) { std::setvbuf(stdout, nullptr, _IONBF, 0); const std::string what = argc > 1 ? argv[1] : "all"; @@ -584,7 +631,11 @@ int main(int argc, char** argv) { failures += run_attention("te282", 16, 4, 256, true, {282}, 42); failures += run_attention("dit76", 24, 24, 128, false, {76, 200, 1}, 43); } - } catch (const std::exception& e) { + if (what == "bench") { + failures += run_attn_bench("dit256", 24, 24, 128, false, 266, 50); + failures += run_attn_bench("dit1024", 24, 24, 128, false, 4113, 5); + failures += run_attn_bench("te", 32, 8, 256, true, 4096, 5); + } } catch (const std::exception& e) { std::printf("EXCEPTION: %s\n", e.what()); return 2; } catch (...) { From 9269bd1a8fee4f8efb507dfbf658f31b07aaac8c Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:10:27 +0800 Subject: [PATCH 04/12] sm75: tile the VAE mid attention over query rows, with compile-time lane counts The VAE mid attention read K and V for every query row: at a 1024^2 latent (16384 tokens) that is 8 TB of reads and 5.4 s of the decode. The kernel now keeps 32 query rows per block in registers (four per warp) and stages the K and V tiles once per block in shared memory, so each tile is read L/32 times. The lane count is a template parameter: with a runtime D the per-row arrays landed in local memory and the kernel was 20x slower. Measured at 1024^2 (RTX 2080): VAE attention 5.4 s -> 1.6 s, decode 7.6 s -> 3.3 s, one-step total 17.4 -> 14.0 s. The decoded PNG is bit-identical to the previous build (same key order in the online softmax). sm75_check vae still matches its host reference. --- src/kernels/sm75/vae_sm75.cu | 176 +++++++++++++++++++++++++++-------- 1 file changed, 136 insertions(+), 40 deletions(-) diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index f010986..a87528e 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -128,47 +128,87 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest } } // ---------------------------------------------------------------- mid-block attention (single head, fp32) -// One warp per query row: the lane holds D/32 elements of q, v and the output; the score is a warp reduction and the -// softmax runs online, so the L x L matrix never exists. D is a multiple of 32 (384 or 512). -constexpr int kAttnWarps = 4; +// One warp per row read K and V from global memory in full for every query row: at a 1024^2 latent (16384 tokens) that +// is 8 TB of reads and 5.4 s on an RTX 2080. This version keeps 32 query rows per block in registers (four per warp) +// and stages the K and V tiles once per block in shared memory, so every tile is read L/32 times instead of L times. +// D is a multiple of 32 (384 or 512); the online softmax keeps the same key order, so the output bits are unchanged. +constexpr int kAttnWarps = 8; +constexpr int kAttnRowsPerWarp = 4; +constexpr int kAttnRows = kAttnWarps * kAttnRowsPerWarp; // 32 query rows per block +constexpr int kAttnTile = 8; // keys per staged tile -__global__ void attention_f32_sm75(const float* __restrict__ q, const float* __restrict__ k, - const float* __restrict__ v, float* __restrict__ o, int L, int D, int q_rows, - int k_rows, int v_rows, std::int64_t q_batch, std::int64_t k_batch, - std::int64_t v_batch, std::int64_t o_batch, float scale) { +// LANES is D/32 at compile time, so the per-row register arrays keep their indices in registers; a runtime lane count +// would place them in local memory and cost an order of magnitude. +template +__global__ void attention_f32_sm75_t(const float* __restrict__ q, const float* __restrict__ k, + const float* __restrict__ v, float* __restrict__ o, int L, int q_rows, + int k_rows, int v_rows, std::int64_t q_batch, std::int64_t k_batch, + std::int64_t v_batch, std::int64_t o_batch, float scale) { + constexpr int kD = LANES * 32; + extern __shared__ __align__(16) float smem[]; const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - const int row = blockIdx.x * kAttnWarps + warp; - if (row >= L) return; - const int nlanes = D / 32; + const int tid = threadIdx.x; const int b = blockIdx.y; - const float* qrow = q + static_cast(b) * q_batch + static_cast(row) * q_rows; - float qreg[16]; - float oacc[16]; - for (int i = 0; i < nlanes; ++i) { - qreg[i] = qrow[lane + 32 * i]; - oacc[i] = 0.f; + const int row0 = blockIdx.x * kAttnRows + warp * kAttnRowsPerWarp; + float qreg[kAttnRowsPerWarp][LANES]; + float oacc[kAttnRowsPerWarp][LANES]; + float m[kAttnRowsPerWarp], s[kAttnRowsPerWarp]; +#pragma unroll + for (int r = 0; r < kAttnRowsPerWarp; ++r) { + const bool valid = row0 + r < L; + const float* qrow = q + static_cast(b) * q_batch + static_cast(row0 + r) * q_rows; + m[r] = -1e30f; + s[r] = 0.f; +#pragma unroll + for (int i = 0; i < LANES; ++i) { + qreg[r][i] = valid ? qrow[lane + 32 * i] : 0.f; + oacc[r][i] = 0.f; + } } - float m = -1e30f, s = 0.f; - for (int j = 0; j < L; ++j) { - const float* krow = k + static_cast(b) * k_batch + static_cast(j) * k_rows; - float dot = 0.f; - for (int i = 0; i < nlanes; ++i) dot += qreg[i] * krow[lane + 32 * i]; + const float* kbase = k + static_cast(b) * k_batch; + const float* vbase = v + static_cast(b) * v_batch; + float* ktile = smem; + float* vtile = smem + kAttnTile * kD; + for (int j0 = 0; j0 < L; j0 += kAttnTile) { + for (int idx = tid; idx < kAttnTile * kD; idx += kAttnWarps * 32) { + const int jj = idx / kD, d = idx % kD; + const int key = j0 + jj; + const bool valid = key < L; + ktile[idx] = valid ? kbase[static_cast(key) * k_rows + d] : 0.f; + vtile[idx] = valid ? vbase[static_cast(key) * v_rows + d] : 0.f; + } + __syncthreads(); + const int jn = min(kAttnTile, L - j0); + for (int jj = 0; jj < jn; ++jj) { + const float* krow = ktile + jj * kD; + const float* vrow = vtile + jj * kD; #pragma unroll - for (int off = 16; off > 0; off >>= 1) dot += __shfl_xor_sync(0xffffffffu, dot, off); - dot *= scale; - const float mn = fmaxf(m, dot); - const float alpha = (m > -1e29f) ? expf(m - mn) : 0.f; - const float p = expf(dot - mn); - s = s * alpha + p; - const float* vrow = v + static_cast(b) * v_batch + static_cast(j) * v_rows; - for (int i = 0; i < nlanes; ++i) oacc[i] = oacc[i] * alpha + p * vrow[lane + 32 * i]; - m = mn; + for (int r = 0; r < kAttnRowsPerWarp; ++r) { + float dot = 0.f; +#pragma unroll + for (int i = 0; i < LANES; ++i) dot += qreg[r][i] * krow[lane + 32 * i]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) dot += __shfl_xor_sync(0xffffffffu, dot, off); + dot *= scale; + const float mn = fmaxf(m[r], dot); + const float alpha = (m[r] > -1e29f) ? expf(m[r] - mn) : 0.f; + const float p = expf(dot - mn); + s[r] = s[r] * alpha + p; +#pragma unroll + for (int i = 0; i < LANES; ++i) oacc[r][i] = oacc[r][i] * alpha + p * vrow[lane + 32 * i]; + m[r] = mn; + } + } + __syncthreads(); + } +#pragma unroll + for (int r = 0; r < kAttnRowsPerWarp; ++r) { + if (row0 + r >= L) continue; + const float inv = 1.f / s[r]; + float* orow = o + static_cast(b) * o_batch + static_cast(row0 + r) * kD; + for (int i = 0; i < LANES; ++i) orow[lane + 32 * i] = oacc[r][i] * inv; } - const float inv = 1.f / s; - float* orow = o + static_cast(b) * o_batch + static_cast(row) * D; - for (int i = 0; i < nlanes; ++i) orow[lane + 32 * i] = oacc[i] * inv; } - } // namespace void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, @@ -254,16 +294,72 @@ void vae_conv2d_down(const Tensor& x, const Tensor& w, const Tensor& bias, const y.ptr(), H, W, C, K, P, Q, kb_count); MI_CUDA_CHECK_LAUNCH(); } +// Generic fallback (any D that is a multiple of 32): one warp per query row, the lane holds D/32 elements. Slow at long +// sequences, used only for head dimensions the register-tiled kernel does not instantiate (the checks' tiny shapes). +__global__ void attention_f32_sm75_dyn(const float* __restrict__ q, const float* __restrict__ k, + const float* __restrict__ v, float* __restrict__ o, int L, int D, int q_rows, + int k_rows, int v_rows, std::int64_t q_batch, std::int64_t k_batch, + std::int64_t v_batch, std::int64_t o_batch, float scale) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int row = blockIdx.x * 4 + warp; + if (row >= L) return; + const int nlanes = D / 32; + const int b = blockIdx.y; + const float* qrow = q + static_cast(b) * q_batch + static_cast(row) * q_rows; + float qreg[16]; + float oacc[16]; + for (int i = 0; i < nlanes; ++i) { + qreg[i] = qrow[lane + 32 * i]; + oacc[i] = 0.f; + } + float m = -1e30f, s = 0.f; + for (int j = 0; j < L; ++j) { + const float* krow = k + static_cast(b) * k_batch + static_cast(j) * k_rows; + float dot = 0.f; + for (int i = 0; i < nlanes; ++i) dot += qreg[i] * krow[lane + 32 * i]; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) dot += __shfl_xor_sync(0xffffffffu, dot, off); + dot *= scale; + const float mn = fmaxf(m, dot); + const float alpha = (m > -1e29f) ? expf(m - mn) : 0.f; + const float p = expf(dot - mn); + s = s * alpha + p; + const float* vrow = v + static_cast(b) * v_batch + static_cast(j) * v_rows; + for (int i = 0; i < nlanes; ++i) oacc[i] = oacc[i] * alpha + p * vrow[lane + 32 * i]; + m = mn; + } + const float inv = 1.f / s; + float* orow = o + static_cast(b) * o_batch + static_cast(row) * D; + for (int i = 0; i < nlanes; ++i) orow[lane + 32 * i] = oacc[i] * inv; +} + void vae_attention_fp32(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& o, float scale, cudaStream_t stream) { const int L = static_cast(q.dim(1)), D = static_cast(q.dim(2)); const int batch = static_cast(q.dim(0)); - const dim3 grid(static_cast((L + kAttnWarps - 1) / kAttnWarps), static_cast(batch)); - attention_f32_sm75<<>>(q.ptr(), k.ptr(), v.ptr(), - o.ptr(), L, D, static_cast(q.stride[1]), - static_cast(k.stride[1]), - static_cast(v.stride[1]), q.stride[0], k.stride[0], - v.stride[0], o.stride[0], scale); + const dim3 grid(static_cast((L + kAttnRows - 1) / kAttnRows), static_cast(batch)); + const std::size_t smem = static_cast(2) * kAttnTile * D * sizeof(float); + const auto launch = [&](auto* kernel) { + if (smem > 48 * 1024) { + MI_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem))); + } + kernel<<>>(q.ptr(), k.ptr(), v.ptr(), + o.ptr(), L, static_cast(q.stride[1]), + static_cast(k.stride[1]), static_cast(v.stride[1]), + q.stride[0], k.stride[0], v.stride[0], o.stride[0], scale); + }; + if (D == 384) { + launch(&attention_f32_sm75_t<12>); + } else if (D == 512) { + launch(&attention_f32_sm75_t<16>); + } else { + const dim3 g(static_cast((L + 3) / 4), static_cast(batch)); + attention_f32_sm75_dyn<<>>(q.ptr(), k.ptr(), v.ptr(), o.ptr(), + L, D, static_cast(q.stride[1]), static_cast(k.stride[1]), + static_cast(v.stride[1]), q.stride[0], k.stride[0], + v.stride[0], o.stride[0], scale); + } MI_CUDA_CHECK_LAUNCH(); } From ca9466cfff116a1e788eaab27bb36e42c482628f Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:12:23 +0800 Subject: [PATCH 05/12] sm75: two output pixels per thread in the VAE conv kernel Nsight Compute showed the conv kernel L1/TEX-bound at 93 % with 40 % compute throughput: the weight tile was reloaded for every output pixel. Each thread now computes two pixels eight rows apart, so each weight float4 serves two pixels and the shared-memory traffic per FMA halves. t=1024^2 VAE decode 3.30 -> 2.19 s, bit-identical output. --- src/kernels/sm75/vae_sm75.cu | 95 ++++++++++++++++++++---------------- 1 file changed, 53 insertions(+), 42 deletions(-) diff --git a/src/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index a87528e..36e716c 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -18,9 +18,9 @@ namespace { // ---------------------------------------------------------------- conv2d (NHWC, KRSC filters, stride 1, pad R/2) // One block: an 8x8 output tile of 32 output channels. Input channels are walked in groups of 8 through shared memory // (the input patch and the weights they multiply), each thread accumulating 2 pixels x 16 channels. -constexpr int kTW = 8, kTH = 8, kTK = 32, kTC = 8; +constexpr int kTW = 8, kTH = 16, kTK = 32, kTC = 8; constexpr int kPX = kTW + 2, kPY = kTH + 2; -constexpr int kConvThreads = 128; +constexpr int kConvThreads = 128; // each thread computes two output pixels with the same weights __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __restrict__ w, const float* __restrict__ bias, const float* __restrict__ residual, @@ -32,12 +32,15 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest const int ox0 = blockIdx.x * kTW, oy0 = blockIdx.y * kTH; const int pad = (R - 1) / 2; const int tid = threadIdx.x; - const int pixel = tid >> 1; // 0..63 + const int pixel = tid >> 1; // 0..63; the second pixel is eight rows below const int cg = (tid & 1) * 16; // first / second 16 output channels const int py = pixel / kTW, px = pixel % kTW; - float4 acc[16]; + float4 acc[2][16]; #pragma unroll - for (int ch = 0; ch < 16; ++ch) acc[ch] = make_float4(0.f, 0.f, 0.f, 0.f); + for (int q = 0; q < 2; ++q) { +#pragma unroll + for (int ch = 0; ch < 16; ++ch) acc[q][ch] = make_float4(0.f, 0.f, 0.f, 0.f); + } for (int c0 = 0; c0 < C; c0 += kTC) { for (int i = tid; i < kPY * kPX * kTC; i += kConvThreads) { @@ -57,21 +60,25 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest : 0.f; } __syncthreads(); - // Four input channels at a time: one float4 shared load per (channel, tap) feeds four FMAs into the float4 - // accumulator, which is four times fewer shared load and arithmetic instructions than the scalar loop. - const int taps = R * R; + // Four input channels at a time: one float4 shared load per (channel, tap) feeds four FMAs into each of the two + // pixels' float4 accumulators, so the weight tile is loaded for two output pixels at once. for (int cc4 = 0; cc4 < kTC; cc4 += 4) { if (c0 + cc4 + 3 < C) { for (int r = 0; r < R; ++r) { for (int s = 0; s < R; ++s) { - const float4 xv = *reinterpret_cast(&xs[((py + r) * kPX + (px + s)) * kTC + cc4]); + const float4 xv0 = *reinterpret_cast(&xs[((py + r) * kPX + (px + s)) * kTC + cc4]); + const float4 xv1 = *reinterpret_cast(&xs[((py + 8 + r) * kPX + (px + s)) * kTC + cc4]); #pragma unroll for (int ch = 0; ch < 16; ++ch) { const float4 wv = *reinterpret_cast(&ws[((cg + ch) * 9 + r * R + s) * kTC + cc4]); - acc[ch].x += xv.x * wv.x; - acc[ch].y += xv.y * wv.y; - acc[ch].z += xv.z * wv.z; - acc[ch].w += xv.w * wv.w; + acc[0][ch].x += xv0.x * wv.x; + acc[0][ch].y += xv0.y * wv.y; + acc[0][ch].z += xv0.z * wv.z; + acc[0][ch].w += xv0.w * wv.w; + acc[1][ch].x += xv1.x * wv.x; + acc[1][ch].y += xv1.y * wv.y; + acc[1][ch].z += xv1.z * wv.z; + acc[1][ch].w += xv1.w * wv.w; } } } @@ -79,54 +86,58 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest for (int t = cc4; t < kTC && c0 + t < C; ++t) { for (int r = 0; r < R; ++r) { for (int s = 0; s < R; ++s) { - const float xv = xs[((py + r) * kPX + (px + s)) * kTC + t]; - const float xw = xv; + const float xv0 = xs[((py + r) * kPX + (px + s)) * kTC + t]; + const float xv1 = xs[((py + 8 + r) * kPX + (px + s)) * kTC + t]; #pragma unroll for (int ch = 0; ch < 16; ++ch) { const float v = ws[((cg + ch) * 9 + r * R + s) * kTC + t]; - acc[ch].x += xw * v; + acc[0][ch].x += xv0 * v; + acc[1][ch].x += xv1 * v; } } } } } } - (void)taps; __syncthreads(); } - float h[16]; #pragma unroll - for (int ch = 0; ch < 16; ++ch) h[ch] = (acc[ch].x + acc[ch].y) + (acc[ch].z + acc[ch].w); - const int ox = ox0 + px, oy = oy0 + py; - if (ox >= W || oy >= H) return; - const std::int64_t out_base = ((static_cast(n) * H + oy) * W + ox) * K; - if ((K & 3) == 0 && kb * kTK + cg + 16 <= K) { - const int kch = kb * kTK + cg; + for (int q = 0; q < 2; ++q) { + const int ox = ox0 + px, oy = oy0 + py + q * 8; + if (ox >= W || oy >= H) continue; + float h[16]; #pragma unroll - for (int ch = 0; ch < 16; ch += 4) { - const float4 b = *reinterpret_cast(&bias[kch + ch]); - float4 v = make_float4(h[ch] + b.x, h[ch + 1] + b.y, h[ch + 2] + b.z, h[ch + 3] + b.w); - if (residual != nullptr) { - const float4 r4 = *reinterpret_cast(&residual[out_base + kch + ch]); - v.x += r4.x; - v.y += r4.y; - v.z += r4.z; - v.w += r4.w; + for (int ch = 0; ch < 16; ++ch) h[ch] = (acc[q][ch].x + acc[q][ch].y) + (acc[q][ch].z + acc[q][ch].w); + const std::int64_t out_base = ((static_cast(n) * H + oy) * W + ox) * K; + if ((K & 3) == 0 && kb * kTK + cg + 16 <= K) { + const int kch = kb * kTK + cg; +#pragma unroll + for (int ch = 0; ch < 16; ch += 4) { + const float4 b = *reinterpret_cast(&bias[kch + ch]); + float4 v = make_float4(h[ch] + b.x, h[ch + 1] + b.y, h[ch + 2] + b.z, h[ch + 3] + b.w); + if (residual != nullptr) { + const float4 r4 = *reinterpret_cast(&residual[out_base + kch + ch]); + v.x += r4.x; + v.y += r4.y; + v.z += r4.z; + v.w += r4.w; + } + *reinterpret_cast(&y[out_base + kch + ch]) = v; } - *reinterpret_cast(&y[out_base + kch + ch]) = v; - } - } else { + } else { #pragma unroll - for (int ch = 0; ch < 16; ++ch) { - const int kch = kb * kTK + cg + ch; - if (kch >= K) break; - float v = h[ch] + bias[kch]; - if (residual != nullptr) v += residual[out_base + kch]; - y[out_base + kch] = v; + for (int ch = 0; ch < 16; ++ch) { + const int kch = kb * kTK + cg + ch; + if (kch >= K) break; + float v = h[ch] + bias[kch]; + if (residual != nullptr) v += residual[out_base + kch]; + y[out_base + kch] = v; + } } } } + // ---------------------------------------------------------------- mid-block attention (single head, fp32) // One warp per row read K and V from global memory in full for every query row: at a 1024^2 latent (16384 tokens) that // is 8 TB of reads and 5.4 s on an RTX 2080. This version keeps 32 query rows per block in registers (four per warp) From bbd0316a46f903cd83c16e944e1a19f3cb28dc9b Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:20:38 +0800 Subject: [PATCH 06/12] docs: quantized tiers on sm_75, their gaps and the plan New docs/sm75-quantization.md: what reference/high/turbo require, the measured RTX 2080 numbers for the three tiers with quantized DiT files (0.12 s/step INT8, 0.14 s/step INT4 at 256^2), the gaps (no INT4 writer in quantize, high needs a quantized DiT on Turing because there is no 8-bit attention, LoRA on quantized tiers unverified) and a five-point plan (INT4 output, Turing q8_fp16 attention, a tier rule fallback, LoRA on quantized tiers, docs/app). docs/sm75.md links it and states the quantized-file requirement. --- docs/sm75-quantization.md | 97 +++++++++++++++++++++++++++++++++++++++ docs/sm75.md | 9 +++- 2 files changed, 105 insertions(+), 1 deletion(-) create mode 100644 docs/sm75-quantization.md diff --git a/docs/sm75-quantization.md b/docs/sm75-quantization.md new file mode 100644 index 0000000..dbd3b98 --- /dev/null +++ b/docs/sm75-quantization.md @@ -0,0 +1,97 @@ +# Quantized models on sm_75: status and plan + +The mainline engine runs three quality tiers. On sm_75 only the bf16 `reference` tier is *available without extra +files* today; the quantized tiers work when the models are already quantized, and the pieces needed to make them +first-class (and to produce the files) are listed below with a plan. + +## What the tiers mean + +| Tier | Weights | Attention | What the loaded DiT must be | +|---|---|---|---| +| `reference` | bf16, exact | exact | a bf16 DiT and text encoder | +| `high` | INT8 ConvRot (W8A8), optional quantized text encoder | 8-bit (`q8_fp16`/`q8_fp8`) | an INT8 ConvRot DiT **file**, or a bf16 DiT *and* 8-bit attention | +| `turbo` | INT4 ConvRot (`convrot_w4a4`) or NVFP4 | 8-bit | an INT4/NVFP4 DiT file (NVFP4 needs an RTX 50) | + +`planner::resolve_tier` (src/planner/planner.cpp) enforces that: `high` with exact attention on a bf16 DiT is "the +Reference configuration" and is refused, and on sm_75 the missing 8-bit attention kernels make `exact_when_no_q8` pick +exact attention, so `high` needs an INT8 DiT file. `turbo` only accepts a 4-bit file. + +## What already works on sm_75 + +Measured on the RTX 2080 8 GiB with the SFT bundle (`balanced`), this branch, after the attention and VAE work: + +| Configuration | DiT file | Denoise, 256^2, 8 steps | Denoise, 1024^2, 2 steps | Device peak | +|---|---|---|---|---| +| `reference` | bf16 bundle | 7.50 s (0.94 s/step) | 8.60 s (4.3 s/step) | 6.37 GiB | +| `high` | INT8 ConvRot, 5.04 GB | **0.99 s (0.12 s/step)** | **5.18 s (2.54 s/step)** | 5.27 GiB | +| `turbo` | INT4 ConvRot, 2.83 GB | 1.14 s (0.14 s/step) | 5.16 s (2.57 s/step) | 3.06 GiB | + +- The INT8 GEMMs are CUTLASS Sm75 m8n8k16 kernels with the sm_80 epilogues (dequantize+bias, gated residual, + SiTU-GLU) under `src/kernels/sm75/int8_gemm_sm75.cu`; one tile configuration (64x128x64) fits Turing's 64 KiB of + shared memory, so forcing another tile is refused (`check_sm75_tile`). +- INT4 weights stay packed and are unpacked into the INT8 GEMM (`quant4/w4a8_decode.cu`), so `turbo` is correct on + Turing (about 1.5x the INT8 step in the planner's price, not a native 4-bit MMA). +- Exact attention runs with every tier; the quantized tiers do not need 8-bit attention on sm_75. +- `moss-infer quantize` runs on an 8 GiB Turing card: the 5B DiT became a 5.04 GB INT8 ConvRot file in 49 s + (252 GEMMs quantized, 0 kept bf16). +- The VAE stays fp32 SIMT either way (see `docs/sm75.md`). + +## Gaps + +**G1 - no INT4 writer.** `moss-infer quantize --format` accepts only `int8` and `nvfp4` +(src/tools/cli/cmd_quantize.cpp). `convrot_w4a4` files (the only DiT format `turbo` accepts) therefore come from +elsewhere. The pieces to write them exist: `src/weights/quant/convrot_w4a4.{h,cpp}` quantizes W4 with the ConvRot +groups and `quant_group_size`, and the loader already reads the format. The file metadata is a +`_quantization_metadata` JSON holding `format_version` and per-layer +`{"format": "convrot_w4a4", "convrot_groupsize": 256, "quant_group_size": 64}` (checked in a community file). + +**G2 - `high` with the bf16 bundle is refused on sm_75.** The tier is defined as INT8 GEMMs *plus* 8-bit attention; +sm_75 has no 8-bit attention, so the resolver falls back to exact attention and then refuses the combination. A +Turing user has to quantize the DiT first (or download an INT8 file) to use `high`. + +**G3 - no 8-bit attention kernel.** `q8_fp16`/`q8_fp8` are refused on sm_75 `--attention`. Besides blocking G2, this +is the mainline's long-sequence accelerator: at 1024^2 the DiT's exact attention is 1.05 s per step +(7.3 TFLOPS with the new kernel), which an INT8 QK mainloop should roughly halve. + +**G4 - LoRA on quantized DiTs is unverified on sm_75.** Merged adapters on INT8 block GEMMs go through +`merge_requantize_int8` (src/models/lora/lora_merge.cpp); on sm_75 every LoRA is merged (no side branches), and an +INT4 DiT has no merge path at all, so `--tier turbo --lora` should refuse clearly instead of silently ignoring the +adapter. Needs tests and, if missing, a message. + +**G5 - the app has to offer it.** MossStudio decides what to download and which tiers to show; a Turing card should be +able to run `quantize` (engine CLI) or fetch INT8/INT4 files, and the tier selector should enable `high`/`turbo` when +one is loaded. + +## Plan + +**P1 - INT4 output in `quantize` (small, independent).** Add `--format int4`: quantize each block GEMM with +`weights::quant::convrot_w4a4`, write packed W4 weights plus the per-layer metadata above (group 256, quant group 64 by +default, `--group` as for INT8), keep the non-block tensors bf16. Then `moss-infer generate --dit --tier turbo` +works with a file produced on the same card. Test: the produced file's size (expected ~2.9 GB), `inspect`, a 256^2 and +a 1024^2 generation, and compare its output with the community file's. + +**P2 - Turing 8-bit attention (`q8_fp16`).** A new kernel in `src/kernels/sm75/`: INT8 QK with the Turing m8n8k16 MMA +over the ConvRot-quantized q/k (the quantizers are already portable: `quant/rotate_quantize.cu`), fp16 P·V with the +existing m16n8k8 MMA, exact softmax in fp32. Wire it into `attention_supported`, `MossDit`'s attention dispatch and +`resolve_tier` so `--attention q8_fp16` is accepted on cc 7.5 and `high` with a bf16 DiT resolves to it. Expected: +attention 14.3 ms -> 8-9 ms at 1024^2 (0.5 s per step saved) and `high` works without a quantized file. + +**P3 - or, if P2 waits: a sm_75 tier rule.** Let `high` on cc 7.5 quantize the block GEMMs at load +(`LoadOptions::int8_block_gemms` currently only follows an INT8 file) and keep exact attention, documented as a +sm_75-specific definition of `high`. Small (planner + load options), but it changes what `high` means on Turing, so +P2 is the cleaner long-term answer. + +**P4 - LoRA on quantized tiers.** Test `--tier high --lora` on sm_75 (merge + requantize), and make `--tier turbo +--lora` refuse with a message that names the missing merge path. If the requantize path turns out to need sm_80 +kernels, port the host merge (`quant::quantize_convrot_rows` is already used by the sm_80 path). + +**P5 - docs and app.** Update `docs/sm75.md`'s limits with the outcome, add the INT4 production step and the measured +tier table here, and tell MossStudio which files/tiers a cc 7.5 device can use. + +## Test matrix + +- `sm75_check int8` (existing) plus a `quantize` -> load -> generate round trip for INT8 and INT4 at 256^2 and 1024^2. +- `high`/`turbo` with `--attention exact` and (after P2) `q8_fp16`: same prompt and seed, check the report's tier line + and `quality.gpu_arch`. +- LoRA: `--tier reference` (merged, works today), `--tier high --lora` (P4), `--tier turbo --lora` (clear refusal). +- Regression: `sm75_check` all cases, and the OFF build (`MOSSINFER_SM75=OFF`) still compiles and passes. diff --git a/docs/sm75.md b/docs/sm75.md index 38f5e1f..0eff99e 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -46,6 +46,12 @@ its CUDA graphs are not used. Every block path produces the same image, so this "(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. +`high` and `turbo` need a quantized DiT file; the bf16 bundle itself only runs `reference`. `moss-infer quantize` +produces INT8 ConvRot files on Turing (5.04 GB for the 5B DiT in 49 s), and a community `convrot_w4a4` file runs as it +is. Measured on an RTX 2080 8 GiB, 256^2: 0.12 s/step (INT8, 5.3 GiB peak) and 0.14 s/step (INT4, 3.1 GiB) against +0.94 s/step for bf16. `docs/sm75-quantization.md` has the tier table, the INT4 writer gap and the plan for 8-bit +attention on Turing. + ## Performance Measured on an RTX 2080 Ti 11 GiB, MOSS-Image SFT, `balanced` mode, warm runs (models resident, prompt cached), the @@ -111,7 +117,8 @@ and float4 accumulators (four input channels per FMA) with vectorised float4 out ## Limits on sm_75 - The 8-bit attention kernels (`q8_fp16`, `q8_fp8`): the tier default resolves to `exact`; passing an 8-bit mode - explicitly is refused. + explicitly is refused. `docs/sm75-quantization.md` (P2) plans the Turing `q8_fp16` kernel, which would also let + `tier high` run on a bf16 DiT. - Native INT4 MMA and NVFP4: INT4 files run through the INT8 unpack fallback (correct and bit-exact; the planner prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. - The fused DiT blocks and their CUDA graphs (the blocks run op by op). From 3684f2ce5a9b2a700fdfd44388557b96a2dc833d Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:33:55 +0800 Subject: [PATCH 07/12] quantize: write the Turbo tier's convrot_w4a4 (INT4) format moss-infer quantize --format int4 produces the file the turbo tier loads: packed 4-bit weights ([N, K/2], even column the low nibble), a per-row f32 scale, 256-column ConvRot rotation groups and the file-level _quantization_metadata the loader reads (convrot_w4a4, convrot_groupsize 256, quant_group_size 64). The scale rule is the format's clipping-ratio search (--recipe mse, the default) or absmax (--recipe rtn); LoRA merging and --keep-bf16 work as for int8. Verified on the RTX 2080: the 5B DiT becomes a 2.83 GB file (55 s with rtn, 283 s with mse); against a published community convrot_w4a4 file the dequantized weights match at cosine 0.9993 and the bf16 tensors are bit-identical, and the produced file runs tier turbo (0.16 s/step at 256^2, 3.9 GiB peak). --- docs/sm75-quantization.md | 24 ++++---- src/tools/cli/cmd_quantize.cpp | 28 ++++++++-- src/weights/quant/quant_convert.cpp | 87 ++++++++++++++++++++++++++++- src/weights/quant/quant_convert.h | 11 ++++ 4 files changed, 131 insertions(+), 19 deletions(-) diff --git a/docs/sm75-quantization.md b/docs/sm75-quantization.md index dbd3b98..e9401ea 100644 --- a/docs/sm75-quantization.md +++ b/docs/sm75-quantization.md @@ -34,16 +34,21 @@ Measured on the RTX 2080 8 GiB with the SFT bundle (`balanced`), this branch, af - Exact attention runs with every tier; the quantized tiers do not need 8-bit attention on sm_75. - `moss-infer quantize` runs on an 8 GiB Turing card: the 5B DiT became a 5.04 GB INT8 ConvRot file in 49 s (252 GEMMs quantized, 0 kept bf16). +- `moss-infer quantize --format int4` (added for P1) writes the Turbo tier's `convrot_w4a4` format: 2.83 GB in 55 s + with `--recipe rtn`, 283 s with the default `--recipe mse` (its clipping-ratio search). Against a published community + `convrot_w4a4` file the quantized weights match at cosine 0.9993 (dequantized; the codes differ only by the 256-column + rotation GEMM's rounding), the bf16 tensors are bit-identical and the file loads and runs `tier turbo` on the RTX + 2080 (0.16 s/step at 256^2, 3.9 GiB peak). - The VAE stays fp32 SIMT either way (see `docs/sm75.md`). ## Gaps -**G1 - no INT4 writer.** `moss-infer quantize --format` accepts only `int8` and `nvfp4` -(src/tools/cli/cmd_quantize.cpp). `convrot_w4a4` files (the only DiT format `turbo` accepts) therefore come from -elsewhere. The pieces to write them exist: `src/weights/quant/convrot_w4a4.{h,cpp}` quantizes W4 with the ConvRot -groups and `quant_group_size`, and the loader already reads the format. The file metadata is a -`_quantization_metadata` JSON holding `format_version` and per-layer -`{"format": "convrot_w4a4", "convrot_groupsize": 256, "quant_group_size": 64}` (checked in a community file). +**G1 - INT4 writer (done, P1).** `moss-infer quantize --format int4` now writes the Turbo tier's `convrot_w4a4` +format: packed W4 (`[N, K/2]`, even column the low nibble), a per-row f32 scale, 256-column rotation groups, and the +file-level `_quantization_metadata` (`{"format_version": "1.0", "layers": {: {"format": "convrot_w4a4", +"convrot_groupsize": 256, "quant_group_size": 64}}}`). The scale rule is comfy-kitchen's clipping-ratio search +(`--recipe mse`, the default) or absmax (`--recipe rtn`); LoRA merging and `--keep-bf16` work as for INT8. K must be a +multiple of 256. **G2 - `high` with the bf16 bundle is refused on sm_75.** The tier is defined as INT8 GEMMs *plus* 8-bit attention; sm_75 has no 8-bit attention, so the resolver falls back to exact attention and then refuses the combination. A @@ -64,11 +69,8 @@ one is loaded. ## Plan -**P1 - INT4 output in `quantize` (small, independent).** Add `--format int4`: quantize each block GEMM with -`weights::quant::convrot_w4a4`, write packed W4 weights plus the per-layer metadata above (group 256, quant group 64 by -default, `--group` as for INT8), keep the non-block tensors bf16. Then `moss-infer generate --dit --tier turbo` -works with a file produced on the same card. Test: the produced file's size (expected ~2.9 GB), `inspect`, a 256^2 and -a 1024^2 generation, and compare its output with the community file's. +**P1 - INT4 output in `quantize` (done).** `--format int4 --recipe mse|rtn` writes a `convrot_w4a4` DiT the same card +can run with `--tier turbo`; the produced file is 2.83 GB and matches a published file's weights at cosine 0.9993. **P2 - Turing 8-bit attention (`q8_fp16`).** A new kernel in `src/kernels/sm75/`: INT8 QK with the Turing m8n8k16 MMA over the ConvRot-quantized q/k (the quantizers are already portable: `quant/rotate_quantize.cu`), fp16 P·V with the diff --git a/src/tools/cli/cmd_quantize.cpp b/src/tools/cli/cmd_quantize.cpp index d77b9ad..7766821 100644 --- a/src/tools/cli/cmd_quantize.cpp +++ b/src/tools/cli/cmd_quantize.cpp @@ -27,16 +27,18 @@ constexpr const char* kUsage = "usage: moss-infer quantize -o \n" " [--fuse-qkv] [--keep-bf16 ] [--lora [--lora-scale ]]... [--group ]\n" " [--comfy-attention-config] [--threads ]\n" - " [--format int8|nvfp4] [--recipe rtn|mse|gptq|gptq-mse] [--keep ] [--calib-model ]\n" + " [--format int8|int4|nvfp4] [--recipe rtn|mse|gptq|gptq-mse] [--keep ] [--calib-model ]\n" " --keep-bf16 comma-separated globs over GEMM names, e.g. 'layers.*.proj.ffn_down,layers.0.*'\n" " --lora kohya (diffusion_model..lora_{down,up}.weight + alpha), PEFT (lora_A/lora_B) or moss-edit-lora\n" " (bare .lora_A / .lora_B, alphas in adapter_config) LoRA,\n" " merged in fp32 before quantizing; --lora-scale applies to the preceding --lora (default 1)\n" " --group ConvRot group of newly quantized GEMMs: 4, 16, 64, 256 (default) or 1024\n" " --threads CPU threads, 0 = all cores (default)\n" - " --format int8 (default): INT8 ConvRot; nvfp4: NVFP4 block GEMMs (RTX 50), the --keep units INT8 ConvRot,\n" + " --format int8 (default): INT8 ConvRot; int4: ConvRot W4A4 (the Turbo tier's convrot_w4a4 file,\n" + " 4-bit weights + per-row scale, rotated per 256 columns); nvfp4: NVFP4 block GEMMs (RTX 50),\n" " q/k/v fused, from bf16 weights\n" - " --recipe nvfp4 weights: gptq-mse (default), gptq, mse or rtn (bit-exact to comfy-kitchen); gptq*\n" + " --recipe int4 weights: mse (default, the format's clipping-ratio search) or rtn (absmax, faster);\n" + " nvfp4 weights: gptq-mse (default), gptq, mse or rtn (bit-exact to comfy-kitchen); gptq*\n" " capture calibration Hessians with the engine on the GPU (3 built-in prompts)\n" " --keep units kept INT8: g0.9 (default, 42 units), g0.95 (21), none, or e.g. 'b37.down,b26.o'\n" " --calib-model the bundle (tokenizer, text encoder) gptq* calibrate with; default: the input's bundle when the\n" @@ -97,7 +99,7 @@ int cmd_quantize(const std::vector& args) { const bool has_value = i + 1 < args.size(); if (a == "--format" && has_value) { format = args[++i]; - if (format != "int8" && format != "nvfp4") return bad_value(a, format, "int8 or nvfp4"); + if (format != "int8" && format != "int4" && format != "nvfp4") return bad_value(a, format, "int8, int4 or nvfp4"); } else if (a == "--recipe" && has_value) { recipe_arg = args[++i]; if (!quant::parse_nvfp4_recipe(recipe_arg)) return bad_value(a, recipe_arg, "rtn, mse, gptq or gptq-mse"); @@ -143,10 +145,24 @@ int cmd_quantize(const std::vector& args) { std::cerr << kUsage; return 2; } - if (format != "nvfp4" && (!recipe_arg.empty() || !keep_arg.empty() || !calib_model.empty() || !hessians_out.empty() || !hessians_in.empty())) { - std::cerr << "moss-infer quantize: --recipe, --keep, --calib-model and --hessians[-out] need --format nvfp4\n" << kUsage; + if (format == "int8" && (!recipe_arg.empty() || !keep_arg.empty() || !calib_model.empty() || !hessians_out.empty() || !hessians_in.empty())) { + std::cerr << "moss-infer quantize: --recipe, --keep, --calib-model and --hessians[-out] need --format nvfp4 or int4\n" << kUsage; return 2; } + if (format == "int4") { + options.int4 = true; + if (!recipe_arg.empty()) { + if (recipe_arg != "rtn" && recipe_arg != "mse") { + std::cerr << "moss-infer quantize: --format int4 takes --recipe rtn or mse\n" << kUsage; + return 2; + } + options.w4_search = recipe_arg == "mse" ? quant::W4ScaleSearch::kMse : quant::W4ScaleSearch::kAbsmax; + } + if (keep_arg.empty() == false || !calib_model.empty() || !hessians_out.empty() || !hessians_in.empty()) { + std::cerr << "moss-infer quantize: --keep, --calib-model and --hessians[-out] need --format nvfp4\n" << kUsage; + return 2; + } + } if (format == "nvfp4") { quant::Nvfp4Options fp4; if (!recipe_arg.empty()) fp4.recipe = *quant::parse_nvfp4_recipe(recipe_arg); diff --git a/src/weights/quant/quant_convert.cpp b/src/weights/quant/quant_convert.cpp index 1444bcb..3d5bccb 100644 --- a/src/weights/quant/quant_convert.cpp +++ b/src/weights/quant/quant_convert.cpp @@ -330,6 +330,7 @@ struct OutModule { std::vector parts; // 1, or q/k/v QuantSpec spec; std::int64_t n = 0, k = 0; + bool int4 = false; // convrot_w4a4 output: packed [n, k/2] + per-row scale, described by _quantization_metadata // Indices into the tensor list (scale/scale_2/config/bias: SIZE_MAX if absent; scale_2: nvfp4 only). std::size_t weight = 0, scale = 0, scale_2 = 0, config = 0, bias = 0; }; @@ -637,6 +638,7 @@ std::vector plan_modules(const QuantModel& model, const ConvertOption if (keep) fail("--keep-bf16 matches '" + l.name + "', which the input stores quantized (the bf16 weights are gone)"); if (loras.count(l.name)) fail("a LoRA targets '" + l.name + "', which the input stores quantized; merge into the bf16 weights instead"); if (fp4 != nullptr) fail("'" + l.name + "' is stored " + describe(l.spec) + "; NVFP4 output is converted from the bf16 weights"); + if (options.int4) fail("'" + l.name + "' is stored " + describe(l.spec) + "; INT4 output is converted from the bf16 weights"); return l.spec; } if (fp4 != nullptr) { @@ -684,6 +686,7 @@ std::vector plan_modules(const QuantModel& model, const ConvertOption om.n = l.n; } om.k = om.parts[0]->k; + om.int4 = options.int4 && om.spec.quantized(); mods.push_back(std::move(om)); } if (fp4 != nullptr) { @@ -794,6 +797,11 @@ ConvertStats convert(const fs::path& input, const fs::path& output, const Conver util::check_user_path(output, "output"); for (const LoraInput& li : options.loras) util::check_user_path(li.file, "LoRA"); if (!valid_convrot_group(options.group)) fail("--group must be a power of 4 in [4, 1024]"); + if (options.int4) { + if (options.nvfp4) fail("--format int4 and nvfp4 are mutually exclusive"); + if (options.policy) fail("a spec policy cannot be combined with int4 output"); + if (options.group != 256) fail("--format int4 rotates in 256-column groups; drop --group"); + } if (options.threads < 0) fail("--threads must be >= 0"); for (const LoraInput& li : options.loras) { if (!std::isfinite(li.strength)) fail("LoRA strength of " + util::path_utf8(li.file) + " is not finite"); @@ -853,7 +861,11 @@ ConvertStats convert(const fs::path& input, const fs::path& output, const Conver for (OutModule& om : mods) { om.scale = om.scale_2 = om.config = om.bias = kNone; spec_counts[describe(om.spec)] += static_cast(om.parts.size()); - if (om.spec.format == QuantFormat::kNvfp4) { + if (om.int4) { + if (om.k % 256 != 0) fail("'" + om.name + "': K=" + std::to_string(om.k) + " is not a multiple of 256, which convrot_w4a4 needs"); + om.weight = add(om.name + ".weight", DType::kI8, {om.n, om.k / 2}); + om.scale = add(om.name + ".weight_scale", DType::kF32, {om.n}); + } else if (om.spec.format == QuantFormat::kNvfp4) { if (om.k % kNvfp4Block != 0) fail("'" + om.name + "': K=" + std::to_string(om.k) + " is not a multiple of 16, which nvfp4 needs"); om.weight = add(om.name + ".weight", DType::kU8, {om.n, om.k / 2}); om.scale = add(om.name + ".weight_scale", DType::kF8E4M3, {nvfp4_scale_rows(om.n), nvfp4_scale_cols(om.k / kNvfp4Block)}); @@ -1014,7 +1026,21 @@ ConvertStats convert(const fs::path& input, const fs::path& output, const Conver kept_json += "]"; std::string specs; for (const auto& [name, count] : spec_counts) specs += (specs.empty() ? "" : ", ") + name + " x" + std::to_string(count); - if (fp4 == nullptr) { + if (fp4 == nullptr && options.int4) { + // convrot_w4a4 keeps its format in a file-level header (the community layout), not in per-tensor comfy_quant. + std::string layers_json = "{"; + bool first_layer = true; + for (const OutModule& om : mods) { + if (!om.int4) continue; + layers_json += (first_layer ? "" : ", ") + quote(om.name) + + ": {\"format\": \"convrot_w4a4\", \"convrot_groupsize\": 256, \"quant_group_size\": 64}"; + first_layer = false; + } + layers_json += "}"; + md["_quantization_metadata"] = "{\"format_version\": \"1.0\", \"layers\": " + layers_json + "}"; + md["mossinfer_converter"] = "moss-infer quantize " + std::to_string(kConverterVersion) + " (" + specs + + "; convrot_w4a4 4-bit weights, per-row f32 scale)"; + } else if (fp4 == nullptr) { md["mossinfer_converter"] = "moss-infer quantize " + std::to_string(kConverterVersion) + " (" + specs + (gguf ? "; GGUF input dequantized to fp32, then quantized round-to-nearest or rounded to bf16)" : "; bf16 inputs quantized round-to-nearest, quantized inputs copied or decoded exactly)"); @@ -1155,6 +1181,63 @@ ConvertStats convert(const fs::path& input, const fs::path& output, const Conver const std::int64_t k = om.k; if (written[mi]) continue; unit_start(); + if (om.int4) { + // convrot_w4a4: bf16 rows (LoRA merged in fp32), then the offline W4 quantizer rotates, scales and packs them. + std::vector w16(static_cast(om.n * k)); + std::vector bias_bytes; + std::int64_t row0 = 0; + for (const QuantLinear* p : om.parts) { + const auto lf = loras.find(p->name); + const std::vector* lfs = lf == loras.end() ? nullptr : &lf->second; + if (lfs) { + lora_used.insert(p->name); + ++stats.lora_modules; + } + const std::vector src = readers.read(p->weight); + note_source(p->name + ".weight", src); + const auto* w16src = reinterpret_cast(src.data()); + const std::uint32_t gtype = p->encoding == Encoding::kGguf ? model.gguf_types.at(p->weight.name) : 0; + const std::uint64_t row_bytes = p->weight.nbytes / static_cast(p->n); + std::atomic bad{false}; + parallel_for(p->n, threads, [&](std::int64_t b, std::int64_t e) { + std::vector rows(static_cast(kBatch * k)); + std::vector acc(lfs ? static_cast(k) : 0); + for (std::int64_t r0 = b; r0 < e; r0 += kBatch) { + const std::int64_t nb = std::min(kBatch, e - r0); + if (p->encoding == Encoding::kGguf) { + weights::dequantize_gguf_rows(gtype, src.data() + static_cast(r0) * row_bytes, nb, k, rows.data()); + } else { + for (std::int64_t i = 0; i < nb * k; ++i) rows[static_cast(i)] = bf16_to_float(w16src[r0 * k + i]); + } + if (lfs) { + for (std::int64_t r = 0; r < nb; ++r) { + for (const LoraFactor& f : *lfs) lora_merge_row(rows.data() + r * k, k, f.up.data() + (r0 + r) * f.rank, f.down.data(), f.rank, f.coef, acc.data()); + } + } + if (!to_bf16_finite(rows.data(), static_cast(nb * k), w16.data() + (row0 + r0) * k)) bad = true; + } + }); + if (bad) fail("'" + p->name + "' holds non-finite weights" + (lfs ? " after the LoRA merge" : "")); + if (p->bias) { + const std::vector b = readers.read(*p->bias); + note_source(p->name + ".bias", b); + bias_bytes.insert(bias_bytes.end(), b.begin(), b.end()); + } + row0 += p->n; + } + ConvRotW4A4Options w4; + w4.search = options.w4_search; + w4.threads = threads; + const ConvRotW4A4Weight out = quantize_convrot_w4a4(w16.data(), om.n, k, w4); + writer.write(tensors[om.weight], out.q.data(), out.q.size()); + writer.write(tensors[om.scale], out.scale.data(), out.scale.size() * sizeof(float)); + if (om.bias != kNone) writer.write(tensors[om.bias], bias_bytes.data(), bias_bytes.size()); + stats.quantized += static_cast(om.parts.size()); + written[mi] = 1; + if (++done % 24 == 0) log(" " + std::to_string(done) + "/" + std::to_string(mods.size()) + " GEMMs"); + unit_done(); + continue; + } if (om.spec.format == QuantFormat::kNvfp4) { write_nvfp4(mi, quantize_nvfp4_module(om, module_f32(om), nullptr)); continue; diff --git a/src/weights/quant/quant_convert.h b/src/weights/quant/quant_convert.h index 0f72e0e..f5b0adc 100644 --- a/src/weights/quant/quant_convert.h +++ b/src/weights/quant/quant_convert.h @@ -47,6 +47,7 @@ #include "util/regular_file.h" #include "weights/quant/nvfp4_recipes.h" +#include "weights/quant/convrot_w4a4.h" #include "weights/quant/quant_model.h" // mi::quant rather than mi::weights: the format is shared by the reader, the converter and the GEMM kernels. @@ -120,6 +121,16 @@ struct LoraInput { struct ConvertOptions { int group = 256; // ConvRot group of newly quantized GEMMs + // convrot_w4a4 output (`--format int4`): block GEMMs become packed 4-bit weights ([N, K/2], even column the low + // nibble, values -7..7) with a per-row f32 scale, rotated per 256 columns (the Turbo tier's file format; its + // activation quant group is 64). Described by a file-level `_quantization_metadata`; no per-tensor comfy_quant is + // written. `group` is fixed at 256 (the rotation width); keep_bf16, LoRA merging and the other options work as for + // int8. Needs K divisible by 256. + bool int4 = false; + // The clipping-ratio search of the format's reference implementation (comfy-kitchen): a published file's weights + // match this at cosine 0.9993 (the codes differ only by the rotation GEMM's rounding). --recipe rtn skips the search + // (absmax), which is faster to produce and slightly worse. + W4ScaleSearch w4_search = W4ScaleSearch::kMse; // --recipe mse (default) or rtn bool fuse_qkv = false; // write layers.N.proj.qkv instead of proj.q/k/v // Also write layers.N.comfy_attention.config (a ComfyUI runtime preference, not part of the weight format) so the // output is byte-identical to Comfy-Org's files. Off by default; files load in ComfyUI either way. From 49a6723c6f1c3fe2718f39c2efdbb080a07a5e0a Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:35:44 +0800 Subject: [PATCH 08/12] sm75: LoRA on the quantized tiers (high merges, turbo refuses) Verified on the RTX 2080: reference + --lora merges into bf16; high + --lora merges into the INT8 ConvRot weights (merge_requantize_int8) and reports the requantization warning; turbo + --lora used to fail with the generic merge error, whose advice (lora_placement side_branch) does not exist on Turing. The pipeline now refuses that combination up front on sm_7x and names what works (tier high or reference). --- docs/sm75-quantization.md | 14 +++++++------- src/pipeline/pipeline.cpp | 8 ++++++++ 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/docs/sm75-quantization.md b/docs/sm75-quantization.md index e9401ea..b7c8f0a 100644 --- a/docs/sm75-quantization.md +++ b/docs/sm75-quantization.md @@ -58,10 +58,11 @@ Turing user has to quantize the DiT first (or download an INT8 file) to use `hig is the mainline's long-sequence accelerator: at 1024^2 the DiT's exact attention is 1.05 s per step (7.3 TFLOPS with the new kernel), which an INT8 QK mainloop should roughly halve. -**G4 - LoRA on quantized DiTs is unverified on sm_75.** Merged adapters on INT8 block GEMMs go through -`merge_requantize_int8` (src/models/lora/lora_merge.cpp); on sm_75 every LoRA is merged (no side branches), and an -INT4 DiT has no merge path at all, so `--tier turbo --lora` should refuse clearly instead of silently ignoring the -adapter. Needs tests and, if missing, a message. +**G4 - LoRA on quantized DiTs (done, P4).** On sm_75 every adapter is merged (the side branches are sm_80+ kernels). +Verified on the RTX 2080: `reference + --lora` merges into bf16; `high + --lora` merges into the INT8 ConvRot weights +through `merge_requantize_int8` and reports "merged into INT8 (requantization error)" per file; `turbo + --lora` is +refused up front with a message naming the reason (INT4 ConvRot has no merge path; Turing has no side branches) and +suggesting `tier high` or `tier reference`. **G5 - the app has to offer it.** MossStudio decides what to download and which tiers to show; a Turing card should be able to run `quantize` (engine CLI) or fetch INT8/INT4 files, and the tier selector should enable `high`/`turbo` when @@ -83,9 +84,8 @@ attention 14.3 ms -> 8-9 ms at 1024^2 (0.5 s per step saved) and `high` works wi sm_75-specific definition of `high`. Small (planner + load options), but it changes what `high` means on Turing, so P2 is the cleaner long-term answer. -**P4 - LoRA on quantized tiers.** Test `--tier high --lora` on sm_75 (merge + requantize), and make `--tier turbo ---lora` refuse with a message that names the missing merge path. If the requantize path turns out to need sm_80 -kernels, port the host merge (`quant::quantize_convrot_rows` is already used by the sm_80 path). +**P4 - LoRA on quantized tiers (done).** `high + --lora` works on sm_75 (host requantize merge, warning as on sm_80); +`turbo + --lora` refuses with the reason and the working alternatives. **P5 - docs and app.** Update `docs/sm75.md`'s limits with the outcome, add the INT4 production step and the measured tier table here, and tell MossStudio which files/tiers a cc 7.5 device can use. diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 6fa4ea4..677590f 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -943,6 +943,14 @@ Pipeline::Impl::CallModels Pipeline::Impl::setup_models(const Request& req, Resu // failing. Keyed on the device, not on the kernels' availability: an sm_80+ build without them (MOSSINFER_ONLY) // still refuses, as before. const bool exact_when_no_q8 = kernels::sm75::turing(); + // sm_75 has no LoRA side-branch kernels, so run() forces every adapter to merge (Pipeline::run, above); INT4 + // ConvRot weights have no merge path (merges need bf16 or INT8 ConvRot), and the tier's own advice ("use + // lora_placement side_branch") does not exist on Turing: refuse here with what does work. + if (lora_requested && exact_when_no_q8 && src.precision == planner::Precision::kInt4) { + throw std::invalid_argument( + "runtime LoRAs cannot run on the turbo tier on sm_7x: its INT4 ConvRot weights have no merge path and Turing " + "has no LoRA side-branch kernels; use tier high (INT8 ConvRot) or tier reference"); + } r.tier = planner::resolve_tier(m.options.tier, m.options.attention, src.precision, supported, tes.precision != qwen35_text::Precision::kBF16, src.precision == planner::Precision::kNvfp4 && kernels::nvfp4::available(), From 3f6bd6942b5f0c9248c83c5b9a9d003f600bedde Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:44:31 +0800 Subject: [PATCH 09/12] docs: P2 design notes for the Turing INT8-QK attention Records the dispatch points (attention_q8.cu's three functions), the workspace layout, the m8n8k16 fragment mapping from CUTLASS's sm75 header and the fact that the exact kernel's skeleton carries over, so the kernel work can start from the document. --- docs/sm75-quantization.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/sm75-quantization.md b/docs/sm75-quantization.md index b7c8f0a..ea6258e 100644 --- a/docs/sm75-quantization.md +++ b/docs/sm75-quantization.md @@ -79,6 +79,23 @@ existing m16n8k8 MMA, exact softmax in fp32. Wire it into `attention_supported`, `resolve_tier` so `--attention q8_fp16` is accepted on cc 7.5 and `high` with a bf16 DiT resolves to it. Expected: attention 14.3 ms -> 8-9 ms at 1024^2 (0.5 s per step saved) and `high` works without a quantized file. +**P2 design notes (from the code, for the implementation).** The whole 8-bit path funnels through three functions in +`src/kernels/attention_q8/attention_q8.cu` (`supported`, `attention_q8_workspace_bytes`, `attention_q8`), so a Turing +implementation only has to answer those three for `kmInt8Fp16`/`kInt8Fp16F32Acc`; `moss_dit` and `resolve_tier` need no +change beyond `supported()` turning true (tier `high` then picks `q8_fp16` on a bf16 DiT, because `kTierAttention` is +`{q8_fp8, q8_fp16}` and picks the first supported one). Workspace on sm_75 should use its own layout: +`q_i8 [total_q][heads][128]`, `k_i8 [total_k][kv][128]`, `v_f16`, `q_scale [total_q][heads]`, `k_scale [total_k][kv]`. +The kernel can reuse the exact kernel's skeleton (`attention_sm75.cu`: 64-query CTA, 4 warps, 16 rows per warp, K/V +tiles in shared, online softmax): only the score mainloop changes. Turing's INT8 MMA is `m8n8k16` +(`mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.s8.s32`, see `third_party/cutlass/include/cutlass/arch/mma_sm75.h`): +A = 4 consecutive int8 per lane (row = lane/4, dims = 16*(chunk) + 4*(lane%4)), B = 4 consecutive int8 per lane +(key = lane/4, same dims), C = 2 int32 per lane (row = lane/4, keys = 8*subtile + 2*(lane%4), +1). With two row groups +per warp (rows r and r+8) this maps exactly onto the existing `s[t][0..3]` register block, so only the fragment loads, +the mma and the dequantization (`s = acc * q_scale[row] * k_scale[key] * scale`) are new. `q` is quantized per row and +`k` per row (SageAttention smooths K by its per-sequence mean; optional here); `v` becomes fp16 and the P·V path stays +`m16n8k8`. The fused QK-norm/RoPE pre-pass of `attention_q8` is not worth porting: `dit_block_quant.cpp` already has a +separate-kernel path (`fused_norm_rope` false, same bits), so the sm_75 q8 branch needs that condition forced off for +Turing. **P3 - or, if P2 waits: a sm_75 tier rule.** Let `high` on cc 7.5 quantize the block GEMMs at load (`LoadOptions::int8_block_gemms` currently only follows an INT8 file) and keep exact attention, documented as a sm_75-specific definition of `high`. Small (planner + load options), but it changes what `high` means on Turing, so From db8408212ccd9eb88efb5d3e299a48fd1c42e1ce Mon Sep 17 00:00:00 2001 From: CloverIris Date: Sun, 8 Nov 2026 13:50:06 +0800 Subject: [PATCH 10/12] sm75: run tier high on a bf16 DiT with a Turing q8_fp16 attention kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New src/kernels/sm75/attention_q8_sm75.cu: per-row INT8 q/k (a quantize pre-pass, v converted to fp16), QK^T with Turing's m8n8k16 INT8 MMA, then the exact kernel's fp32 online softmax and fp16 P·V. attention_q8.cu's three entry points (supported, attention_q8_workspace_bytes, attention_q8) dispatch to it on cc 7.5, and dit_block_quant keeps the separate qk_norm_rope_moss on Turing (the fused pre-pass is sm_80-only). Verified on the RTX 2080: tier high now runs on the bf16 bundle ('DiT bf16, q8_fp16 attention'), the 256^2 output matches the exact path at cosine 0.9986, and 1024^2 works. No speedup: Turing's INT8 MMA m8n8k16 does the same MACs per instruction as the fp16 m16n8k8 it replaces (m16n8k32 is the 2x part on Ampere), so 8-bit attention is a bandwidth path here; documented in docs/sm75-quantization.md and docs/sm75.md. --- docs/sm75-quantization.md | 22 +- docs/sm75.md | 6 +- src/kernels/attention_q8/attention_q8.cu | 17 ++ src/kernels/sm75/attention_q8_sm75.cu | 348 +++++++++++++++++++++++ src/kernels/sm75/sm75.h | 7 + src/models/moss_dit/dit_block_quant.cpp | 5 +- 6 files changed, 393 insertions(+), 12 deletions(-) create mode 100644 src/kernels/sm75/attention_q8_sm75.cu diff --git a/docs/sm75-quantization.md b/docs/sm75-quantization.md index ea6258e..7b7b8dd 100644 --- a/docs/sm75-quantization.md +++ b/docs/sm75-quantization.md @@ -54,9 +54,15 @@ multiple of 256. sm_75 has no 8-bit attention, so the resolver falls back to exact attention and then refuses the combination. A Turing user has to quantize the DiT first (or download an INT8 file) to use `high`. -**G3 - no 8-bit attention kernel.** `q8_fp16`/`q8_fp8` are refused on sm_75 `--attention`. Besides blocking G2, this -is the mainline's long-sequence accelerator: at 1024^2 the DiT's exact attention is 1.05 s per step -(7.3 TFLOPS with the new kernel), which an INT8 QK mainloop should roughly halve. +**G3 - 8-bit attention (done, P2).** `src/kernels/sm75/attention_q8_sm75.cu` implements `q8_fp16` on Turing: +per-row INT8 q/k (a `q8_quantize_rows` pre-pass, v converted to fp16), QK^T with the m8n8k16 INT8 MMA, the exact +kernel's fp32 online softmax and fp16 P·V. It is wired through `attention_q8.cu`'s three entry points, so +`tier high` now runs on a **bf16 DiT** (`tier high: DiT bf16, q8_fp16 attention`) — verified on the RTX 2080 at +256^2 (output cosine 0.9986 against the exact path) and 1024^2. Speed: at 256^2 and 1024^2 it is within noise of the +exact path (0.96 vs 0.94 s/step; 4.39 vs 4.30 s/step), because Turing's INT8 MMA (`m8n8k16`) does the same MACs per +instruction as the fp16 `m16n8k8` it replaces — Ampere's `m16n8k32` is what doubles INT8 throughput, and Turing does not +have it. The 8-bit mode is therefore a memory-footprint/bandwidth path on Turing, not a compute accelerator; +`q8_fp8` stays refused (no e4m3 MMA). **G4 - LoRA on quantized DiTs (done, P4).** On sm_75 every adapter is merged (the side branches are sm_80+ kernels). Verified on the RTX 2080: `reference + --lora` merges into bf16; `high + --lora` merges into the INT8 ConvRot weights @@ -73,11 +79,11 @@ one is loaded. **P1 - INT4 output in `quantize` (done).** `--format int4 --recipe mse|rtn` writes a `convrot_w4a4` DiT the same card can run with `--tier turbo`; the produced file is 2.83 GB and matches a published file's weights at cosine 0.9993. -**P2 - Turing 8-bit attention (`q8_fp16`).** A new kernel in `src/kernels/sm75/`: INT8 QK with the Turing m8n8k16 MMA -over the ConvRot-quantized q/k (the quantizers are already portable: `quant/rotate_quantize.cu`), fp16 P·V with the -existing m16n8k8 MMA, exact softmax in fp32. Wire it into `attention_supported`, `MossDit`'s attention dispatch and -`resolve_tier` so `--attention q8_fp16` is accepted on cc 7.5 and `high` with a bf16 DiT resolves to it. Expected: -attention 14.3 ms -> 8-9 ms at 1024^2 (0.5 s per step saved) and `high` works without a quantized file. +**P2 - Turing 8-bit attention (`q8_fp16`) (done).** `attention_q8_sm75.cu` (per-row INT8 q/k, m8n8k16 QK^T, the +exact kernel's softmax and fp16 P·V) is wired into `attention_q8.cu`; `tier high` runs on a bf16 DiT and matches the +exact path at cosine 0.9986. The expected 1.5-2x attention speedup did not materialize: Turing's INT8 MMA is +`m8n8k16`, the same MACs per instruction as the fp16 `m16n8k8` it replaces (Ampere's `m16n8k32` is the 2x part), so +8-bit attention is a bandwidth/footprint path on this architecture, not a compute one. **P2 design notes (from the code, for the implementation).** The whole 8-bit path funnels through three functions in `src/kernels/attention_q8/attention_q8.cu` (`supported`, `attention_q8_workspace_bytes`, `attention_q8`), so a Turing diff --git a/docs/sm75.md b/docs/sm75.md index 0eff99e..9432e3e 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -116,9 +116,9 @@ and float4 accumulators (four input channels per FMA) with vectorised float4 out ## Limits on sm_75 -- The 8-bit attention kernels (`q8_fp16`, `q8_fp8`): the tier default resolves to `exact`; passing an 8-bit mode - explicitly is refused. `docs/sm75-quantization.md` (P2) plans the Turing `q8_fp16` kernel, which would also let - `tier high` run on a bf16 DiT. +- 8-bit attention: `q8_fp16` runs (a Turing INT8-QK kernel, `tier high` accepts a bf16 DiT with it); `q8_fp8` is + refused (no e4m3 MMA on Turing). On Turing the INT8 MMA does the same MACs per instruction as the fp16 one, so + `q8_fp16` is a bandwidth/footprint path rather than a speedup (see `docs/sm75-quantization.md`). - Native INT4 MMA and NVFP4: INT4 files run through the INT8 unpack fallback (correct and bit-exact; the planner prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. - The fused DiT blocks and their CUDA graphs (the blocks run op by op). diff --git a/src/kernels/attention_q8/attention_q8.cu b/src/kernels/attention_q8/attention_q8.cu index b8a6013..45c17bb 100644 --- a/src/kernels/attention_q8/attention_q8.cu +++ b/src/kernels/attention_q8/attention_q8.cu @@ -1,4 +1,5 @@ #include "kernels/attention_q8/attention_q8.h" +#include "kernels/sm75/sm75.h" #include #include @@ -166,6 +167,8 @@ double e4m3_rn(double x) { } // namespace bool supported(Variant v) { + // sm_75 has its own INT8-QK kernel (kernels/sm75/attention_q8_sm75.cu) for the fp16-accumulation variants. + if (kernels::sm75::turing()) return v == Variant::kInt8Fp16 || v == Variant::kInt8Fp16F32Acc; int device = 0; MI_CUDA_CHECK(cudaGetDevice(&device)); return support(device, v) == Support::kOk; @@ -195,12 +198,26 @@ float effective_v_range(const Options& o) { std::size_t attention_q8_workspace_bytes(const Options& o, std::int64_t total_q, std::int64_t total_k, int batch, int max_seqlen, int num_heads, int num_kv_heads) { + if (kernels::sm75::turing()) return kernels::sm75::q8_workspace_bytes(total_q, total_k, num_heads, num_kv_heads); return layout(o.variant, total_q, total_k, batch, max_seqlen, num_heads, num_kv_heads).total; } void attention_q8(const Options& o, 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, float scale, void* workspace, std::size_t workspace_bytes, cudaStream_t stream, const QkNormRope* norm_rope) { + if (kernels::sm75::turing()) { + if (o.variant != Variant::kInt8Fp16 && o.variant != Variant::kInt8Fp16F32Acc) { + throw std::runtime_error("attention_q8: the sm_75 build runs the fp16-accumulation variants only"); + } + if (norm_rope != nullptr) { + throw std::invalid_argument( + "attention_q8: the sm_75 build has no fused norm/rope pre-pass; run qk_norm_rope_moss first"); + } + check_qkvo(q, "q"); + kernels::sm75::q8_flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, scale, workspace, workspace_bytes, + stream); + return; + } check_qkvo(q, "q"); check_qkvo(k, "k"); check_qkvo(v, "v"); diff --git a/src/kernels/sm75/attention_q8_sm75.cu b/src/kernels/sm75/attention_q8_sm75.cu new file mode 100644 index 0000000..95d533a --- /dev/null +++ b/src/kernels/sm75/attention_q8_sm75.cu @@ -0,0 +1,348 @@ +// 8-bit attention on sm_75 (`q8_fp16` / `q8_fp16_f32acc`): INT8 QK^T with Turing's m8n8k16 MMA and fp16 P·V with the +// existing m16n8k8 path. q and k are quantized per row (amax / 127) in a pre-pass; k is not mean-smoothed (the engine's +// q8 is lossy by definition, and the smoothing only buys INT8 range). v becomes fp16. The tiling, online softmax and +// P·V reuse the exact kernel's (attention_sm75.cu): 64 queries per CTA, 4 warps, 16 rows per warp, 32 keys per tile. +#include +#include +#include + +#include +#include +#include +#include + +#include "core/cuda_check.h" +#include "kernels/sm75/sm75.h" + +namespace mi::kernels::sm75 { +namespace { + +constexpr int kQ8D = 128; // q8 attention is head_dim 128 only (the DiT's joint attention) +constexpr int kQ8BN = 32; +constexpr int kQ8Warps = 4; +constexpr int kQ8BM = kQ8Warps * 16; +constexpr int kQ8Stride = kQ8D + 16; // int8 K tile row stride (4-byte aligned, pads the bank pattern) +constexpr int kQ8Chunks = kQ8D / 16; // m8n8k16 in-dimension chunks (QK^T) +constexpr int kQ8PvChunks = kQ8D / 8; // m16n8k8 in-dimension chunks (P·V) +constexpr float kLog2e = 1.4426950408889634f; + +// m8n8k16 int8 -> int32 MMA (Turing). a = 4 int8 of one query row, b = 4 int8 of one key, d = 2 int32. +__device__ __forceinline__ void mma_m8n8k16(int& d0, int& d1, std::uint32_t a, std::uint32_t b) { + asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%0,%1};\n" + : "+r"(d0), "+r"(d1) + : "r"(a), "r"(b)); +} + +__device__ __forceinline__ void mma_m16n8k8(float (&d)[4], std::uint32_t a0, std::uint32_t a1, std::uint32_t b) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a0), "r"(a1), "r"(b)); +} + +__device__ __forceinline__ std::uint32_t pack_half2(float lo, float hi) { + union { + __half2 h; + std::uint32_t u; + } c; + c.h = __halves2half2(__float2half_rn(lo), __float2half_rn(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__ std::uint32_t load_u32(const void* p) { + return *reinterpret_cast(p); +} + +__device__ __forceinline__ float block_amax(float x) { + __shared__ float sm[kQ8Warps * 8]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) x = fmaxf(x, __shfl_xor_sync(0xffffffffu, x, off)); + if (lane == 0) sm[warp] = x; + __syncthreads(); + float m = 0.f; + for (int w = 0; w < 4; ++w) m = fmaxf(m, sm[w]); + __syncthreads(); + return m; +} + +// One CTA per (token, head): scale = amax / 127, then q = round(x / scale) clamped to [-127, 127]. +__global__ void q8_quantize_rows(const __nv_bfloat16* __restrict__ x, std::int8_t* __restrict__ y, + float* __restrict__ scale, int rows, int heads, std::int64_t row_stride, + std::int64_t head_stride) { + const int idx = blockIdx.x; + const int token = idx / heads, head = idx % heads; + const __nv_bfloat16* src = x + static_cast(token) * row_stride + static_cast(head) * head_stride; + float amax = 0.f; + for (int i = threadIdx.x; i < kQ8D; i += blockDim.x) amax = fmaxf(amax, fabsf(__bfloat162float(src[i]))); + amax = block_amax(amax); + const float s = amax > 0.f ? amax / 127.f : 1.f; + if (threadIdx.x == 0) scale[idx] = s; + const float inv = 1.f / s; + std::int8_t* dst = y + (static_cast(token) * heads + head) * kQ8D; + for (int i = threadIdx.x; i < kQ8D; i += blockDim.x) { + int v = static_cast(__float2int_rn(__bfloat162float(src[i]) * inv)); + v = v < -127 ? -127 : (v > 127 ? 127 : v); + dst[i] = static_cast(v); + } + (void)rows; +} + +// One CTA per (token, head): v bf16 -> fp16. +__global__ void q8_convert_v(const __nv_bfloat16* __restrict__ v, __half* __restrict__ out, std::int64_t row_stride, + std::int64_t head_stride, int heads) { + const int idx = blockIdx.x; + const int token = idx / heads, head = idx % heads; + const __nv_bfloat16* src = v + static_cast(token) * row_stride + static_cast(head) * head_stride; + __half* dst = out + (static_cast(token) * heads + head) * kQ8D; + for (int i = threadIdx.x; i < kQ8D; i += blockDim.x) dst[i] = __float2half_rn(__bfloat162float(src[i])); +} + +struct Q8Params { + const std::int8_t* q = nullptr; + const std::int8_t* k = nullptr; + const __half* v = nullptr; + const float* q_scale = nullptr; + const float* k_scale = nullptr; + __nv_bfloat16* o = nullptr; + const int* cu_q = nullptr; + const int* cu_k = nullptr; + int heads = 0, kv_heads = 0, total_q = 0; + std::int64_t q_row = 0, q_head = 0, k_row = 0, k_head = 0, v_row = 0, v_head = 0, o_row = 0, o_head = 0; + float scale = 1.f; +}; + +__global__ void attention_q8_kernel(const Q8Params p) { + __shared__ __align__(16) std::int8_t ks[kQ8BN * kQ8Stride]; + __shared__ __align__(16) __half vt[kQ8D * (kQ8BN + 8)]; + __shared__ float kscale[kQ8BN]; + const int batch = blockIdx.y, head = blockIdx.z; + const int q_start = p.cu_q[batch], q_end = p.cu_q[batch + 1]; + const int k_start = p.cu_k[batch], k_end = p.cu_k[batch + 1]; + const int len_q = q_end - q_start, len_k = k_end - k_start; + const int row_base = static_cast(blockIdx.x) * kQ8BM; + if (row_base >= len_q) return; + const int rows = min(kQ8BM, len_q - row_base); + const int kv_head = head / (p.heads / p.kv_heads); + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; + const int group = lane >> 2, tig = lane & 3; + const int r0 = warp * 16 + group, r1 = r0 + 8; + + // Q fragments: 4 int8 of one query row per chunk; two row groups per lane-pair. + std::uint32_t qf[2][kQ8Chunks]; + const std::int64_t q_head_off = static_cast(head) * p.q_head; + const auto load_q = [&](int row, int chunk) -> std::uint32_t { + if (row >= rows) return 0u; + return load_u32(p.q + static_cast(q_start + row_base + row) * p.q_row + q_head_off + chunk * 16 + tig * 4); + }; +#pragma unroll + for (int i = 0; i < kQ8Chunks; ++i) { + qf[0][i] = load_q(r0, i); + qf[1][i] = load_q(r1, i); + } + const float sq0 = r0 < rows ? p.q_scale[(q_start + row_base + r0) * p.heads + head] : 1.f; + const float sq1 = r1 < rows ? p.q_scale[(q_start + row_base + r1) * p.heads + head] : 1.f; + + float o[kQ8PvChunks][4]; +#pragma unroll + for (int i = 0; i < kQ8PvChunks; ++i) o[i][0] = o[i][1] = o[i][2] = o[i][3] = 0.f; + float m0 = -1e30f, m1 = -1e30f, l0 = 0.f, l1 = 0.f; + + for (int j0 = 0; j0 < len_k; j0 += kQ8BN) { + const int n = min(kQ8BN, len_k - j0); + for (int idx = tid; idx < kQ8BN * kQ8D; idx += kQ8Warps * 32) { + const int j = idx / kQ8D, d = idx % kQ8D; + const int key = j0 + j; + const bool valid = key < len_k; + ks[j * kQ8Stride + d] = + valid ? p.k[static_cast(k_start + key) * p.k_row + static_cast(kv_head) * p.k_head + d] : 0; + vt[d * (kQ8BN + 8) + j] = + valid ? p.v[static_cast(k_start + key) * p.v_row + static_cast(kv_head) * p.v_head + d] + : __float2half_rn(0.f); + } + for (int idx = tid; idx < kQ8BN; idx += kQ8Warps * 32) { + kscale[idx] = j0 + idx < len_k ? p.k_scale[(k_start + j0 + idx) * p.kv_heads + kv_head] : 0.f; + } + __syncthreads(); + constexpr int kSubtiles = kQ8BN / 8; + float s[kSubtiles][4]; + const int jn = (n + 7) / 8; +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + int a0 = 0, a1 = 0, b0 = 0, b1 = 0; + const int key = t * 8 + group; +#pragma unroll + for (int i = 0; i < kQ8Chunks; ++i) { + const std::uint32_t kb = load_u32(&ks[key * kQ8Stride + i * 16 + tig * 4]); + mma_m8n8k16(a0, a1, qf[0][i], kb); + mma_m8n8k16(b0, b1, qf[1][i], kb); + } + const float sk0 = kscale[t * 8 + 2 * tig] * p.scale; + const float sk1 = kscale[t * 8 + 2 * tig + 1] * p.scale; + s[t][0] = t < jn ? static_cast(a0) * sq0 * sk0 : -1e30f; + s[t][1] = t < jn ? static_cast(a1) * sq0 * sk1 : -1e30f; + s[t][2] = t < jn ? static_cast(b0) * sq1 * sk0 : -1e30f; + s[t][3] = t < jn ? static_cast(b1) * sq1 * sk1 : -1e30f; + } + float tile_max0 = -1e30f, tile_max1 = -1e30f; +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + const int key0 = j0 + t * 8 + tig * 2; + for (int e = 0; e < 4; ++e) { + const int key = key0 + (e & 1); + if (key >= len_k) s[t][e] = -1e30f; + } + tile_max0 = fmaxf(tile_max0, fmaxf(s[t][0], s[t][1])); + tile_max1 = fmaxf(tile_max1, fmaxf(s[t][2], s[t][3])); + } + tile_max0 = fmaxf(tile_max0, __shfl_xor_sync(0xffffffffu, tile_max0, 1)); + tile_max0 = fmaxf(tile_max0, __shfl_xor_sync(0xffffffffu, tile_max0, 2)); + tile_max1 = fmaxf(tile_max1, __shfl_xor_sync(0xffffffffu, tile_max1, 1)); + tile_max1 = fmaxf(tile_max1, __shfl_xor_sync(0xffffffffu, tile_max1, 2)); + const float m_new0 = fmaxf(m0, tile_max0), m_new1 = fmaxf(m1, tile_max1); + const float alpha0 = (m0 <= -1e29f) ? 0.f : exp2f((m0 - m_new0) * kLog2e); + const float alpha1 = (m1 <= -1e29f) ? 0.f : exp2f((m1 - m_new1) * kLog2e); + const auto prob = [](float sv, float mv) { return sv <= -1e29f ? 0.f : exp2f((sv - mv) * kLog2e); }; + float tile_sum0 = 0.f, tile_sum1 = 0.f; +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + const float p0 = prob(s[t][0], m_new0), p1 = prob(s[t][1], m_new0); + const float p2 = prob(s[t][2], m_new1), p3 = prob(s[t][3], m_new1); + tile_sum0 += p0 + p1; + tile_sum1 += p2 + p3; + s[t][0] = p0; + s[t][1] = p1; + s[t][2] = p2; + s[t][3] = p3; + } + tile_sum0 += __shfl_xor_sync(0xffffffffu, tile_sum0, 1); + tile_sum0 += __shfl_xor_sync(0xffffffffu, tile_sum0, 2); + tile_sum1 += __shfl_xor_sync(0xffffffffu, tile_sum1, 1); + tile_sum1 += __shfl_xor_sync(0xffffffffu, tile_sum1, 2); + l0 = l0 * alpha0 + tile_sum0; + l1 = l1 * alpha1 + tile_sum1; + m0 = m_new0; + m1 = m_new1; +#pragma unroll + for (int i = 0; i < kQ8PvChunks; ++i) { + o[i][0] *= alpha0; + o[i][1] *= alpha0; + o[i][2] *= alpha1; + o[i][3] *= alpha1; + } +#pragma unroll + for (int t = 0; t < kSubtiles; ++t) { + const std::uint32_t pa0 = pack_half2(s[t][0], s[t][1]); + const std::uint32_t pa1 = pack_half2(s[t][2], s[t][3]); +#pragma unroll + for (int i = 0; i < kQ8PvChunks; ++i) { + const std::uint32_t b = load_u32(&vt[(i * 8 + group) * (kQ8BN + 8) + t * 8 + tig * 2]); + mma_m16n8k8(o[i], pa0, pa1, b); + } + } + __syncthreads(); + } + + 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 < kQ8PvChunks; ++i) { + const std::int64_t idx = static_cast(q_start + row_base + r0) * p.o_row + + static_cast(head) * p.o_head + i * 8 + tig * 2; + *reinterpret_cast<__nv_bfloat162*>(&p.o[idx]) = __floats2bfloat162_rn(o[i][0] * inv0, o[i][1] * inv0); + } + } + if (r1 < rows) { +#pragma unroll + for (int i = 0; i < kQ8PvChunks; ++i) { + const std::int64_t idx = static_cast(q_start + row_base + r1) * p.o_row + + static_cast(head) * p.o_head + i * 8 + tig * 2; + *reinterpret_cast<__nv_bfloat162*>(&p.o[idx]) = __floats2bfloat162_rn(o[i][2] * inv1, o[i][3] * inv1); + } + } +} + +} // namespace + +std::size_t q8_workspace_bytes(std::int64_t total_q, std::int64_t total_k, int num_heads, int num_kv_heads) { + const auto pad = [](std::size_t b) { return (b + 255u) & ~std::size_t{255}; }; + const std::size_t q8 = static_cast(total_q * num_heads * kQ8D); + const std::size_t k8 = static_cast(total_k * num_kv_heads * kQ8D); + const std::size_t v16 = k8 * 2; + const std::size_t qs = static_cast(total_q * num_heads) * sizeof(float); + const std::size_t ks = static_cast(total_k * num_kv_heads) * sizeof(float); + return pad(q8) + pad(k8) + pad(v16) + pad(qs) + pad(ks); +} + +void q8_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, float scale, void* workspace, std::size_t workspace_bytes, + cudaStream_t stream) { + const std::int64_t total_q = q.dim(0); + if (total_q == 0) return; + const std::int64_t total_k = k.dim(0); + const int heads = static_cast(q.dim(1)); + const int kv_heads = static_cast(k.dim(1)); + const int batch = static_cast(cu_seqlens_q.dim(0) - 1); + const auto pad = [](std::size_t b) { return (b + 255u) & ~std::size_t{255}; }; + std::size_t off = 0; + std::int8_t* q_i8 = static_cast(workspace) + off; + off += pad(static_cast(total_q * heads * kQ8D)); + std::int8_t* k_i8 = static_cast(workspace) + off; + off += pad(static_cast(total_k * kv_heads * kQ8D)); + __half* v16 = reinterpret_cast<__half*>(static_cast(workspace) + off); + off += pad(static_cast(total_k * kv_heads * kQ8D * 2)); + float* q_scale = reinterpret_cast(static_cast(workspace) + off); + off += pad(static_cast(total_q * heads * sizeof(float))); + float* k_scale = reinterpret_cast(static_cast(workspace) + off); + if (workspace_bytes < q8_workspace_bytes(total_q, total_k, heads, kv_heads)) { + throw std::invalid_argument("sm75::q8_flash_fwd: workspace too small"); + } + const __nv_bfloat16* qb = static_cast(q.data); + const __nv_bfloat16* kb = static_cast(k.data); + const __nv_bfloat16* vb = static_cast(v.data); + q8_quantize_rows<<(total_q * heads), 128, 0, stream>>>(qb, q_i8, q_scale, static_cast(total_q), heads, + q.stride[0], q.stride[1]); + MI_CUDA_CHECK_LAUNCH(); + q8_quantize_rows<<(total_k * kv_heads), 128, 0, stream>>>(kb, k_i8, k_scale, static_cast(total_k), + kv_heads, k.stride[0], k.stride[1]); + MI_CUDA_CHECK_LAUNCH(); + q8_convert_v<<(total_k * kv_heads), 128, 0, stream>>>(vb, v16, v.stride[0], v.stride[1], kv_heads); + MI_CUDA_CHECK_LAUNCH(); + Q8Params p; + p.q = q_i8; + p.k = k_i8; + p.v = v16; + p.q_scale = q_scale; + p.k_scale = k_scale; + p.o = static_cast<__nv_bfloat16*>(out.data); + p.cu_q = cu_seqlens_q.ptr(); + p.cu_k = cu_seqlens_k.ptr(); + p.heads = heads; + p.kv_heads = kv_heads; + p.total_q = static_cast(total_q); + p.q_row = static_cast(heads) * kQ8D; + p.q_head = kQ8D; + p.k_row = static_cast(kv_heads) * kQ8D; + p.k_head = kQ8D; + p.v_row = static_cast(kv_heads) * kQ8D; + p.v_head = kQ8D; + p.o_row = out.stride[0]; + p.o_head = out.stride[1]; + p.scale = scale; + const dim3 grid(static_cast((max_seqlen + kQ8BM - 1) / kQ8BM), static_cast(batch), + static_cast(heads)); + attention_q8_kernel<<>>(p); + MI_CUDA_CHECK_LAUNCH(); +} + +} // namespace mi::kernels::sm75 diff --git a/src/kernels/sm75/sm75.h b/src/kernels/sm75/sm75.h index 3cfef53..2befd91 100644 --- a/src/kernels/sm75/sm75.h +++ b/src/kernels/sm75/sm75.h @@ -86,4 +86,11 @@ void flash_fwd(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& const Tensor& cu_seqlens_k, int max_seqlen, bool causal, float scale, void* workspace, std::size_t workspace_bytes, cudaStream_t stream); +// 8-bit attention (`q8_fp16`/`q8_fp16_f32acc`), the sm_75 counterpart of attention_q8::attention_q8's fp16 variants: +// head_dim 128, non-causal, varlen/GQA; q and k are quantized per row (int8) and v becomes fp16 (attention_q8_sm75.cu, +// Turing's m8n8k16 INT8 MMA for QK^T, the exact kernel's fp16 P·V). Lossy by design, like the sm_80 variants. +std::size_t q8_workspace_bytes(std::int64_t total_q, std::int64_t total_k, int num_heads, int num_kv_heads); +void q8_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, float scale, void* workspace, std::size_t workspace_bytes, + cudaStream_t stream); } // namespace mi::kernels::sm75 diff --git a/src/models/moss_dit/dit_block_quant.cpp b/src/models/moss_dit/dit_block_quant.cpp index fe62138..1d5bde7 100644 --- a/src/models/moss_dit/dit_block_quant.cpp +++ b/src/models/moss_dit/dit_block_quant.cpp @@ -9,6 +9,7 @@ #include "kernels/attention/flash_attention.h" #include "kernels/attention_q8/attention_q8.h" +#include "kernels/sm75/sm75.h" #include "kernels/basic/elementwise.h" #include "kernels/basic/norm.h" #include "kernels/fused/fused_gemm.h" @@ -429,8 +430,10 @@ void run_block_quant(const BlockTensors& w, const QuantBlockBuffers& qa, const B // cu_seqlens_q == cu_seqlens_k and nothing reads the normed k first: no reference slots (kv is then the slotted // region, with cu_seqlens_k) and no K/V tap (it copies the normed, roped k). Otherwise qk_norm_rope_moss runs before // the tap and attention_q8 gets no norm_rope. + // sm_75's 8-bit attention has no fused norm/rope pre-pass (its own pre-pass quantizes rows); the separate + // qk_norm_rope_moss below gives the same bits, so keep it there. const bool fused_norm_rope = attention != AttentionMode::kExact && slots == 0 && run.kv_tap == nullptr && - m >= kFusedNormRopeMinTokens; + m >= kFusedNormRopeMinTokens && !kernels::sm75::turing(); if (!fused_norm_rope) kernels::fused::qk_norm_rope_moss(qkv, static_cast(kH), w.q_norm, w.k_norm, rope, stream); if (run.kv_tap != nullptr) { tap_reference_kv(qkv, run.kv_tap, stream); From b62b1cf00776a28e925024ca84925cfae7200c36 Mon Sep 17 00:00:00 2001 From: xycld Date: Sun, 8 Nov 2026 04:08:37 -0500 Subject: [PATCH 11/12] sm75: review fixes on top of PR #13 - attention_q8: the Turing INT8-QK kernel is built for 75-real only, so Volta (also sm75::turing()) no longer selects it; q8 stays unsupported there, as before. The Turing path now runs the same operand checks as the sm_80 kernels (k, v, out, cu_seqlens, head_dim 128, kv_heads, scale, workspace size and alignment). - Tier high/turbo on sm_7x: an unset attention stays exact. Turing's q8_fp16 is no faster than exact there (PR's own figures), so it runs only when asked for. - Session runs force LoRA placement merge on sm_7x as Pipeline::run does (one helper, force_lora_merge_on_sm7x); before, a session LoRA reached the sm_80-only side branches and failed mid-run. - Merged LoRAs on sm_7x: remerge_lora and rebase_merges convert the rewritten device-resident weights back to fp16, instead of dropping the DiT on a strength change (which bypassed reset_arena and left bf16 bits in fp16 views on a stack switch). A stack merging a module for the first time reloads (its resident bits are already fp16). - VAE encoder on sm_7x: encoder_conv2d goes through conv2d (Turing and Volta dispatch, its checks); conv2d_down shares the CUTLASS path's operand checks. - sm75 convert kernels: vectorized only when rows are aligned for it; the float4 loop also dropped the last 2-3 columns when cols % 4 >= 2. - quantize --format int4: the converter metadata names convrot_w4a4 layers as such. - docs/sm75.md: the attention default, the re-merge behaviour, the Numerics heading, one encoder timing. --- docs/sm75.md | 28 ++++++++++-------- src/kernels/attention_q8/attention_q8.cu | 30 ++++++++++++++----- src/kernels/sm75/convert.cu | 37 ++++++++++++++---------- src/kernels/vae/encoder_conv2d.cu | 23 ++++++++++----- src/models/moss_dit/moss_dit.cpp | 26 +++++++++++++++-- src/pipeline/pipeline.cpp | 32 ++++++++++---------- src/pipeline/pipeline_impl.h | 4 +++ src/pipeline/session.cpp | 1 + src/planner/planner.cpp | 4 +-- src/planner/planner.h | 4 +-- src/weights/quant/quant_convert.cpp | 2 +- tests/unit/test_planner.cpp | 4 +-- 12 files changed, 128 insertions(+), 67 deletions(-) diff --git a/docs/sm75.md b/docs/sm75.md index 9432e3e..c462457 100644 --- a/docs/sm75.md +++ b/docs/sm75.md @@ -31,7 +31,7 @@ presets stay Ampere-and-newer. Without the option the engine refuses compute cap | 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) | +| 8-bit attention (`q8_fp16`, `q8_fp8`) | SageAttention kernels | `q8_fp16`: a Turing INT8-QK kernel (sm_75 only, not Volta), run when `attention=q8_fp16` is given; `q8_fp8` is refused. An unset `attention` resolves to `exact` with the reason saying so (the tier is otherwise unchanged) | | VAE decode | TF32/bf16 CUTLASS implicit-GEMM convolutions, TF32 attention | fp32 SIMT convolutions and a fp32 attention kernel (reference precision; slower, see below) | | GatedDeltaNet (text encoder) | chunked kernel (~85 KiB of shared memory) | a per-token recurrent kernel that fits 33 KiB; the same rule, fp32 | | Other kernels (norm, RoPE, elementwise, samplers, group norm, latent preview, …) | compiled for sm_80+ | compiled for sm_75 and sm_80+ from the same source (pure SIMT) | @@ -41,10 +41,11 @@ 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. 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. +`reference` (bf16 bundle), `high` (INT8 ConvRot DiT) and `turbo` (INT4 ConvRot DiT) all run. On sm_75 an unset +`attention` resolves to `exact` and the report's `quality.reason` says "(no default 8-bit attention kernel on this +device)": Turing's `q8_fp16` is no faster than exact there, so it runs only when `attention=q8_fp16` is given (it also +lets `tier high` run on a bf16 DiT). A caller may still pass `attention=exact` explicitly. NVFP4 DiTs are refused as +on every non-Blackwell GPU. `high` and `turbo` need a quantized DiT file; the bf16 bundle itself only runs `reference`. `moss-infer quantize` produces INT8 ConvRot files on Turing (5.04 GB for the 5B DiT in 49 s), and a community `convrot_w4a4` file runs as it @@ -82,8 +83,10 @@ time before the vectorisation below) and the sm75 conversion kernels. At 1024^2 The kernels this backend owns were optimised from that profile: `conv2d_f32_sm75` now uses float4 shared-memory loads and float4 accumulators (four input channels per FMA) with vectorised float4 output stores (VAE decode 0.37 -> 0.20 s at 256^2 and 10-12 -> 7.8 s at 1024^2; encode 0.39 -> 0.19 s), `gemm_out_f32` loads float4 and stores packed bf16 -(176 -> 93 ms per four-step run), and `bf16_to_f16_rows_scaled` uses packed pair loads/stores. -## Numerics The exact attention kernel is next: it kept its 64-query tile but now loads the Q fragments straight from global memory instead of staging them in shared memory, which lets three blocks share an SM instead of one, and it uses eight warps for the text encoder's head_dim 256 (four for the DiT's 128): the 1024^2 DiT attention call went from 32.8 to 14.3 ms (3.2 to 7.3 TFLOPS) and the text encoder's from 29.4 to 17.3 ms (7.9 TFLOPS). +(176 -> 93 ms per four-step run), and `bf16_to_f16_rows_scaled` uses packed pair loads/stores (both fall back to one +column per thread for rows that are not aligned for it). The exact attention kernel is next: it kept its 64-query tile but now loads the Q fragments straight from global memory instead of staging them in shared memory, which lets three blocks share an SM instead of one, and it uses eight warps for the text encoder's head_dim 256 (four for the DiT's 128): the 1024^2 DiT attention call went from 32.8 to 14.3 ms (3.2 to 7.3 TFLOPS) and the text encoder's from 29.4 to 17.3 ms (7.9 TFLOPS). + +## Numerics - 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 @@ -116,19 +119,20 @@ and float4 accumulators (four input channels per FMA) with vectorised float4 out ## Limits on sm_75 -- 8-bit attention: `q8_fp16` runs (a Turing INT8-QK kernel, `tier high` accepts a bf16 DiT with it); `q8_fp8` is +- 8-bit attention: `q8_fp16` runs when asked for (a Turing INT8-QK kernel; `tier high` accepts a bf16 DiT with it; + Volta has none); `q8_fp8` is refused (no e4m3 MMA on Turing). On Turing the INT8 MMA does the same MACs per instruction as the fp16 one, so `q8_fp16` is a bandwidth/footprint path rather than a speedup (see `docs/sm75-quantization.md`). - Native INT4 MMA and NVFP4: INT4 files run through the INT8 unpack fallback (correct and bit-exact; the planner prices it at about 1.5× a native INT4 step); NVFP4 files are refused on every non-Blackwell GPU. - The fused DiT blocks and their CUDA graphs (the blocks run op by op). - img2img, inpainting, reference images and the hires fix run: the VAE **encoder** uses the same fp32 SIMT - convolutions as the decoder on Turing (a stride-2 downsample kernel included). It is not fast (about 5 s for a - 256x256 encode, 16 s for a 384x384 hires second pass), but it works. + convolutions as the decoder on Turing (a stride-2 downsample kernel included); its cost is under Performance above. - Runtime LoRAs are **merged into the weights at load** instead of the sm_80+ fused side branches: the merge kernels are plain SIMT, run before the bf16-to-fp16 conversion, and the report says `merged` for their adapters. A strength - change drops the DiT and reloads it (the resident weights are fp16, so an in-place re-merge would mix dtypes), which - costs a load. Control adapters still need the fused LoRA path and are refused. + change or a switch between stacks re-merges in place from the pristine bf16 copies and converts the rewritten + resident weights to fp16 again; a stack that merges a module the loaded DiT never merged reloads it. Control adapters + still need the fused LoRA path and are refused. - 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. diff --git a/src/kernels/attention_q8/attention_q8.cu b/src/kernels/attention_q8/attention_q8.cu index 45c17bb..e924eca 100644 --- a/src/kernels/attention_q8/attention_q8.cu +++ b/src/kernels/attention_q8/attention_q8.cu @@ -1,4 +1,5 @@ #include "kernels/attention_q8/attention_q8.h" +#include "kernels/sm70/sm70.h" #include "kernels/sm75/sm75.h" #include @@ -166,9 +167,15 @@ double e4m3_rn(double x) { } // namespace +namespace { +// Turing (cc 7.5) has its own INT8-QK kernel (kernels/sm75/attention_q8_sm75.cu) for the fp16-accumulation variants. +// Built for 75-real only (its m8n8k16 INT8 and m16n8k8 MMA need sm_75): Volta, which also runs the sm_75 family's +// other dispatchers, has none. +bool turing_q8_kernel() { return kernels::sm75::turing() && !kernels::sm70::volta(); } +} // namespace + bool supported(Variant v) { - // sm_75 has its own INT8-QK kernel (kernels/sm75/attention_q8_sm75.cu) for the fp16-accumulation variants. - if (kernels::sm75::turing()) return v == Variant::kInt8Fp16 || v == Variant::kInt8Fp16F32Acc; + if (turing_q8_kernel()) return v == Variant::kInt8Fp16 || v == Variant::kInt8Fp16F32Acc; int device = 0; MI_CUDA_CHECK(cudaGetDevice(&device)); return support(device, v) == Support::kOk; @@ -198,14 +205,15 @@ float effective_v_range(const Options& o) { std::size_t attention_q8_workspace_bytes(const Options& o, std::int64_t total_q, std::int64_t total_k, int batch, int max_seqlen, int num_heads, int num_kv_heads) { - if (kernels::sm75::turing()) return kernels::sm75::q8_workspace_bytes(total_q, total_k, num_heads, num_kv_heads); + if (turing_q8_kernel()) return kernels::sm75::q8_workspace_bytes(total_q, total_k, num_heads, num_kv_heads); return layout(o.variant, total_q, total_k, batch, max_seqlen, num_heads, num_kv_heads).total; } void attention_q8(const Options& o, 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, float scale, void* workspace, std::size_t workspace_bytes, cudaStream_t stream, const QkNormRope* norm_rope) { - if (kernels::sm75::turing()) { + const bool turing = turing_q8_kernel(); + if (turing) { if (o.variant != Variant::kInt8Fp16 && o.variant != Variant::kInt8Fp16F32Acc) { throw std::runtime_error("attention_q8: the sm_75 build runs the fp16-accumulation variants only"); } @@ -213,10 +221,6 @@ void attention_q8(const Options& o, const Tensor& q, const Tensor& k, const Tens throw std::invalid_argument( "attention_q8: the sm_75 build has no fused norm/rope pre-pass; run qk_norm_rope_moss first"); } - check_qkvo(q, "q"); - kernels::sm75::q8_flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, scale, workspace, workspace_bytes, - stream); - return; } check_qkvo(q, "q"); check_qkvo(k, "k"); @@ -236,6 +240,16 @@ void attention_q8(const Options& o, const Tensor& q, const Tensor& k, const Tens require(cu_seqlens_q.dim(0) - 1 <= kMaxGridYz, "batch must be <= 65535 (grid z)"); require(max_seqlen > 0, "max_seqlen must be positive"); require(std::isfinite(scale) && scale > 0, "scale must be finite and positive"); + if (turing) { // the same operand contract as the sm_80 kernels; its own workspace layout + const std::size_t need = kernels::sm75::q8_workspace_bytes(total_q, total_k, static_cast(heads), + static_cast(kv_heads)); + require(total_q == 0 || (workspace != nullptr && reinterpret_cast(workspace) % kAlign == 0 && + workspace_bytes >= need), + "workspace too small or not 256-byte aligned (need " + std::to_string(need) + " bytes)"); + kernels::sm75::q8_flash_fwd(q, k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen, scale, workspace, workspace_bytes, + stream); + return; + } if (norm_rope != nullptr) { const auto vec128 = [](const Tensor& t) { return t.dtype == DType::kBF16 && t.device == Device::kCUDA && t.rank == 1 && t.dim(0) == kHeadDim && diff --git a/src/kernels/sm75/convert.cu b/src/kernels/sm75/convert.cu index be293aa..aebcaf1 100644 --- a/src/kernels/sm75/convert.cu +++ b/src/kernels/sm75/convert.cu @@ -16,6 +16,8 @@ namespace mi::kernels::sm75 { namespace { +bool aligned(const void* p, std::uintptr_t a) { return reinterpret_cast(p) % a == 0; } + constexpr int kThreads = 256; constexpr int kMaxBlocks = 65535; constexpr float kF16Max = 65504.f; @@ -107,18 +109,19 @@ namespace { // 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) { + std::int64_t dst_stride, std::int64_t pair_cols) { __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; float amax = 0.f; - std::int64_t i = 2 * static_cast(threadIdx.x); - for (; i + 1 < cols; i += 2 * static_cast(blockDim.x)) { + // Pairs over [0, pair_cols) (0 when the rows are not 4-byte aligned), then one column per thread for the rest. + const std::int64_t tid = threadIdx.x, threads = blockDim.x; + for (std::int64_t i = 2 * tid; i < pair_cols; i += 2 * threads) { const float2 f = __bfloat1622float2(*reinterpret_cast(s + i)); amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); } - for (; i < cols; i += blockDim.x) amax = fmaxf(amax, fabsf(__bfloat162float(s[i]))); + for (std::int64_t i = pair_cols + tid; i < cols; i += threads) 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(); @@ -133,18 +136,17 @@ __global__ void bf16_to_f16_rows_scaled(const __nv_bfloat16* __restrict__ src, _ } __syncthreads(); const float scale = row_scale; - std::int64_t j = 2 * static_cast(threadIdx.x); - for (; j + 1 < cols; j += 2 * static_cast(blockDim.x)) { + for (std::int64_t j = 2 * tid; j < pair_cols; j += 2 * threads) { const float2 f = __bfloat1622float2(*reinterpret_cast(s + j)); *reinterpret_cast<__half2*>(d + j) = __floats2half2_rn(f.x * scale, f.y * scale); } - for (; j < cols; j += blockDim.x) d[j] = __float2half_rn(__bfloat162float(s[j]) * scale); + for (std::int64_t j = pair_cols + tid; j < cols; j += threads) d[j] = __float2half_rn(__bfloat162float(s[j]) * 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) { + std::int64_t acc_stride, std::int64_t vec_cols) { 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]; @@ -155,9 +157,10 @@ __global__ void gemm_out_f32(const float* __restrict__ acc, const float* __restr o[c] = v; } } else { - // Four columns per iteration: one float4 accumulator load and two packed bf16 stores. - std::int64_t c = 4 * static_cast(threadIdx.x); - for (; c + 3 < cols; c += 4 * static_cast(blockDim.x)) { + // Four columns per iteration over [0, vec_cols) (0 when the rows are not aligned for it): one float4 accumulator + // load and two packed bf16 stores. Then one column per thread for the rest. + const std::int64_t tid = threadIdx.x, threads = blockDim.x; + for (std::int64_t c = 4 * tid; c < vec_cols; c += 4 * threads) { const float4 a4 = *reinterpret_cast(a + c); float v0 = a4.x * sc, v1 = a4.y * sc, v2 = a4.z * sc, v3 = a4.w * sc; if (bias != nullptr) { @@ -171,7 +174,7 @@ __global__ void gemm_out_f32(const float* __restrict__ acc, const float* __restr *reinterpret_cast<__nv_bfloat162*>(o + c) = __floats2bfloat162_rn(v0, v1); *reinterpret_cast<__nv_bfloat162*>(o + c + 2) = __floats2bfloat162_rn(v2, v3); } - for (; c < cols; c += blockDim.x) { + for (std::int64_t c = vec_cols + tid; c < cols; c += threads) { float v = a[c] * sc; if (bias != nullptr) v += static_cast(bias[c]); o[c] = __float2bfloat16_rn(v); @@ -184,8 +187,10 @@ void convert_rows_bf16_to_f16_scaled(void* dst, float* inv_scale, const void* sr std::int64_t cols, std::int64_t src_stride, std::int64_t dst_stride, cudaStream_t stream) { if (rows <= 0 || cols <= 0) return; + const bool pairs = aligned(src, 4) && aligned(dst, 4) && src_stride % 2 == 0 && dst_stride % 2 == 0; bf16_to_f16_rows_scaled<<(rows), kThreads, 0, stream>>>( - static_cast(src), static_cast<__half*>(dst), inv_scale, cols, src_stride, dst_stride); + static_cast(src), static_cast<__half*>(dst), inv_scale, cols, src_stride, dst_stride, + pairs ? cols & ~std::int64_t{1} : 0); MI_CUDA_CHECK_LAUNCH(); } @@ -197,11 +202,13 @@ void convert_gemm_out_f32(void* out, bool out_f32, const float* acc, const float if (out_f32) { gemm_out_f32<<>>(acc, inv_scale, static_cast(out), static_cast(bias), cols, out_stride, - acc_stride); + acc_stride, 0); } else { + const bool vec = aligned(acc, 16) && acc_stride % 4 == 0 && aligned(out, 4) && out_stride % 2 == 0 && + (bias == nullptr || aligned(bias, 4)); gemm_out_f32<__nv_bfloat16, __nv_bfloat16><<>>( acc, inv_scale, static_cast<__nv_bfloat16*>(out), static_cast(bias), cols, out_stride, - acc_stride); + acc_stride, vec ? cols & ~std::int64_t{3} : 0); } MI_CUDA_CHECK_LAUNCH(); } diff --git a/src/kernels/vae/encoder_conv2d.cu b/src/kernels/vae/encoder_conv2d.cu index 596e91f..478b1aa 100644 --- a/src/kernels/vae/encoder_conv2d.cu +++ b/src/kernels/vae/encoder_conv2d.cu @@ -136,9 +136,10 @@ void check_cutlass(cutlass::Status st, const char* op, const char* what) { throw std::runtime_error(std::string(op) + ": CUTLASS " + what + " failed: " + cutlassGetStatusString(st)); } -template -void run(const char* name, const Shape& s, const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, - const Tensor& y, void* workspace, std::size_t workspace_bytes, cudaStream_t stream) { +// The operand rules every path shares (the CUTLASS kernels and the sm_7x SIMT ones, which load bias and residual as +// float4). +void check_operands(const char* name, const Shape& s, const Tensor& x, const Tensor& w, const Tensor& bias, + const Tensor* residual, const Tensor& y) { require(bias.dtype == DType::kF32 && bias.numel() == s.k && bias.device == Device::kCUDA, std::string(name) + ": bias must be f32 [K] on CUDA"); if (residual) @@ -152,6 +153,12 @@ void run(const char* name, const Shape& s, const Tensor& x, const Tensor& w, con std::string(name) + ": x/w/y/residual need 16-byte and bias 32-byte alignment"); require(x.nbytes() < (std::size_t{1} << 31) && y.nbytes() < (std::size_t{1} << 31), std::string(name) + ": tensors must be smaller than 2 GiB (CUTLASS 32-bit offsets)"); +} + +template +void run(const char* name, const Shape& s, const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, + const Tensor& y, void* workspace, std::size_t workspace_bytes, cudaStream_t stream) { + check_operands(name, s, x, w, bias, residual, y); auto args = make_args(s, x, w, bias, residual, y); check_cutlass(Cfg::Op::can_implement(args), name, "can_implement"); require(Cfg::Op::get_workspace_size(args) <= workspace_bytes, std::string(name) + ": workspace too small"); @@ -172,9 +179,9 @@ std::size_t encoder_conv2d_workspace_bytes(const Tensor& x, const Tensor& w, Con void encoder_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, const Tensor& y, void* workspace, std::size_t workspace_bytes, ConvMath math, cudaStream_t stream) { - if (kernels::sm75::turing()) { // fp32 SIMT conv: Turing has no tensor cores for this path - check("encoder_conv2d", x, w, y, 1); - kernels::sm75::vae_conv2d(x, w, bias, residual, y, stream); + if (kernels::sm75::turing()) { // sm_7x: the decoder's fp32 SIMT conv (conv2d dispatches Turing and Volta) + check_operands("encoder_conv2d", check("encoder_conv2d", x, w, y, 1), x, w, bias, residual, y); + conv2d(x, w, bias, residual, y, workspace, workspace_bytes, stream); return; } if (math == ConvMath::kTf32) { @@ -193,8 +200,8 @@ std::size_t conv2d_down_workspace_bytes(const Tensor& x, const Tensor& w, const void conv2d_down(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor& y, void* workspace, std::size_t workspace_bytes, ConvMath math, cudaStream_t stream) { - if (kernels::sm75::turing()) { - check("conv2d_down", x, w, y, 2); + if (kernels::sm75::turing()) { // sm_7x: fp32 SIMT (kernels/sm75/vae_sm75.cu, also built for sm_70) + check_operands("conv2d_down", check("conv2d_down", x, w, y, 2), x, w, bias, nullptr, y); kernels::sm75::vae_conv2d_down(x, w, bias, y, stream); return; } diff --git a/src/models/moss_dit/moss_dit.cpp b/src/models/moss_dit/moss_dit.cpp index 2695b84..eebadb5 100644 --- a/src/models/moss_dit/moss_dit.cpp +++ b/src/models/moss_dit/moss_dit.cpp @@ -447,6 +447,16 @@ void MossDit::remerge_into(const lora::PristineWeights& pristine, const lora::Lo namespace { std::atomic g_host_merges{false}; + +// sm_7x: the resident bf16 linears were converted to fp16 at load (after the load's merges, which run in bf16). A +// re-merge writes bf16 bits into a device-resident target again (its MergeTarget keeps the load's kBF16 view), so they +// are converted back in place. Streamed blocks merge into their pinned host copy, which stays bf16 and is converted on +// every refill (BlockWeights::Pass::acquire), and INT8 targets are never converted. +void refresh_fp16(const Tensor& w) { + if (!kernels::sm75::fp16_weights() || w.data == nullptr || w.device != Device::kCUDA) return; + if (w.dtype != DType::kBF16) return; + kernels::sm75::convert_bf16_to_f16(w.data, w.data, w.numel(), nullptr); +} } // namespace void MossDit::set_host_merges_for_tests(bool host) { g_host_merges.store(host); } @@ -504,6 +514,9 @@ bool MossDit::rebase_merges(const lora::LoraStack* stack, const std::string& mem if (t.w.data == nullptr) return false; // a block module: only a load can merge it added.emplace_back(module, t); } + // sm_7x: a module merged for the first time needs its pristine bf16 bits, but its resident weight already holds the + // load's fp16 conversion: only a load can merge it. + if (!added.empty() && kernels::sm75::fp16_weights()) return false; MI_CUDA_CHECK(cudaDeviceSynchronize()); // no forward may read the weights while they change // Transactional: the pristine copies of the new modules and every module's new bits are computed into owned host // buffers first; a failure there (a non-finite merge, a bad factor, out of memory) drops the new copies and leaves the @@ -549,7 +562,11 @@ bool MossDit::rebase_merges(const lora::LoraStack* stack, const std::string& mem for (const std::string& module : new_copies) pristine_->drop(module); throw; } - for (const MergeStaging& st : staged) st.store(); + for (std::size_t i = 0; i < staged.size(); ++i) { // staged[i] is all[i]'s + staged[i].store(); + refresh_fp16(all[i].second.w); + } + if (kernels::sm75::fp16_weights()) MI_CUDA_CHECK(cudaStreamSynchronize(nullptr)); if (!memo_key.empty() && !memo_hit) { // The new stack's merged bits, when they and their pristine copies fit the cap (a stack merging more is merged again // when it comes back). @@ -617,11 +634,16 @@ void MossDit::remerge_lora(const lora::LoraStack& stack) { // The pristine bits back into the target, then the merge on the GPU. pristine_->restore_bf16(module, staging.w(), staging.bias()); staging.store(); - if (merge_on_device(stack, module, t)) continue; + if (merge_on_device(stack, module, t)) { + refresh_fp16(t.w); + continue; + } } remerge_into(*pristine_, stack, module, t, staging); staging.store(); + refresh_fp16(t.w); } + if (kernels::sm75::fp16_weights()) MI_CUDA_CHECK(cudaStreamSynchronize(nullptr)); } std::size_t MossDit::pristine_bytes(const lora::LoraStack& stack, const quant::QuantIndex* quant) { diff --git a/src/pipeline/pipeline.cpp b/src/pipeline/pipeline.cpp index 677590f..dcbd06b 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -938,11 +938,16 @@ Pipeline::Impl::CallModels Pipeline::Impl::setup_models(const Request& req, Resu // Quality tier from what is loaded, before anything heavy (its refusals throw std::invalid_argument here, ahead of // the LoRA files too). The lossy accelerators are an explicit speed opt-in: the tier is resolved without them. const std::filesystem::path te_file = tes.model ? tes.path : std::filesystem::path(); - const auto supported = [](planner::AttentionMode mode) { return moss_dit::attention_supported(dit_attention(mode)); }; - // sm_75: no 8-bit attention kernel exists for Turing, so the caller's unset attention resolves to exact instead of - // failing. Keyed on the device, not on the kernels' availability: an sm_80+ build without them (MOSSINFER_ONLY) - // still refuses, as before. + // sm_7x: the caller's unset attention resolves to exact instead of an 8-bit default. Volta has no 8-bit attention + // kernel; Turing's q8_fp16 (kernels/sm75/attention_q8_sm75.cu) runs when asked for, but is no faster than exact + // there (Turing's INT8 MMA does the fp16 MMA's MACs per instruction), so it is never the default. Keyed on the + // device, not on the kernels' availability: an sm_80+ build without them (MOSSINFER_ONLY) still refuses, as before. const bool exact_when_no_q8 = kernels::sm75::turing(); + const bool attention_given = m.options.attention.has_value(); + const auto supported = [&](planner::AttentionMode mode) { + if (exact_when_no_q8 && !attention_given && planner::is_q8(mode)) return false; + return moss_dit::attention_supported(dit_attention(mode)); + }; // sm_75 has no LoRA side-branch kernels, so run() forces every adapter to merge (Pipeline::run, above); INT4 // ConvRot weights have no merge path (merges need bf16 or INT8 ConvRot), and the tier's own advice ("use // lora_placement side_branch") does not exist on Turing: refuse here with what does work. @@ -1071,12 +1076,8 @@ Pipeline::Impl::CallModels Pipeline::Impl::setup_models(const Request& req, Resu // Merged adapters: re-merged from the pristine weights kept at load, in place (no reload). // A model merged without pristine copies (an edit adapter's blocks) is reloaded instead (its key cleared). if (m.dit && m.dit_merge_key == files && !m.dit->can_remerge()) m.dit_merge_key.clear(); - if (kernels::sm75::turing() && !m.lora_stack->merge_modules().empty()) { - // The merge is in bf16 at load, and the weights are converted to fp16 afterwards: re-merging in place would - // mix the two, so drop the DiT and let the next call reload (its load re-merges, then re-converts). - m.dit_merge_key.clear(); - m.dit.reset(); - } else if (m.dit && m.dit_merge_key == files && !m.lora_stack->merge_modules().empty()) { + // On sm_7x the re-merge converts the rewritten resident weights back to fp16 (MossDit::remerge_lora). + if (m.dit && m.dit_merge_key == files && !m.lora_stack->merge_modules().empty()) { const auto t0m = Clock::now(); try { m.dit->remerge_lora(*m.lora_stack); @@ -1517,6 +1518,11 @@ void Pipeline::Impl::explain_control(const Request::Control& control, Result& r) // Pipeline::run, and by a Session for its options and every run's spec): control adapters (a LoRA with a replacement // embedding, fused into the block GEMMs). The VAE encoder runs on Turing (fp32 SIMT), and runtime LoRAs are merged at // load (Pipeline::run sets the placement). +void detail::force_lora_merge_on_sm7x(Request& req) { + // The load-time merge runs in bf16 and the fp16 conversion follows it, as for the bf16 bundle. + if (kernels::sm75::turing() && !req.loras.empty()) req.lora_placement = lora::Placement::kMerge; +} + void detail::refuse_unsupported_on_sm7x(const Request& req) { if (!kernels::sm75::turing()) return; // The VAE encoder runs on Turing (fp32 SIMT convs, the same kernels as the decoder), so img2img, inpainting, @@ -1530,11 +1536,7 @@ void detail::refuse_unsupported_on_sm7x(const Request& req) { Result Pipeline::run(const Request& request, const GenerateControl& control, bool preload) { Impl& m = *impl_; Request req = request; - // Turing has no fused LoRA side-branch kernels (sm_80+ CUTLASS), so a runtime LoRA is merged into the weights at - // load instead (its load-time merge then feeds the fp16 conversion the same way as the bf16 bundle). - if (kernels::sm75::turing() && !req.loras.empty() && req.lora_placement != lora::Placement::kMerge) { - req.lora_placement = lora::Placement::kMerge; - } + force_lora_merge_on_sm7x(req); check_user_paths(m.options, req); validate(req); refuse_unsupported_on_sm7x(req); diff --git a/src/pipeline/pipeline_impl.h b/src/pipeline/pipeline_impl.h index bd61867..1040691 100644 --- a/src/pipeline/pipeline_impl.h +++ b/src/pipeline/pipeline_impl.h @@ -221,6 +221,10 @@ inline void check_user_paths(const PipelineOptions& o, const Request& r) { // On a Volta or Turing device: throws NotSupportedError for what still needs sm_80+ kernels (pipeline.cpp). No-op elsewhere. void refuse_unsupported_on_sm7x(const Request& req); +// On a Volta or Turing device, runtime LoRAs are merged into the weights at load: the side branches are fused into +// sm_80+ CUTLASS GEMMs. Sets req.lora_placement to kMerge there (Pipeline::run and every Session run); no-op elsewhere. +void force_lora_merge_on_sm7x(Request& req); + inline moss_dit::SkipSet skip_set(const std::vector& layers) { moss_dit::SkipSet s; for (const int l : layers) s.set(static_cast(l)); diff --git a/src/pipeline/session.cpp b/src/pipeline/session.cpp index 7b50593..eb7b109 100644 --- a/src/pipeline/session.cpp +++ b/src/pipeline/session.cpp @@ -139,6 +139,7 @@ Request run_request(const SessionOptions& o, const RunSpec& spec, const Guidance q.slg_start = spec.slg_start; q.slg_stop = spec.slg_stop; q.loras = spec.loras; + force_lora_merge_on_sm7x(q); // no LoRA side branches on sm_7x q.control.adapter = o.control_adapter; q.control.lora_strength = spec.control_lora_strength; q.control.images = spec.control_images; diff --git a/src/planner/planner.cpp b/src/planner/planner.cpp index b86ae83..a05083d 100644 --- a/src/planner/planner.cpp +++ b/src/planner/planner.cpp @@ -928,7 +928,7 @@ TierDecision resolve_tier(Tier requested, std::optional attention return d; } // High and Turbo: the tier's default attention (kTierAttention) unless the caller gave one; never a silent fallback - // (exact_when_no_q8 makes the absence explicit in the reason instead: sm_75 has no 8-bit attention kernel). + // (exact_when_no_q8 makes the absence explicit in the reason instead: sm_7x has no default 8-bit attention kernel). const char* const tier = requested == Tier::kTurbo ? "turbo" : "high"; if (attention) { d.attention = *attention; @@ -952,7 +952,7 @@ TierDecision resolve_tier(Tier requested, std::optional attention d.reason = std::string(requested == Tier::kTurbo ? "Turbo tier" : "High tier") + ": " + weights + ", " + attention_name(d.attention) + " attention" + (exact_when_no_q8 && !attention && d.attention == AttentionMode::kExact - ? " (no 8-bit attention kernel on this device)" + ? " (no default 8-bit attention kernel on this device)" : ""); return d; } diff --git a/src/planner/planner.h b/src/planner/planner.h index 71f40a4..74331f3 100644 --- a/src/planner/planner.h +++ b/src/planner/planner.h @@ -131,8 +131,8 @@ inline constexpr AttentionMode kTierAttention[] = {AttentionMode::kQ8Fp8, Attent // Throws std::invalid_argument (naming the conflict) for a refused combination. An empty `supported` runs every kernel. // te_quantized: the text encoder is an INT8 / W4A8 file, a lossy dimension as an INT8 DiT is (reference refused; the // exact-attention configuration is a High-tier one). nvfp4_native: the NVFP4 GEMMs run here (Capabilities::nvfp4_native). -// exact_when_no_q8: the device has no 8-bit attention kernel (sm_75), so an unset `attention` resolves to exact with a -// reason instead of refusing the tier (the default would otherwise be a silent, absent fallback). +// exact_when_no_q8: the device has no default 8-bit attention kernel (sm_7x), so an unset `attention` resolves to exact +// with a reason instead of refusing the tier (the default would otherwise be a silent, absent fallback). TierDecision resolve_tier(Tier tier, std::optional attention, Precision dit_precision, const AttentionSupported& supported, bool te_quantized = false, bool nvfp4_native = false, bool exact_when_no_q8 = false); diff --git a/src/weights/quant/quant_convert.cpp b/src/weights/quant/quant_convert.cpp index 3d5bccb..9819571 100644 --- a/src/weights/quant/quant_convert.cpp +++ b/src/weights/quant/quant_convert.cpp @@ -860,7 +860,7 @@ ConvertStats convert(const fs::path& input, const fs::path& output, const Conver std::map spec_counts; // describe(spec) -> GEMMs written with it for (OutModule& om : mods) { om.scale = om.scale_2 = om.config = om.bias = kNone; - spec_counts[describe(om.spec)] += static_cast(om.parts.size()); + spec_counts[om.int4 ? std::string("convrot_w4a4 g256") : describe(om.spec)] += static_cast(om.parts.size()); if (om.int4) { if (om.k % 256 != 0) fail("'" + om.name + "': K=" + std::to_string(om.k) + " is not a multiple of 256, which convrot_w4a4 needs"); om.weight = add(om.name + ".weight", DType::kI8, {om.n, om.k / 2}); diff --git a/tests/unit/test_planner.cpp b/tests/unit/test_planner.cpp index acf8d10..a546c1a 100644 --- a/tests/unit/test_planner.cpp +++ b/tests/unit/test_planner.cpp @@ -1179,11 +1179,11 @@ TEST(Planner, TurboTier) { EXPECT_EQ(d.attention, AttentionMode::kQ8Fp8); EXPECT_EQ(d.reason, "Turbo tier: INT4 ConvRot GEMMs, q8_fp8 attention"); EXPECT_THROW(resolve_tier(Tier::kTurbo, O{}, Precision::kInt4, no_q8), std::invalid_argument); // no fallback - // sm_75: no 8-bit attention kernel exists, so an unset attention may resolve to exact (explicitly said so). + // sm_7x: no default 8-bit attention kernel, so an unset attention may resolve to exact (explicitly said so). { const TierDecision t = resolve_tier(Tier::kTurbo, O{}, Precision::kInt4, no_q8, false, false, /*exact_when_no_q8=*/true); EXPECT_EQ(t.attention, AttentionMode::kExact); - EXPECT_NE(t.reason.find("no 8-bit attention kernel"), std::string::npos) << t.reason; + EXPECT_NE(t.reason.find("no default 8-bit attention kernel"), std::string::npos) << t.reason; const TierDecision h = resolve_tier(Tier::kHigh, O{}, Precision::kInt8, no_q8, false, false, /*exact_when_no_q8=*/true); EXPECT_EQ(h.attention, AttentionMode::kExact); } From 03dd7d9103c85a0fa7dc20321133b9270d4c74a3 Mon Sep 17 00:00:00 2001 From: xycld Date: Sun, 8 Nov 2026 04:22:09 -0500 Subject: [PATCH 12/12] sm75: a memoized LoRA stack re-merges in place; the session's plan request merges too rebase_merges on sm_7x reloads only when an added module has no pristine copy (the memo keeps copies of modules a later stack restored). The Session's plan request goes through force_lora_merge_on_sm7x like its runs, so the setup before the first run builds the stack the runs use. --- src/models/moss_dit/moss_dit.cpp | 10 +++++++--- src/pipeline/session.cpp | 1 + 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/models/moss_dit/moss_dit.cpp b/src/models/moss_dit/moss_dit.cpp index eebadb5..03eca26 100644 --- a/src/models/moss_dit/moss_dit.cpp +++ b/src/models/moss_dit/moss_dit.cpp @@ -514,9 +514,13 @@ bool MossDit::rebase_merges(const lora::LoraStack* stack, const std::string& mem if (t.w.data == nullptr) return false; // a block module: only a load can merge it added.emplace_back(module, t); } - // sm_7x: a module merged for the first time needs its pristine bf16 bits, but its resident weight already holds the - // load's fp16 conversion: only a load can merge it. - if (!added.empty() && kernels::sm75::fp16_weights()) return false; + // sm_7x: a module with no pristine copy yet would take it from its resident weight, which already holds the load's + // fp16 conversion: only a load can merge it. One the memo kept a copy of merges from that copy. + if (kernels::sm75::fp16_weights()) { + for (const auto& [module, t] : added) { + if (!pristine_ || !pristine_->has(module)) return false; + } + } MI_CUDA_CHECK(cudaDeviceSynchronize()); // no forward may read the weights while they change // Transactional: the pristine copies of the new modules and every module's new bits are computed into owned host // buffers first; a failure there (a non-finite merge, a bad factor, out of memory) drops the new copies and leaves the diff --git a/src/pipeline/session.cpp b/src/pipeline/session.cpp index eb7b109..6731195 100644 --- a/src/pipeline/session.cpp +++ b/src/pipeline/session.cpp @@ -380,6 +380,7 @@ Session::Session(Pipeline& pipeline, JobGate& gate, SessionOptions options) { q.control.adapter = options.control_adapter; if (q.control.adapter.empty()) q.control.images.clear(); if (tok::python_strip(q.prompt).empty()) q.prompt = "preload"; + force_lora_merge_on_sm7x(q); // the plan's models match the runs' (run_request) check_user_paths(pipeline.impl_->options, q); pipeline::validate(q); refuse_unsupported_on_sm7x(q);