[PR #112] [MERGED] Fix int8 attention rejecting batch-one inputs on PyTorch 2.9 and older #153

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/112
Author: @0xDELUXA
Created: 8/12/2026
Status: ✅ Merged
Merged: 8/12/2026
Merged by: @comfyanonymous

Base: main ← Head: amd/hip-int8-attention-size-one-strides

📝 Commits (1)

2d5bf49 Exempt size-one extents from the int8 attention stride checks

📊 Changes

5 files changed (+75 additions, -24 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu (+16 -8)
📝 comfy_kitchen/backends/cuda/sage_attention/quant_v_int8.cu (+9 -4)
📝 comfy_kitchen/backends/hip/sage_attention/quant_qk_int8.hip (+13 -7)
📝 comfy_kitchen/backends/hip/sage_attention/quant_v_int8.hip (+9 -5)
📝 tests/test_int8_attention.py (+28 -0)

📄 Description

Correctness fix for the int8 attention host-side guards, in both the CUDA and the HIP backend. Reported in #110 on gfx1151 with torch 2.9.1+rocm7.2.1, where every MiniMax H3 workflow fails at the first attention call with sage int8 attention: Q/K base pointers and B/H/N strides must preserve 4-element alignment. Updating PyTorch resolves the issue.

Cause

PyTorch 2.9.1 normalizes DLPack strides for size-one dimensions to 1. MiniMax H3 passes q with shape [1, 56, 372, 128] and strides (7168, 128, 7168, 1), but the DLPack capsule becomes (1, 128, 7168, 1). The guard sees stride_b = 1 and incorrectly rejects it as misaligned.

A size-one dimension is only indexed at zero, so its stride cannot affect addressing or alignment. This affects torch 2.9.x and older.

Change

The Q/K and V int8 quantization guards now skip stride-alignment checks for dimensions with extent < 2. Alignment checks for real extents and base pointers remain unchanged.

Verification

HIP: Run-verified on gfx1200; compile-verified on gfx1100, gfx1151, and gfx1200.
Normalized DLPack strides are accepted and produce bit-identical results.
Genuine misalignment is still rejected.
V broadcast with stride_b == 0 remains supported.
Tests: 61 passed, 5 skipped (one new regression test).
CUDA: Compile-verified by CI. Not hardware-tested due to no NVIDIA GPU/nvcc. The edited guards compile cleanly with -std=c++17 -Wall -Wextra, and standalone tests matched HIP behavior.

_{🔄 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/112 **Author:** [@0xDELUXA](https://github.com/0xDELUXA) **Created:** 8/12/2026 **Status:** ✅ Merged **Merged:** 8/12/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `amd/hip-int8-attention-size-one-strides` --- ### 📝 Commits (1) - [`2d5bf49`](https://github.com/Comfy-Org/comfy-kitchen/commit/2d5bf49fb110de5a733c95f9adf99f1cf2a26e65) Exempt size-one extents from the int8 attention stride checks ### 📊 Changes **5 files changed** (+75 additions, -24 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu` (+16 -8) 📝 `comfy_kitchen/backends/cuda/sage_attention/quant_v_int8.cu` (+9 -4) 📝 `comfy_kitchen/backends/hip/sage_attention/quant_qk_int8.hip` (+13 -7) 📝 `comfy_kitchen/backends/hip/sage_attention/quant_v_int8.hip` (+9 -5) 📝 `tests/test_int8_attention.py` (+28 -0) </details> ### 📄 Description Correctness fix for the int8 attention host-side guards, in both the CUDA and the HIP backend. Reported in #110 on gfx1151 with `torch 2.9.1+rocm7.2.1`, where every MiniMax H3 workflow fails at the first attention call with `sage int8 attention: Q/K base pointers and B/H/N strides must preserve 4-element alignment`. Updating PyTorch resolves the issue. ### Cause PyTorch 2.9.1 normalizes DLPack strides for size-one dimensions to `1`. MiniMax H3 passes `q` with shape `[1, 56, 372, 128]` and strides `(7168, 128, 7168, 1)`, but the DLPack capsule becomes `(1, 128, 7168, 1)`. The guard sees `stride_b = 1` and incorrectly rejects it as misaligned. A size-one dimension is only indexed at zero, so its stride cannot affect addressing or alignment. This affects torch 2.9.x and older. ### Change The Q/K and V int8 quantization guards now skip stride-alignment checks for dimensions with extent `< 2`. Alignment checks for real extents and base pointers remain unchanged. ### Verification - HIP: Run-verified on gfx1200; compile-verified on gfx1100, gfx1151, and gfx1200. - Normalized DLPack strides are accepted and produce bit-identical results. - Genuine misalignment is still rejected. - V broadcast with `stride_b == 0` remains supported. - Tests: `61 passed, 5 skipped` (one new regression test). - CUDA: Compile-verified by CI. Not hardware-tested due to no NVIDIA GPU/`nvcc`. The edited guards compile cleanly with `-std=c++17 -Wall -Wextra`, and standalone tests matched HIP behavior. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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