Try to optimize sdvquant implementation a bit.#50
Open
comfyanonymous wants to merge 1 commit into
Open
Conversation
📝 WalkthroughWalkthroughThis PR adds an AWQ W4A16 weight-dequantization CUDA kernel with Nanobind bindings, env-gated Python caching infrastructure, and a direct-CUDA dispatch fast path. It also lowers the AWQ GEMV/GEMM threshold from 8 to 4, parameterizes the SVDQuant W4A4 GEMM kernel by warp-grid shape, and adds a transposed LoRA-down path and a grouped-fuse inference fast path in the Python tensor layer. ChangesAWQ W4A16 Dequantization Kernel and Routing
SVDQuant W4A4 Warp-Grid GEMM, LoRA-Down, and Grouped-Fuse Fast Path
Sequence DiagramssequenceDiagram
rect rgba(100, 180, 255, 0.5)
note over _awq_forward,gemv_awq_w4a16: AWQ W4A16 dispatch (direct-CUDA path)
end
participant _awq_forward
participant _direct_cuda_backend
participant gemv_awq_w4a16 as cuda_backend.gemv_awq_w4a16
participant _awq_w4a16_cached_matmul_with_optional_bias
participant _awq_w4a16_dequant_weight
participant _C_awq_w4a16 as _C.awq_w4a16 (fused small-M)
_awq_forward->>_direct_cuda_backend: input_tensor, qdata, wscales, wzeros
_direct_cuda_backend-->>_awq_forward: cuda_backend or None
_awq_forward->>gemv_awq_w4a16: x, qweight, wscales, wzeros, bias, group_size
gemv_awq_w4a16->>_awq_w4a16_cached_matmul_with_optional_bias: try cache hit
alt cache hit
_awq_w4a16_cached_matmul_with_optional_bias-->>gemv_awq_w4a16: out, bias_applied=True
else small-m dequant path
gemv_awq_w4a16->>_awq_w4a16_dequant_weight: qweight → fp16/bf16 weight (cached or temp)
_awq_w4a16_dequant_weight-->>gemv_awq_w4a16: weight
gemv_awq_w4a16-->>gemv_awq_w4a16: matmul → out, bias_applied
else fused small-M kernel
gemv_awq_w4a16->>_C_awq_w4a16: x2d, qweight, wscales, wzeros, group_size
_C_awq_w4a16-->>gemv_awq_w4a16: out, bias_applied=False
end
gemv_awq_w4a16-->>_awq_forward: final output (+ bias if not applied)
sequenceDiagram
rect rgba(180, 130, 255, 0.5)
note over svdquant_w4a4_grouped_linear,_w4a4_forward_prepared: SVDQuant W4A4 grouped-fuse fast path
end
participant svdquant_w4a4_grouped_linear
participant _svdquant_cached_fused_group
participant _w4a4_forward_prepared
participant quantize_svdquant_w4a4
participant scaled_mm_svdquant_w4a4
participant _svdquant_split_fused_output
svdquant_w4a4_grouped_linear->>_svdquant_cached_fused_group: weights list + input m
_svdquant_cached_fused_group-->>svdquant_w4a4_grouped_linear: fused_qdata, fused_params, splits
svdquant_w4a4_grouped_linear->>_w4a4_forward_prepared: x, fused_qdata, fused_params, cuda_backend
_w4a4_forward_prepared->>quantize_svdquant_w4a4: x, smooth_factor, proj_down, pad_size
quantize_svdquant_w4a4-->>_w4a4_forward_prepared: q_act, ascales, lora_act
_w4a4_forward_prepared->>scaled_mm_svdquant_w4a4: q_act, fused_qdata, ascales, lora_act
scaled_mm_svdquant_w4a4-->>_w4a4_forward_prepared: fused_out
_w4a4_forward_prepared-->>svdquant_w4a4_grouped_linear: fused_out
svdquant_w4a4_grouped_linear->>_svdquant_split_fused_output: fused_out, splits
_svdquant_split_fused_output-->>svdquant_w4a4_grouped_linear: [out_0, out_1, ..., out_n]
🚥 Pre-merge checks | ✅ 2✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 4
🤖 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 981-1005: The issue is a race condition where the weight is cached
in qweight._ck_awq_dequant_cache immediately after enqueueing the async
_C.awq_dequant_w4a16 kernel, but the cache can be accessed by another CUDA
stream before the kernel finishes executing. To fix this, store a CUDA event
with the cache tuple when it is created (capturing the event after the async
kernel is enqueued), then modify the code that reads from this cache (in
_awq_w4a16_cached_dequant_weight) to synchronize on that event before using the
cached weight. Alternatively, key the cache by both the cache_key and the
producer stream to prevent cross-stream reuse until the first fill is complete.
In `@comfy_kitchen/backends/cuda/dlpack_bindings.cpp`:
- Around line 728-749: The awq_dequant_w4a16 function validates the data types
of wscales, wzeros, and weight tensors but fails to validate the dtype of
qweight, which should be byte-packed (int8 or uint8). Add a validation check
after the existing dtype validation using svdquant_dtype_code to ensure
qweight.dtype() is int8 or uint8, throwing a runtime_error if it is not. This
check must occur before launch_awq_dequant_w4a16_kernel is called to prevent the
CUDA kernel from misinterpreting non-byte-packed data as packed uint4 values.
In `@comfy_kitchen/backends/cuda/ops/awq_w4a16.cu`:
- Around line 615-625: The group index `g` is computed once from `k` but is used
for both the low nibble at position `k` and the high nibble at position `k + 1`.
When `group_size` is odd, the high nibble may belong to a different quantization
group but incorrectly uses the previous group's scale and zero values. Fix this
by either validating that `group_size` is even before launching all AWQ
byte-pair kernels, or by computing separate group indices for both nibble
positions (compute `g` for `k` and a separate group index for `k + 1`) in the
dequantization logic in the store_fp32_pair call.
In `@comfy_kitchen/tensor/svdquant_w4a4.py`:
- Around line 295-306: The grouped CUDA fuse path does not validate whether
shared quantization is possible before calling _svdquant_cached_fused_group(),
which internally assumes shared quantization. Add a validation check using
svdquant_w4a4_can_share_quant() on the weights parameter before entering the if
block that checks cuda_backend and _svdquant_grouped_fuse_enabled(). This
ensures that all weights have compatible quantization parameters (same
smooth_factor and proj_down) before attempting the fused operation, preventing
silent incorrect fusion of incompatible weights.
🪄 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: 133e9f5e-ea54-408c-a639-88c0b9cd570c
📒 Files selected for processing (8)
comfy_kitchen/backends/cuda/__init__.pycomfy_kitchen/backends/cuda/dlpack_bindings.cppcomfy_kitchen/backends/cuda/ops/awq_w4a16.cucomfy_kitchen/backends/cuda/ops/scaled_mm_svdquant_w4a4.cucomfy_kitchen/tensor/awq_w4a16.pycomfy_kitchen/tensor/svdquant_w4a4.pytests/test_awq_w4a16.pytests/test_svdquant_w4a4.py
Comment on lines
+981
to
+1005
| if hasattr(_C, "awq_dequant_w4a16"): | ||
| weight = torch.empty(n, k, dtype=compute_dtype, device=qweight.device) | ||
| stream_ptr = torch.cuda.current_stream(qweight.device).cuda_stream | ||
| _C.awq_dequant_w4a16( | ||
| _wrap_for_dlpack(qweight.view(torch.uint8)), | ||
| _wrap_for_dlpack(wscales), | ||
| _wrap_for_dlpack(wzeros), | ||
| _wrap_for_dlpack(weight), | ||
| group_size, | ||
| stream_ptr, | ||
| ) | ||
| else: | ||
| x16 = qweight.view(torch.uint8).to(torch.int16) | ||
| lo = (x16 & 0xF).to(torch.int8) | ||
| hi = ((x16 >> 4) & 0xF).to(torch.int8) | ||
| nibbles = torch.stack([lo, hi], dim=-1).reshape(n, k).to(compute_dtype) | ||
| weight = ( | ||
| (nibbles.view(n, k // group_size, group_size) - 8.0) | ||
| * wscales.t().unsqueeze(-1) | ||
| + wzeros.t().unsqueeze(-1) | ||
| ).view(n, k) | ||
|
|
||
| if cache: | ||
| weight = weight.contiguous() | ||
| qweight._ck_awq_dequant_cache = (cache_key, weight) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Do not publish the AWQ dequant cache before the producer stream is safe.
_C.awq_dequant_w4a16 enqueues an async kernel on stream_ptr, but line 1005 exposes weight via qweight._ck_awq_dequant_cache immediately. A second request on another CUDA stream can hit _awq_w4a16_cached_dequant_weight and enqueue mm reading weight before the dequant kernel finishes. Store a CUDA event with the cache and wait on reuse, or key the cache by producer stream until the first fill is complete.
🤖 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 981 - 1005, The issue
is a race condition where the weight is cached in qweight._ck_awq_dequant_cache
immediately after enqueueing the async _C.awq_dequant_w4a16 kernel, but the
cache can be accessed by another CUDA stream before the kernel finishes
executing. To fix this, store a CUDA event with the cache tuple when it is
created (capturing the event after the async kernel is enqueued), then modify
the code that reads from this cache (in _awq_w4a16_cached_dequant_weight) to
synchronize on that event before using the cached weight. Alternatively, key the
cache by both the cache_key and the producer stream to prevent cross-stream
reuse until the first fill is complete.
Comment on lines
+728
to
+749
| if (qweight.ndim() != 2 || wscales.ndim() != 2 || | ||
| wzeros.ndim() != 2 || weight.ndim() != 2) { | ||
| throw std::runtime_error("awq_dequant_w4a16: all tensors must be 2D"); | ||
| } | ||
| const int N = static_cast<int>(qweight.shape(0)); | ||
| const int K = static_cast<int>(qweight.shape(1)) * 2; | ||
| if (group_size <= 0 || K % group_size != 0) { | ||
| throw std::runtime_error("awq_dequant_w4a16: K must be divisible by group_size"); | ||
| } | ||
| const int dtype_code = svdquant_dtype_code(wscales.dtype()); | ||
| if (dtype_code != 1 && dtype_code != 2) { | ||
| throw std::runtime_error("awq_dequant_w4a16: only fp16 (1) and bf16 (2) outputs supported"); | ||
| } | ||
| if (svdquant_dtype_code(wzeros.dtype()) != dtype_code || | ||
| svdquant_dtype_code(weight.dtype()) != dtype_code) { | ||
| throw std::runtime_error("awq_dequant_w4a16: wscales, wzeros, and weight dtypes must match"); | ||
| } | ||
| if (wscales.shape(0) != K / group_size || wscales.shape(1) != N || | ||
| wzeros.shape(0) != K / group_size || wzeros.shape(1) != N || | ||
| weight.shape(0) != N || weight.shape(1) != K) { | ||
| throw std::runtime_error("awq_dequant_w4a16: expected qweight (N,K/2), scales/zeros (K/G,N), weight (N,K)"); | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate qweight is byte-packed before launching CUDA.
The wrapper validates the fp output tensors but never checks qweight.dtype(). A fp16/fp32 CUDA tensor with the expected 2D shape would pass these checks, then launch_awq_dequant_w4a16_kernel would interpret its raw bytes as packed uint4 and return garbage. Reject anything except int8/uint8 before the shape-only contract lets the nibble gremlin in.
🤖 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/dlpack_bindings.cpp` around lines 728 - 749, The
awq_dequant_w4a16 function validates the data types of wscales, wzeros, and
weight tensors but fails to validate the dtype of qweight, which should be
byte-packed (int8 or uint8). Add a validation check after the existing dtype
validation using svdquant_dtype_code to ensure qweight.dtype() is int8 or uint8,
throwing a runtime_error if it is not. This check must occur before
launch_awq_dequant_w4a16_kernel is called to prevent the CUDA kernel from
misinterpreting non-byte-packed data as packed uint4 values.
Comment on lines
+615
to
+625
| const int k = kh * 2; | ||
| const int g = k / G; | ||
| const int idx = n * K_half + kh; | ||
|
|
||
| const uint8_t byte = static_cast<uint8_t>(qweight[idx]); | ||
| const float scale = to_fp32(wscales[g * N + n]); | ||
| const float zero = to_fp32(wzeros[g * N + n]); | ||
| store_fp32_pair( | ||
| weight + n * K + k, | ||
| (float(byte & 0xF) - 8.0f) * scale + zero, | ||
| (float((byte >> 4) & 0xF) - 8.0f) * scale + zero); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Keep byte-paired nibbles inside the correct quantization group.
g is computed from k once and then used for both k and k + 1. Since the public validation only requires K % group_size == 0, an odd group_size lets the high nibble at a group boundary use the previous group's scale/zero, so CUDA dequant can silently diverge from the reference path. Either reject odd group_size before launching all AWQ byte-pair kernels, or compute separate group indices here. Tiny nibble, big quibble.
🐛 Possible kernel-side fix if odd groups should be supported
- const int g = k / G;
+ const int g_lo = k / G;
+ const int g_hi = (k + 1) / G;
const int idx = n * K_half + kh;
const uint8_t byte = static_cast<uint8_t>(qweight[idx]);
- const float scale = to_fp32(wscales[g * N + n]);
- const float zero = to_fp32(wzeros[g * N + n]);
+ const float scale_lo = to_fp32(wscales[g_lo * N + n]);
+ const float zero_lo = to_fp32(wzeros[g_lo * N + n]);
+ const float scale_hi = to_fp32(wscales[g_hi * N + n]);
+ const float zero_hi = to_fp32(wzeros[g_hi * N + n]);
store_fp32_pair(
weight + n * K + k,
- (float(byte & 0xF) - 8.0f) * scale + zero,
- (float((byte >> 4) & 0xF) - 8.0f) * scale + zero);
+ (float(byte & 0xF) - 8.0f) * scale_lo + zero_lo,
+ (float((byte >> 4) & 0xF) - 8.0f) * scale_hi + zero_hi);🤖 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/awq_w4a16.cu` around lines 615 - 625, The
group index `g` is computed once from `k` but is used for both the low nibble at
position `k` and the high nibble at position `k + 1`. When `group_size` is odd,
the high nibble may belong to a different quantization group but incorrectly
uses the previous group's scale and zero values. Fix this by either validating
that `group_size` is even before launching all AWQ byte-pair kernels, or by
computing separate group indices for both nibble positions (compute `g` for `k`
and a separate group index for `k + 1`) in the dequantization logic in the
store_fp32_pair call.
Comment on lines
+295
to
+306
| if cuda_backend is not None and len(weights) > 1 and _svdquant_grouped_fuse_enabled(): | ||
| fused_weight, fused_bias, splits = _svdquant_cached_fused_group(weights, biases) | ||
| out = _w4a4_forward_prepared( | ||
| fused_weight, | ||
| fused_bias, | ||
| orig_shape=orig_shape, | ||
| m=m, | ||
| x_main=x_main, | ||
| lora_x=lora_x, | ||
| cuda_backend=cuda_backend, | ||
| ) | ||
| return _svdquant_split_fused_output(out, splits) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Validate shared quantization before taking the grouped CUDA fuse path.
Line 295 enters the fast path before svdquant_w4a4_can_share_quant(...), then _svdquant_cached_fused_group() fuses with assume_shared_quant=True. That skips the caller’s validate_shared_quant/assume_shared_quant contract, so CUDA can silently fuse weights with different smooth_factor/proj_down using the first tensor’s params. Quite the sneaky gremlin: fast fuse, wrong muse.
🐛 Proposed fix
cuda_backend = _direct_cuda_backend(x_main, qdata)
- if cuda_backend is not None and len(weights) > 1 and _svdquant_grouped_fuse_enabled():
- fused_weight, fused_bias, splits = _svdquant_cached_fused_group(weights, biases)
- out = _w4a4_forward_prepared(
- fused_weight,
- fused_bias,
- orig_shape=orig_shape,
- m=m,
- x_main=x_main,
- lora_x=lora_x,
- cuda_backend=cuda_backend,
- )
- return _svdquant_split_fused_output(out, splits)
-
if not svdquant_w4a4_can_share_quant(
weights, validate=validate_shared_quant, trust=assume_shared_quant,
):
raise ValueError("SVDQuant weights do not share quantization parameters")
+
+ if cuda_backend is not None and len(weights) > 1 and _svdquant_grouped_fuse_enabled():
+ fused_weight, fused_bias, splits = _svdquant_cached_fused_group(weights, biases)
+ out = _w4a4_forward_prepared(
+ fused_weight,
+ fused_bias,
+ orig_shape=orig_shape,
+ m=m,
+ x_main=x_main,
+ lora_x=lora_x,
+ cuda_backend=cuda_backend,
+ )
+ return _svdquant_split_fused_output(out, splits)🤖 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/svdquant_w4a4.py` around lines 295 - 306, The grouped
CUDA fuse path does not validate whether shared quantization is possible before
calling _svdquant_cached_fused_group(), which internally assumes shared
quantization. Add a validation check using svdquant_w4a4_can_share_quant() on
the weights parameter before entering the if block that checks cuda_backend and
_svdquant_grouped_fuse_enabled(). This ensures that all weights have compatible
quantization parameters (same smooth_factor and proj_down) before attempting the
fused operation, preventing silent incorrect fusion of incompatible weights.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.