[PR #237] sm86 INT8 GEMM tuning: measured tile chooser, per-shape cfg table, and path logging #247

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/237
Author: @Many0therFunctions
Created: 10/10/2026
Status: 🔄 Open

Base: main ← Head: sm86-int8-gemm-tuning-v2

📝 Commits (3)

c424a72 sm86 INT8 GEMM tuning: measured tile chooser, per-shape cfg table, path logging
20769ed Address review: table loader gates and merge-safe sweeps
6fb6599 Ignore a cfg table whose JSON root is not an object

📊 Changes

8 files changed (+2094 additions, -1 deletions)

View changed files

➕ comfy_kitchen/_int8_cfg_cache.py (+233 -0)
➕ comfy_kitchen/a6000_int8_cfg_table.json (+1093 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+108 -0)
📝 comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu (+82 -1)
📝 comfy_kitchen/sage_attention.py (+51 -0)
📝 pyproject.toml (+1 -0)
➕ samples/int8_autotune_sweep.py (+399 -0)
➕ tests/test_int8_cfg_cache.py (+127 -0)

📄 Description

Fresh, squashed replacement for #142, rebased onto v0.2.37 (which brought the banded stream-K of #215 that changed what the sm86 picks should be) and re-measured. All ten review findings from #142 are folded in (details at the bottom).

What this does

INT8 inference on sm86 (RTX 3090 / A6000 / A40) was slower than it should be, and it was hard to tell why. Three things fix that:

1. sm86 tile chooser for the fused CUTLASS INT8 GEMM (cutlass_gemm_int8.cu)

select_fused_int8_config's thresholds were fitted on Ada and Blackwell. On GA102 (84 SMs, 6 MB L2) the generic path loses up to 43% on real diffusion shapes: 1797x2048x8192 (cfg12, −43%), 274x8192x2048 (−29%), 274x2048x8192 (−17%). The new branch, active only when the device is sm86, was fitted on a 14-config sweep over 45 shapes (LTX 2.5, MiniMax H3, Wan 2.2 14B, and a mid-M bridge series; M from 274 to 161920) on an A6000 at its 300 W cap, and checked with an interleaved sustained-load head-to-head of the choosers, since short bursts run above the capped clock and disagree with sustained numbers on small shapes:

m > 2048 and cfg13 bands         -> 13   banded stream-K wins all 22 banding shapes, 2-16% over cfg0
m > 2048, not banding            -> generic tree
1024 < m <= 2048: k <= 2048      -> n >= 16384 ? 0 : 1
                  k >  2048      -> n >=  8192 ? 0 : 1    (narrow N, tall K: cfg1, the tree's cfg12 is 43% slower)
 512 < m <= 1024                 -> 12
       m <=  512                 -> (n >= 8192 && k <= 4096) ? 3 : 2   (the tree's stream-K pick is 17-29% slower)

Sustained gap to the fastest config, average / worst over the 45 shapes: 0.5% / 5.7% for this chooser vs 3.6% / 43% for the generic path on the same card. Upstream's wave_guard and the #215 L2 reroute stay active; they agree with every pick on the swept shapes.

COMFY_KITCHEN_FORCE_CUTLASS_INT8_CONFIG=<0-13> forces one config for A/B runs without a rebuild. It is validated as a whole string ("1junk" and " 1" mean unset), checked before wave_guard and the reroute, never substituted by the launcher's fallback list, and is what cutlass_int8_selected_config reports.

2. Per-shape config table (_int8_cfg_cache.py, comfy_kitchen/a6000_int8_cfg_table.json, samples/int8_autotune_sweep.py)

int8_linear looks the exact (M, N, K) up in the packaged table before the heuristic and dispatches the swept winner directly; off-table shapes fall through to the chooser. The table is sm-gated (written for sm 8.6, ignored on any other device) and ships as package data. COMFY_KITCHEN_INT8_CFG_CACHE=<path> overrides it; a missing path warns and falls through to the packaged copy; malformed entries are skipped with a warning. samples/int8_autotune_sweep.py regenerates the table in process through _C.benchmark_cutlass_int8_dequant_config (no rebuild, no env vars) and writes it atomically. Entries whose sustained winner differed from the burst sweep carry a dated revalidated block; best_ms always comes from the same sweep as per_cfg_ms.

3. Path logging (cuda/__init__.py, sage_attention.py)

COMFY_KITCHEN_LOG_INT8_PATH=1 prints one line per unique configuration saying which path ran: fused CUTLASS, cuBLAS INT8 GEMM + separate dequant fallback, convrot rotate-matmul fallback, or the INT8 attention kernel with its launch parameters (dedup key = every printed field + path + device). Off by default, no overhead when off. "No int8 speedup" reports on Ampere usually turn out to be the cuBLAS fallback running silently.

Testing

A6000 (sm86): tests/test_int8.py, test_int8_input_act.py, test_int8_residual.py, test_gemm_bias_contract.py: 192 passed, 4 skipped, 1 failed — TestVectorAlignment::test_fp16_linear_realigns_bias_and_residual_scale, which is pre-existing and unrelated: its 512x2048x512 shape is 8 tiles, below the fp16 CUTLASS kernel's 96-tile minimum, so fp16_linear falls back to torch.nn.functional.linear on every GPU and the test ends up comparing cuBLAS against itself with two bias alignments (cuBLAS picks algorithms by pointer alignment; one-ulp fp16 differences). A shape the CUTLASS path serves (e.g. 2048x2048x4096) would make the test check what it means to.
Cache wiring verified end to end: swept shapes dispatch with their table config, off-table shapes use the chooser, a mismatched SM ignores the table, a forced config is exactly what runs.
ruff check . clean.

Ada/Blackwell behaviour is unchanged: the chooser branch only runs on sm86, and the table and logging are opt-in.

Review findings from #142, all included

Package the table as package-data · reset _cache_sm with the other cache state · test the original bias argument for the no-bias cache path · validate the complete forced-config string · cache SM detection per device · keep revalidated results separate from the sweep timings · describe the cuBLAS fallback as INT8 GEMM + separate dequant · never apply wave_guard to a forced config · include every field and the attention path in the once-only log key · have the sweep write the packaged table.

_{🔄 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/237 **Author:** [@Many0therFunctions](https://github.com/Many0therFunctions) **Created:** 10/10/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm86-int8-gemm-tuning-v2` --- ### 📝 Commits (3) - [`c424a72`](https://github.com/Comfy-Org/comfy-kitchen/commit/c424a727395c9961acbc086f617812d3cb1f7929) sm86 INT8 GEMM tuning: measured tile chooser, per-shape cfg table, path logging - [`20769ed`](https://github.com/Comfy-Org/comfy-kitchen/commit/20769ede86d12144deb0f2792e9b47ad187c5df1) Address review: table loader gates and merge-safe sweeps - [`6fb6599`](https://github.com/Comfy-Org/comfy-kitchen/commit/6fb659971bbf2a9f33fa637a8f3603167df72385) Ignore a cfg table whose JSON root is not an object ### 📊 Changes **8 files changed** (+2094 additions, -1 deletions) <details> <summary>View changed files</summary> ➕ `comfy_kitchen/_int8_cfg_cache.py` (+233 -0) ➕ `comfy_kitchen/a6000_int8_cfg_table.json` (+1093 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+108 -0) 📝 `comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu` (+82 -1) 📝 `comfy_kitchen/sage_attention.py` (+51 -0) 📝 `pyproject.toml` (+1 -0) ➕ `samples/int8_autotune_sweep.py` (+399 -0) ➕ `tests/test_int8_cfg_cache.py` (+127 -0) </details> ### 📄 Description Fresh, squashed replacement for #142, rebased onto v0.2.37 (which brought the banded stream-K of #215 that changed what the sm86 picks should be) and re-measured. All ten review findings from #142 are folded in (details at the bottom). ## What this does INT8 inference on sm86 (RTX 3090 / A6000 / A40) was slower than it should be, and it was hard to tell why. Three things fix that: **1. sm86 tile chooser for the fused CUTLASS INT8 GEMM** (`cutlass_gemm_int8.cu`) `select_fused_int8_config`'s thresholds were fitted on Ada and Blackwell. On GA102 (84 SMs, 6 MB L2) the generic path loses up to 43% on real diffusion shapes: 1797x2048x8192 (cfg12, −43%), 274x8192x2048 (−29%), 274x2048x8192 (−17%). The new branch, active only when the device is sm86, was fitted on a 14-config sweep over 45 shapes (LTX 2.5, MiniMax H3, Wan 2.2 14B, and a mid-M bridge series; M from 274 to 161920) on an A6000 at its 300 W cap, and checked with an interleaved sustained-load head-to-head of the choosers, since short bursts run above the capped clock and disagree with sustained numbers on small shapes: ``` m > 2048 and cfg13 bands -> 13 banded stream-K wins all 22 banding shapes, 2-16% over cfg0 m > 2048, not banding -> generic tree 1024 < m <= 2048: k <= 2048 -> n >= 16384 ? 0 : 1 k > 2048 -> n >= 8192 ? 0 : 1 (narrow N, tall K: cfg1, the tree's cfg12 is 43% slower) 512 < m <= 1024 -> 12 m <= 512 -> (n >= 8192 && k <= 4096) ? 3 : 2 (the tree's stream-K pick is 17-29% slower) ``` Sustained gap to the fastest config, average / worst over the 45 shapes: **0.5% / 5.7%** for this chooser vs 3.6% / 43% for the generic path on the same card. Upstream's `wave_guard` and the #215 L2 reroute stay active; they agree with every pick on the swept shapes. `COMFY_KITCHEN_FORCE_CUTLASS_INT8_CONFIG=<0-13>` forces one config for A/B runs without a rebuild. It is validated as a whole string (`"1junk"` and `" 1"` mean unset), checked before `wave_guard` and the reroute, never substituted by the launcher's fallback list, and is what `cutlass_int8_selected_config` reports. **2. Per-shape config table** (`_int8_cfg_cache.py`, `comfy_kitchen/a6000_int8_cfg_table.json`, `samples/int8_autotune_sweep.py`) `int8_linear` looks the exact `(M, N, K)` up in the packaged table before the heuristic and dispatches the swept winner directly; off-table shapes fall through to the chooser. The table is sm-gated (written for sm 8.6, ignored on any other device) and ships as package data. `COMFY_KITCHEN_INT8_CFG_CACHE=<path>` overrides it; a missing path warns and falls through to the packaged copy; malformed entries are skipped with a warning. `samples/int8_autotune_sweep.py` regenerates the table in process through `_C.benchmark_cutlass_int8_dequant_config` (no rebuild, no env vars) and writes it atomically. Entries whose sustained winner differed from the burst sweep carry a dated `revalidated` block; `best_ms` always comes from the same sweep as `per_cfg_ms`. **3. Path logging** (`cuda/__init__.py`, `sage_attention.py`) `COMFY_KITCHEN_LOG_INT8_PATH=1` prints one line per unique configuration saying which path ran: fused CUTLASS, cuBLAS INT8 GEMM + separate dequant fallback, convrot rotate-matmul fallback, or the INT8 attention kernel with its launch parameters (dedup key = every printed field + path + device). Off by default, no overhead when off. "No int8 speedup" reports on Ampere usually turn out to be the cuBLAS fallback running silently. ## Testing - A6000 (sm86): `tests/test_int8.py`, `test_int8_input_act.py`, `test_int8_residual.py`, `test_gemm_bias_contract.py`: 192 passed, 4 skipped, 1 failed — `TestVectorAlignment::test_fp16_linear_realigns_bias_and_residual_scale`, which is pre-existing and unrelated: its 512x2048x512 shape is 8 tiles, below the fp16 CUTLASS kernel's 96-tile minimum, so `fp16_linear` falls back to `torch.nn.functional.linear` on every GPU and the test ends up comparing cuBLAS against itself with two bias alignments (cuBLAS picks algorithms by pointer alignment; one-ulp fp16 differences). A shape the CUTLASS path serves (e.g. 2048x2048x4096) would make the test check what it means to. - Cache wiring verified end to end: swept shapes dispatch with their table config, off-table shapes use the chooser, a mismatched SM ignores the table, a forced config is exactly what runs. - `ruff check .` clean. Ada/Blackwell behaviour is unchanged: the chooser branch only runs on sm86, and the table and logging are opt-in. ## Review findings from #142, all included Package the table as package-data · reset `_cache_sm` with the other cache state · test the original `bias` argument for the no-bias cache path · validate the complete forced-config string · cache SM detection per device · keep `revalidated` results separate from the sweep timings · describe the cuBLAS fallback as INT8 GEMM + separate dequant · never apply `wave_guard` to a forced config · include every field and the attention path in the once-only log key · have the sweep write the packaged table. --- <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:59 +00:00

bsakaguchi referenced this issue

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

[PR #247] [MERGED] Add sm100 support to windows nvidia builds. #251