[PR #239] [MERGED] Performance improvements for qwen 3.x LLM (Part 1) #246

New issue

Closed

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

bsakaguchi commented

2026-10-26 22:21:58 +00:00

Owner

📋 Pull Request Information

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

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

📝 Commits (10+)

3075219 Fuse scaled residuals into single-row INT8 GEMV
86fc2fa Factor out the single-row INT8 GEMV launch
008b88e Use aligned 16-byte loads in single-row INT8 GEMV
79458c4 Factor out the ConvRot row quantizer
e74f445 Fuse ConvRot input quantization into single-row INT8 GEMV
6ee300e Scope CUTLASS workspace to each stream-ordered launch
016d15b Add universal GQA decode with in-kernel split reduction
7d7a691 Merge causal step rows into prefix attention on CUDA HIP and torch
1fe78e7 Add bounded paged INT8 decode cache
cce5e24 Add deferred chain replay for DeltaNet decode

📊 Changes

27 files changed (+3635 additions, -277 deletions)

View changed files

📝 comfy_kitchen/__init__.py (+16 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+2 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+41 -10)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+405 -37)
➕ comfy_kitchen/backends/cuda/int8_decode_bindings.cpp (+47 -0)
📝 comfy_kitchen/backends/cuda/ops/cutlass_gemm_common.cuh (+9 -34)
📝 comfy_kitchen/backends/cuda/ops/flash_decode.cu (+278 -20)
📝 comfy_kitchen/backends/cuda/ops/gated_delta.cu (+418 -0)
📝 comfy_kitchen/backends/cuda/ops/int8_linear.cu (+306 -69)
➕ comfy_kitchen/backends/cuda/sage_attention/int8_decode.cu (+67 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_i8_cuda.cuh (+14 -14)
📝 comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu (+33 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/quant_v_int8.cu (+23 -2)
📝 comfy_kitchen/backends/hip/__init__.py (+134 -0)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+351 -2)
📝 comfy_kitchen/backends/hip/launchers.h (+21 -8)
📝 comfy_kitchen/backends/hip/ops/flash_decode.hip (+193 -78)
📝 comfy_kitchen/backends/hip/ops/gated_delta.hip (+407 -0)
📝 comfy_kitchen/flash_attention.py (+155 -1)
📝 comfy_kitchen/gated_delta.py (+222 -0)

...and 7 more files

📄 Description

Improve fusions and kernels for qwen 3.x LLM.

Faster INT8 GEMV with fused quantization and residuals.
GQA decode attention and causal step merging.
Paged INT8 KV cache with rollback support.
Deferred DeltaNet chain replay.

The bigger fish:

Deferred DeltaNet

Previously, MTP verification saved a full recurrent-state snapshot after each speculative token, then restored the accepted snapshot on rejection. Now it saves small token intermediates. The next call replays only the accepted updates and writes the committed state once. Each recurrent state is 64 KiB per head (128 × 128 FP32). Avoiding several snapshots per head, per layer, is substantial. On CUDA, it also moves the working state from shared memory into registers across four cooperating blocks, shortening reductions and reducing shared-memory traffic.

Before: verify → save full states → restore accepted state
After: verify → save token data → replay accepted updates next call

This trades extra arithmetic for fewer large memory transfers. With MTP off, the snapshot-saving benefit disappears; the CUDA register/reduction improvements remain.

GQA decode

A dedicated decode kernel handles the short query and causal verification rows directly, without constructing a full attention mask.
Grouped query heads reuse KV reads where the kernel groups them.
Splitting the KV prefix gives more GPU blocks work, rather than leaving most of the GPU idle on a tiny decode query.
On CUDA, the last finishing split combines the results inside the same launch, removing a separate reduction-kernel launch.

_{🔄 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/239 **Author:** [@rattus128](https://github.com/rattus128) **Created:** 10/10/2026 **Status:** ✅ Merged **Merged:** 10/11/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `prs/qwen-3x-general` --- ### 📝 Commits (10+) - [`3075219`](https://github.com/Comfy-Org/comfy-kitchen/commit/307521960767f99dfd3088292388b1b60a120699) Fuse scaled residuals into single-row INT8 GEMV - [`86fc2fa`](https://github.com/Comfy-Org/comfy-kitchen/commit/86fc2faf19c4bc53e37963ab6b0a4e4f9cbdaaf9) Factor out the single-row INT8 GEMV launch - [`008b88e`](https://github.com/Comfy-Org/comfy-kitchen/commit/008b88e68045bb12dad75f9e7570b2454f7883b5) Use aligned 16-byte loads in single-row INT8 GEMV - [`79458c4`](https://github.com/Comfy-Org/comfy-kitchen/commit/79458c451917e958748f77155ee5f190a1748410) Factor out the ConvRot row quantizer - [`e74f445`](https://github.com/Comfy-Org/comfy-kitchen/commit/e74f4454af98de374e98c73b54dd2c713cbfb6ff) Fuse ConvRot input quantization into single-row INT8 GEMV - [`6ee300e`](https://github.com/Comfy-Org/comfy-kitchen/commit/6ee300e5d8f8fcd959e38bd447ffb988ed8c2be6) Scope CUTLASS workspace to each stream-ordered launch - [`016d15b`](https://github.com/Comfy-Org/comfy-kitchen/commit/016d15b83c761771498272915c7abe4bf8e1c72f) Add universal GQA decode with in-kernel split reduction - [`7d7a691`](https://github.com/Comfy-Org/comfy-kitchen/commit/7d7a6918cb6c5d590030892f0929949bf260b219) Merge causal step rows into prefix attention on CUDA HIP and torch - [`1fe78e7`](https://github.com/Comfy-Org/comfy-kitchen/commit/1fe78e7f76e28d86dbf5d2969cbab2fda956cd2b) Add bounded paged INT8 decode cache - [`cce5e24`](https://github.com/Comfy-Org/comfy-kitchen/commit/cce5e24b59e505f1aabe6d9f4c3ce5b84375bd66) Add deferred chain replay for DeltaNet decode ### 📊 Changes **27 files changed** (+3635 additions, -277 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/__init__.py` (+16 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+2 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+41 -10) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+405 -37) ➕ `comfy_kitchen/backends/cuda/int8_decode_bindings.cpp` (+47 -0) 📝 `comfy_kitchen/backends/cuda/ops/cutlass_gemm_common.cuh` (+9 -34) 📝 `comfy_kitchen/backends/cuda/ops/flash_decode.cu` (+278 -20) 📝 `comfy_kitchen/backends/cuda/ops/gated_delta.cu` (+418 -0) 📝 `comfy_kitchen/backends/cuda/ops/int8_linear.cu` (+306 -69) ➕ `comfy_kitchen/backends/cuda/sage_attention/int8_decode.cu` (+67 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/qk_int_sv_i8_cuda.cuh` (+14 -14) 📝 `comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu` (+33 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/quant_v_int8.cu` (+23 -2) 📝 `comfy_kitchen/backends/hip/__init__.py` (+134 -0) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+351 -2) 📝 `comfy_kitchen/backends/hip/launchers.h` (+21 -8) 📝 `comfy_kitchen/backends/hip/ops/flash_decode.hip` (+193 -78) 📝 `comfy_kitchen/backends/hip/ops/gated_delta.hip` (+407 -0) 📝 `comfy_kitchen/flash_attention.py` (+155 -1) 📝 `comfy_kitchen/gated_delta.py` (+222 -0) _...and 7 more files_ </details> ### 📄 Description Improve fusions and kernels for qwen 3.x LLM. - Faster INT8 GEMV with fused quantization and residuals. - GQA decode attention and causal step merging. - Paged INT8 KV cache with rollback support. - Deferred DeltaNet chain replay. The bigger fish: ---------------------------------------------------------------------------------------------------- ### Deferred DeltaNet Previously, MTP verification saved a full recurrent-state snapshot after each speculative token, then restored the accepted snapshot on rejection. Now it saves small token intermediates. The next call replays only the accepted updates and writes the committed state once. Each recurrent state is 64 KiB per head (128 × 128 FP32). Avoiding several snapshots per head, per layer, is substantial. On CUDA, it also moves the working state from shared memory into registers across four cooperating blocks, shortening reductions and reducing shared-memory traffic. Before: verify → save full states → restore accepted state After: verify → save token data → replay accepted updates next call This trades extra arithmetic for fewer large memory transfers. With MTP off, the snapshot-saving benefit disappears; the CUDA register/reduction improvements remain. ---------------------------------------------------------------------------------------------------- ### GQA decode - A dedicated decode kernel handles the short query and causal verification rows directly, without constructing a full attention mask. - Grouped query heads reuse KV reads where the kernel groups them. - Splitting the KV prefix gives more GPU blocks work, rather than leaving most of the GPU idle on a tiny decode query. - On CUDA, the last finishing split combines the results inside the same launch, removing a separate reduction-kernel launch. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:21:58 +00:00