Comfy-Org / Comfy-Org/comfy-kitchen
retune int8 GEMM tiles for gfx1103 and vectorize rowwise quantize
- Dominant language
- Python
- Stars
- 220
- Forks
- 91
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 12
Description
[ops/quantize_int8.hip](https://github.com/user-attachments/files/30875595/quantize_int8.hip.txt) should be useful for all cards.
[ops/gemm_int8.hip](https://github.com/user-attachments/files/30875597/gemm_int8.hip.txt) changed tiles for gfx1103. This should ideally be automatically selected based on arch, but implementing it requires modifying more files, so it hasn't been implemented yet.
I locally compiled the wheel file, tested it with the Anima model, and confirmed it ran faster than before the modifications. However, my own compiled wheel file (before modifications) itself is slower than the official compilation, and after modifications it is as fast as the official compilation.
The test results are based on the following two scripts.
[tile_scan.py](https://github.com/user-attachments/files/30875762/tile_scan.py)
[bench_int8.py](https://github.com/user-attachments/files/30875761/bench_int8.py)
Change 1 — GEMM tile dispatch rework (ops/gemm_int8.hip)
Replaced the old three-tier shape heuristic
M >= 512 && N >= 128 && K > N -> 256x128 tile
M >= 96 && N >= 96 -> 128x128 tile
otherwise -> 64x64 tile
with a K-based heuristic
K >= 8192 -> 128x128 tile, BKB=64 (deep K: larger tile amortizes per-tile global reads)
K < 1024 -> 64x64 tile, BKB=32 (shallow K: smaller LDS footprint -> higher occupancy)
otherwise -> 64x64 tile, BKB=64
Why: on a 12-CU iGPU the large 256x128/128x128 blocks (256 threads, 128 accumulator VGPRs) leave too few resident warps per SIMD to hide DRAM latency (200–400 ns). The small 64x64 tile wins 22/28 scanned cases — its high occupancy hides latency, and the 2 MiB L2 already covers the cross-block reuse a bigger tile would otherwise buy with extra global traffic. Deep K is the exception: with enough K-steps, the larger tile's reduction in global reads outweighs the occupancy loss. A BKB=32 K-tile further helps shallow K (K ≤ 640: −13% to −25%) at the cost of more K-loop syncs, which is why it is only used for small K.
Also added a COMFY_KITCHEN_INT8_TILE environment hook (12 tile configs, off by default) so the heuristic can be re-scanned on other hardware without rebuilding.
Change 2 — vectorized rowwise quantize (ops/quantize_int8.hip)
The rowwise quantizer was the dominant cost for many shapes — e.g. Anima05 (8192×2048 bf16) spent 2.3 ms quantizing vs 0.74 ms in the GEMM, and it was 1.7–2.4x slower than triton's. The old kernel walked each row with scalar 2/4-byte load_in calls at stride 256, twice (absmax pass + write pass).
The new kernel gives each thread whole 16-byte chunks of a row (8 f16/bf16 or 4 f32): both passes use uint4 vector loads, and the int8 writeback is packed into uint2/uint32 stores. The old scalar kernel is kept as a fallback for rows that do not split into whole chunks (K % elems != 0) or whose base pointer is not 16-byte aligned (strided torch views).
Contributor guide
Research direction
Start with ops/gemm_int8.hip and ops/quantize_int8.hip, then run tile_scan.py and bench_int8.py to reproduce the reported measurements. Review the gfx1103 tile dispatch and rowwise quantization paths, including their fallback conditions. Done means the optimized paths preserve fallback behavior and benchmarking confirms the intended gains on the relevant hardware.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- hpc, performance
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100