[PR #63] int4_tensorwise: W4A4 layout (+ConvRot) with fused CUDA kernels #118

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/63
Author: @HK416-TYPED
Created: 7/6/2026
Status: 🔄 Open

Base: main ← Head: feat/int4-tensorwise-convrot

📝 Commits (9)

83d4566 feat: int4_tensorwise W4A4 (+ConvRot) eager reference and layout
5d800c2 cuda: fused int4_tensorwise activation quantize + s4 MMA dispatch
ba6ab49 cuda: rank-one accumulation fast path for the int4_tensorwise GEMM
1c77543 cuda: widen the rank-one GEMM M-tile (kMUnroll 1 -> 4)
005cf9a cuda: int4 rank-one tuning sweep + same-activation quantize memo
32349ec cuda: ldmatrix fragment loads + 48B smem row stride for the rank-one GEMM
1ef96a5 style: lowercase GEMM dim locals in int4_linear (ruff N806)
3e94fd4 review: keep int4 GEMM scales in the out_dtype lane; raise MIN_SM to 8.0
be4a39e eager: warn once when int4_tensorwise runs the slow reference path on CUDA

📊 Changes

11 files changed (+1823 additions, -27 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+127 -4)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+47 -1)
➕ comfy_kitchen/backends/cuda/ops/int4_tensorwise.cu (+315 -0)
📝 comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu (+171 -22)
📝 comfy_kitchen/backends/cuda/ops/svdquant_utils.cuh (+12 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+72 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+363 -0)
📝 comfy_kitchen/tensor/__init__.py (+3 -0)
➕ comfy_kitchen/tensor/int4.py (+296 -0)
➕ tests/test_int4_tensorwise.py (+416 -0)

📄 Description

Adds TensorWiseINT4Layout — a 4-bit downward extension of int8_tensorwise. Weights are packed signed int4 pairs (low-nibble-first, [N, K/2] int8 container) with fp32 [N, 1] per-row scales; activations are rotated + quantized to int4 dynamically inside the kernel.

Eager: quantize_int4_rowwise / quantize_int4_convrot_weight / int4_linear reference path, exact pack/unpack roundtrip, works on CPU.

CUDA (SM >= 8.0 for the MMA path, layout MIN_SM (7,5)):

fused activation kernel: FHT rotation (smem, radix-4) + rowwise int4 quantize + pack in one pass
s4·s4 MMA GEMM with rank-one epilogue dequant: per-row/per-col scales factor out of the K loop, so the main loop is pure cp.async + mma.m16n8k64.s4 with carried int32 accumulators (|acc| <= 49·K << 2^31) — no in-loop CUDA-core dequant
ldmatrix.x4/x2 fragment loads over a 48-byte-stride smem layout (bank-conflict-free); both operands share the ldmatrix k-permutation so the dot product is unchanged
capacity-1 activation-quantize memo: q/k (and add_q/add_k) projections receive the same activation object, so the second layer reuses the rotated+quantized activation (object identity + _version guarded, inference-mode safe)

Per-layer, L4 (SM 8.9, native int4 tensor cores), M=4608, vs the compiled int8 path:

 3072x3072    0.54 ms   (int8: 1.18)
 3072->12288  2.60 ms   (int8: 4.54)
12288->3072   2.30 ms   (int8: 4.04)

E2E Qwen-Image-Edit-2511 mixed checkpoint (80% int4 / 20% int8 layers, 20 steps, 1024², L4, same session): 4.22 s/it vs 4.29 s/it for full int8, at 75% of the file size (14.9 GB vs 20 GB, bf16 is 40 GB). Krea-2 Turbo/Raw: 1.33×/1.22× vs bf16. Published checkpoints for testing:

