[PR #222] Add SM120 H3 QKV RMS RoPE and attention preparation fusion #233

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/222
Author: @Tokha233
Created: 9/30/2026
Status: 🔄 Open

Base: main ← Head: cuda/h3-qkv-quant

📝 Commits (5)

5304d03 Add SM120 H3 QKV RMS RoPE and attention preparation fusion
e070b2a Release H3 QKV scratch before declined native fallback
240e76a Retain quantization source attribution and updated test count
bbfee43 Fix H3 QKV fallback memory and source builds
c8e9ad7 Merge main and retain indexed operator exports

📊 Changes

18 files changed (+1913 additions, -1 deletions)

View changed files

📝 MANIFEST.in (+2 -0)
📝 comfy_kitchen/__init__.py (+2 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+3 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+112 -0)
➕ comfy_kitchen/backends/cuda/ops/h3_qkv/anchor_v.cu (+210 -0)
➕ comfy_kitchen/backends/cuda/ops/h3_qkv/mma.cuh (+94 -0)
➕ comfy_kitchen/backends/cuda/ops/h3_qkv/qkv_quant.cu (+386 -0)
➕ comfy_kitchen/backends/cuda/ops/h3_qkv/runtime.cuh (+21 -0)
➕ comfy_kitchen/backends/cuda/ops/h3_qkv/sample_k.cu (+198 -0)
➕ comfy_kitchen/h3_qkv.py (+176 -0)
➕ docs/benchmarks/h3-qkv-benchmark.json (+272 -0)
➕ docs/benchmarks/h3-qkv-image-probes.json (+10 -0)
➕ docs/benchmarks/h3-qkv-sampler.json (+148 -0)
➕ docs/benchmarks/h3-qkv-sdist-check.json (+5 -0)
➕ docs/benchmarks/h3-qkv.md (+50 -0)
📝 pyproject.toml (+1 -1)
➕ tests/benchmark_h3_qkv.py (+51 -0)
➕ tests/test_h3_qkv.py (+172 -0)

📄 Description

H3 currently writes the full BF16 QKV projection, then rereads it for RMS/RoPE and INT8 attention preparation. This adds an opt-in h3_qkv_prequantize operation combining the projection and preparation while retaining the current BF16 rounding, nine-current-key anchor selection, and all Q/K/V tokens.

The initial fused path targets SM120, H3 hidden 5376, 56×128 heads, partial RoPE 96, and 8192–200000 packed rows. Smaller/strided/other CUDA inputs use the existing chain. The result is the existing PrequantizedInt8Attention container; dense attention is unchanged. No cache or sparse approximation.

Validation on RTX 5090 D v2 / Torch 2.12.0+cu130 / CUDA 13.0.88:

24 tests pass, 1 skips (the isolated container exposes one GPU), including all six packed tensors and attention output matching exactly, finite magnitude cases, strided inputs, nondefault streams, CUDA Graphs, real Inductor, input immutability, and invalid epsilon/rank.
Eager full preparation chain, 8 alternating groups × 4 calls: 16.73% less time at 87142 rows after the review fixes (50.770336 → 42.275801 ms). Small shapes remain on the fallback. Raw data and runnable benchmark: docs/benchmarks/h3-qkv.md.
Full current-ComfyUI sampler experiment, 4 alternating formal pairs, Larry v4 INT8 / 8 steps / 768×512 / 124 frames: 15.803847 → 15.366207 s (−2.7692%), identical video/audio latent hashes in all 8 formal requests, unchanged peak Torch allocation.

Ready for review of the standalone model-specific API and integration boundary. The measured consumer explicitly selects Kitchen INT8 attention; it is not a generic ComfyUI dispatch implementation. The single-operation percentage above covers sampling only. The later combined validation also checked decoded RGB/PCM; no continuous-service throughput claim. These percentages cannot be added to other overlapping fusions.

Prepared with AI assistance, manually exercised on the stated GPU; DCO signed off. Source notices retained.

Latest combined validation and per-PR fixes: https://github.com/Tokha233/ComfyUI-H3-SpeedKit/blob/main/docs/UPSTREAM-RECOVERY-TESTS-20261001.zh-CN.md . Combined gains overlap and must not be added to the single-operation timings above.

_{🔄 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/222 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/h3-qkv-quant` --- ### 📝 Commits (5) - [`5304d03`](https://github.com/Comfy-Org/comfy-kitchen/commit/5304d03f8ebe3dd1ae9074234921e5f27e8e95e0) Add SM120 H3 QKV RMS RoPE and attention preparation fusion - [`e070b2a`](https://github.com/Comfy-Org/comfy-kitchen/commit/e070b2ad915f98ca302492be124ce2f1905a73ed) Release H3 QKV scratch before declined native fallback - [`240e76a`](https://github.com/Comfy-Org/comfy-kitchen/commit/240e76a5a21a2f726fe93e20b981138d86341c40) Retain quantization source attribution and updated test count - [`bbfee43`](https://github.com/Comfy-Org/comfy-kitchen/commit/bbfee43be7d9091c2e63cc4cf282b5c41c116af2) Fix H3 QKV fallback memory and source builds - [`c8e9ad7`](https://github.com/Comfy-Org/comfy-kitchen/commit/c8e9ad762959e933ffc7cbcbe4644769d08a1ae7) Merge main and retain indexed operator exports ### 📊 Changes **18 files changed** (+1913 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `MANIFEST.in` (+2 -0) 📝 `comfy_kitchen/__init__.py` (+2 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+3 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+112 -0) ➕ `comfy_kitchen/backends/cuda/ops/h3_qkv/anchor_v.cu` (+210 -0) ➕ `comfy_kitchen/backends/cuda/ops/h3_qkv/mma.cuh` (+94 -0) ➕ `comfy_kitchen/backends/cuda/ops/h3_qkv/qkv_quant.cu` (+386 -0) ➕ `comfy_kitchen/backends/cuda/ops/h3_qkv/runtime.cuh` (+21 -0) ➕ `comfy_kitchen/backends/cuda/ops/h3_qkv/sample_k.cu` (+198 -0) ➕ `comfy_kitchen/h3_qkv.py` (+176 -0) ➕ `docs/benchmarks/h3-qkv-benchmark.json` (+272 -0) ➕ `docs/benchmarks/h3-qkv-image-probes.json` (+10 -0) ➕ `docs/benchmarks/h3-qkv-sampler.json` (+148 -0) ➕ `docs/benchmarks/h3-qkv-sdist-check.json` (+5 -0) ➕ `docs/benchmarks/h3-qkv.md` (+50 -0) 📝 `pyproject.toml` (+1 -1) ➕ `tests/benchmark_h3_qkv.py` (+51 -0) ➕ `tests/test_h3_qkv.py` (+172 -0) </details> ### 📄 Description H3 currently writes the full BF16 QKV projection, then rereads it for RMS/RoPE and INT8 attention preparation. This adds an opt-in `h3_qkv_prequantize` operation combining the projection and preparation while retaining the current BF16 rounding, nine-current-key anchor selection, and all Q/K/V tokens. The initial fused path targets SM120, H3 hidden 5376, 56×128 heads, partial RoPE 96, and 8192–200000 packed rows. Smaller/strided/other CUDA inputs use the existing chain. The result is the existing `PrequantizedInt8Attention` container; dense attention is unchanged. No cache or sparse approximation. Validation on RTX 5090 D v2 / Torch 2.12.0+cu130 / CUDA 13.0.88: - **24 tests pass, 1 skips** (the isolated container exposes one GPU), including all six packed tensors and attention output matching exactly, finite magnitude cases, strided inputs, nondefault streams, CUDA Graphs, real Inductor, input immutability, and invalid epsilon/rank. - Eager full preparation chain, 8 alternating groups × 4 calls: **16.73% less time** at 87142 rows after the review fixes (50.770336 → 42.275801 ms). Small shapes remain on the fallback. Raw data and runnable benchmark: `docs/benchmarks/h3-qkv.md`. - Full current-ComfyUI sampler experiment, 4 alternating formal pairs, Larry v4 INT8 / 8 steps / 768×512 / 124 frames: **15.803847 → 15.366207 s (−2.7692%)**, identical video/audio latent hashes in all 8 formal requests, unchanged peak Torch allocation. Ready for review of the standalone model-specific API and integration boundary. The measured consumer explicitly selects Kitchen INT8 attention; it is not a generic ComfyUI dispatch implementation. The single-operation percentage above covers sampling only. The later combined validation also checked decoded RGB/PCM; no continuous-service throughput claim. These percentages cannot be added to other overlapping fusions. Prepared with AI assistance, manually exercised on the stated GPU; DCO signed off. Source notices retained. Latest combined validation and per-PR fixes: https://github.com/Tokha233/ComfyUI-H3-SpeedKit/blob/main/docs/UPSTREAM-RECOVERY-TESTS-20261001.zh-CN.md . Combined gains overlap and must not be added to the single-operation timings above. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

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

bsakaguchi referenced this issue

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

[PR #233] Optimize NVFP4 and MXFP8 quantization in Triton backend #243