[PR #172] Use int64 offsets in the Triton INT8 GEMM kernels #192

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/172
Author: @neuregex
Created: 9/11/2026
Status: 🔄 Open

Base: main ← Head: fix/int8-gemm-int64-offsets

📝 Commits (1)

525c1ff Use int64 offsets in the Triton INT8 GEMM kernels

📊 Changes

2 files changed (+42 additions, -4 deletions)

View changed files

📝 comfy_kitchen/backends/triton/quantization.py (+6 -4)
📝 tests/test_int8.py (+36 -0)

📄 Description

Fixes #136

Both Triton INT8 GEMM kernels (_int8_matmul_dequant_kernel and _int8_matmul_dequant_per_row_kernel) compute the output offset stride_cm * offs_am in int32, so it wraps once m * n >= 2**31 and the kernel hits an illegal memory access. This casts offs_am/offs_bn to int64 in both kernels, as @mjungnickel18 suggested in the issue.

Test: test_triton_int8_linear_output_offsets_past_int32 (per-tensor and per-row weight scale). With m=32769, n=65536, k=16 the last row starts at offset 2**31; it checks the last two rows against a float32 reference. Needs ~6 GiB of free VRAM, marked slow.

Checked on an RTX 3090 Ti (WSL2, torch 2.11.0+cu128, triton 3.6.0):

before: illegal memory access in both kernels; after: 0 mismatches on the rows past 2**31
tests/test_int8.py: 64 passed, 39 skipped, no new failures

Throughput (the open question in the issue): median int64/int32 time over 30 interleaved batches, with an identical int32 copy as control (within ±0.4%):

H3 QKV shape (m=8192, n=21504, k=5376): 0.994 per-tensor, 0.995 per-row
DiT MLP shape (m=16384, n=3072, k=3072): 0.992 per-tensor, 1.014 per-row

So gating the int64 path on m * n doesn't seem necessary.

_{🔄 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/172 **Author:** [@neuregex](https://github.com/neuregex) **Created:** 9/11/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/int8-gemm-int64-offsets` --- ### 📝 Commits (1) - [`525c1ff`](https://github.com/Comfy-Org/comfy-kitchen/commit/525c1ff1ed397543295024400aa4d7f59bd5c3b9) Use int64 offsets in the Triton INT8 GEMM kernels ### 📊 Changes **2 files changed** (+42 additions, -4 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/triton/quantization.py` (+6 -4) 📝 `tests/test_int8.py` (+36 -0) </details> ### 📄 Description Fixes #136 Both Triton INT8 GEMM kernels (`_int8_matmul_dequant_kernel` and `_int8_matmul_dequant_per_row_kernel`) compute the output offset `stride_cm * offs_am` in int32, so it wraps once `m * n >= 2**31` and the kernel hits an illegal memory access. This casts `offs_am`/`offs_bn` to int64 in both kernels, as @mjungnickel18 suggested in the issue. **Test:** `test_triton_int8_linear_output_offsets_past_int32` (per-tensor and per-row weight scale). With m=32769, n=65536, k=16 the last row starts at offset 2**31; it checks the last two rows against a float32 reference. Needs ~6 GiB of free VRAM, marked `slow`. **Checked** on an RTX 3090 Ti (WSL2, torch 2.11.0+cu128, triton 3.6.0): - before: illegal memory access in both kernels; after: 0 mismatches on the rows past 2**31 - `tests/test_int8.py`: 64 passed, 39 skipped, no new failures **Throughput** (the open question in the issue): median int64/int32 time over 30 interleaved batches, with an identical int32 copy as control (within ±0.4%): - H3 QKV shape (m=8192, n=21504, k=5376): 0.994 per-tensor, 0.995 per-row - DiT MLP shape (m=16384, n=3072, k=3072): 0.992 per-tensor, 1.014 per-row So gating the int64 path on `m * n` doesn't seem necessary. --- <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:26 +00:00

bsakaguchi referenced this issue

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

[PR #192] [MERGED] Add zero_pad and strided NDHWC views to conv3d and GroupNorm #209

bsakaguchi referenced this issue

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

[PR #202] [MERGED] HIP port of #192: zero_pad and strided NDHWC views for conv3d and GroupNorm #216