[PR #235] [CLOSED] Move int8 attention conversions off the XU pipe on Ampere #244

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/235
Author: @Many0therFunctions
Created: 10/9/2026
Status: ❌ Closed

Base: main ← Head: int8-attention-xu-pipe

📝 Commits (1)

4f6ed43 Move int8 attention conversions off the XU pipe on Ampere

📊 Changes

3 files changed (+68 additions, -8 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/sage_attention/attn_utils.cuh (+15 -8)
📝 comfy_kitchen/backends/cuda/sage_attention/math.cuh (+22 -0)
📝 comfy_kitchen/backends/cuda/sage_attention/mma.cuh (+31 -0)

📄 Description

What

In the fused softmax of qk_int_sv_i8_attn_kernel, every score goes through an I2F (int32 score → fp32), an ex2 and a paired F2I + cvt.pack.sat.u8 (fp32 probability → u8), plus a __uint2float_rn on each row sum. On Ampere all of those issue on the XU pipe (16/clk/SM). Nsight Compute on an A6000 (GA102) at a MiniMax-H3 shape (B1 H56 S53730 D128) shows that pipe as the kernel's busiest, with the tensor pipe stalling behind it:

kernel	tensor	XU	FMA	ALU	LSU	top stalls (cycles/issue)	ms
int8 before	52%	56%	6%	8%	18%	wait 2.4, math_pipe_throttle 2.4, mio 1.7	547
int8 after	63%	33%	12%	16%	21%	math_pipe_throttle 2.5, wait 1.3	457
flash bf16 (torch SDPA)	82%	12%	13%	10%	24%	math_pipe_throttle 3.4, wait 1.4	745

This PR replaces the two conversions with FP32/integer arithmetic on the FMA/ALU pipes, which were nearly idle:

math::i32_to_f32_exact(x) = __int_as_float(x + 0x4B400000) - 1.5 * 2^23. Inside [2^23, 2^24) the FP32 ulp is exactly 1, so the integer add lands exactly on float(x) and the subtraction is exact (Sterbenz). Valid for |x| ≤ 2^22; the S8×S8 dot over ≤ 256 products is bounded by 2^22, and the dp4a row sum by 1020.
mma::pack_u8x4_bounded(a, b, c, d) = __float_as_uint(p + 2^23) per value, then three PRMTs to gather the low mantissa bytes. The +2^23 FADD rounds to nearest with ties to even, same as cvt.rni, and leaves the integer in the low byte. Valid for 0 ≤ p ≤ 255, which holds at both call sites because the tile max carries S_U8_OFFSET = log2(255). Masked scores (−inf bias) give exp2(−inf) = 0 under both forms; a row masked across a whole tile has a NaN exponent that cvt packs as 0 and this form as 0xFF, but update_mdo multiplies that tile's contribution by tile_scale = 0 either way, so the output is unchanged.

After the change the XU pipe carries only the ex2 itself: static SASS of the D128 instantiation goes I2F 421 → 261, F2I 193 → 65, I2IP 96 → 32 (the rest is the once-per-CTA partial-tile tail), with +288 FADD, +160 IADD3, +96 PRMT.

Architecture gate

Both helpers switch on __CUDA_ARCH__ < 890. From sm_89 on, ptxas already emits the packed I2FP.F32 / F2IP.U8.F32 forms, which issue on the FMA pipe two values at a time, so the arithmetic form only adds instructions there: measured +4% slower on an RTX 5000 Ada (668 → 695 ms at the same shape) when ungated. With the gate, Ada compiles to the original code (bit-identical, 661 ms = before). Hopper/Blackwell are unmeasured and keep the original conversions as well.

Results (A6000, sm_80 cubin from the default arch manifest)

B1 H56 S53730 D128: 553 → 466 ms (−16%, 1.56× bf16 flash SDPA)
B1 H56 S80661 D128: 1238 → 1039 ms (−16%, 1.68× flash)
Output bit-identical to the previous kernel on D64/D128/D256, a partial tile (S=1000) and both masked paths (unmasked, prepared-key mask, dense mask), ties-to-even checked in both directions.
tests/test_int8_attention.py: 419 passed; the two test_int8_attention_long_sequence_and_partial_tile[0.0 / -0.088] failures are pre-existing (the torch 2.13 flash reference returns NaN for scale <= 0; ck.int8_attention is finite there and within 0.017 nrmse of the MATH reference).

Three files, no API or Python changes. Orthogonal to #234 (tile selection / launch bounds on Ampere): the two apply cleanly on top of each other since this only touches the per-score math inside the kernel.

_{🔄 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/235 **Author:** [@Many0therFunctions](https://github.com/Many0therFunctions) **Created:** 10/9/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `int8-attention-xu-pipe` --- ### 📝 Commits (1) - [`4f6ed43`](https://github.com/Comfy-Org/comfy-kitchen/commit/4f6ed4336c60a120f289aff235225b548cb3347a) Move int8 attention conversions off the XU pipe on Ampere ### 📊 Changes **3 files changed** (+68 additions, -8 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/sage_attention/attn_utils.cuh` (+15 -8) 📝 `comfy_kitchen/backends/cuda/sage_attention/math.cuh` (+22 -0) 📝 `comfy_kitchen/backends/cuda/sage_attention/mma.cuh` (+31 -0) </details> ### 📄 Description ## What In the fused softmax of `qk_int_sv_i8_attn_kernel`, every score goes through an `I2F` (int32 score → fp32), an `ex2` and a paired `F2I` + `cvt.pack.sat.u8` (fp32 probability → u8), plus a `__uint2float_rn` on each row sum. On Ampere all of those issue on the XU pipe (16/clk/SM). Nsight Compute on an A6000 (GA102) at a MiniMax-H3 shape (B1 H56 S53730 D128) shows that pipe as the kernel's busiest, with the tensor pipe stalling behind it: | kernel | tensor | XU | FMA | ALU | LSU | top stalls (cycles/issue) | ms | |---|---|---|---|---|---|---|---| | int8 before | 52% | **56%** | 6% | 8% | 18% | wait 2.4, math_pipe_throttle 2.4, mio 1.7 | 547 | | int8 after | 63% | 33% | 12% | 16% | 21% | math_pipe_throttle 2.5, wait 1.3 | 457 | | flash bf16 (torch SDPA) | 82% | 12% | 13% | 10% | 24% | math_pipe_throttle 3.4, wait 1.4 | 745 | This PR replaces the two conversions with FP32/integer arithmetic on the FMA/ALU pipes, which were nearly idle: - `math::i32_to_f32_exact(x)` = `__int_as_float(x + 0x4B400000) - 1.5 * 2^23`. Inside [2^23, 2^24) the FP32 ulp is exactly 1, so the integer add lands exactly on float(x) and the subtraction is exact (Sterbenz). Valid for |x| ≤ 2^22; the S8×S8 dot over ≤ 256 products is bounded by 2^22, and the dp4a row sum by 1020. - `mma::pack_u8x4_bounded(a, b, c, d)` = `__float_as_uint(p + 2^23)` per value, then three `PRMT`s to gather the low mantissa bytes. The `+2^23` FADD rounds to nearest with ties to even, same as `cvt.rni`, and leaves the integer in the low byte. Valid for 0 ≤ p ≤ 255, which holds at both call sites because the tile max carries `S_U8_OFFSET = log2(255)`. Masked scores (−inf bias) give `exp2(−inf) = 0` under both forms; a row masked across a whole tile has a NaN exponent that `cvt` packs as 0 and this form as 0xFF, but `update_mdo` multiplies that tile's contribution by `tile_scale = 0` either way, so the output is unchanged. After the change the XU pipe carries only the `ex2` itself: static SASS of the D128 instantiation goes I2F 421 → 261, F2I 193 → 65, I2IP 96 → 32 (the rest is the once-per-CTA partial-tile tail), with +288 FADD, +160 IADD3, +96 PRMT. ## Architecture gate Both helpers switch on `__CUDA_ARCH__ < 890`. From sm_89 on, ptxas already emits the packed `I2FP.F32` / `F2IP.U8.F32` forms, which issue on the FMA pipe two values at a time, so the arithmetic form only adds instructions there: measured **+4% slower** on an RTX 5000 Ada (668 → 695 ms at the same shape) when ungated. With the gate, Ada compiles to the original code (bit-identical, 661 ms = before). Hopper/Blackwell are unmeasured and keep the original conversions as well. ## Results (A6000, sm_80 cubin from the default arch manifest) - B1 H56 S53730 D128: 553 → 466 ms (−16%, 1.56× bf16 flash SDPA) - B1 H56 S80661 D128: 1238 → 1039 ms (−16%, 1.68× flash) - Output **bit-identical** to the previous kernel on D64/D128/D256, a partial tile (S=1000) and both masked paths (unmasked, prepared-key mask, dense mask), ties-to-even checked in both directions. - `tests/test_int8_attention.py`: 419 passed; the two `test_int8_attention_long_sequence_and_partial_tile[0.0 / -0.088]` failures are pre-existing (the torch 2.13 flash reference returns NaN for `scale <= 0`; `ck.int8_attention` is finite there and within 0.017 nrmse of the MATH reference). Three files, no API or Python changes. Orthogonal to #234 (tile selection / launch bounds on Ampere): the two apply cleanly on top of each other since this only touches the per-score math inside the kernel. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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