[PR #248] [MERGED] Performance improvements for qwen 3.x (Part 2) #255

New issue

Closed

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

bsakaguchi commented

2026-10-26 22:22:09 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/248
Author: @rattus128
Created: 10/13/2026
Status: ✅ Merged
Merged: 10/13/2026
Merged by: @comfyanonymous

Base: main ← Head: prs/qwen-3x-general-2

📝 Commits (10+)

beb01d5 Round W4A8 and W6A8 decode onto the INT8 grid without conversions
ae30597 Add clustered ConvRot quantization with fused input activations
a84bfa9 Add quantized activation API with INT8 dispatch
5126939 Add the resident decode-layout hook
ba5b11b Support W4A8 and W6A8 through the activation API
dc09a05 Extract the canonical W4A8 chunked GEMM path
c03f18f Support packed W4A8 and W6A8 weights in chunked GEMM
fd835d5 Extract shared W4A8 and W6A8 code packing helpers
b836cae Add packed W4A8 and W6A8 small-M MMA decode
2f2dd58 Fuse packed W4A8 and W6A8 input activations and residuals

📊 Changes

17 files changed (+2024 additions, -157 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/__init__.py (+275 -74)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+208 -34)
📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+281 -0)
📝 comfy_kitchen/backends/cuda/ops/w4a8_gemm.cu (+526 -11)
📝 comfy_kitchen/backends/eager/__init__.py (+4 -0)
📝 comfy_kitchen/backends/eager/w4a8_int8.py (+46 -34)
📝 comfy_kitchen/backends/hip/__init__.py (+16 -0)
📝 comfy_kitchen/backends/triton/w4a8_int8.py (+17 -0)
📝 comfy_kitchen/constraints.py (+9 -0)
📝 comfy_kitchen/tensor/base.py (+33 -0)
📝 comfy_kitchen/tensor/int8.py (+16 -0)
📝 comfy_kitchen/tensor/w4a8_int8.py (+103 -4)
➕ comfy_kitchen/tensor/w4a8_packed.py (+111 -0)
📝 tests/test_constraints.py (+40 -0)
📝 tests/test_tensor.py (+41 -0)
➕ tests/test_w4a8_mma.py (+274 -0)
📝 tests/test_w4a8_requant.py (+24 -0)

📄 Description

Part 2 of qwen 3.x LLM performance improvements.

The main change is w4a8 and w6a8 M<8 memory-bandwidth optimized GEMV kernels and weight relayout facilities.

Interleaving the weights and their scales allows for significantly less memory access churn by linearizing access and allows a higher steady state read of weights. Some kernels to relayout weights are added accordingly.

These kernels use the tensor-cores instead of dp4a, which is similar performance for M=1 but gives significant optimization for MTP structures up the M=8.

The API for linear_input_act and switching between kernels for quants is put behind an abstraction in CK.

Add the usual suspect fusions.

Theres a super arcane fp32 quantization optimization in there too.

_{🔄 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/248 **Author:** [@rattus128](https://github.com/rattus128) **Created:** 10/13/2026 **Status:** ✅ Merged **Merged:** 10/13/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `prs/qwen-3x-general-2` --- ### 📝 Commits (10+) - [`beb01d5`](https://github.com/Comfy-Org/comfy-kitchen/commit/beb01d503d2eb1727d986ada889159e506d3a3b0) Round W4A8 and W6A8 decode onto the INT8 grid without conversions - [`ae30597`](https://github.com/Comfy-Org/comfy-kitchen/commit/ae30597b287c06f526484953fad023a8d0d158ba) Add clustered ConvRot quantization with fused input activations - [`a84bfa9`](https://github.com/Comfy-Org/comfy-kitchen/commit/a84bfa92d136d84be513f8f4f7426779e4bd0feb) Add quantized activation API with INT8 dispatch - [`5126939`](https://github.com/Comfy-Org/comfy-kitchen/commit/51269396e4fe2591cfdd46f2e5bb103583899320) Add the resident decode-layout hook - [`ba5b11b`](https://github.com/Comfy-Org/comfy-kitchen/commit/ba5b11b7a1fb339e8be3078aecaeb6831e195bdc) Support W4A8 and W6A8 through the activation API - [`dc09a05`](https://github.com/Comfy-Org/comfy-kitchen/commit/dc09a05a71637fc5dbef2a2ed1952c7f75944ca8) Extract the canonical W4A8 chunked GEMM path - [`c03f18f`](https://github.com/Comfy-Org/comfy-kitchen/commit/c03f18f2b22bfbca6e471aaec73d124be1176edb) Support packed W4A8 and W6A8 weights in chunked GEMM - [`fd835d5`](https://github.com/Comfy-Org/comfy-kitchen/commit/fd835d5a1e304b337c21e8704ad1419c1b3195f7) Extract shared W4A8 and W6A8 code packing helpers - [`b836cae`](https://github.com/Comfy-Org/comfy-kitchen/commit/b836caea78f4e433259c4bfaf343ec5a28397c00) Add packed W4A8 and W6A8 small-M MMA decode - [`2f2dd58`](https://github.com/Comfy-Org/comfy-kitchen/commit/2f2dd580cd13e06c96cbbd1e37a4bf642af02aaa) Fuse packed W4A8 and W6A8 input activations and residuals ### 📊 Changes **17 files changed** (+2024 additions, -157 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/__init__.py` (+275 -74) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+208 -34) 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+281 -0) 📝 `comfy_kitchen/backends/cuda/ops/w4a8_gemm.cu` (+526 -11) 📝 `comfy_kitchen/backends/eager/__init__.py` (+4 -0) 📝 `comfy_kitchen/backends/eager/w4a8_int8.py` (+46 -34) 📝 `comfy_kitchen/backends/hip/__init__.py` (+16 -0) 📝 `comfy_kitchen/backends/triton/w4a8_int8.py` (+17 -0) 📝 `comfy_kitchen/constraints.py` (+9 -0) 📝 `comfy_kitchen/tensor/base.py` (+33 -0) 📝 `comfy_kitchen/tensor/int8.py` (+16 -0) 📝 `comfy_kitchen/tensor/w4a8_int8.py` (+103 -4) ➕ `comfy_kitchen/tensor/w4a8_packed.py` (+111 -0) 📝 `tests/test_constraints.py` (+40 -0) 📝 `tests/test_tensor.py` (+41 -0) ➕ `tests/test_w4a8_mma.py` (+274 -0) 📝 `tests/test_w4a8_requant.py` (+24 -0) </details> ### 📄 Description Part 2 of qwen 3.x LLM performance improvements. The main change is w4a8 and w6a8 M<8 memory-bandwidth optimized GEMV kernels and weight relayout facilities. Interleaving the weights and their scales allows for significantly less memory access churn by linearizing access and allows a higher steady state read of weights. Some kernels to relayout weights are added accordingly. These kernels use the tensor-cores instead of dp4a, which is similar performance for M=1 but gives significant optimization for MTP structures up the M=8. The API for linear_input_act and switching between kernels for quants is put behind an abstraction in CK. Add the usual suspect fusions. Theres a super arcane fp32 quantization optimization in there too. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:22:09 +00:00