[PR #221] Fuse indexed modulation, RMSNorm and ConvRot INT8 quantization #231

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/221
Author: @Tokha233
Created: 9/30/2026
Status: 🔄 Open

Base: main ← Head: cuda/indexed-norm-convrot

📝 Commits (3)

c058e49 Fuse indexed modulation with RMSNorm and ConvRot quantization
f66dda5 Release unused fallback buffers and verify native benchmark dispatch
34394a5 Merge main and retain indexed operator exports

📊 Changes

12 files changed (+1008 additions, -0 deletions)

View changed files

📝 NOTICE (+8 -0)
📝 comfy_kitchen/__init__.py (+2 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+41 -0)
➕ comfy_kitchen/backends/cuda/ops/indexed_norm_convrot.cu (+287 -0)
➕ comfy_kitchen/indexed_norm.py (+116 -0)
➕ comfy_kitchen/licenses/PYTORCH_LICENSE.txt (+84 -0)
➕ docs/benchmarks/norm-benchmark.json (+210 -0)
➕ docs/benchmarks/norm.md (+34 -0)
📝 pyproject.toml (+1 -0)
➕ tests/benchmark_norm.py (+94 -0)
➕ tests/test_indexed_norm.py (+130 -0)

📄 Description

Fuses weighted RMSNorm, compact indexed BF16 scale/shift modulation, ConvRot256 and rowwise INT8 quantization. This removes the materialized norm/mod intermediates. The new opt-in indexed_norm_convrot API returns INT8 data + FP32 row scales; existing rms_adaln is unchanged.

Initial native scope is SM120/K5376 and the tested PyTorch RMSNorm reduction (git 7661cd9). Other Torch revisions, devices and layouts use their native RMSNorm/modulation/quantization chain. Each intermediate BF16 rounding is preserved. Invalid indices select zero shift/scale without host synchronization. Inputs are not modified; fake/custom-op support Inductor and CUDA Graphs. PyTorch BSD license is packaged.

RTX 5090 D v2 / Torch 2.12.0+cu130: 34 tests passed (CPU/CUDA, empty tables, invalid INT32 indices, strides, magnitude extremes, direct native dispatch, streams, Graphs, real Inductor). Full CUDA extension rebuilt. Benchmark baseline is native RMSNorm + six actual contiguous modulation segments + Kitchen ConvRot256, not a large indexed-gather baseline.

Tokens / K5376	Original chain	Fused	Time reduction
14,850	see raw samples	see raw samples	68.83%
32,700	1.97962 ms	0.57689 ms	70.86%
87,142	6.23443 ms	1.49892 ms	75.96%
90,461	6.47000 ms	1.55566 ms	75.96%

All Q bytes and FP32 scale values match. Eight alternating groups / ten CUDA Graph replays per sample. These are preparation-chain timings; they are not full DiT speedups. Raw samples and reproduction: docs/benchmarks/norm.md, tests/benchmark_norm.py. An explicit INT8 consumer was also exercised in the combined full-sampler validation linked below. The generic ComfyUI consumer is not part of this PR; no isolated whole-model speedup is claimed for this operation.

Ready for API/reduction-compatibility review as a standalone opt-in library operation. The narrow Torch gate is intentional until additional reductions are qualified. Prepared with AI assistance; measurements were executed on the stated hardware. DCO signed off.

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/221 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/indexed-norm-convrot` --- ### 📝 Commits (3) - [`c058e49`](https://github.com/Comfy-Org/comfy-kitchen/commit/c058e49cb68397b5f30be2cdca51900d20c67a16) Fuse indexed modulation with RMSNorm and ConvRot quantization - [`f66dda5`](https://github.com/Comfy-Org/comfy-kitchen/commit/f66dda592a0796a12a96b3adbf8d61e51113e58d) Release unused fallback buffers and verify native benchmark dispatch - [`34394a5`](https://github.com/Comfy-Org/comfy-kitchen/commit/34394a516ce5ed5e7bc89effe43980c35d0bc66a) Merge main and retain indexed operator exports ### 📊 Changes **12 files changed** (+1008 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `NOTICE` (+8 -0) 📝 `comfy_kitchen/__init__.py` (+2 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+41 -0) ➕ `comfy_kitchen/backends/cuda/ops/indexed_norm_convrot.cu` (+287 -0) ➕ `comfy_kitchen/indexed_norm.py` (+116 -0) ➕ `comfy_kitchen/licenses/PYTORCH_LICENSE.txt` (+84 -0) ➕ `docs/benchmarks/norm-benchmark.json` (+210 -0) ➕ `docs/benchmarks/norm.md` (+34 -0) 📝 `pyproject.toml` (+1 -0) ➕ `tests/benchmark_norm.py` (+94 -0) ➕ `tests/test_indexed_norm.py` (+130 -0) </details> ### 📄 Description Fuses weighted RMSNorm, compact indexed BF16 scale/shift modulation, ConvRot256 and rowwise INT8 quantization. This removes the materialized norm/mod intermediates. The new opt-in `indexed_norm_convrot` API returns INT8 data + FP32 row scales; existing `rms_adaln` is unchanged. Initial native scope is SM120/K5376 and the tested PyTorch RMSNorm reduction (git 7661cd9). Other Torch revisions, devices and layouts use their native RMSNorm/modulation/quantization chain. Each intermediate BF16 rounding is preserved. Invalid indices select zero shift/scale without host synchronization. Inputs are not modified; fake/custom-op support Inductor and CUDA Graphs. PyTorch BSD license is packaged. RTX 5090 D v2 / Torch 2.12.0+cu130: 34 tests passed (CPU/CUDA, empty tables, invalid INT32 indices, strides, magnitude extremes, direct native dispatch, streams, Graphs, real Inductor). Full CUDA extension rebuilt. Benchmark baseline is native RMSNorm + **six actual contiguous modulation segments** + Kitchen ConvRot256, not a large indexed-gather baseline. | Tokens / K5376 | Original chain | Fused | Time reduction | |---|---:|---:|---:| | 14,850 | see raw samples | see raw samples | 68.83% | | 32,700 | 1.97962 ms | 0.57689 ms | 70.86% | | 87,142 | 6.23443 ms | 1.49892 ms | 75.96% | | 90,461 | 6.47000 ms | 1.55566 ms | 75.96% | All Q bytes and FP32 scale values match. Eight alternating groups / ten CUDA Graph replays per sample. These are preparation-chain timings; they are not full DiT speedups. Raw samples and reproduction: `docs/benchmarks/norm.md`, `tests/benchmark_norm.py`. An explicit INT8 consumer was also exercised in the combined full-sampler validation linked below. The generic ComfyUI consumer is not part of this PR; no isolated whole-model speedup is claimed for this operation. Ready for API/reduction-compatibility review as a standalone opt-in library operation. The narrow Torch gate is intentional until additional reductions are qualified. Prepared with AI assistance; measurements were executed on the stated hardware. DCO signed off. 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:54 +00:00

[PR #231] Tune SM120 V quantization for large D128 inputs #240