int4_tensorwise: W4A4 layout (+ConvRot) with fused CUDA kernels#63
int4_tensorwise: W4A4 layout (+ConvRot) with fused CUDA kernels#63HK416-TYPED wants to merge 9 commits into
Conversation
Downward extension of the int8_tensorwise protocol to signed INT4: per-row fp32 scale (absmax/7, division at the source dtype), packed 2-per-byte low-nibble-first int8 storage [N, K/2], optional group-wise regular-Hadamard ConvRot with online activation rotation in int4_linear. Dequantize un-rotates back to the original basis and requantize_kwargs preserves convrot so LoRA-offload requantization stays correct. Eager reference only; CUDA kernel to follow.
Adds ops/int4_tensorwise.cu: fused group-256 ConvRot FHT (same butterfly as int8_linear.cu) + per-row symmetric int4 quantize (absmax/7, [-7,7]) packed low-nibble-first, with the per-row scale broadcast into the (K/64, M_pad) svdquant ascales layout. The cuda int4_linear wrapper then runs the existing scaled_mm_svdquant_w4a4 s4 MMA GEMM with rank-1 scales (an exact special case of its per-group layout, empty LoRA), so int4 layers ride native INT4 tensor cores on SM 8.0-8.9. The FHT output is rounded once through the source dtype to reproduce the eager reference's rotated values (eager rotates via one fp32-accumulated matmul); residual cuda/eager divergence is boundary code flips at ~1e-3 relative, covered by a parity test matrix.
The svdquant MMA kernel dequantizes per K-group inside the main loop (int32 -> fp32 conversion + 16 FMAs per MMA), which caps it near 15% of s4 peak at diffusion shapes. int4_tensorwise scales are rank-1, so a new kRankOne kernel variant carries raw int32 accumulators across the whole K loop inside the MMA itself (|acc| <= 49*K << 2^31) and collapses dequantization into the epilogue: one per-row and one per-col scale. GEMM time drops ~35% at qwen shapes and accuracy improves (exact integer accumulation; batch parity rel 0.030 -> 0.004). The quantize kernel takes single-group ascales buffers (rank-one needs no broadcast) and gets 16-byte vectorized loads + uchar4 packed stores. Existing svdquant dispatches are untouched (kRankOne=false).
The default 32-row CTA rereads every B tile M/32 times — at qwen shapes (M~4600) that is ~144 full weight-matrix rereads and the dominant GEMM traffic. The rank-one variant frees the per-group scale registers, so it can afford a 128-row CTA (kMUnroll=4), quartering B rereads. kMUnroll is now a template parameter defaulting to the existing geometry; all non-rank-one instantiations are unchanged. L4 (SM 8.9) per-layer at M=4608, vs the compiled int8 path: 3072x3072: 0.94ms vs 1.18ms 3072->12288: 3.06ms vs 4.54ms 12288->3072: 3.58ms vs 4.04ms E2E qwen-edit-2511 int4-mixed: 90.1s/image vs same-session wheel int8 90.4-93.1s (eager baseline was 330s).
Static profiling (perf counters unavailable in the container; cuobjdump res-usage instead): the rank-one instantiation uses 128 regs / 24 KB smem / 256 threads -> 2 CTAs/SM (33% occupancy, register-bound) on SM 8.9. Occupancy-vs-reread sweep via COMFY_KITCHEN_INT4_R1_CONFIG: 64M- and 96M-row tiles raise occupancy but lose to L2-resident B rereads; 4-stage cp.async is flat (latency already covered). 128M/8w/3st stays default. LOCAL:0 for all int4 kernels (no off-chip spills); the K loop is pure cp.async + MMA with carried int32 accumulators. Adds a capacity-1 activation-quantize memo: q/k (and add_q/add_k) projections receive the same activation tensor object, so the second layer reuses the rotated+quantized activation (object identity + torch._version guarded; in-place mutation invalidates).
…GEMM Replaces the rank-one path's per-thread scattered 4B smem fragment loads (16 for A + 8 for B per K group) with one ldmatrix.x4 (A tile) and one ldmatrix.x2 (B chunk) per fragment. Both operands switch to the ldmatrix k-permutation together, so the MMA dot product — and bit-exactness vs the eager reference — is unchanged (full 33-case parity matrix passes). A 48-byte smem row stride staggers ldmatrix's 8 row addresses per phase across distinct bank groups (32B stride aliases rows r/r+4); staging pads OOB rows with zeros so the fragment-load guards drop out. L4 (SM 8.9) per-layer, M=4608, vs the compiled int8 path: 3072x3072: 0.54ms vs 1.18ms (was 0.88) 3072->12288: 2.60ms vs 4.54ms (was 2.90) 12288->3072: 2.30ms vs 4.04ms (was 3.50) Non-rank-one instantiations keep the 32B stride and manual loads, byte for byte.
📝 WalkthroughWalkthroughThis PR adds INT4 tensor-wise quantization and dequantization across eager and CUDA backends, including ConvRot support, a CUDA fused quantization kernel, a rank-one SVDQuant W4A4 matmul path, a ChangesINT4 Quantization and Rank-One Matmul
Sequence Diagram(s)sequenceDiagram
participant Caller as linear/mm/addmm
participant Layout as TensorWiseINT4Layout handler
participant FastPath as int4_linear op
participant Fallback as dequantize + dense op
Caller->>Layout: call with quantized weight
alt orientation matches fast path
Layout->>FastPath: torch.ops.comfy_kitchen.int4_linear(...)
FastPath-->>Layout: output
else fallback
Layout->>Fallback: dequantize(weight) then dense op
Fallback-->>Layout: output
end
Layout-->>Caller: result
sequenceDiagram
participant Python as int4_linear
participant Bind as int4_tensorwise_quantize
participant Launcher as launch_int4_tensorwise_quantize_kernel
participant GEMM as svdquant_scaled_mm_w4a4
Python->>Bind: quantize activations (rotate optional)
Bind->>Launcher: launch kernel, dispatch by dtype/rotate
Launcher-->>Python: packed q, ascales
Python->>GEMM: scaled_mm_svdquant_w4a4(rank_one=True)
GEMM-->>Python: output tensor
Possibly related issues
🚥 Pre-merge checks | ✅ 2✅ Passed checks (2 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy_kitchen/backends/cuda/__init__.py`:
- Around line 1502-1543: The tensorwise int4 scale buffers are being allocated
from x.dtype but later consumed by svdquant_scaled_mm_w4a4 as out_dtype, so fix
the dtype mismatch in the rank-1 path around the q_x/ascales/wscales setup.
Update the quantization flow to either cast the activation and weight scale
inputs to out_dtype before building ascales and wscales, or explicitly validate
that x.dtype matches out_dtype and fail fast when it does not. Make sure the
change is applied in the same code path that calls scaled_mm_svdquant_w4a4 so
the CUDA launcher always reads both scale lanes with the expected type.
In `@comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu`:
- Around line 809-841: Remove the rank-one dispatch that instantiates
LAUNCH_GEMM_R1 with ST=4 in DISPATCH_R1, since svdquant_scaled_mm_w4a4_kernel
with that template combination exceeds the per-block shared memory limit. Keep
the other r1_config cases intact, and either route case 1 to a supported ST
value or rework the tile/config selection in LAUNCH_GEMM_R1 so the chosen
svdquant_scaled_mm_w4a4_kernel instantiation stays under 48 KiB.
In `@comfy_kitchen/tensor/int4.py`:
- Line 63: The INT4 fast-matmul gate is too low for the actual kernel path used
by int4_linear. Update MIN_SM_VERSION in the int4.py gate and the nearby
explanatory note so the advertised minimum matches scaled_mm_svdquant_w4a4’s
real requirement of SM 8.0; if there is no separate SM75 fallback in
int4_linear, raise the gate from 7.5 to 8.0 and keep the dispatch logic
consistent with the existing parity test behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 56af49be-4bd8-4915-8dff-35004c0bce37
📒 Files selected for processing (11)
comfy_kitchen/backends/cuda/CMakeLists.txtcomfy_kitchen/backends/cuda/__init__.pycomfy_kitchen/backends/cuda/dlpack_bindings.cppcomfy_kitchen/backends/cuda/ops/int4_tensorwise.cucomfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cucomfy_kitchen/backends/cuda/ops/svdquant_utils.cuhcomfy_kitchen/backends/eager/__init__.pycomfy_kitchen/backends/eager/quantization.pycomfy_kitchen/tensor/__init__.pycomfy_kitchen/tensor/int4.pytests/test_int4_tensorwise.py
| #define LAUNCH_GEMM_R1(OutType, MU, WM, ST) \ | ||
| do { \ | ||
| constexpr int kBlockM_R1 = (MU) * 16 * (WM); \ | ||
| const dim3 grid_r1((N + kBlockN - 1) / kBlockN, \ | ||
| (M + kBlockM_R1 - 1) / kBlockM_R1); \ | ||
| const dim3 block_r1((WM) * kWarpsN * 32); \ | ||
| svdquant_scaled_mm_w4a4_kernel<OutType, false, false, false, false, false, \ | ||
| true, (MU), (WM), (ST)> \ | ||
| <<<grid_r1, block_r1, 0, stream>>>( \ | ||
| reinterpret_cast<const int8_t*>(act), \ | ||
| reinterpret_cast<const int8_t*>(wgt), \ | ||
| reinterpret_cast<const OutType*>(ascales), \ | ||
| reinterpret_cast<const OutType*>(wscales), \ | ||
| reinterpret_cast<const OutType*>(lora_act_in), \ | ||
| reinterpret_cast<const OutType*>(lora_up), \ | ||
| reinterpret_cast<const OutType*>(bias), \ | ||
| reinterpret_cast<OutType*>(out), \ | ||
| M, N, K, R); \ | ||
| } while (0) | ||
| #define DISPATCH_R1(OutType) \ | ||
| do { \ | ||
| switch (r1_config) { \ | ||
| case 1: LAUNCH_GEMM_R1(OutType, 4, 2, 4); break; \ | ||
| case 2: LAUNCH_GEMM_R1(OutType, 2, 2, 3); break; \ | ||
| case 3: LAUNCH_GEMM_R1(OutType, 3, 2, 3); break; \ | ||
| default: LAUNCH_GEMM_R1(OutType, 4, 2, 3); break; \ | ||
| } \ | ||
| } while (0) | ||
| if (out_dtype_code == 2 /* bf16 */) { | ||
| DISPATCH_R1(__nv_bfloat16); | ||
| } else if (out_dtype_code == 1 /* fp16 */) { | ||
| DISPATCH_R1(__half); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Estimate static shared usage per R1 config (kBlockN=128, stride=48, sizeof(OutType)=2).
for cfg in "0 4 2 3" "1 4 2 4" "2 2 2 3" "3 3 2 3"; do
set -- $cfg; id=$1; MU=$2; WM=$3; ST=$4
kBlockM=$((MU*16*WM)); kBlockN=128; stride=48
A=$((ST*kBlockM*stride)); B=$((ST*kBlockN*stride)); WS=$((kBlockN*2))
echo "config $id: A=$A B=$B WS=$WS total=$((A+B+WS)) bytes (cap=49152)"
doneRepository: Comfy-Org/comfy-kitchen
Length of output: 412
🏁 Script executed:
#!/bin/bash
# Inspect the R1 kernel body and shared-memory declarations around the launch site.
sed -n '640,860p' comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu | cat -nRepository: Comfy-Org/comfy-kitchen
Length of output: 13346
🏁 Script executed:
#!/bin/bash
# Inspect the kernel's shared-memory declarations and launch shape.
rg -n "smem_(A|B|WS|LoraA|LoraB)|extern __shared__|__shared__" comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cuRepository: Comfy-Org/comfy-kitchen
Length of output: 1613
🌐 Web query:
CUDA per-block static shared memory limit sm_90 sm_100 48 KiB static shared memory opt-in
💡 Result:
In CUDA architectures including sm_90 (Hopper) and sm_100 (Blackwell), the limit for statically allocated shared memory remains capped at 48 KiB per thread block to maintain architectural compatibility [1][2][3]. While these architectures support much larger total shared memory capacities per streaming multiprocessor (SM)—for instance, 228 KiB on sm_90 and sm_100—kernels that require more than 48 KiB of shared memory per block must use dynamic shared memory rather than statically sized arrays [1][2][3]. To use shared memory beyond the 48 KiB limit, you must perform the following steps: 1. Declare the shared memory as an extern array (e.g., extern shared int my_array[];) [4]. 2. Specify the additional memory size at kernel launch in the third parameter of the execution configuration (<<<grid_dim, block_dim, dynamic_smem_bytes>>>) [4]. 3. Explicitly opt-in by setting the cudaFuncAttributeMaxDynamicSharedMemorySize attribute for the kernel using cudaFuncSetAttribute [2][5]: // Example opt-in for dynamic shared memory size_t max_bytes = 98304; // e.g., 96 KiB cudaFuncSetAttribute(kernel_function, cudaFuncAttributeMaxDynamicSharedMemorySize, max_bytes); This requirement ensures that developers explicitly manage and account for the larger shared memory allocations, as exceeding the 48 KiB threshold may impact kernel occupancy and scheduling [6][2].
Citations:
- 1: https://docs.nvidia.com/cuda/archive/13.0.3/hopper-tuning-guide/index.html
- 2: https://ai-infrastructure.net/shared-memory-tiling/
- 3: https://docs.nvidia.com/cuda/archive/12.6.3/hopper-tuning-guide/index.html
- 4: https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/cpp-language-extensions.html
- 5: https://forums.developer.nvidia.com/t/shared-memory-size-per-thread-block/74858
- 6: https://docs.nvidia.com/cuda/archive/13.1.0/cuda-programming-guide/02-basics/writing-cuda-kernels.html
🏁 Script executed:
#!/bin/bash
# Locate the kernel template parameters and confirm which shared arrays are instantiated in rank-one mode.
rg -n "template <|kFuseLora|kSharedScale|kRankOne|svdquant_scaled_mm_w4a4_kernel<" comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cuRepository: Comfy-Org/comfy-kitchen
Length of output: 1881
Drop the ST=4 rank-one dispatch — this instantiation uses 49,408 bytes of static shared memory, which is 256 bytes over the 48 KiB per-block cap. The other R1 configs stay under the limit, so this one needs to go or the tiles need to be reworked. Tiny tweak, big relief.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu` around lines 809
- 841, Remove the rank-one dispatch that instantiates LAUNCH_GEMM_R1 with ST=4
in DISPATCH_R1, since svdquant_scaled_mm_w4a4_kernel with that template
combination exceeds the per-block shared memory limit. Keep the other r1_config
cases intact, and either route case 1 to a supported ST value or rework the
tile/config selection in LAUNCH_GEMM_R1 so the chosen
svdquant_scaled_mm_w4a4_kernel instantiation stays under 48 KiB.
The rank-one GEMM launcher reinterprets ascales/wscales as out_dtype, so a bf16 activation with an fp16 out_dtype read scale bytes as the wrong type. Cast both buffers to out_dtype after quantization (the quantize kernel still writes ascales in x.dtype) and add a mismatched-dtype parity test. The fused CUDA path is built on cp.async, which Turing lacks — advertise MIN_SM (8, 0) instead of (7, 5) so SM 7.5 refuses the layout cleanly instead of failing at kernel launch.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
comfy_kitchen/tensor/int4.py (1)
267-273: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCheck the transpose flag before you dequantize — no need to un-pack a byte you're about to toss.
On the fallback path (
transposedis False),input_tensoris dequantized at Lines 267‑268, then the handler bails totorch.mm(*dequantize_args(args))at Line 273, which re-dequantizes the input from scratch. For a large quantized activation that's a wasted un-pack. Thelinearhandler already checks the orientation gate before touching the input; mirroring that order here (and inaddmmat Lines 289‑293) keeps things consistent and skips the throwaway work.♻️ Proposed reorder for the mm handler
- if isinstance(input_tensor, QuantizedTensor): - input_tensor = input_tensor.dequantize() - - if not getattr(weight._params, "transposed", False): - # Per-row scales belong to the rows of the logical RHS, not output - # columns, so a directly-quantized RHS cannot ride int4_linear. - return torch.mm(*dequantize_args(args), **dequantize_args(kwargs)) + if not getattr(weight._params, "transposed", False): + # Per-row scales belong to the rows of the logical RHS, not output + # columns, so a directly-quantized RHS cannot ride int4_linear. + return torch.mm(*dequantize_args(args), **dequantize_args(kwargs)) + + if isinstance(input_tensor, QuantizedTensor): + input_tensor = input_tensor.dequantize()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@comfy_kitchen/tensor/int4.py` around lines 267 - 273, The fallback branch in the int4 mm path dequantizes input_tensor before checking weight._params.transposed, but that work is discarded when the function returns torch.mm via dequantize_args. Reorder the logic in the mm handler so the transposed gate is checked first, and only dequantize input_tensor on the int4_linear path; apply the same ordering cleanup in the addmm handler for consistency. Reference the mm/addmm flow and the transposed check on weight._params to keep the fix easy to locate.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@comfy_kitchen/tensor/int4.py`:
- Around line 267-273: The fallback branch in the int4 mm path dequantizes
input_tensor before checking weight._params.transposed, but that work is
discarded when the function returns torch.mm via dequantize_args. Reorder the
logic in the mm handler so the transposed gate is checked first, and only
dequantize input_tensor on the int4_linear path; apply the same ordering cleanup
in the addmm handler for consistency. Reference the mm/addmm flow and the
transposed check on weight._params to keep the fix easy to locate.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 88db77c3-ec57-4076-82dc-880caa1a5353
📒 Files selected for processing (3)
comfy_kitchen/backends/cuda/__init__.pycomfy_kitchen/tensor/int4.pytests/test_int4_tensorwise.py
|
So cool ,thanks for the PR. Int4 is really needed on cards below 50s. I wonder if video model like ltx2.3 and bernini can work too. Can you provide a convert script too,like comfy-quants |
|
Too slow ,Krea-2-Turbo-int4-tensorwise-mixed int4 4.27s/it, but int8 1s/it. 4070 313tor212cu130 |
|
I will investigate. On the L4 test card I used, the operator performed
well, faster than int8. There might be something I missed that caused the
worse results on the consumer-grade graphics card.
zwukong ***@***.***> 于 2026年7月6日周一 21:28写道：
… *zwukong* left a comment (Comfy-Org/comfy-kitchen#63)
<#63 (comment)>
Too slow , int4 4.27s/it, but int8 1s/it. 4070 313tor212cu130
—
Reply to this email directly, view it on GitHub
<#63?email_source=notifications&email_token=CAONRYVCUP2AIRBQQBLGAUD5DOSPTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBZGMZTIOBZG42KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KYZTPN52GK4S7MNWGSY3L#issuecomment-4893348974>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/CAONRYRFUIH2VOS6UAIKMID5DOSPTAVCNFSNUABGKJSXA33TNF2G64TZHMYTCMJZGIYDKNZYGY5US43TOVSTWNBYGE2TENZRGU4THILWAI>
.
Triage notifications, keep track of coding agent tasks and review pull
requests on the go with GitHub Mobile for iOS
<https://github.com/notifications/mobile/ios/CAONRYTM2EG7BP7NVRLDI4T5DOSPTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBZGMZTIOBZG42KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KUZTPN52GK4S7NFXXG>
and Android
<https://github.com/notifications/mobile/android/CAONRYUJO2BI4J6B4OWXIWD5DOSPTA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBZGMZTIOBZG42KM4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2K4ZTPN52GK4S7MFXGI4TPNFSA>.
Download it today!
You are receiving this because you authored the thread.Message ID:
***@***.***>
|
… CUDA A missing compiled extension silently drops int4_tensorwise to the eager reference (~4x slower than the int8 kernels), which reads as 'int4 is slow' rather than 'the kernels are not installed' (reported on a 4070 in PR Comfy-Org#63: 4.27 s/it eager vs the expected ~1 s/it). Warn once with the actual reason: extension missing vs fused-path constraints unmet.
… CUDA A missing compiled extension silently drops int4_tensorwise to the eager reference (~4x slower than the int8 kernels), which reads as 'int4 is slow' rather than 'the kernels are not installed' (reported on a 4070 in PR Comfy-Org#63: 4.27 s/it eager vs the expected ~1 s/it). Warn once with the actual reason: extension missing vs fused-path constraints unmet.
f905170 to
be4a39e
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy_kitchen/backends/eager/quantization.py`:
- Around line 1328-1351: The eager CUDA warning in the int4_tensorwise
quantization path is describing the wrong fused-path gate, so update the warning
text in the quantization function that checks x.is_cuda and logs via
logging.warning to either stay generic or match the real kernel constraints
enforced by the CUDA launch path (min_compute_capability=(8, 0) and the actual K
divisibility requirement, including the ConvRot case). Keep the existing
kernel-detection logic around comfy_kitchen.backends.cuda._C and
_int4_eager_cuda_warned, but change the reason/message so it no longer mentions
SM >= 8.0, K % 64 == 0 unless that is truly the active precondition.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 930b6611-cfe7-481c-967f-3a412cda9f15
📒 Files selected for processing (1)
comfy_kitchen/backends/eager/quantization.py
| global _int4_eager_cuda_warned | ||
| if x.is_cuda and not _int4_eager_cuda_warned: | ||
| _int4_eager_cuda_warned = True | ||
| import logging | ||
| try: | ||
| from comfy_kitchen.backends.cuda import _C as _CUDA_EXT | ||
| _has_kernels = hasattr(_CUDA_EXT, "int4_tensorwise_quantize") | ||
| except Exception: | ||
| _has_kernels = False | ||
| if _has_kernels: | ||
| reason = ( | ||
| "the compiled kernels are present but this call does not meet the " | ||
| "fused path's constraints (SM >= 8.0, K % 64 == 0)" | ||
| ) | ||
| else: | ||
| reason = ( | ||
| "the compiled comfy_kitchen CUDA extension with the int4 kernels was " | ||
| "not found — build it from this checkout " | ||
| "(python setup.py build_ext --inplace) or install a wheel that ships it" | ||
| ) | ||
| logging.warning( | ||
| "int4_tensorwise is running on the SLOW eager reference path on a CUDA " | ||
| "device (roughly 4x slower than the int8 kernels): %s.", reason | ||
| ) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm the fused int4 path constraints referenced by the eager warning.
rg -nP -C4 '(K\s*%\s*64|% 64|sm_?\d|getDeviceCapability|device_capability|8\.0|>= *80|major\s*>=\s*8)' \
comfy_kitchen/backends/cuda/ops/int4_tensorwise.cu \
comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cu \
comfy_kitchen/backends/cuda/__init__.pyRepository: Comfy-Org/comfy-kitchen
Length of output: 4968
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the eager warning site and the CUDA int4 dispatch/gating logic.
sed -n '1290,1375p' comfy_kitchen/backends/eager/quantization.py
printf '\n--- CUDA init (int4-related) ---\n'
rg -n -C 4 'int4_tensorwise|scaled_mm_svdquant_w4a4|K % 64|major >= 8|device_capability|get_device_capability|supports.*int4|register.*int4|dispatch.*int4' \
comfy_kitchen/backends/cuda/__init__.py \
comfy_kitchen/backends/cuda/ops/int4_tensorwise.cu \
comfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cuRepository: Comfy-Org/comfy-kitchen
Length of output: 21207
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the CUDA-side constraint definitions around int4_tensorwise.
sed -n '1455,1545p' comfy_kitchen/backends/cuda/__init__.py
printf '\n--- constraints block ---\n'
sed -n '1908,1938p' comfy_kitchen/backends/cuda/__init__.py
printf '\n--- int4 kernel launch checks ---\n'
sed -n '251,310p' comfy_kitchen/backends/cuda/ops/int4_tensorwise.cuRepository: Comfy-Org/comfy-kitchen
Length of output: 8225
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the exact constraint block for int4_linear and the nearby fused path rules.
sed -n '1918,1955p' comfy_kitchen/backends/cuda/__init__.py
printf '\n--- eager fallback warning site ---\n'
sed -n '1328,1355p' comfy_kitchen/backends/eager/quantization.pyRepository: Comfy-Org/comfy-kitchen
Length of output: 3481
The eager warning names the wrong gate. int4_tensorwise only checks min_compute_capability=(8, 0) here, and the CUDA launch path actually enforces K % 128 == 0 (or K % 256 == 0 with ConvRot). Hard-coding SM >= 8.0, K % 64 == 0 can send users on a wild goose chase; keep the warning generic or report the real kernel preconditions.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy_kitchen/backends/eager/quantization.py` around lines 1328 - 1351, The
eager CUDA warning in the int4_tensorwise quantization path is describing the
wrong fused-path gate, so update the warning text in the quantization function
that checks x.is_cuda and logs via logging.warning to either stay generic or
match the real kernel constraints enforced by the CUDA launch path
(min_compute_capability=(8, 0) and the actual K divisibility requirement,
including the ConvRot case). Keep the existing kernel-detection logic around
comfy_kitchen.backends.cuda._C and _int4_eager_cuda_warned, but change the
reason/message so it no longer mentions SM >= 8.0, K % 64 == 0 unless that is
truly the active precondition.
Could you provide the logs? This speed difference seems to be caused by incorrect compilation of CUDA C files, resulting in a rollback to the eager path. |
|
I did not see any difference logs , Just copy your comfy-kitchen to replace the old one in site-packages, and use your forked comfyui. |
You need to recompile kitchen to generate the _c kernel file to utilize the fast int4 kernel; otherwise, it will fall back to the slow path without operator acceleration. |
|
Recompiling the C file using the latest commits and applying the accelerated kernel, hybrid quantization can reduce the size by 60% and achieve the same or even slightly faster speeds as pure int8. |
|
I see， but how to generate int4 models, if it is easy then the method is very useful, or not. Nunchaku is dead because of too complicated to adapt to new model. Video model is the more sensitive to 4 bit, maybe you should try bernini or ltx2.3. Nunchaku failed to generate wan2.2, which i think is the most reason to die |
My comfy quants repository provides scripts for converting other models into convolution int4 mix models. It uses heuristic search to degrade and rank sensitive layers, theoretically applicable to most models. Alternatively, you can directly use my pre-configured quantization profiles for specific models. Additionally, if you require more model quantization support for Nunchaku Int4, you can contact me or my GitHub account @Yidhar. A potential issue with Nunchaku Int4 is that my rewritten CUDA operators don't achieve the same speed as the original versions. |
|
Sounds good. I really hope int4 can work as easily as int8 covrot does. Really easy to convert to int8 covrot |
2 participants
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_linearreference path, exact pack/unpack roundtrip, works on CPU.CUDA (SM >= 8.0 for the MMA path, layout MIN_SM (7,5)):
cp.async+mma.m16n8k64.s4with carried int32 accumulators (|acc| <= 49·K << 2^31) — no in-loop CUDA-core dequantldmatrix.x4/x2fragment loads over a 48-byte-stride smem layout (bank-conflict-free); both operands share the ldmatrix k-permutation so the dot product is unchanged_versionguarded, inference-mode safe)Per-layer, L4 (SM 8.9, native int4 tensor cores), M=4608, vs the compiled int8 path:
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-loaderwith this branch