Basic convrot int4 support.#65
Conversation
Should work on 30 and 40 series. 50 series doesn't have int4 ops so it will use int8 ops. 20 series needs special code for int4 and should fallback to int8 at the moment but is untested.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds ConvRot W4A4 int4 quantization, dequantization, and linear support across eager PyTorch, CUDA, and tensor-layout backends. It also extends public exports, backend bindings, CUTLASS strided dequant routing, and validation coverage. ChangesConvRot W4A4 Feature
🚥 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: 2
🤖 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/convrot_w4a4.py`:
- Around line 156-171: The int4 accumulation in int4_linear currently happens in
the low-precision out_dtype because x_qdata and weight are cast before the
matmul, which can lose precision for large K and diverge from the CUDA path.
Update int4_linear so the unpacked operands are accumulated in fp32 during x_int
@ w_int.t(), then apply x_scale, weight_scale, and bias, and only cast the final
tensor to out_dtype at the end; use int4_linear and convrot_w4a4_linear as the
main points to locate the change.
- Around line 24-25: The power-of-4 validation in the regular Hadamard size
check relies on math.log in the convrot_w4a4 module, which can fail for large
valid powers due to float precision. Update the size validation logic in the
Regular Hadamard check to use an integer-only test instead of math.log(size, 4)
% 1, ideally with a bit-length/parity-based power-of-4 check, while keeping the
existing ValueError path and message in the same validation block.
🪄 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: 1834f6a8-172e-46f7-b57c-dc05638ce144
📒 Files selected for processing (10)
comfy_kitchen/__init__.pycomfy_kitchen/backends/cuda/CMakeLists.txtcomfy_kitchen/backends/cuda/__init__.pycomfy_kitchen/backends/cuda/dlpack_bindings.cppcomfy_kitchen/backends/cuda/ops/convrot_w4a4.cucomfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cucomfy_kitchen/backends/eager/__init__.pycomfy_kitchen/backends/eager/convrot_w4a4.pycomfy_kitchen/tensor/__init__.pycomfy_kitchen/tensor/convrot_w4a4.py
| if size < 4 or (size & (size - 1)) != 0 or math.log(size, 4) % 1 != 0: | ||
| raise ValueError(f"Regular Hadamard size must be a power of 4, got {size}") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Demonstrate the float fragility across powers of 4 and confirm bit_length parity is exact.
python - <<'PY'
import math
for e in range(2, 40, 2): # even exponents -> powers of 4
size = 1 << e
print(size,
"log%1=", math.log(size, 4) % 1,
"bitlen_ok=", (size.bit_length() - 1) % 2 == 0)
PYRepository: Comfy-Org/comfy-kitchen
Length of output: 818
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the surrounding implementation and any size constraints.
git ls-files comfy_kitchen/backends/eager/convrot_w4a4.py
wc -l comfy_kitchen/backends/eager/convrot_w4a4.py
sed -n '1,120p' comfy_kitchen/backends/eager/convrot_w4a4.py
# Search for callers/constraints on the Hadamard size.
rg -n "Regular Hadamard size|convrot_w4a4|Hadamard|power of 4|size < 4" comfy_kitchen -SRepository: Comfy-Org/comfy-kitchen
Length of output: 11072
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the other implementation that uses the same guard and the nearby constraint docs.
wc -l comfy_kitchen/tensor/int8_utils.py comfy_kitchen/backends/eager/__init__.py comfy_kitchen/backends/cuda/__init__.py
sed -n '1,80p' comfy_kitchen/tensor/int8_utils.py
sed -n '210,260p' comfy_kitchen/backends/eager/__init__.py
sed -n '2230,2275p' comfy_kitchen/backends/cuda/__init__.py
# Find where the group size is documented as a constraint on the public APIs.
rg -n "convrot_groupsize|group size|power of 4|Regular orthogonal Hadamard" comfy_kitchen/tensor comfy_kitchen/backends -SRepository: Comfy-Org/comfy-kitchen
Length of output: 19922
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Compare the current float-based guard against an exact integer criterion
# over a wide range of powers of 4 and nearby non-powers.
python3 - <<'PY'
import math
def float_guard(size: int) -> bool:
return size < 4 or (size & (size - 1)) != 0 or math.log(size, 4) % 1 != 0
def exact_guard(size: int) -> bool:
return size < 4 or (size & (size - 1)) != 0 or ((size.bit_length() - 1) % 2 != 0)
mismatches = []
for e in range(2, 4000): # plenty of powers of 4 and very large ints
size = 1 << e
if float_guard(size) != exact_guard(size):
mismatches.append(("pow2", e, size, math.log(size, 4) % 1, float_guard(size), exact_guard(size)))
break
if e % 2 == 0 and e > 2:
# test immediate neighbors too
for delta in (-1, 1):
n = size + delta
if float_guard(n) != exact_guard(n):
mismatches.append(("neighbor", e, n, math.log(n, 4) % 1 if n > 0 else None, float_guard(n), exact_guard(n)))
break
if mismatches:
break
print("mismatches:", mismatches[:3] if mismatches else "none")
PYRepository: Comfy-Org/comfy-kitchen
Length of output: 244
Use an integer-only power-of-4 check here. math.log(size, 4) % 1 can reject valid large powers of 4 once float precision slips; for example, 2^58 = 4^29 trips the test. A bit-length parity check is exact and keeps the Hadamard power on the rails, no log-jam needed.
🤖 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/convrot_w4a4.py` around lines 24 - 25, The
power-of-4 validation in the regular Hadamard size check relies on math.log in
the convrot_w4a4 module, which can fail for large valid powers due to float
precision. Update the size validation logic in the Regular Hadamard check to use
an integer-only test instead of math.log(size, 4) % 1, ideally with a
bit-length/parity-based power-of-4 check, while keeping the existing ValueError
path and message in the same validation block.
| def int4_linear( | ||
| x_qdata: torch.Tensor, | ||
| weight: torch.Tensor, | ||
| x_scale: torch.Tensor, | ||
| weight_scale: torch.Tensor, | ||
| bias: torch.Tensor | None, | ||
| out_dtype: torch.dtype, | ||
| ) -> torch.Tensor: | ||
| x_int = _unpack_int4_row_major(x_qdata).to(dtype=out_dtype) | ||
| w_int = _unpack_int4_row_major(weight).to(dtype=out_dtype) | ||
| out = x_int @ w_int.t() | ||
| out = out * x_scale.reshape(-1, 1).to(device=out.device, dtype=out_dtype) | ||
| out = out * weight_scale.reshape(1, -1).to(device=out.device, dtype=out_dtype) | ||
| if bias is not None: | ||
| out = out + bias.to(device=out.device, dtype=out_dtype).reshape(1, -1) | ||
| return out |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
INT4 GEMM accumulates in the low‑precision output dtype — expect precision to depart in a hurry.
out_dtype is passed as x.dtype from convrot_w4a4_linear (Line 195), which for real models is typically bfloat16/float16. Casting the unpacked int4 operands to that dtype before x_int @ w_int.t() means the accumulation runs in bf16/fp16. bf16 represents integers exactly only up to 256, so for any K >= 256 (and convrot_groupsize defaults to 256) the summed products silently lose bits, producing wrong reference results that won't match the int32-accumulating CUDA int4 MMA.
Accumulate in fp32 and cast only the final result:
🐛 Proposed fix: fp32 accumulation, cast at the end
- x_int = _unpack_int4_row_major(x_qdata).to(dtype=out_dtype)
- w_int = _unpack_int4_row_major(weight).to(dtype=out_dtype)
- out = x_int @ w_int.t()
- out = out * x_scale.reshape(-1, 1).to(device=out.device, dtype=out_dtype)
- out = out * weight_scale.reshape(1, -1).to(device=out.device, dtype=out_dtype)
- if bias is not None:
- out = out + bias.to(device=out.device, dtype=out_dtype).reshape(1, -1)
- return out
+ x_int = _unpack_int4_row_major(x_qdata).to(torch.float32)
+ w_int = _unpack_int4_row_major(weight).to(torch.float32)
+ out = x_int @ w_int.t()
+ out = out * x_scale.reshape(-1, 1).to(device=out.device, dtype=torch.float32)
+ out = out * weight_scale.reshape(1, -1).to(device=out.device, dtype=torch.float32)
+ if bias is not None:
+ out = out + bias.to(device=out.device, dtype=torch.float32).reshape(1, -1)
+ return out.to(out_dtype)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def int4_linear( | |
| x_qdata: torch.Tensor, | |
| weight: torch.Tensor, | |
| x_scale: torch.Tensor, | |
| weight_scale: torch.Tensor, | |
| bias: torch.Tensor | None, | |
| out_dtype: torch.dtype, | |
| ) -> torch.Tensor: | |
| x_int = _unpack_int4_row_major(x_qdata).to(dtype=out_dtype) | |
| w_int = _unpack_int4_row_major(weight).to(dtype=out_dtype) | |
| out = x_int @ w_int.t() | |
| out = out * x_scale.reshape(-1, 1).to(device=out.device, dtype=out_dtype) | |
| out = out * weight_scale.reshape(1, -1).to(device=out.device, dtype=out_dtype) | |
| if bias is not None: | |
| out = out + bias.to(device=out.device, dtype=out_dtype).reshape(1, -1) | |
| return out | |
| def int4_linear( | |
| x_qdata: torch.Tensor, | |
| weight: torch.Tensor, | |
| x_scale: torch.Tensor, | |
| weight_scale: torch.Tensor, | |
| bias: torch.Tensor | None, | |
| out_dtype: torch.dtype, | |
| ) -> torch.Tensor: | |
| x_int = _unpack_int4_row_major(x_qdata).to(torch.float32) | |
| w_int = _unpack_int4_row_major(weight).to(torch.float32) | |
| out = x_int @ w_int.t() | |
| out = out * x_scale.reshape(-1, 1).to(device=out.device, dtype=torch.float32) | |
| out = out * weight_scale.reshape(1, -1).to(device=out.device, dtype=torch.float32) | |
| if bias is not None: | |
| out = out + bias.to(device=out.device, dtype=torch.float32).reshape(1, -1) | |
| return out.to(out_dtype) |
🤖 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/convrot_w4a4.py` around lines 156 - 171, The
int4 accumulation in int4_linear currently happens in the low-precision
out_dtype because x_qdata and weight are cast before the matmul, which can lose
precision for large K and diverge from the CUDA path. Update int4_linear so the
unpacked operands are accumulated in fp32 during x_int @ w_int.t(), then apply
x_scale, weight_scale, and bias, and only cast the final tensor to out_dtype at
the end; use int4_linear and convrot_w4a4_linear as the main points to locate
the change.
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/cuda/__init__.py`:
- Around line 640-645: The helper _int4_int8_weight_chunk_cols has a redundant m
<= 128 branch because both paths return the same chunk width. Simplify the logic
by collapsing that branch in _int4_int8_weight_chunk_cols so it only keeps the n
<= 2560 special case and then falls back to the shared min(n,
_INT4_INT8_WEIGHT_CHUNK_N) result, unless you intend to introduce a distinct
m-dependent width later.
🪄 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: 51443c89-87a1-4c2d-aa77-60d2826d7501
📒 Files selected for processing (8)
comfy_kitchen/backends/cuda/__init__.pycomfy_kitchen/backends/cuda/dlpack_bindings.cppcomfy_kitchen/backends/cuda/ops/convrot_w4a4.cucomfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cucomfy_kitchen/backends/eager/__init__.pycomfy_kitchen/backends/eager/convrot_w4a4.pycomfy_kitchen/tensor/convrot_w4a4.pytests/test_convrot_w4a4.py
| def _int4_int8_weight_chunk_cols(m: int, n: int) -> int: | ||
| if n <= 2560: | ||
| return n | ||
| if m <= 128: | ||
| return min(n, _INT4_INT8_WEIGHT_CHUNK_N) | ||
| return min(n, _INT4_INT8_WEIGHT_CHUNK_N) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm no other callers rely on m-based branching and check for an intended large-M chunk constant.
rg -nP '_int4_int8_weight_chunk_cols|_INT4_INT8_WEIGHT_CHUNK' comfy_kitchen -C2Repository: Comfy-Org/comfy-kitchen
Length of output: 1766
🏁 Script executed:
#!/bin/bash
sed -n '630,705p' comfy_kitchen/backends/cuda/__init__.py
printf '\n----\n'
sed -n '780,1010p' comfy_kitchen/backends/cuda/__init__.pyRepository: Comfy-Org/comfy-kitchen
Length of output: 13655
Collapse the redundant m <= 128 branch
_int4_int8_weight_chunk_cols returns min(n, _INT4_INT8_WEIGHT_CHUNK_N) in both m paths, and there’s no separate large-M chunk width defined. Collapse the branch, or wire in a real m-dependent width if that was the plan — this split is all sizzle, no steak.
♻️ Collapse the redundant branch
def _int4_int8_weight_chunk_cols(m: int, n: int) -> int:
if n <= 2560:
return n
- if m <= 128:
- return min(n, _INT4_INT8_WEIGHT_CHUNK_N)
return min(n, _INT4_INT8_WEIGHT_CHUNK_N)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _int4_int8_weight_chunk_cols(m: int, n: int) -> int: | |
| if n <= 2560: | |
| return n | |
| if m <= 128: | |
| return min(n, _INT4_INT8_WEIGHT_CHUNK_N) | |
| return min(n, _INT4_INT8_WEIGHT_CHUNK_N) | |
| def _int4_int8_weight_chunk_cols(m: int, n: int) -> int: | |
| if n <= 2560: | |
| return n | |
| return min(n, _INT4_INT8_WEIGHT_CHUNK_N) |
🤖 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/__init__.py` around lines 640 - 645, The helper
_int4_int8_weight_chunk_cols has a redundant m <= 128 branch because both paths
return the same chunk width. Simplify the logic by collapsing that branch in
_int4_int8_weight_chunk_cols so it only keeps the n <= 2560 special case and
then falls back to the shared min(n, _INT4_INT8_WEIGHT_CHUNK_N) result, unless
you intend to introduce a distinct m-dependent width later.
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 `@tests/test_convrot_w4a4.py`:
- Around line 184-185: Add test coverage in the convrot W4A4 tests for automatic
INT4→INT8 fallback routing, since the current parametrization only verifies
explicit INT8 execution via linear_dtype in test_convrot_w4a4. Introduce a case
that requests INT4 while mocking native INT4 support as unavailable, then assert
the result matches the INT8 fallback path; make sure this applies to the
relevant parametrized sections around convrot_groupsize and the other referenced
test blocks so the fallback behavior is exercised for the targeted GPU routing.
🪄 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: a19e42f2-576f-4843-be03-ed1d6c84c397
📒 Files selected for processing (2)
comfy_kitchen/backends/cuda/__init__.pytests/test_convrot_w4a4.py
| @pytest.mark.parametrize("linear_dtype", ["int4", "int8"]) | ||
| @pytest.mark.parametrize("convrot_groupsize", [16, 64, 256]) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Add coverage for automatic INT4→INT8 fallback routing.
Selecting linear_dtype="int8" only tests explicit INT8 execution. Add a test that requests INT4 while mocking native-INT4 support as unavailable, then verifies the INT8 fallback result. Otherwise, the architecture fallback is four-bit folklore rather than tested fact.
As per PR objectives, 20- and 50-series GPUs must route INT4 requests through INT8 fallback.
Also applies to: 215-216, 245-246
🤖 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 `@tests/test_convrot_w4a4.py` around lines 184 - 185, Add test coverage in the
convrot W4A4 tests for automatic INT4→INT8 fallback routing, since the current
parametrization only verifies explicit INT8 execution via linear_dtype in
test_convrot_w4a4. Introduce a case that requests INT4 while mocking native INT4
support as unavailable, then assert the result matches the INT8 fallback path;
make sure this applies to the relevant parametrized sections around
convrot_groupsize and the other referenced test blocks so the fallback behavior
is exercised for the targeted GPU routing.
1 participant
Should work on 30 and 40 series. 50 series doesn't have int4 ops so it will use int8 ops. 20 series needs special code for int4 and should fallback to int8 at the moment but is untested.