Repository navigation
sm75: VAE encoder, runtime LoRAs, quantized tiers and a q8_fp16 attention kernel on Turing - #13
Merged
Merged
Conversation
- encoder_conv2d and conv2d_down dispatch to fp32 SIMT kernels on cc<8, with a new stride-2 3x3 downsample kernel in kernels/sm75/vae_sm75.cu. img2img, inpainting, reference images and the hires fix now work on RTX 20xx instead of being refused with not_supported. - Runtime LoRAs are merged into the weights at load on sm7x: the fused side-branch GEMMs are sm_80+ CUTLASS, but the merge kernels are plain SIMT, so the stack is forced to Placement::kMerge before the bf16-to-fp16 conversion. A strength change drops the DiT and reloads it (the resident weights are fp16, so an in-place re-merge would mix dtypes) instead of re-merging. - lora_merge.cu joins the portable sm75 sources; docs/sm75.md updated (only control adapters remain refused). Verified on the RTX 2080: --init-image, --reference and --hires produce images; --lora reports its 6 adapters as merged and changes the image; the reference-tier apple renders correctly with and without a LoRA.
CloverIris
force-pushed
the
sm75-lora-vae
branch
from
November 8, 2026 04:31
5bda750 to
c9806ee
Compare
Profiling on the RTX 2080 (Nsight Systems) showed the VAE convolutions as the largest single kernel (52.6 % of the img2img GPU time, 10.7 ms per launch) and the sm75 GEMM output/bias conversion at 176 ms per four-step generation. - conv2d_f32_sm75: float4 shared-memory loads and float4 accumulators (four input channels per FMA), vectorized float4 output stores with the bias and residual added in the same order. VAE decode 0.37 -> 0.20 s at 256^2 (10-12 -> 7.8 s at 1024^2), VAE encode 0.39 -> 0.19 s. - gemm_out_f32: float4 accumulator loads and packed bf16 stores, 176 -> 93 ms per four-step run. - bf16_to_f16_rows_scaled: pair-wise bf16 loads and packed half2 stores in both passes.
CloverIris
force-pushed
the
sm75-lora-vae
branch
from
November 8, 2026 04:32
c9806ee to
d466475
Compare
…ps for head_dim 256 The exact attention kernel staged the 64x64 Q tile in shared memory, which left one block per SM and four warps running. The Q fragments are now loaded straight from global memory (bf16 -> fp16 per lane), the block's shared memory is only the K and V tiles (about 19 KiB), and three blocks fit an SM. head_dim 256 (the text encoder's causal attention) uses eight warps per block. Measured with tools/sm75_check bench (RTX 2080): the DiT's 1024^2 attention call 32.8 -> 14.3 ms (3.2 -> 7.3 TFLOPS), the 256^2 call 0.23 -> 0.13 ms, the text encoder's 4096-token causal attention 29.4 -> 17.3 ms (7.9 TFLOPS). End to end at 1024^2 the denoise step drops from 6.4 to 5.2 s (two steps: 28.2 -> 23.6 s); 256^2 is unchanged within noise. tools/sm75_check gains a bench command for these shapes.
…ane counts The VAE mid attention read K and V for every query row: at a 1024^2 latent (16384 tokens) that is 8 TB of reads and 5.4 s of the decode. The kernel now keeps 32 query rows per block in registers (four per warp) and stages the K and V tiles once per block in shared memory, so each tile is read L/32 times. The lane count is a template parameter: with a runtime D the per-row arrays landed in local memory and the kernel was 20x slower. Measured at 1024^2 (RTX 2080): VAE attention 5.4 s -> 1.6 s, decode 7.6 s -> 3.3 s, one-step total 17.4 -> 14.0 s. The decoded PNG is bit-identical to the previous build (same key order in the online softmax). sm75_check vae still matches its host reference.
Nsight Compute showed the conv kernel L1/TEX-bound at 93 % with 40 % compute throughput: the weight tile was reloaded for every output pixel. Each thread now computes two pixels eight rows apart, so each weight float4 serves two pixels and the shared-memory traffic per FMA halves. t=1024^2 VAE decode 3.30 -> 2.19 s, bit-identical output.
New docs/sm75-quantization.md: what reference/high/turbo require, the measured RTX 2080 numbers for the three tiers with quantized DiT files (0.12 s/step INT8, 0.14 s/step INT4 at 256^2), the gaps (no INT4 writer in quantize, high needs a quantized DiT on Turing because there is no 8-bit attention, LoRA on quantized tiers unverified) and a five-point plan (INT4 output, Turing q8_fp16 attention, a tier rule fallback, LoRA on quantized tiers, docs/app). docs/sm75.md links it and states the quantized-file requirement.
moss-infer quantize --format int4 produces the file the turbo tier loads: packed 4-bit weights ([N, K/2], even column the low nibble), a per-row f32 scale, 256-column ConvRot rotation groups and the file-level _quantization_metadata the loader reads (convrot_w4a4, convrot_groupsize 256, quant_group_size 64). The scale rule is the format's clipping-ratio search (--recipe mse, the default) or absmax (--recipe rtn); LoRA merging and --keep-bf16 work as for int8. Verified on the RTX 2080: the 5B DiT becomes a 2.83 GB file (55 s with rtn, 283 s with mse); against a published community convrot_w4a4 file the dequantized weights match at cosine 0.9993 and the bf16 tensors are bit-identical, and the produced file runs tier turbo (0.16 s/step at 256^2, 3.9 GiB peak).
Verified on the RTX 2080: reference + --lora merges into bf16; high + --lora merges into the INT8 ConvRot weights (merge_requantize_int8) and reports the requantization warning; turbo + --lora used to fail with the generic merge error, whose advice (lora_placement side_branch) does not exist on Turing. The pipeline now refuses that combination up front on sm_7x and names what works (tier high or reference).
Records the dispatch points (attention_q8.cu's three functions), the workspace layout, the m8n8k16 fragment mapping from CUTLASS's sm75 header and the fact that the exact kernel's skeleton carries over, so the kernel work can start from the document.
New src/kernels/sm75/attention_q8_sm75.cu: per-row INT8 q/k (a quantize pre-pass, v converted to fp16), QK^T
with Turing's m8n8k16 INT8 MMA, then the exact kernel's fp32 online softmax and fp16 P·V. attention_q8.cu's
three entry points (supported, attention_q8_workspace_bytes, attention_q8) dispatch to it on cc 7.5, and
dit_block_quant keeps the separate qk_norm_rope_moss on Turing (the fused pre-pass is sm_80-only).
Verified on the RTX 2080: tier high now runs on the bf16 bundle ('DiT bf16, q8_fp16 attention'), the 256^2 output
matches the exact path at cosine 0.9986, and 1024^2 works. No speedup: Turing's INT8 MMA m8n8k16 does the same
MACs per instruction as the fp16 m16n8k8 it replaces (m16n8k32 is the 2x part on Ampere), so 8-bit attention is a
bandwidth path here; documented in docs/sm75-quantization.md and docs/sm75.md.
CloverIris
changed the title
- attention_q8: the Turing INT8-QK kernel is built for 75-real only, so Volta (also sm75::turing()) no longer selects it; q8 stays unsupported there, as before. The Turing path now runs the same operand checks as the sm_80 kernels (k, v, out, cu_seqlens, head_dim 128, kv_heads, scale, workspace size and alignment). - Tier high/turbo on sm_7x: an unset attention stays exact. Turing's q8_fp16 is no faster than exact there (PR's own figures), so it runs only when asked for. - Session runs force LoRA placement merge on sm_7x as Pipeline::run does (one helper, force_lora_merge_on_sm7x); before, a session LoRA reached the sm_80-only side branches and failed mid-run. - Merged LoRAs on sm_7x: remerge_lora and rebase_merges convert the rewritten device-resident weights back to fp16, instead of dropping the DiT on a strength change (which bypassed reset_arena and left bf16 bits in fp16 views on a stack switch). A stack merging a module for the first time reloads (its resident bits are already fp16). - VAE encoder on sm_7x: encoder_conv2d goes through conv2d (Turing and Volta dispatch, its checks); conv2d_down shares the CUTLASS path's operand checks. - sm75 convert kernels: vectorized only when rows are aligned for it; the float4 loop also dropped the last 2-3 columns when cols % 4 >= 2. - quantize --format int4: the converter metadata names convrot_w4a4 layers as such. - docs/sm75.md: the attention default, the re-merge behaviour, the Numerics heading, one encoder timing.
…quest merges too rebase_merges on sm_7x reloads only when an added module has no pristine copy (the memo keeps copies of modules a later stack restored). The Session's plan request goes through force_lora_merge_on_sm7x like its runs, so the setup before the first run builds the stack the runs use.
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ten commits that make the sm_75 (Turing) build a full engine, verified on an RTX 2080 8 GiB with the SFT bundle.
1. The VAE encoder runs on Turing (img2img, inpainting, reference images, hires fix)
encoder_conv2d/conv2d_downdispatch to fp32 SIMT kernels on cc<8, including a new stride-2 3x3 downsamplekernel; the
not_supportedrefusals for these features are gone.2. Runtime LoRAs on Turing (merged instead of side branches)
mergeon cc<8 (the side-branch GEMMs are sm_80+ CUTLASS); the merge kernelsare plain SIMT and run before the bf16 -> fp16 conversion. A strength change reloads the DiT instead of re-merging
into fp16 weights.
turbo + --loranow refuses up front with thereason and the working alternatives.
3. Performance (Nsight Systems + Nsight Compute on the RTX 2080)
L1/TEX-bound). Decode 0.37 -> 0.20 s at 256^2, 10-12 -> 2.19 s at 1024^2; encode 0.39 -> 0.19 s.
to local memory). 5.4 -> 0.24 s at a 1024^2 latent, bit-identical output.
warps for head_dim 256. 1024^2 call 32.8 -> 14.3 ms (7.3 TFLOPS); text encoder 29.4 -> 17.3 ms.
gemm_out_f32and the bf16->fp16 row conversions vectorized.4. Quantized tiers
moss-infer quantize --format int4writes the Turbo tier'sconvrot_w4a4format (packed 4-bit weights, per-rowscale, 256-column rotation groups,
_quantization_metadata); the scale rule is comfy-kitchen's clipping search(
--recipe mse, the default) or absmax. The produced 2.83 GB file matches a published community file's dequantizedweights at cosine 0.9993.
q8_fp16attention kernel (src/kernels/sm75/attention_q8_sm75.cu): per-row INT8 q/k, m8n8k16 QK^T, theexact kernel's fp32 softmax and fp16 P·V, wired through
attention_q8.cu, sotier highruns on a bf16 DiT.It matches the exact path at cosine 0.9986. No speedup on Turing: INT8
m8n8k16has the same MACs per instructionas fp16
m16n8k8(Ampere'sm16n8k32is the 2x part), so this is a bandwidth path on this architecture.Verified on the RTX 2080 8 GiB
reference(bf16)high(INT8 ConvRot 5.04 GB)turbo(INT4 convrot_w4a4 2.83 GB)highon a bf16 DiT (q8_fp16)Also verified:
--init-image,--reference,--hires, merged LoRA on bf16 and INT8 DiTs,quantizeINT8 and INT4round trips,
sm75_check allgreen, and theMOSSINFER_SM75=OFFbuild unchanged.docs/sm75.mdanddocs/sm75-quantization.mdcarry the details and the measured numbers.