sm75: VAE encoder, runtime LoRAs, quantized tiers and a q8_fp16 attention kernel on Turing

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

docs/sm75-quantization.md

@@ -0,0 +1,122 @@
+# 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).
+- `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 - 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": {<module>: {"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
+Turing user has to quantize the DiT first (or download an INT8 file) to use `high`.
+
+**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
+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
+one is loaded.
+
+## Plan
+
+**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`) (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
+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
+P2 is the cleaner long-term answer.
+
+**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.
+
+## 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.

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,17 @@ 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
+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
 
@@ -59,9 +66,26 @@ 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 (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
@@ -85,27 +109,30 @@ 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.
+- 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: 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); 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 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.