[PR #142] [CLOSED] sm86 INT8 GEMM tuning: measured tile heuristic, per-shape cfg cache, and path logging #171

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/142
Author: @Many0therFunctions
Created: 8/29/2026
Status: ❌ Closed

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

📝 Commits (3)

ea91848 sm86 INT8 GEMM tuning: measured tile heuristic, per-shape cfg cache, and path logging
f39b616 sm86: do not let wave_guard re-decide the measured tile pick
43c31a3 Address review: forced config wins over wave_guard, accurate fallback diagnostics, packaged sweep output

📊 Changes

7 files changed (+1771 additions, -2 deletions)

View changed files

➕ comfy_kitchen/_int8_cfg_cache.py (+196 -0)
➕ comfy_kitchen/a6000_int8_cfg_table.json (+917 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+109 -1)
📝 comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu (+170 -1)
📝 comfy_kitchen/sage_attention.py (+51 -0)
➕ int8_autotune_sweep.py (+327 -0)
📝 pyproject.toml (+1 -0)

📄 Description

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 heuristic for the fused CUTLASS INT8 GEMM (cutlass_gemm_int8.cu)

The existing select_fused_int8_config thresholds were fitted on Ada and Blackwell. On GA102 (84 SMs, 6MB L2, GDDR6) they pick StreamK configs for large-M shapes, and StreamK's partial-sum atomics thrash the small L2 once N is also wide. I measured this with a 14-config sweep over 38 real model shapes on an A6000 (LTX 2.5 and MiniMax H3 layer shapes, M from 274 to 80666): for the worst shapes (e.g. 25900x16384x4096, 80666x28672x5376) the Ada/Blackwell heuristic's pick loses to cfg0 by 44-50%. The new sm86 branch picks winners from the measured table, average regret 0.004% across the sweep, and was re-validated on 2026-09-26 with the card at its 300 W power cap (rule = measured winner on 34/38 shapes, 1.0% average regret; details in the comments below).

Also adds COMFY_KITCHEN_FORCE_CUTLASS_INT8_CONFIG=<int> for A/B testing a single config without rebuilding, and when forcing, the launcher does not silently fall back to a different tile.

2. Per-shape config cache (_int8_cfg_cache.py, int8_autotune_sweep.py, a6000_int8_cfg_table.json)

int8_autotune_sweep.py benchmarks every config against shapes you care about (LTX 2.5 and MiniMax H3 shapes are included) and writes a JSON table. int8_linear consults the table before the heuristic and uses the benchmarked winner for exact shape matches. The cache is SM-gated: a table written for sm86 is ignored on any other GPU, so a shared repo can't misselect. When no table exists, behavior is exactly current upstream.

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

COMFY_KITCHEN_LOG_INT8_PATH=1 prints one line per unique shape telling you which path actually ran: fused CUTLASS, cuBLAS fake-quant fallback, convrot variants, or the INT8 attention kernel with its launch parameters. Off by default, no overhead when off. This exists because "no int8 speedup" reports on Ampere usually turn out to be the cuBLAS fallback running silently, which means int8 tensor cores never execute at all.

Numbers from the sweep (A6000, 300W PL)

shape	old heuristic	new heuristic	speedup
25900x16384x4096	26.27 ms (cfg13)	16.98 ms (cfg0)	1.55x
80666x28672x5376	186.8 ms (cfg13)	122.8 ms (cfg0)	1.52x
8192x16384x4096	8.26 ms (cfg13)	5.60 ms (cfg0)	1.47x
1024x4096x4096	0.171 ms (cfg0)	0.162 ms (cfg12)	1.06x
274x2048x2048	0.036 ms (cfg2)	0.020 ms (cfg7)	1.8x

Full per-config timings for all 38 shapes are in the committed comfy_kitchen/a6000_int8_cfg_table.json.

Testing

Rebased branch, pytest tests/test_int8.py tests/test_int8_attention.py: 176 passed, 2 skipped, 2 failed — the 2 failures are test_int8_attention_long_sequence_and_partial_tile[0.0 / -0.088], where the reference scaled_dot_product_attention returns NaN on torch 2.13.0+cu132 (FLASH backend with scale=0.0 / negative scale); unrelated to this PR, see the comment below
Cache wiring verified end to end: exact swept shapes dispatch with their cached config, off-table shapes fall through to the heuristic, mismatched SM ignores the table
ruff check clean on all touched files

Existing Ada/Blackwell behavior is untouched: the new heuristic branch only activates on sm86, and the cache and logging are opt-in.

_{🔄 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/142 **Author:** [@Many0therFunctions](https://github.com/Many0therFunctions) **Created:** 8/29/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `sm86-int8-gemm-tuning` --- ### 📝 Commits (3) - [`ea91848`](https://github.com/Comfy-Org/comfy-kitchen/commit/ea91848ee85647d28b854d11fa1c924afa8969c7) sm86 INT8 GEMM tuning: measured tile heuristic, per-shape cfg cache, and path logging - [`f39b616`](https://github.com/Comfy-Org/comfy-kitchen/commit/f39b6167c3c2936278b90749452b949e38e42e19) sm86: do not let wave_guard re-decide the measured tile pick - [`43c31a3`](https://github.com/Comfy-Org/comfy-kitchen/commit/43c31a3ec1f9e826acbffe8535e218e71836e439) Address review: forced config wins over wave_guard, accurate fallback diagnostics, packaged sweep output ### 📊 Changes **7 files changed** (+1771 additions, -2 deletions) <details> <summary>View changed files</summary> ➕ `comfy_kitchen/_int8_cfg_cache.py` (+196 -0) ➕ `comfy_kitchen/a6000_int8_cfg_table.json` (+917 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+109 -1) 📝 `comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu` (+170 -1) 📝 `comfy_kitchen/sage_attention.py` (+51 -0) ➕ `int8_autotune_sweep.py` (+327 -0) 📝 `pyproject.toml` (+1 -0) </details> ### 📄 Description ## 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 heuristic for the fused CUTLASS INT8 GEMM** (`cutlass_gemm_int8.cu`) The existing `select_fused_int8_config` thresholds were fitted on Ada and Blackwell. On GA102 (84 SMs, 6MB L2, GDDR6) they pick StreamK configs for large-M shapes, and StreamK's partial-sum atomics thrash the small L2 once N is also wide. I measured this with a 14-config sweep over 38 real model shapes on an A6000 (LTX 2.5 and MiniMax H3 layer shapes, M from 274 to 80666): for the worst shapes (e.g. 25900x16384x4096, 80666x28672x5376) the Ada/Blackwell heuristic's pick loses to cfg0 by 44-50%. The new sm86 branch picks winners from the measured table, average regret 0.004% across the sweep, and was re-validated on 2026-09-26 with the card at its 300 W power cap (rule = measured winner on 34/38 shapes, 1.0% average regret; details in the comments below). Also adds `COMFY_KITCHEN_FORCE_CUTLASS_INT8_CONFIG=<int>` for A/B testing a single config without rebuilding, and when forcing, the launcher does not silently fall back to a different tile. **2. Per-shape config cache** (`_int8_cfg_cache.py`, `int8_autotune_sweep.py`, `a6000_int8_cfg_table.json`) `int8_autotune_sweep.py` benchmarks every config against shapes you care about (LTX 2.5 and MiniMax H3 shapes are included) and writes a JSON table. `int8_linear` consults the table before the heuristic and uses the benchmarked winner for exact shape matches. The cache is SM-gated: a table written for sm86 is ignored on any other GPU, so a shared repo can't misselect. When no table exists, behavior is exactly current upstream. **3. Path logging** (`cuda/__init__.py`, `sage_attention.py`) `COMFY_KITCHEN_LOG_INT8_PATH=1` prints one line per unique shape telling you which path actually ran: fused CUTLASS, cuBLAS fake-quant fallback, convrot variants, or the INT8 attention kernel with its launch parameters. Off by default, no overhead when off. This exists because "no int8 speedup" reports on Ampere usually turn out to be the cuBLAS fallback running silently, which means int8 tensor cores never execute at all. ## Numbers from the sweep (A6000, 300W PL) | shape | old heuristic | new heuristic | speedup | |---|---|---|---| | 25900x16384x4096 | 26.27 ms (cfg13) | 16.98 ms (cfg0) | 1.55x | | 80666x28672x5376 | 186.8 ms (cfg13) | 122.8 ms (cfg0) | 1.52x | | 8192x16384x4096 | 8.26 ms (cfg13) | 5.60 ms (cfg0) | 1.47x | | 1024x4096x4096 | 0.171 ms (cfg0) | 0.162 ms (cfg12) | 1.06x | | 274x2048x2048 | 0.036 ms (cfg2) | 0.020 ms (cfg7) | 1.8x | Full per-config timings for all 38 shapes are in the committed `comfy_kitchen/a6000_int8_cfg_table.json`. ## Testing - Rebased branch, `pytest tests/test_int8.py tests/test_int8_attention.py`: 176 passed, 2 skipped, 2 failed — the 2 failures are `test_int8_attention_long_sequence_and_partial_tile[0.0 / -0.088]`, where the *reference* `scaled_dot_product_attention` returns NaN on torch 2.13.0+cu132 (FLASH backend with `scale=0.0` / negative scale); unrelated to this PR, see the comment below - Cache wiring verified end to end: exact swept shapes dispatch with their cached config, off-table shapes fall through to the heuristic, mismatched SM ignores the table - `ruff check` clean on all touched files Existing Ada/Blackwell behavior is untouched: the new heuristic branch only activates on sm86, and the cache and logging are opt-in. --- <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:11 +00:00