[PR #267] Add long key sequence tests for int8 attention #271

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/267
Author: @potoior
Created: 10/22/2026
Status: 🔄 Open

Base: main ← Head: long-sequence-int8-attention-tests

📝 Commits (1)

d57a90c Add long key sequence tests for int8 attention

📊 Changes

1 file changed (+50 additions, -0 deletions)

View changed files

📝 tests/test_int8_attention.py (+50 -0)

📄 Description

Follow-up to #178 with verification and regression coverage.

Context

#178 reported an OverflowError from int8_attention past 299,520 key rows (comfy-kitchen 0.2.31). On current master the boundary no longer exists: commit b54f851 ("Support larger shapes in cuda int8 attention") rewrote the stride checks, split within-head offsets from batch/head bases, and gave the launcher a 64-bit Offset variant for large base strides. That commit was never linked to the issue, so the fix went unnoticed there.

Verified on main (NVIDIA Thor, sm_110, torch 2.11.0+cu130), comparing against scaled_dot_product_attention with the usual nrmse < 0.03 gate:

shape (B : H_kv : key rows)	result
1 : 56 : 299520 (old working boundary)	OK, nrmse 0.0175
1 : 56 : 299521 (old failing boundary, CTA_K-unaligned)	OK, nrmse 0.0175
1 : 56 : 314437 (the 50 s 720p MiniMax H3 case from the issue)	OK, nrmse 0.0175
2 : 64 : 320000 (multi-batch, per-term uint32 products near the limit)	OK, nrmse 0.0175
1 : 64 : 540000 (head-base products past uint32 → 64-bit `Offset` variant)	OK, nrmse 0.0177

The last row is the interesting one: the 64-bit Offset head-base path had no test coverage at all before this change.

Tests added (tests/test_int8_attention.py)

test_int8_attention_long_key_sequences — the issue's boundary shapes including the CTA_K-unaligned one
test_int8_attention_long_key_sequences_multi_batch — multi-batch long key sequences
test_int8_attention_wide_head_base_offsets — a shape that forces the launcher's 64-bit Offset variant

All three skip without a CUDA int8-attention extension, are marked slow, and skip when free VRAM is insufficient (the wide-offset case needs ~25 GB).

Remaining limits (documented here for the issue): within-head offsets are still int32 by design, so kv_len * head_dim > 2^31 (≈ 16.7 M rows at head_dim 128) still raises a clean overflow_error — unreachable on current hardware. The HIP backend passes all strides as int64 and is not affected.

Tests run: the 5 new cases pass; 30 existing comparison tests in the same file still pass.

_{🔄 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/267 **Author:** [@potoior](https://github.com/potoior) **Created:** 10/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `long-sequence-int8-attention-tests` --- ### 📝 Commits (1) - [`d57a90c`](https://github.com/Comfy-Org/comfy-kitchen/commit/d57a90c2992b08279514aa76bde4bdbb082988cb) Add long key sequence tests for int8 attention ### 📊 Changes **1 file changed** (+50 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `tests/test_int8_attention.py` (+50 -0) </details> ### 📄 Description Follow-up to #178 with verification and regression coverage. **Context** #178 reported an `OverflowError` from `int8_attention` past 299,520 key rows (comfy-kitchen 0.2.31). On current master the boundary no longer exists: commit b54f851 ("Support larger shapes in cuda int8 attention") rewrote the stride checks, split within-head offsets from batch/head bases, and gave the launcher a 64-bit `Offset` variant for large base strides. That commit was never linked to the issue, so the fix went unnoticed there. **Verified on main** (NVIDIA Thor, sm_110, torch 2.11.0+cu130), comparing against `scaled_dot_product_attention` with the usual nrmse < 0.03 gate: | shape (B : H_kv : key rows) | result | |---|---| | 1 : 56 : 299520 (old working boundary) | OK, nrmse 0.0175 | | 1 : 56 : 299521 (old failing boundary, CTA_K-unaligned) | OK, nrmse 0.0175 | | 1 : 56 : 314437 (the 50 s 720p MiniMax H3 case from the issue) | OK, nrmse 0.0175 | | 2 : 64 : 320000 (multi-batch, per-term uint32 products near the limit) | OK, nrmse 0.0175 | | 1 : 64 : 540000 (head-base products past uint32 → 64-bit `Offset` variant) | OK, nrmse 0.0177 | The last row is the interesting one: the 64-bit `Offset` head-base path had no test coverage at all before this change. **Tests added** (`tests/test_int8_attention.py`) - `test_int8_attention_long_key_sequences` — the issue's boundary shapes including the CTA_K-unaligned one - `test_int8_attention_long_key_sequences_multi_batch` — multi-batch long key sequences - `test_int8_attention_wide_head_base_offsets` — a shape that forces the launcher's 64-bit `Offset` variant All three skip without a CUDA int8-attention extension, are marked `slow`, and skip when free VRAM is insufficient (the wide-offset case needs ~25 GB). **Remaining limits** (documented here for the issue): within-head offsets are still int32 by design, so `kv_len * head_dim > 2^31` (≈ 16.7 M rows at head_dim 128) still raises a clean `overflow_error` — unreachable on current hardware. The HIP backend passes all strides as int64 and is not affected. **Tests run**: the 5 new cases pass; 30 existing comparison tests in the same file still pass. --- <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:22:20 +00:00

bsakaguchi referenced this issue

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

[PR #271] sol_attn: caller-supplied block lists (block-sparse attention) #275