[PR #80] [CLOSED] Optimize quantize_int8_rowwise_convrot64 launch config (up to 4.9x on the generic path) #126

New issue

Closed

opened 2026-10-26 22:20:38 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:20:38 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/80
Author: @TheLegendOfKitty
Created: 7/20/2026
Status: ❌ Closed

Base: main ← Head: optimize-convrot64-quantize-launch

📝 Commits (6)

c17a1dc Right-size convrot64 quantize launch; chunk 2048<K<=4096 to bound smem
b41a280 Add direct test coverage for quantize_int8_rowwise_convrot64
38469fe Extract shared load+rotate helper in chunked kernel; add K=2304 test
a9f9ae2 Trim comment density to match this file's/test suite's existing style
4343fa9 Add K=768/1280/1792 launch-config test coverage
c957c76 Trim stale comments; drop unneeded chunked-kernel smem attr and barrier

📊 Changes

2 files changed (+373 additions, -1 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+214 -1)
➕ tests/test_convrot64_quantize.py (+159 -0)

📄 Description

The generic launch path of quantize_int8_rowwise_convrot64 used 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.

_{🔄 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/80 **Author:** [@TheLegendOfKitty](https://github.com/TheLegendOfKitty) **Created:** 7/20/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `optimize-convrot64-quantize-launch` --- ### 📝 Commits (6) - [`c17a1dc`](https://github.com/Comfy-Org/comfy-kitchen/commit/c17a1dc2aa2a206c339d8a1b71ade87eb256f534) Right-size convrot64 quantize launch; chunk 2048<K<=4096 to bound smem - [`b41a280`](https://github.com/Comfy-Org/comfy-kitchen/commit/b41a28030ca62d33e366a201f47c687dd9561d8e) Add direct test coverage for quantize_int8_rowwise_convrot64 - [`38469fe`](https://github.com/Comfy-Org/comfy-kitchen/commit/38469fe726f4ed444c4e6c133b16b906f1531497) Extract shared load+rotate helper in chunked kernel; add K=2304 test - [`a9f9ae2`](https://github.com/Comfy-Org/comfy-kitchen/commit/a9f9ae2f7c566ad6b1defc6251b9fb9eb438707e) Trim comment density to match this file's/test suite's existing style - [`4343fa9`](https://github.com/Comfy-Org/comfy-kitchen/commit/4343fa9c857010d09784c7bc920ea1e9dbf13cc5) Add K=768/1280/1792 launch-config test coverage - [`c957c76`](https://github.com/Comfy-Org/comfy-kitchen/commit/c957c76f6620e83ffae796c28f644258eb917c3b) Trim stale comments; drop unneeded chunked-kernel smem attr and barrier ### 📊 Changes **2 files changed** (+373 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+214 -1) ➕ `tests/test_convrot64_quantize.py` (+159 -0) </details> ### 📄 Description The generic launch path of `quantize_int8_rowwise_convrot64` used 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. --- <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:38 +00:00