Comfy-Org / Comfy-Org/comfy-kitchen

INT8 GEMM kernels index the output in int32: illegal memory access once m*n >= 2**31 (MiniMax H3 at high resolution x length)

Open
#136 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
220
Forks
91
Avg merge
1d 7h
Merged PRs (30d)
12

Description

## Summary

Both Triton INT8 GEMM kernels in `comfy_kitchen/backends/triton/quantization.py`
(`_int8_matmul_dequant_kernel` and `_int8_matmul_dequant_per_row_kernel`) compute the
output address with **int32** arithmetic. As soon as `m * n >= 2**31` the offset wraps and
the kernel faults with `CUDA error: an illegal memory access was encountered`, which
aborts the whole ComfyUI process.

```python
offs_am = (pid_m * block_m + tl.arange(0, block_m)) % m # int32
offs_bn = (pid_n * block_n + tl.arange(0, block_n)) % n # int32
...
c_ptrs = c_ptr + stride_cm * offs_am[:, None] + stride_cn * offs_bn[None, :]
```

`stride_cm` is `n`, so the largest store offset is `(m-1)*n + (n-1)` and it crosses
`2**31` exactly when `m*n` does. `m` is the packed sequence length and `n` the projection
width, which is why the failure depends on **resolution x length** rather than on either
one alone.

This is what takes MiniMax H3 down on a 24 GB card above a certain size. It is independent
of dynamic VRAM, pinned memory and the attention backend.

## Minimal reproduction (no model, no weights)

One `m` per process, because the CUDA context is unusable after the fault:

```python
import sys, torch
from comfy_kitchen.backends.triton.quantization import int8_linear

m = int(sys.argv[1])
n = int(sys.argv[2]) if len(sys.argv) > 2 else 21504
k = 5376
x = torch.randn(m, k, device='cuda', dtype=torch.bfloat16)
w = torch.randint(-127, 127, (n, k), device='cuda', dtype=torch.int8)
ws = torch.rand(n, device='cuda', dtype=torch.float32) * 0.01
out = int8_linear(x, w, ws, out_dtype=torch.bfloat16)
torch.cuda.synchronize()
print("OK", tuple(out.shape))
```

`n = 21504, k = 5376` are the real MiniMax H3 QKV projection dimensions (the same shapes
as in #98).

Measured boundary, three different `n`:

| n | last good m | first failing m | m*n at the boundary |
|---|---|---|---|
| 21504 | 99,865 | 99,900 | 2.147e9 |
| 10752 | 199,000 | 200,500 | 2.147e9 |
| 5376 | 399,000 | 400,500 | 2.147e9 |

The crossing sits at `m*n = 2**31` in all three cases. Below the boundary the results are
numerically fine (~1 % against a bf16 reference, i.e. plain quantization error), so this is
a hard fault at a fixed index limit and not a gradual accuracy problem.

## Stack trace

With `CUDA_LAUNCH_BLOCKING=1` the fault is pinned to the kernel launch itself rather than
to a later synchronize:

```
File "comfy_kitchen\backends\triton\quantization.py", line 1078, in int8_linear
_int8_matmul_dequant_per_row_kernel[grid](
File "triton\runtime\autotuner.py", line 164, in _bench
return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
File "triton\runtime\jit.py", line 744, in run
kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, ...)
File "triton\backends\nvidia\driver.py", line 752, in __call__
self.launch(gridX, gridY, gridZ, stream, function, ...)
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
```

Without it the error surfaces one frame later, in `triton/testing.py do_bench ->
di.synchronize()`, which makes it look like an autotuner problem. It is not — the
autotuner is simply the first thing that launches the kernel.

## Real-world impact (MiniMax H3, RTX 3090 24 GB)

Text-to-video with `minimax_h3_fl2va_pruned_int8_convrot` and `--enable-triton-backend`:

| size | pixels x frames | before | with the fix below |
|---|---|---|---|
| 1344x768, 243 frames (10 s) | 251 M | ok | ok |
| 1344x768, 294 frames (12 s) | 303 M | **crash after 43 s** | ok, 410 s |
| 1344x768, 362 frames (15 s) | 374 M | **crash after 43 s** | ok, 620 s |

1024x576 at 15 s — more frames, fewer pixels — always worked, which matches the `m*n` law
rather than a frame-count or VRAM limit.

Ruled out here by direct test: `--disable-dynamic-vram`, `--disable-pinned-memory`,
comfy-kitchen 0.2.30 vs 0.2.31 (the kernel file is byte-identical between the two), and the
w4a8 and GGUF weight variants — w4a8 goes through the same `int8_linear` and fails at the
same shapes.

## Suggested fix

Casting the offsets to int64 in both kernels is enough:

```python
offs_am = ((pid_m * block_m + tl.arange(0, block_m)) % m).to(tl.int64)
offs_bn = ((pid_n * block_n + tl.arange(0, block_n)) % n).to(tl.int64)
```

With that applied locally, `m = 200,000, n = 21504` runs clean and the output is
numerically correct (~1 % against the reference, the same as below the old boundary), and
both H3 sizes above complete end to end. I have not measured the throughput cost of the
wider index arithmetic; if it turns out to matter, gating the int64 path on
`m * n >= 2**31` would leave the fast path untouched.

Related: #98 (same tall-M H3 shapes, but about config selection and performance) and #64
(same error text on AMD gfx1151, different backend).

## Environment

- RTX 3090 24 GB, driver 595.97, Windows 11
- comfy-kitchen 0.2.31 (and 0.2.30, identical kernel file)
- ComfyUI 0.32.0, torch 2.10.0+cu128, triton-windows 3.6.0.post26, Python 3.12.11

Contributor guide

Open the contributing guide

Research direction

Start in comfy_kitchen/backends/triton/quantization.py by reading _int8_matmul_dequant_kernel, _int8_matmul_dequant_per_row_kernel, and the int8_linear entry point. Run the minimal reproduction with a shape crossing m*n = 2**31, then verify both kernels complete without an illegal memory access and retain the stated numerical accuracy at the reported H3 sizes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.