Comfy-Org / Comfy-Org/comfy-kitchen

int8_attention: OverflowError past 299,520 key rows (int32 strides), reachable on long-video workloads

Open
#178 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

`int8_attention` raises `OverflowError: sage_sdpa: tensor strides exceed int32 range; reduce batch/seq/head dimensions` once the key sequence passes **299,520 rows** at 56 heads × 128, with batch size 1.

This is reachable on ordinary long-video work rather than an exotic shape. It came up generating a 50-second MiniMax H3 video at 720p: the packed sequence is 314,437 rows, and the call that fails is the one where the text and reference rows attend to the whole timeline.

## Where the boundary is

Bisected on a 5090:

```
largest working key rows: 299520 (first failing: 299521)
k.stride(0) at 299520 = 2,146,959,360 int32 max = 2,147,483,647
```

The interesting part: **299,520 works and 299,521 fails, but neither stride exceeds int32.** At 299,521 the batch stride is 2,147,425,792, still under the limit — so something internal (padding to `cta_k`, or a derived stride) crosses it first. 299,520 is exactly `64 × 4680`, i.e. `CTA_K`-aligned, which fits that guess but I have not verified it.

Also worth noting: **batch is 1 here**, so `stride(0)` is never used to index anything. If the check is on the batch stride specifically, skipping it when `batch == 1` might be a cheap fix for the common single-batch inference case.

## Repro

Needs about 14 GB of VRAM, since the overflow is a function of element count.

```python
import torch, comfy_kitchen

DEV, HEADS, HD = "cuda", 56, 128
for rows in (299520, 299521):
q = torch.randn(1, HEADS, 512, HD, device=DEV, dtype=torch.bfloat16)
k = torch.randn(1, HEADS, rows, HD, device=DEV, dtype=torch.bfloat16)
v = torch.randn(1, HEADS, rows, HD, device=DEV, dtype=torch.bfloat16)
print("%d rows, k.stride(0)=%d" % (rows, k.stride(0)), end=" -> ")
try:
comfy_kitchen.int8_attention(q, k, v, scale=HD ** -0.5)
print("ok")
except Exception as e:
print("%s: %s" % (type(e).__name__, e))
del q, k, v
torch.cuda.empty_cache()
```

## What it means in practice

The limit is a **row count**, so it arrives sooner the larger the frame — for MiniMax H3 at 24 fps, roughly 101 s at 832×480 but only 47 s at 720p. Video models are exactly the workload that reaches it, and a hard raise means a long render dies at the last step rather than degrading.

For what it is worth, the error is a good one: it refuses rather than computing something wrong, and `OverflowError` is specific enough to catch. We now do catch it, remember the width, and send those rows to PyTorch attention while everything narrower keeps using Kitchen — so this is not blocking us. But an attention backend quietly having a maximum sequence length is worth documenting even if it is not raised, and int64 offsets (or the `batch == 1` shortcut) would remove it.

Related: #177, which is about the masked path's cost on the same workload.

## Environment

`comfy-kitchen` 0.2.31 · RTX 5090, driver 616.86 · torch 2.9.1+cu130 · Python 3.10 · Windows 11

---

## Related: the memory cost is not obvious either

Mentioning it here since it is the same call path. `_int8_attention_cuda` allocates `q_int8`,
`k_int8`, `v_int8` and their scales on top of the bf16 tensors it is handed. On one block of the
same workload — 4096 queries against a 58,000-key window, 56 heads of 128 — that is roughly:

```
k_int8 416 MB
v_int8 416 MB (padded to cta_k)
q_int8 29 MB
output 59 MB
------
~920 MB where PyTorch SDPA allocates the 59 MB output alone
```

Which is a fair trade for 3.4× and not a bug. But it does mean a workload can be fast enough and
still not fit, and the failure mode is a hard `OutOfMemoryError` mid-render rather than anything
recoverable from inside the kernel. A note in the docs that the int8 path trades memory for speed
would have saved us a render; we now catch it and fall back to PyTorch for the affected calls.

Contributor guide

Open the contributing guide

Research direction

Start at the `int8_attention` and `_int8_attention_cuda` entry points, then run the supplied 299520/299521-row reproduction to inspect the failing stride or derived offset. Compare the int32 checks with the CTA_K padding behavior and the batch-size-one case; done means the boundary is explained and the chosen fix or limitation is covered by a regression test or documentation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
43/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.