Optimize quantize_int8_rowwise_convrot64 launch config (up to 4.9x on the generic path) - #80
Open
TheLegendOfKitty wants to merge 6 commits into
Open
Conversation
quantize_int8_rowwise_convrot64's generic launch branch (K not in
{256, 2560, 6144}, M>1) always used a fixed 1024-thread block
regardless of K. At K=1024 that leaves 768 of 1024 threads idle every
iteration (only (K/256)*64 = 256 of them ever have `active` true),
and the full-row shared-memory buffer (row_buf, K floats) both grows
with K and forces roughly one resident block per SM at larger K.
Device-track profiling on an RTX 3090 (SM86) showed the kernel
achieving only 9-35% of the GPU's peak memory bandwidth depending on
K, despite being clearly memory-bound (~7 FLOP/byte vs the card's
~38 FLOP/byte ridge point).
Three changes, all confined to the launcher's generic branch (the
M==1 path and the K in {256, 2560, 6144} special cases are untouched):
1. K <= 2048: right-size the launch to
threads = min(1024, (K/256)*64), reusing the existing single-pass
kernel body unchanged -- only the launcher's block-size selection
changes (new template instantiations at 128/192/256/320/384/448/
512 threads, dispatched via a switch on K/256, same pattern the
existing 256/640/768-thread special cases already use).
2. 2048 < K <= 4096: a new two-pass "chunked" kernel
(quantize_int8_rowwise_convrot64_chunked_kernel) that never
buffers more than one 2048-element chunk's rotated values at a
time (16KB shared memory total at 512 threads, independent of K).
Pass 1 rotates each chunk and folds its local max into the row's
abs_max (discarding the rotated values); pass 2 re-reads each
chunk (hopefully L2-resident) and re-rotates it -- deterministic
fp32 math, so this reproduces pass 1's values exactly -- fusing
the quantize step into the last butterfly stage so pass 2 needs no
buffer either.
3. K > 4096: kept on the ORIGINAL, unmodified single-pass kernel at a
fixed 1024 threads -- verified byte-identical dispatch to stock.
The chunked kernel from (2) was initially applied to all K > 2048,
but measured a reproducible ~8% regression at K=8192 (M=9216):
with 9216 blocks in flight, pass 2's re-read no longer reliably
hits L2 (the resident-rows working set exceeds the RTX 3090's 6MB
L2 at that row count), so the recompute becomes a genuine second
DRAM read instead of an L2 hit, and the occupancy gain isn't
enough to pay for that. Gated to K<=4096, where the smaller
per-row footprint keeps the working set L2-resident and the trade
stays net-positive. See the PR description for the full numbers;
occupancy/L2-aware chunk sizing that adapts to M is noted there as
possible future work.
Verified bitwise-identical (qdata and scale, torch.equal) against the
unmodified kernel across all real and boundary shapes -- K in
{1024, 2048, 3072, 3584, 4096, 8192} (3072/3584 specifically exercise
the chunked kernel's partial-last-chunk / inactive-lane path, since
3072/256=12 and 3584/256=14 groups are not multiples of the chunked
kernel's 8-groups-per-chunk width), M in {1, 9216, 9217}, fp16/fp32/
bf16, and an adversarial suite (all-zero rows, near-dtype-max rows
that overflow to NaN mid-butterfly, denormal-heavy rows, single-huge-
outlier rows) -- 64/64 cases, zero mismatches. Full pytest suite
(with the CUTLASS submodule initialized) passes the same 23
pre-existing failures / 101 skipped as the unmodified tree (fp8/
nvfp4/mxfp8 and torch.compile-integration issues on this box,
unrelated to this file), plus all new tests in the companion commit.
Preserves PR Comfy-Org#66's finite-max clamp (finite_absmax_for_int8_scale)
unchanged. PR Comfy-Org#65 (ConvRot int4/W4A4 support) also touched this file,
but only an unrelated int64-overflow fix in the plain (non-ConvRot)
quantize_int8_rowwise_kernel -- no overlap with the convrot64 kernel
or launcher this commit changes.
Signed-off-by: Parsa Adl-Tabatabai <61601745+TheLegendOfKitty@users.noreply.github.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01YUhW5zPPgvPHxPJYcb11ri
This kernel previously had no direct tests (only indirect coverage through int8_linear / TensorWiseINT8Layout convrot weight-quantize tests). Added alongside the launch-config fix in the previous commit to pin down the kernel's contract independent of that change: - shape/dtype matrix across the launcher's special cases, the right-sized single-pass branch, the chunked branch, and the K>4096 branch (K in 256/512/1024/1536/2048/2560/3072/3584/4096/6144/8192, M in 1/4/63/9216, bf16/fp16/fp32) -- K=3072 and K=3584 specifically target the chunked kernel's partial-last-chunk path (12 and 14 groups respectively, not multiples of its 8-groups-per-chunk width, so some lanes are masked `active=false` in the final chunk iteration, unlike K=4096's exact 2x8-group split) - determinism (two calls on the same input are bitwise identical), including across the resident/chunked and chunked/K>4096 boundaries - all-zero-row scale floor (1e-30), including a degenerate row mixed in among normal rows - near-dtype-max rows stay finite (exercises PR Comfy-Org#66's clamp on the shapes/launch configs this change touches) - CUDA vs the eager backend's reference (dense Hadamard matmul rotation) within a documented tolerance -- not bitwise, since the eager path's dense-matmul rotation sums the same terms in a different order than the CUDA kernel's explicit 4-stage radix-4 butterfly, so floating-point non-associativity puts a small, bounded fraction of int8 codes on the other side of a rounding boundary (measured ~2-4% of codes, max +/-2, stable across K) 154 tests, all passing; ruff clean. Signed-off-by: Parsa Adl-Tabatabai <61601745+TheLegendOfKitty@users.noreply.github.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YUhW5zPPgvPHxPJYcb11ri
|
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged. |
|
Warning Review limit reached
Next review available in: 6 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughChangesThe CUDA ConvRot64 INT8 quantization path adds a two-pass chunked kernel for ConvRot64 quantization
Sequence Diagram(s)sequenceDiagram
participant Launcher
participant ConvRotKernel
participant Output
Launcher->>ConvRotKernel: select kernel and thread count from K
ConvRotKernel->>ConvRotKernel: rotate chunks and reduce row scales
ConvRotKernel->>Output: quantize and store int8 values with scales
Suggested reviewers: 🚥 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/cuda/ops/int8_linear.cu`:
- Around line 1110-1183: Extract the duplicated load, four-value rotation,
synchronization, and FHT stages from the two passes into a shared device helper
near the relevant kernel code, such as convrot_load_rotate_group. Update both
pass 1 and pass 2 in the enclosing kernel to call it and continue consuming the
resulting values in buf1 for absmax accumulation and quantization respectively.
Preserve the existing active-group handling, buffer usage, synchronization
points, and numerical operation order.
In `@tests/test_convrot64_quantize.py`:
- Around line 80-97: Update the k parameter list in
TestConvrot64QuantizeDeterminism.test_two_calls_bitwise_equal to include 2304,
ensuring the advertised resident/chunked boundary coverage is actually exercised
while preserving the existing test cases.
🪄 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: b75d2917-5ea3-4a95-bbd3-df8758aae95f
📒 Files selected for processing (2)
comfy_kitchen/backends/cuda/ops/int8_linear.cutests/test_convrot64_quantize.py
Author
|
I have read and agree to the Contributor License Agreement |
comfy-legal
added a commit
to Comfy-Org/comfy-cla
that referenced
this pull request
Jul 20, 2026
Addresses CodeRabbit review on PR Comfy-Org#80 (2 comments): 1. quantize_int8_rowwise_convrot64_chunked_kernel's pass 1 and pass 2 duplicated the exact same load + first-three-FHT-stages sequence verbatim. Extracted into a new device helper, convrot_load_rotate_group (templated on InputType, called identically by both passes, same active-group masking, same buf0/ buf1 usage, same __syncthreads() placement, same operand order) so the two passes literally cannot drift apart -- pass 2's bitwise reproduction of pass 1's rotated values, which the whole two-pass design depends on, now holds by construction rather than by two hand-kept-in-sync copies. Each pass still applies its own final S=64 stage afterward (store_absmax vs the fused quantize), since that step differs between them. 2. The determinism test class's docstring advertised coverage "across the resident/chunked boundary at K=2048/2304" but 2304 was never actually in test_two_calls_bitwise_equal's parametrize list. Added 2304 there and to the shape/dtype matrix (2304/256=9 groups -> an 8+1 chunk split, the most extreme partial-last-chunk case in the suite: only 1 of the chunked kernel's 8 "sub" lanes active in the second chunk iteration, vs 3072's 8+4 and 3584's 8+6). Verified bitwise-identical against the shipped-kernel references: re-ran all previously-verified shapes (50 + the K=3072/3584 set of 14 = 64/64, confirming the refactor didn't change any kernel's output) plus 7 new K=2304 cases (bf16/fp16/fp32, M=9216/M=1, random + all-zero-mixed) -- 71/71 total, zero mismatches. Full test file: 167 tests (up from 154), all passing, ruff clean. Upstream suite: unchanged 23 pre-existing failures / 975 passed / 101 skipped. Standalone perf spot-check at K=4096 (the refactored kernel): 377.45 us/call, matching the pre-refactor measurement (377.5 us/call) within noise -- the __forceinline__ extraction didn't change codegen performance. Signed-off-by: Parsa Adl-Tabatabai <61601745+TheLegendOfKitty@users.noreply.github.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YUhW5zPPgvPHxPJYcb11ri
Comment-only change (proven below), no code-token differences. int8_linear.cu: several comment blocks added across this branch ran 15-41 lines -- narrative design rationale, roofline numbers, and historical "was previously X" framing that belongs in the PR description and commit messages, not the code. Stock's own comment density in this file (and in adaln.cu/apply_rope.cu/etc.) is 1-3 lines per block, with an occasional ~17-22 line header for a genuinely novel kernel (e.g. this file's own pre-existing ConvRot rotation comment at line 201). Trimmed every block we added to that density: the chunked kernel's header is down from 41 to 9 lines (comparable to stock's own longest precedent, not longer), the launcher's three branch comments from 18/18/8 lines to 2-4 lines each. Kept exactly the load-bearing one-liners a maintainer would actually want at the call site: the NaN-fold fmaxf semantics (now next to the fold itself, not buried in a kernel-header essay), and the one-line reason the chunked kernel is gated to K<=4096. tests/test_convrot64_quantize.py: existing test files in this repo (test_int8.py, test_rms_rope.py, test_adaln.py) use exactly one-line docstrings throughout, no exceptions found. This file's docstrings had grown to 6-15 lines each (shape-boundary enumerations, tolerance derivations, K=2304/3072/3584 partial-chunk explanations). Cut every docstring to one line, matching the rest of the suite; the detail that's genuinely useful for a reader is now in the PR description. Proof of comment/docstring-only change: - .cu: a comment-stripping diff (// line comments and /* */ blocks removed via a string/char-literal-aware scanner, blank lines dropped) against the pre-trim version shows zero code-token differences. Sanity-checked against a known-different pair (this file vs the original v0.2.22 base) to confirm the stripper actually detects real changes, not a no-op. - tests/*.py: an AST diff with every docstring's string value blanked out (docstrings are string literals, not comment tokens, so a plain '#'-comment strip does not catch them) shows the two ASTs are identical once docstring text is ignored. Also sanity-checked against a deliberately mutated copy to confirm it flags real changes. Given both proofs show code-identity, the previously-verified bitwise results (71/71 across all shapes including K=2304/3072/3584, full pytest suite unchanged) still hold without re-running them. Did rebuild and run tests/test_convrot64_quantize.py once (167 passed) as a sanity check. Signed-off-by: Parsa Adl-Tabatabai <61601745+TheLegendOfKitty@users.noreply.github.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YUhW5zPPgvPHxPJYcb11ri
Also adds a seeded-stochastic-rounding test class spanning the resident and chunked kernels. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YUhW5zPPgvPHxPJYcb11ri Signed-off-by: Parsa Adl-Tabatabai <61601745+TheLegendOfKitty@users.noreply.github.com>
Chunked kernel's dynamic smem is always under the default 48 KiB limit, so cudaFuncSetAttribute in launch_chunked was dead. The __syncthreads after storing scales[row] was also redundant: block_reduce_max_t's own final barrier already orders the block_smem read before any thread proceeds. Rewords comments that referenced "stock", benchmarked regressions, or PR history into present-tense facts, and makes the <=2048 switch's K=2048 case explicit rather than falling through a "default" that doubled as a comment explaining why it wasn't reachable for other K. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01YUhW5zPPgvPHxPJYcb11ri Signed-off-by: Parsa Adl-Tabatabai <61601745+TheLegendOfKitty@users.noreply.github.com>
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.
The generic launch path of
quantize_int8_rowwise_convrot64used a fixed 1024-thread block and a full-row shared-memory buffer for every K, reaching only 9–35% of peak bandwidth on Ampere (75% of lanes idle at K=1024). This PR right-sizes the block to the actual group count for K≤2048 and adds a two-pass chunked kernel with bounded shared memory for 2048<K≤4096, while K>4096 keeps the existing path unchanged (chunking measured ~8% slower at K=8192 — the second pass's re-read loses L2 residency at high row counts — so it is gated off there). Outputs are bitwise-identical to the current kernel, verified against the shipped v0.2.22 wheel across K∈{1024, 2048, 2304, 3072, 3584, 4096, 8192}, fp32/fp16/bf16, and adversarial inputs (all-zero rows, near-dtype-max NaN cascades — preserving the overflow-clamp semantics from #66), and the kernel gains direct test coverage it previously lacked. Measured on an RTX 3090: 4.9x at K=1024, 2.5x at K=2048, 1.12x at K=4096 standalone; end-to-end on an int8 ConvRot checkpoint (Anima) at 1536x1536 this is 0.842→0.777 s/it eager and 0.797→0.729 s/it with torch.compile.