[PR #36] [MERGED] feat: svdquant w4a4 + awq w4a16 kitchen native (CORE-31) #95

New issue

Closed

opened 2026-10-26 22:20:21 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:21 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/36
Author: @HK416-TYPED
Created: 4/21/2026
Status: ✅ Merged
Merged: 5/29/2026
Merged by: @comfyanonymous

Base: main ← Head: feat/svdquant-w4a4-kitchen-native

📝 Commits (10+)

c9cdad7 Add SVDQuant W4A4 + AWQ W4A16 eager backends and public API
f3206d1 Add kitchen-native CUDA int4 MMA kernels for SVDQuant W4A4
b6a587d Add TensorCoreAWQW4A16Layout for modulation-style W4A16 linears
671a5eb Add CUDA backend for AWQ W4A16 with tiered M-routing
3367e73 Fix ruff lint: lowercase math var locals + minor cleanup
12df742 Fix remaining ruff lint: cuda/init.py + tensor/init.py
8e8456b Add SVDQuant W4A4 tile-packed runtime support
07b4a60 Fix bf16 W4A4 fast-accum CUDA compile
01b1104 Avoid bf16 hfma2 overload in W4A4 fast accum
cf8c5cd Fix SVDQuant W4A4 test ruff lint

📊 Changes

15 files changed (+4481 additions, -1 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+101 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+3 -1)
📝 comfy_kitchen/backends/cuda/__init__.py (+370 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+208 -0)
➕ comfy_kitchen/backends/cuda/ops/awq_w4a16.cu (+617 -0)
➕ comfy_kitchen/backends/cuda/ops/quantize_svdquant_w4a4.cu (+186 -0)
➕ comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu (+764 -0)
➕ comfy_kitchen/backends/cuda/ops/svdquant_utils.cuh (+256 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+46 -0)
➕ comfy_kitchen/backends/eager/awq.py (+108 -0)
➕ comfy_kitchen/backends/eager/svdquant.py (+364 -0)
📝 comfy_kitchen/tensor/__init__.py (+16 -0)
➕ comfy_kitchen/tensor/awq_w4a16.py (+217 -0)
➕ comfy_kitchen/tensor/svdquant_w4a4.py (+595 -0)
➕ tests/test_svdquant_w4a4.py (+630 -0)

📄 Description

Add SVDQuant W4A4 + AWQ W4A16 int4 quantization support for Qwen-Image-Edit and similar SVD-LoRA / AWQ-modulation checkpoints. Built as kitchen-native row-major backends.

Targets sm_80+ (Ampere).

SVDQuant W4A4

comfy_kitchen.quantize_svdquant_w4a4 / scaled_mm_svdquant_w4a4 as torch.library.custom_op entry points with eager pure-PyTorch reference.
TensorCoreSVDQuantW4A4Layout dispatching aten.t / aten.mm / aten.addmm / aten.linear over int4 qdata + proj_down / proj_up / smooth_factor + per-group weight_scale.
CUDA kernels: m16n8k64 int4 MMA GEMM + fused activation quantize over kitchen-native row-major tensors; fp32 accumulator (fp16 overflows at production Qwen scale ≈ 28 — see kernel header).
act_unsigned end-to-end for post-GELU fc2 layers: u4.s4 MMA variant + +0.171875 shift applied at the layer (kernel API stays shift-free).
Signed quantizer clamp to [-7, 7] (skip -8) to match nunchaku's absmax/7 dequant-symmetric contract.
LoRA-up + bias applied in Python via bf16 addmm_ (faster than kernel fusion for R=96; see scaled_mm_svdquant_w4a4 docstring for the precision tradeoff and the cuBLASLt upgrade path if
needed).

AWQ W4A16

comfy_kitchen.gemv_awq_w4a16 as a single op covering both GEMV (M ≤ 8) and GEMM regimes — kernel routes internally.
TensorCoreAWQW4A16Layout dispatching aten.t / aten.mm / aten.addmm / aten.linear over kitchen-native uint4 qdata + per-group wscales / wzeros.
CUDA kernels:
- M ≤ 8: naive 1-thread-per-output (18× faster than eager at M=1, 4× at M=8 on RTX 5090 / Qwen-Image modulation shape N=18432, K=3072).
- 8 < M ≤ 256: fused mma.m16n8k16.f32.bf16.bf16.f32 with cooperative dequant of int4 weight tile to bf16 in shmem. No 113 MB intermediate W workspace. shmem stride padded 8 b16 to break
  the 128-byte ldmatrix bank conflict.
- M > 256: dequantize + cuBLAS bf16 matmul fallback (the MMA kernel's single-thread-per-N-row dequant pass and lack of cp.async pipelining stop scaling past M ≈ 256 on Blackwell; raising the
  limit is a follow-up tuning task).
bias applied externally in the Python wrapper (out.add_(bias)), mirroring scaled_mm_svdquant_w4a4's epilogue contract.

Tests

tests/test_svdquant_w4a4.py — 16 tests: quantizer clamp contract, signed/unsigned MMA dispatch, lora_x separation, cross-shape smoke.

Verification

AWQ: vs nunchaku runtime ops.gemv_awq on a real r96 modulation tensor — max-abs / max-magnitude ≤ 0.65% at all M (within bf16 ULP precision floor).
E2E: Qwen-Image-Edit r96 ComfyUI sampling on the kitchen-native AWQ checkpoint — wall 274 s vs 342 s eager-AWQ (20% faster), PSNR 33.5 dB vs the bf16-dequant baseline. Stratified sample across 6
variant families (balanced / fast / mid / quality × ranks 32/64/96/128 × base/lightning4/lightning8) all sample successfully with the right (steps, cfg) per distillation contract.ontract, signed/unsigned MMA dispatch, lora_x separation, cross-shape smoke.

Model repository

LAXMAYDAY/Qwen_image_edit_2511_int4
Models carrying the awq suffix are models that have been transformed and can be inferred using that branch

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/36 **Author:** [@HK416-TYPED](https://github.com/HK416-TYPED) **Created:** 4/21/2026 **Status:** ✅ Merged **Merged:** 5/29/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `feat/svdquant-w4a4-kitchen-native` --- ### 📝 Commits (10+) - [`c9cdad7`](https://github.com/Comfy-Org/comfy-kitchen/commit/c9cdad744c0931f8396e846d1043993d52a62d84) Add SVDQuant W4A4 + AWQ W4A16 eager backends and public API - [`f3206d1`](https://github.com/Comfy-Org/comfy-kitchen/commit/f3206d11342f4135415e679bd88aabfebd11f5b9) Add kitchen-native CUDA int4 MMA kernels for SVDQuant W4A4 - [`b6a587d`](https://github.com/Comfy-Org/comfy-kitchen/commit/b6a587d6a823a14c976c86b787d27c010c95a366) Add TensorCoreAWQW4A16Layout for modulation-style W4A16 linears - [`671a5eb`](https://github.com/Comfy-Org/comfy-kitchen/commit/671a5eb41892e3bd46384ed3651982b8be9a2bf6) Add CUDA backend for AWQ W4A16 with tiered M-routing - [`3367e73`](https://github.com/Comfy-Org/comfy-kitchen/commit/3367e73e88e80db05cda60653986c3620c4b94d0) Fix ruff lint: lowercase math var locals + minor cleanup - [`12df742`](https://github.com/Comfy-Org/comfy-kitchen/commit/12df74233845d1b59b8be0d4c26420a905d0c916) Fix remaining ruff lint: cuda/__init__.py + tensor/__init__.py - [`8e8456b`](https://github.com/Comfy-Org/comfy-kitchen/commit/8e8456ba1a42c95a40ca09ee4137321a303c304a) Add SVDQuant W4A4 tile-packed runtime support - [`07b4a60`](https://github.com/Comfy-Org/comfy-kitchen/commit/07b4a607b633b527290602763392598b6284ac64) Fix bf16 W4A4 fast-accum CUDA compile - [`01b1104`](https://github.com/Comfy-Org/comfy-kitchen/commit/01b1104a570c00229602409d5a7936f0b6fed627) Avoid bf16 hfma2 overload in W4A4 fast accum - [`cf8c5cd`](https://github.com/Comfy-Org/comfy-kitchen/commit/cf8c5cdd4e4b7563d6785d0048e876369fad1faa) Fix SVDQuant W4A4 test ruff lint ### 📊 Changes **15 files changed** (+4481 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+101 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+3 -1) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+370 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+208 -0) ➕ `comfy_kitchen/backends/cuda/ops/awq_w4a16.cu` (+617 -0) ➕ `comfy_kitchen/backends/cuda/ops/quantize_svdquant_w4a4.cu` (+186 -0) ➕ `comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu` (+764 -0) ➕ `comfy_kitchen/backends/cuda/ops/svdquant_utils.cuh` (+256 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+46 -0) ➕ `comfy_kitchen/backends/eager/awq.py` (+108 -0) ➕ `comfy_kitchen/backends/eager/svdquant.py` (+364 -0) 📝 `comfy_kitchen/tensor/__init__.py` (+16 -0) ➕ `comfy_kitchen/tensor/awq_w4a16.py` (+217 -0) ➕ `comfy_kitchen/tensor/svdquant_w4a4.py` (+595 -0) ➕ `tests/test_svdquant_w4a4.py` (+630 -0) </details> ### 📄 Description Add SVDQuant W4A4 + AWQ W4A16 int4 quantization support for Qwen-Image-Edit and similar SVD-LoRA / AWQ-modulation checkpoints. Built as kitchen-native row-major backends. Targets sm_80+ (Ampere). ## SVDQuant W4A4 - `comfy_kitchen.quantize_svdquant_w4a4` / `scaled_mm_svdquant_w4a4` as `torch.library.custom_op` entry points with eager pure-PyTorch reference. - `TensorCoreSVDQuantW4A4Layout` dispatching `aten.t` / `aten.mm` / `aten.addmm` / `aten.linear` over int4 qdata + `proj_down` / `proj_up` / `smooth_factor` + per-group `weight_scale`. - CUDA kernels: `m16n8k64` int4 MMA GEMM + fused activation quantize over kitchen-native row-major tensors; fp32 accumulator (fp16 overflows at production Qwen scale ≈ 28 — see kernel header). - `act_unsigned` end-to-end for post-GELU `fc2` layers: `u4.s4` MMA variant + `+0.171875` shift applied at the layer (kernel API stays shift-free). - Signed quantizer clamp to `[-7, 7]` (skip `-8`) to match nunchaku's `absmax/7` dequant-symmetric contract. - LoRA-up + bias applied in Python via bf16 `addmm_` (faster than kernel fusion for R=96; see `scaled_mm_svdquant_w4a4` docstring for the precision tradeoff and the cuBLASLt upgrade path if needed). ## AWQ W4A16 - `comfy_kitchen.gemv_awq_w4a16` as a single op covering both GEMV (M ≤ 8) and GEMM regimes — kernel routes internally. - `TensorCoreAWQW4A16Layout` dispatching `aten.t` / `aten.mm` / `aten.addmm` / `aten.linear` over kitchen-native uint4 qdata + per-group `wscales` / `wzeros`. - CUDA kernels: - **M ≤ 8**: naive 1-thread-per-output (18× faster than eager at M=1, 4× at M=8 on RTX 5090 / Qwen-Image modulation shape N=18432, K=3072). - **8 < M ≤ 256**: fused `mma.m16n8k16.f32.bf16.bf16.f32` with cooperative dequant of int4 weight tile to bf16 in shmem. No 113 MB intermediate W workspace. shmem stride padded 8 b16 to break the 128-byte ldmatrix bank conflict. - **M > 256**: dequantize + cuBLAS bf16 matmul fallback (the MMA kernel's single-thread-per-N-row dequant pass and lack of cp.async pipelining stop scaling past M ≈ 256 on Blackwell; raising the limit is a follow-up tuning task). - bias applied externally in the Python wrapper (`out.add_(bias)`), mirroring `scaled_mm_svdquant_w4a4`'s epilogue contract. ## Tests - `tests/test_svdquant_w4a4.py` — 16 tests: quantizer clamp contract, signed/unsigned MMA dispatch, `lora_x` separation, cross-shape smoke. ## Verification - AWQ: vs nunchaku runtime `ops.gemv_awq` on a real r96 modulation tensor — max-abs / max-magnitude ≤ 0.65% at all M (within bf16 ULP precision floor). - E2E: Qwen-Image-Edit r96 ComfyUI sampling on the kitchen-native AWQ checkpoint — wall 274 s vs 342 s eager-AWQ (20% faster), PSNR 33.5 dB vs the bf16-dequant baseline. Stratified sample across 6 variant families (balanced / fast / mid / quality × ranks 32/64/96/128 × base/lightning4/lightning8) all sample successfully with the right `(steps, cfg)` per distillation contract.ontract, signed/unsigned MMA dispatch, lora_x separation, cross-shape smoke. ## Model repository - [LAXMAYDAY/Qwen_image_edit_2511_int4](https://huggingface.co/LAXMAYDAY/qwen-image-edit-2511-int4-comfy) - Models carrying the awq suffix are models that have been transformed and can be inferred using that branch --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:20:21 +00:00