To test them, use ComfyUI from the companion loader branch: HK416-TYPED/ComfyUI feat/int4-tensorwise-loader with this 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/63 **Author:** [@HK416-TYPED](https://github.com/HK416-TYPED) **Created:** 7/6/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/int4-tensorwise-convrot` --- ### 📝 Commits (9) - [`83d4566`](https://github.com/Comfy-Org/comfy-kitchen/commit/83d4566db16accc98a368f9638e006bae42890c7) feat: int4_tensorwise W4A4 (+ConvRot) eager reference and layout - [`5d800c2`](https://github.com/Comfy-Org/comfy-kitchen/commit/5d800c27f682cca7de649c7c8e753f5561317935) cuda: fused int4_tensorwise activation quantize + s4 MMA dispatch - [`ba6ab49`](https://github.com/Comfy-Org/comfy-kitchen/commit/ba6ab492012a4b2c17d6ca33efdb2d48f7effe98) cuda: rank-one accumulation fast path for the int4_tensorwise GEMM - [`1c77543`](https://github.com/Comfy-Org/comfy-kitchen/commit/1c77543dd776401c77130c5eea137ff62c96311f) cuda: widen the rank-one GEMM M-tile (kMUnroll 1 -> 4) - [`005cf9a`](https://github.com/Comfy-Org/comfy-kitchen/commit/005cf9ad1fdc6e64af4b15f6e7aa480592e9e593) cuda: int4 rank-one tuning sweep + same-activation quantize memo - [`32349ec`](https://github.com/Comfy-Org/comfy-kitchen/commit/32349ecc2cc58194d490c51bb0babc690d0fb648) cuda: ldmatrix fragment loads + 48B smem row stride for the rank-one GEMM - [`1ef96a5`](https://github.com/Comfy-Org/comfy-kitchen/commit/1ef96a5827fd624e98eddf218739a190b9eb54d1) style: lowercase GEMM dim locals in int4_linear (ruff N806) - [`3e94fd4`](https://github.com/Comfy-Org/comfy-kitchen/commit/3e94fd4a66d7666778ecafa2dbfc8563fe9184e6) review: keep int4 GEMM scales in the out_dtype lane; raise MIN_SM to 8.0 - [`be4a39e`](https://github.com/Comfy-Org/comfy-kitchen/commit/be4a39e86897a8b98ce3fef636dfb5123b6214f0) eager: warn once when int4_tensorwise runs the slow reference path on CUDA ### 📊 Changes **11 files changed** (+1823 additions, -27 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+127 -4) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+47 -1) ➕ `comfy_kitchen/backends/cuda/ops/int4_tensorwise.cu` (+315 -0) 📝 `comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu` (+171 -22) 📝 `comfy_kitchen/backends/cuda/ops/svdquant_utils.cuh` (+12 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+72 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+363 -0) 📝 `comfy_kitchen/tensor/__init__.py` (+3 -0) ➕ `comfy_kitchen/tensor/int4.py` (+296 -0) ➕ `tests/test_int4_tensorwise.py` (+416 -0) </details> ### 📄 Description Adds `TensorWiseINT4Layout` — a 4-bit downward extension of int8_tensorwise. Weights are packed signed int4 pairs (low-nibble-first, `[N, K/2]` int8 container) with fp32 `[N, 1]` per-row scales; activations are rotated + quantized to int4 dynamically inside the kernel. **Eager**: `quantize_int4_rowwise` / `quantize_int4_convrot_weight` / `int4_linear` reference path, exact pack/unpack roundtrip, works on CPU. **CUDA** (SM >= 8.0 for the MMA path, layout MIN_SM (7,5)): - fused activation kernel: FHT rotation (smem, radix-4) + rowwise int4 quantize + pack in one pass - s4·s4 MMA GEMM with rank-one epilogue dequant: per-row/per-col scales factor out of the K loop, so the main loop is pure `cp.async` + `mma.m16n8k64.s4` with carried int32 accumulators (`|acc| <= 49·K << 2^31`) — no in-loop CUDA-core dequant - `ldmatrix.x4`/`x2` fragment loads over a 48-byte-stride smem layout (bank-conflict-free); both operands share the ldmatrix k-permutation so the dot product is unchanged - capacity-1 activation-quantize memo: q/k (and add_q/add_k) projections receive the same activation object, so the second layer reuses the rotated+quantized activation (object identity + `_version` guarded, inference-mode safe) Per-layer, L4 (SM 8.9, native int4 tensor cores), M=4608, vs the compiled int8 path: ``` 3072x3072 0.54 ms (int8: 1.18) 3072->12288 2.60 ms (int8: 4.54) 12288->3072 2.30 ms (int8: 4.04) ``` E2E Qwen-Image-Edit-2511 mixed checkpoint (80% int4 / 20% int8 layers, 20 steps, 1024², L4, same session): 4.22 s/it vs 4.29 s/it for full int8, at 75% of the file size (14.9 GB vs 20 GB, bf16 is 40 GB). Krea-2 Turbo/Raw: 1.33×/1.22× vs bf16. Published checkpoints for testing: - https://huggingface.co/LAXMAYDAY/Qwen-Image-Edit-2511-int4-tensorwise-mixed - https://huggingface.co/LAXMAYDAY/Krea-2-Turbo-int4-tensorwise-mixed - https://huggingface.co/LAXMAYDAY/Krea-2-Raw-int4-tensorwise-mixed To test them, use ComfyUI from the companion loader branch: [HK416-TYPED/ComfyUI `feat/int4-tensorwise-loader`](https://github.com/HK416-TYPED/ComfyUI/tree/feat/int4-tensorwise-loader) with this branch --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

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

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#118

No description provided.

Rows
Columns