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-quantization.md b/docs/sm75-quantization.md new file mode 100644 index 0000000..7b7b8dd --- /dev/null +++ b/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": {: {"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. diff --git a/docs/sm75.md b/docs/sm75.md index a7fdf2f..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,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. diff --git a/src/kernels/attention_q8/attention_q8.cu b/src/kernels/attention_q8/attention_q8.cu index b8a6013..e924eca 100644 --- a/src/kernels/attention_q8/attention_q8.cu +++ b/src/kernels/attention_q8/attention_q8.cu @@ -1,4 +1,6 @@ #include "kernels/attention_q8/attention_q8.h" +#include "kernels/sm70/sm70.h" +#include "kernels/sm75/sm75.h" #include #include @@ -165,7 +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) { + 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; @@ -195,12 +205,23 @@ 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 (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) { + 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"); + } + 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"); check_qkvo(k, "k"); check_qkvo(v, "v"); @@ -219,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/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/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/src/kernels/sm75/convert.cu b/src/kernels/sm75/convert.cu index 90db435..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,13 +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; - for (std::int64_t i = threadIdx.x; i < cols; i += blockDim.x) amax = fmaxf(amax, fabsf(__bfloat162float(s[i]))); + // 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 (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(); @@ -128,36 +136,61 @@ __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); + 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 (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]; - 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 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) { + 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 (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); } } } - } // namespace void convert_rows_bf16_to_f16_scaled(void* dst, float* inv_scale, const void* src, std::int64_t rows, std::int64_t cols, std::int64_t src_stride, std::int64_t dst_stride, cudaStream_t stream) { if (rows <= 0 || cols <= 0) return; + 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(); } @@ -169,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/sm75/sm75.h b/src/kernels/sm75/sm75.h index 29b6ac1..2befd91 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). @@ -82,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/kernels/sm75/vae_sm75.cu b/src/kernels/sm75/vae_sm75.cu index 14364ef..36e716c 100644 --- a/src/kernels/sm75/vae_sm75.cu +++ b/src/kernels/sm75/vae_sm75.cu @@ -18,24 +18,29 @@ 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, 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; 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; - 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[2][16]; +#pragma unroll + 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) { @@ -55,13 +60,224 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest : 0.f; } __syncthreads(); - for (int cc = 0; cc < kTC; ++cc) { + // 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 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[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; + } + } + } + } 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 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[0][ch].x += xv0 * v; + acc[1][ch].x += xv1 * v; + } + } + } + } + } + } + __syncthreads(); + } + +#pragma unroll + 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) 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; + } + } 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 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 + +// 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 tid = threadIdx.x; + const int b = blockIdx.y; + 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; + } + } + 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 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; + } +} +} // namespace + +void vae_conv2d(const Tensor& x, const Tensor& w, const Tensor& bias, const Tensor* residual, 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)), R = static_cast(w.dim(1)); + const int kb_count = (K + kTK - 1) / kTK; + const dim3 grid(static_cast((W + kTW - 1) / kTW), static_cast((H + kTH - 1) / kTH), + static_cast(x.dim(0) * kb_count)); + conv2d_f32_sm75<<>>( + x.ptr(), w.ptr(), bias.ptr(), residual != nullptr ? residual->ptr() : nullptr, + y.ptr(), H, W, C, K, R, kb_count); + 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 < R; ++r) { - for (int s = 0; s < R; ++s) { - const float xv = xs[((py + r) * kPX + (px + s)) * kTC + cc]; + 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 * R + s) * kTC + cc]; + for (int ch = 0; ch < 16; ++ch) acc[ch] += xv * ws[((cg + ch) * 9 + r * 3 + s) * kDownTC + cc]; } } } @@ -69,29 +285,34 @@ __global__ void conv2d_f32_sm75(const float* __restrict__ x, const float* __rest } const int ox = ox0 + px, oy = oy0 + py; - if (ox >= W || oy >= H) return; + 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) * H + oy) * W + ox) * K + kch; - float v = acc[ch] + bias[kch]; - if (residual != nullptr) v += residual[idx]; - y[idx] = v; + const std::int64_t idx = ((static_cast(n) * P + oy) * Q + ox) * K + kch; + y[idx] = acc[ch] + bias[kch]; } } - -// ---------------------------------------------------------------- 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; - -__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) { +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(); +} +// 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 * kAttnWarps + warp; + const int row = blockIdx.x * 4 + warp; if (row >= L) return; const int nlanes = D / 32; const int b = blockIdx.y; @@ -123,31 +344,33 @@ __global__ void attention_f32_sm75(const float* __restrict__ q, const float* __r 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, - 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)), R = static_cast(w.dim(1)); - const int kb_count = (K + kTK - 1) / kTK; - const dim3 grid(static_cast((W + kTW - 1) / kTW), static_cast((H + kTH - 1) / kTH), - static_cast(x.dim(0) * kb_count)); - conv2d_f32_sm75<<>>( - x.ptr(), w.ptr(), bias.ptr(), residual != nullptr ? residual->ptr() : nullptr, - y.ptr(), H, W, C, K, R, 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)); 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(); } diff --git a/src/kernels/vae/encoder_conv2d.cu b/src/kernels/vae/encoder_conv2d.cu index 5198044..478b1aa 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" @@ -135,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) @@ -151,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"); @@ -163,6 +171,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 +179,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()) { // 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) { conv2d(x, w, bias, residual, y, workspace, workspace_bytes, stream); return; @@ -179,12 +193,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()) { // 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; + } 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/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); diff --git a/src/models/moss_dit/moss_dit.cpp b/src/models/moss_dit/moss_dit.cpp index 2695b84..03eca26 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,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 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 @@ -549,7 +566,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 +638,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 ff21b15..dcbd06b 100644 --- a/src/pipeline/pipeline.cpp +++ b/src/pipeline/pipeline.cpp @@ -938,11 +938,24 @@ 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. + 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(), @@ -1063,6 +1076,7 @@ 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(); + // 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 { @@ -1500,36 +1514,29 @@ 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::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; - 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; + 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..6731195 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; @@ -379,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); 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/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..9819571 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"); @@ -852,8 +860,12 @@ 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()); - if (om.spec.format == QuantFormat::kNvfp4) { + 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}); + 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. 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); } 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 (...) {