[PR #215] [MERGED] Further optimize int8 GEMM (CORE-520) #226

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/215
Author: @kijai
Created: 9/29/2026
Status: ✅ Merged
Merged: 9/30/2026
Merged by: @comfyanonymous

Base: main ← Head: int8-l2-raster

📝 Commits (2)

00a9cbe Band the int8 stream-K tile walk to keep weights in L2
92e9aa6 Optimize for Ampere

📊 Changes

6 files changed (+167 additions, -20 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/__init__.py (+25 -1)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+6 -0)
📝 comfy_kitchen/backends/cuda/ops/cutlass_gemm_common.cuh (+49 -3)
📝 comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu (+26 -16)
📝 tests/test_int8.py (+48 -0)
📝 tests/test_w6a8.py (+13 -0)

📄 Description

The int8 GEMM's tile order re-streams a weight that outgrows L2 on every wave.
On a 5090 this read 20-26 GB per call for H3's fc1 at 32k tokens and dropped SM clocks under the power limit.

The stream-K swizzle now walks N in bands, band width depends on the card's L2:

16 MB or more (Ada, Blackwell): 32 tiles, so a band's weight slice stays resident.
Below that (Ampere: 5-6 MB): no slice can stay resident across waves, so the band is
sized to minimize what one wave reads, sqrt(SMs * tile_m / tile_n) tiles (6-8 on a
3080). Measured on a 3080: within 2% of the best width at any K, where 32 was up to
47% slower.
Small-L2 cards also decode W4A8/W6A8 weights whole instead of in chunks: a chunk that
cannot stay in L2 only adds launches (up to 20% faster at M=512 on a 3080).

Bit-identical output, no new kernel instantiations. End to end gain in MiniMax-H3 workflow is modest ~2%, but free.

_{🔄 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/215 **Author:** [@kijai](https://github.com/kijai) **Created:** 9/29/2026 **Status:** ✅ Merged **Merged:** 9/30/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `int8-l2-raster` --- ### 📝 Commits (2) - [`00a9cbe`](https://github.com/Comfy-Org/comfy-kitchen/commit/00a9cbeefdb5f0e820a7b3d774a2dcbd3545a554) Band the int8 stream-K tile walk to keep weights in L2 - [`92e9aa6`](https://github.com/Comfy-Org/comfy-kitchen/commit/92e9aa602e494dcea0fe1eccf985c9badc2a630a) Optimize for Ampere ### 📊 Changes **6 files changed** (+167 additions, -20 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/__init__.py` (+25 -1) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+6 -0) 📝 `comfy_kitchen/backends/cuda/ops/cutlass_gemm_common.cuh` (+49 -3) 📝 `comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu` (+26 -16) 📝 `tests/test_int8.py` (+48 -0) 📝 `tests/test_w6a8.py` (+13 -0) </details> ### 📄 Description The int8 GEMM's tile order re-streams a weight that outgrows L2 on every wave. On a 5090 this read 20-26 GB per call for H3's fc1 at 32k tokens and dropped SM clocks under the power limit. The stream-K swizzle now walks N in bands, band width depends on the card's L2: - 16 MB or more (Ada, Blackwell): 32 tiles, so a band's weight slice stays resident. - Below that (Ampere: 5-6 MB): no slice can stay resident across waves, so the band is sized to minimize what one wave reads, sqrt(SMs * tile_m / tile_n) tiles (6-8 on a 3080). Measured on a 3080: within 2% of the best width at any K, where 32 was up to 47% slower. - Small-L2 cards also decode W4A8/W6A8 weights whole instead of in chunks: a chunk that cannot stay in L2 only adds launches (up to 20% faster at M=512 on a 3080). Bit-identical output, no new kernel instantiations. End to end gain in MiniMax-H3 workflow is modest ~2%, but free. --- <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:47 +00:00

closed this issue
added the
pull-request
label

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#226

No description provided.

Rows
Columns