[PR #90] [MERGED] Add optimized w4a8 with int8 codebook #136

Closed
opened 2026-10-26 22:20:44 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/90
Author: @kijai
Created: 8/1/2026
Status: ✅ Merged
Merged: 8/6/2026
Merged by: @comfyanonymous

Base: main ← Head: w4a8-int8-codebook


📝 Commits (10+)

  • 1c50520 a4a8
  • 9e8c323 w4a8
  • 543a89d triton/eager
  • 69620df Fixes
  • d0d968e validate quant contracts + fix transpose fallback orientation
  • eaff9a7 fixes
  • f9d2565 harden W4A8: capsule reuse, CUDA/ROCm detection, binding + addmm/scale contracts
  • 7440c2e validate int8 scale/bias lengths + exact output dtype
  • 3580635 reject unsupported out_dtype_code before dtype match
  • d908d58 cast fallback weight to x.dtype in CUDA path

📊 Changes

12 files changed (+1744 additions, -4 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+8 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+354 -4)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+261 -0)
➕ comfy_kitchen/backends/cuda/ops/w4a8_gemm.cu (+165 -0)
📝 comfy_kitchen/backends/eager/__init__.py (+56 -0)
📝 comfy_kitchen/backends/eager/quantization.py (+6 -0)
➕ comfy_kitchen/backends/eager/w4a8_int8.py (+369 -0)
📝 comfy_kitchen/backends/triton/__init__.py (+18 -0)
➕ comfy_kitchen/backends/triton/w4a8_int8.py (+143 -0)
📝 comfy_kitchen/tensor/__init__.py (+11 -0)
➕ comfy_kitchen/tensor/w4a8_int8.py (+352 -0)

📄 Description

Calibration-free 4-bit weight storage that runs on int8 gemm with convrot activations.

AsymW4A8Int8Layout: ConvRot-rotated int4 weights + a per-tensor Lloyd-Max
codebook + fp8 group scales, dequantized to grouped int8 and fed to the int8
CUTLASS GEMM. ~0.073 weight relL2 on real DiT weights (NVFP4 ~0.094), ~0.56
B/elem, ~1.09× int8 speed, no calibration.

Backends:

  • CUDA: chunked fused int4→int8 dequant + strided INT8 GEMM.
  • Triton + pure-torch eager (AMD/CPU): fused int4→int8 dequant, bit-exact with the CUDA kernel

Core PR:

https://github.com/Comfy-Org/ComfyUI/pull/15308

Test model:

https://huggingface.co/Kijai/MiniMax-H3-experimental/blob/main/minimax_h3_fl2va_pruned_w4a8_mixed.safetensors

Example result on a 5090 (nvfp4 is not accelerated here):

https://github.com/user-attachments/assets/3ab99015-300a-4ef2-9051-2701347afac5


🔄 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/90 **Author:** [@kijai](https://github.com/kijai) **Created:** 8/1/2026 **Status:** ✅ Merged **Merged:** 8/6/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `w4a8-int8-codebook` --- ### 📝 Commits (10+) - [`1c50520`](https://github.com/Comfy-Org/comfy-kitchen/commit/1c5052073bd0fd6879acff6aa7e031dff3bdbdae) a4a8 - [`9e8c323`](https://github.com/Comfy-Org/comfy-kitchen/commit/9e8c323f8fb3dcfdc742c58034262d62ca4f9ca2) w4a8 - [`543a89d`](https://github.com/Comfy-Org/comfy-kitchen/commit/543a89d3f21fa346f194e57265f794a3fefceb9c) triton/eager - [`69620df`](https://github.com/Comfy-Org/comfy-kitchen/commit/69620dfbd717956fb51819dc4e8b2f1e8e7ae06b) Fixes - [`d0d968e`](https://github.com/Comfy-Org/comfy-kitchen/commit/d0d968ed2d72a0e6ffe295dd25a0be2d28185166) validate quant contracts + fix transpose fallback orientation - [`eaff9a7`](https://github.com/Comfy-Org/comfy-kitchen/commit/eaff9a7a9f6c6eb4485be697670831ed8ea05e93) fixes - [`f9d2565`](https://github.com/Comfy-Org/comfy-kitchen/commit/f9d25653edd6ca8f686ea352f55873982758d725) harden W4A8: capsule reuse, CUDA/ROCm detection, binding + addmm/scale contracts - [`7440c2e`](https://github.com/Comfy-Org/comfy-kitchen/commit/7440c2eccbe1feec21fe1596f861e0148e75f00d) validate int8 scale/bias lengths + exact output dtype - [`3580635`](https://github.com/Comfy-Org/comfy-kitchen/commit/35806350707886ab2bac97a6edef9cac65748529) reject unsupported out_dtype_code before dtype match - [`d908d58`](https://github.com/Comfy-Org/comfy-kitchen/commit/d908d589fbda44595f6d78d9c47ef42ba4aa8fad) cast fallback weight to x.dtype in CUDA path ### 📊 Changes **12 files changed** (+1744 additions, -4 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+8 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+354 -4) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+261 -0) ➕ `comfy_kitchen/backends/cuda/ops/w4a8_gemm.cu` (+165 -0) 📝 `comfy_kitchen/backends/eager/__init__.py` (+56 -0) 📝 `comfy_kitchen/backends/eager/quantization.py` (+6 -0) ➕ `comfy_kitchen/backends/eager/w4a8_int8.py` (+369 -0) 📝 `comfy_kitchen/backends/triton/__init__.py` (+18 -0) ➕ `comfy_kitchen/backends/triton/w4a8_int8.py` (+143 -0) 📝 `comfy_kitchen/tensor/__init__.py` (+11 -0) ➕ `comfy_kitchen/tensor/w4a8_int8.py` (+352 -0) </details> ### 📄 Description Calibration-free 4-bit weight storage that runs on int8 gemm with convrot activations. `AsymW4A8Int8Layout`: ConvRot-rotated int4 weights + a per-tensor Lloyd-Max codebook + fp8 group scales, dequantized to grouped int8 and fed to the int8 CUTLASS GEMM. ~0.073 weight relL2 on real DiT weights (NVFP4 ~0.094), ~0.56 B/elem, ~1.09× int8 speed, no calibration. Backends: - CUDA: chunked fused int4→int8 dequant + strided INT8 GEMM. - Triton + pure-torch eager (AMD/CPU): fused int4→int8 dequant, bit-exact with the CUDA kernel Core PR: https://github.com/Comfy-Org/ComfyUI/pull/15308 Test model: https://huggingface.co/Kijai/MiniMax-H3-experimental/blob/main/minimax_h3_fl2va_pruned_w4a8_mixed.safetensors Example result on a 5090 (nvfp4 is not accelerated here): https://github.com/user-attachments/assets/3ab99015-300a-4ef2-9051-2701347afac5 --- <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:44 +00:00
Sign in to join this conversation.
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#136
No description provided.