[PR #219] Add prequantized INT8 GEMM with indexed BF16 gate and residual #234

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: cuda/indexed-gate-epilogue

📝 Commits (5)

dc7c739 Add prequantized INT8 GEMM with indexed BF16 gate and residual
5f7290d Fix indexed gate fallback layout and wheel license
f5e8698 Merge main and retain indexed operator exports
d3b3ec2 Avoid unused indexed-gate output allocation on fallback
62609cc Sync indexed gate exports with Kitchen 0.2.40

📊 Changes

14 files changed (+1074 additions, -1 deletions)

View changed files

📝 NOTICE (+7 -0)
📝 comfy_kitchen/__init__.py (+2 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+1 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+42 -0)
➕ comfy_kitchen/backends/cuda/ops/cutlass_gemm_indexed_gate.cu (+223 -0)
➕ comfy_kitchen/backends/cuda/ops/indexed_gate.cuh (+155 -0)
➕ comfy_kitchen/indexed_gate.py (+129 -0)
➕ docs/benchmarks/gate-api-benchmark.json (+210 -0)
➕ docs/benchmarks/indexed-gate-wheel.json (+11 -0)
➕ docs/benchmarks/int8-indexed-gate.md (+52 -0)
➕ docs/benchmarks/source-sha256.json (+11 -0)
📝 pyproject.toml (+1 -1)
➕ tests/benchmark_int8_indexed_gate.py (+72 -0)
➕ tests/test_int8_indexed_gate.py (+158 -0)

📄 Description

This introduces an opt-in prequantized int8_gemm_indexed_gate operation for diffusion branches with a compact gate table. It fuses indexed gate loads and residual addition into the INT8 GEMM epilogue, avoiding the BF16 branch and expanded gate intermediates.

out = ck.int8_gemm_indexed_gate(a_int8, b_int8, x_scale, w_scale,
                              gate_bf16, row_indices_int32, residual_bf16)

The contract intentionally rounds the scaled GEMM result to BF16 before the FP32 FMA with BF16 gate/residual and final BF16 rounding. Existing int8_linear and column-broadcast residual semantics are unchanged. This uses the existing Stream-K implementation and does not duplicate #215's N-banded scheduling work.

RTX 5090 D v2 measurements

Main 19ea55b (Kitchen0.2.36), CUDA13.0.88, PyTorch2.12.0+cu130; pinned upstream CUTLASS. Baseline: existing fused-dequant GEMM + gate gather + torch.addcmul. Candidate: public API including validation/fallback overhead. Eight alternating groups of ten warmed calls, CUDA events, exact BF16 outputs in every case.

M,N,K	Existing chain ms	New API ms	Less time
8192,5376,14336	2.670003	2.363296	11.49%
32700,5376,14336	10.983971	8.985152	18.20%
65536,5376,14336	21.839299	17.755728	18.70%
8192,7168,18944	4.655568	4.014264	13.77%

M1024 cases use the ordinary chain after the fused tile regressed in testing; measured wrapper overhead was 0.5–0.9%. No whole-model speedup or perceptual quality improvement is claimed: consumers must explicitly use this API. Full samples, a reproduction script and the contract are in docs/benchmarks/int8-indexed-gate.md.

API and validation

Initial fused backend: SM120 BF16, K%16=0/N%8=0; fallback for unsupported/noncontiguous inputs.
Gate [G,N], INT32 row indices [M]. Out-of-range indices select zero safely without host synchronization. Empty dimensions/gates supported; inputs are not modified.
Inference-only custom op with fake implementation for torch.compile.
CUDA Graph uses a workspace-free schedule, avoiding first-use Stream-K workspace allocation during capture.
Full extension built; 36 passed, 1 skipped (new indexed-gate tests plus existing residual tests). Tests include independent integer oracle, bounds, dtype/shape validation, noncontiguous fallback, compile, streams and graph replay.
Ruff and whitespace checks pass. CUTLASS visitor attribution and BSD license retained; DCO signed off.

Prepared with AI assistance; measurements were run on the stated hardware. Feedback on whether this should later be integrated into the broader linear API is welcome; this PR keeps the new numerical/indexing contract separate.

_{🔄 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/219 **Author:** [@Tokha233](https://github.com/Tokha233) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `cuda/indexed-gate-epilogue` --- ### 📝 Commits (5) - [`dc7c739`](https://github.com/Comfy-Org/comfy-kitchen/commit/dc7c739fe5fc7502b50206e81fbe26c6fe3930a1) Add prequantized INT8 GEMM with indexed BF16 gate and residual - [`5f7290d`](https://github.com/Comfy-Org/comfy-kitchen/commit/5f7290d1301e6bc7e2109c7c3d35a3d14f7f3de9) Fix indexed gate fallback layout and wheel license - [`f5e8698`](https://github.com/Comfy-Org/comfy-kitchen/commit/f5e869872efde9e49b363f94a0fab615f10088d7) Merge main and retain indexed operator exports - [`d3b3ec2`](https://github.com/Comfy-Org/comfy-kitchen/commit/d3b3ec25f7b986c3ae14df87dade94f194d12225) Avoid unused indexed-gate output allocation on fallback - [`62609cc`](https://github.com/Comfy-Org/comfy-kitchen/commit/62609cc84910d66ae4237e1a23332fa2c71574f3) Sync indexed gate exports with Kitchen 0.2.40 ### 📊 Changes **14 files changed** (+1074 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `NOTICE` (+7 -0) 📝 `comfy_kitchen/__init__.py` (+2 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+1 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+42 -0) ➕ `comfy_kitchen/backends/cuda/ops/cutlass_gemm_indexed_gate.cu` (+223 -0) ➕ `comfy_kitchen/backends/cuda/ops/indexed_gate.cuh` (+155 -0) ➕ `comfy_kitchen/indexed_gate.py` (+129 -0) ➕ `docs/benchmarks/gate-api-benchmark.json` (+210 -0) ➕ `docs/benchmarks/indexed-gate-wheel.json` (+11 -0) ➕ `docs/benchmarks/int8-indexed-gate.md` (+52 -0) ➕ `docs/benchmarks/source-sha256.json` (+11 -0) 📝 `pyproject.toml` (+1 -1) ➕ `tests/benchmark_int8_indexed_gate.py` (+72 -0) ➕ `tests/test_int8_indexed_gate.py` (+158 -0) </details> ### 📄 Description This introduces an opt-in prequantized `int8_gemm_indexed_gate` operation for diffusion branches with a compact gate table. It fuses indexed gate loads and residual addition into the INT8 GEMM epilogue, avoiding the BF16 branch and expanded gate intermediates. ```python out = ck.int8_gemm_indexed_gate(a_int8, b_int8, x_scale, w_scale, gate_bf16, row_indices_int32, residual_bf16) ``` The contract intentionally rounds the scaled GEMM result to BF16 before the FP32 FMA with BF16 gate/residual and final BF16 rounding. Existing `int8_linear` and column-broadcast residual semantics are unchanged. This uses the existing Stream-K implementation and does not duplicate #215's N-banded scheduling work. ### RTX 5090 D v2 measurements Main `19ea55b` (Kitchen0.2.36), CUDA13.0.88, PyTorch2.12.0+cu130; pinned upstream CUTLASS. Baseline: existing fused-dequant GEMM + gate gather + torch.addcmul. Candidate: public API including validation/fallback overhead. Eight alternating groups of ten warmed calls, CUDA events, exact BF16 outputs in every case. | M,N,K | Existing chain ms | New API ms | Less time | |---|---:|---:|---:| | 8192,5376,14336 | 2.670003 | 2.363296 | 11.49% | | 32700,5376,14336 | 10.983971 | 8.985152 | 18.20% | | 65536,5376,14336 | 21.839299 | 17.755728 | 18.70% | | 8192,7168,18944 | 4.655568 | 4.014264 | 13.77% | M1024 cases use the ordinary chain after the fused tile regressed in testing; measured wrapper overhead was 0.5–0.9%. No whole-model speedup or perceptual quality improvement is claimed: consumers must explicitly use this API. Full samples, a reproduction script and the contract are in `docs/benchmarks/int8-indexed-gate.md`. ### API and validation - Initial fused backend: SM120 BF16, K%16=0/N%8=0; fallback for unsupported/noncontiguous inputs. - Gate [G,N], INT32 row indices [M]. Out-of-range indices select zero safely without host synchronization. Empty dimensions/gates supported; inputs are not modified. - Inference-only custom op with fake implementation for torch.compile. - CUDA Graph uses a workspace-free schedule, avoiding first-use Stream-K workspace allocation during capture. - Full extension built; 36 passed, 1 skipped (new indexed-gate tests plus existing residual tests). Tests include independent integer oracle, bounds, dtype/shape validation, noncontiguous fallback, compile, streams and graph replay. - Ruff and whitespace checks pass. CUTLASS visitor attribution and BSD license retained; DCO signed off. Prepared with AI assistance; measurements were run on the stated hardware. Feedback on whether this should later be integrated into the broader linear API is welcome; this PR keeps the new numerical/indexing contract separate. --- <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:51 +00:00

bsakaguchi referenced this issue

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

[PR #234] [MERGED] Optimize int8 attention for ampere #242

bsakaguchi referenced this issue

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

[PR #235] [CLOSED] Move int8 attention conversions off the XU pipe on Ampere #244

bsakaguchi referenced this issue

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