Comfy-Org/comfy-kitchen

Fork 0

mirror of https://github.com/Comfy-Org/comfy-kitchen.git synced 2026-11-04 23:37:42 +00:00

[GH-ISSUE #101] Feature Request: Direct int8 Q/K passthrough between Comfy Kitchen and SageAttention #32

New issue

Closed

opened 2026-10-26 22:19:25 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-26 22:19:25 +00:00

Owner

Originally created by @LuXuxue on GitHub (Aug 9, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/101

1. Summary

When an int8-quantized model is served in ComfyUI (weights quantized by
Comfy Kitchen), the Q/K/V projections emit fp16 activations even though the
GEMM itself ran in int8 (int32 accumulator). SageAttention then re-quantizes
Q and K back to int8 with its own mean_seq + quant_qk_int8 kernels. The
fp16 round-trip is pure overhead: it costs two extra quantization passes
(precision loss), two extra kernel launches and full-size fp16 tensor
read/writes on a memory-bandwidth-limited iGPU.

This request is to add a direct int8 Q/K passthrough: Comfy Kitchen's
int8 GEMM optionally emits per-block-int8 Q/K + scale in the exact layout
SageAttention already consumes, and SageAttention exposes the low-level
"external int8 Q/K + scale" entry point it already has internally as a
first-class API. The fp16 intermediate disappears entirely.

Current:  int8 GEMM ──dequant──▶ fp16 Q/K ──mean+quant──▶ int8 Q/K + scale ──▶ int8 attn
Proposed: int8 GEMM ──block-quant──▶ int8 Q/K + scale ───────────────────────▶ int8 attn

2. Motivation

2.1 The fp16 round-trip is redundant

Comfy Kitchen's int8_linear (all backends) works as:

x (fp16) → rowwise int8 quant → int8 GEMM (WMMA i8×i8→i32)
         → int32 result × (x_scale × w_scale) → fp16 output

SageAttention's int8 path (qk_int8_pv_fp16) then does, on the fp16
activations:

mean_seq(k)      → k_mean (smooth_k)
quant_qk_int8    → per-block int8 Q/K + q_scale/k_scale
qk_int8_sv_bf16_attn → int8 QK WMMA + fp16 PV

The Q/K int8 representation that SageAttention wants is exactly the
representation the GEMM already produced internally (before dequantization).
Nothing is gained by materializing fp16 in between; both precision and
bandwidth are lost.

2.2 Measured cost of the quantization stage

Measured on AMD Radeon 780M (gfx1103, 12 CU), ROCm 7.14, torch 2.12,
SageAttention RDNA3 HIP-native backend, all int8-path self-attention cases in

benchmark_attn.py

(full pipeline vs. per-stage, same-process, warmed up):

Case	dtype	head_dim	seq	full (ms)	mean_seq	quant_qk_int8	Q/K quant total
Anima01	bf16	128	4096	18.17	1.4%	12.4%	13.8%
Anima03	bf16	128	6144	38.35	1.0%	8.5%	9.5%
Anima05	bf16	128	9216	82.68	0.7%	5.8%	6.5%
SDXL01	fp16	64	4096	5.24	1.3%	5.1%	6.4%
SDXL07	fp16	64	6144	11.54	0.9%	3.3%	4.2%
SDXL13	fp16	64	9216	24.73	0.7%	2.3%	3.0%

The quantization stage (mean_seq + quant_qk_int8) costs 3–14% of total
attention time, with the biggest share on head_dim=128 (Anima-class) cases
(6.5–13.8%). The d128 quant_qk_int8 kernel runs ~3.6× slower than its
bandwidth roofline, i.e. there is no cleverness to keep here — the whole stage
is best skipped.

Additional bandwidth (per inference step, Anima01 shape: b=1, h=16, s=4096,
d=128):

Tensor	current (fp16 round-trip)	proposed (int8 passthrough)	saved
Q write (GEMM out)	16 MB fp16	8 MB int8	8 MB
K write (GEMM out)	16 MB fp16	8 MB int8	8 MB
Q read (quant)	16 MB	—	16 MB
K read (mean+quant)	32 MB	—	32 MB
int8 Q/K write	16 MB	16 MB (same)	0
Total per step	96 MB	32 MB	64 MB (~67%)

2.3 Precision

The current chain quantizes Q/K twice (fp16 intermediate is only 11-bit
mantissa and truncates the int32 result; then int8 re-quantization). A direct
int32→int8 block quantization is a single rounding step and is strictly
less lossy. (Optionally fused with sm_scale and smooth_k semantics.)

3. Proposed change

3.1 Comfy Kitchen side

Add an "int8 output mode" to the int8 GEMM path (or a separate
int8_linear_qk entry):

int32 accumulator is quantized per block (Q: 32 rows, K: 16 rows —
SageAttention's MIN_BLK_Q / MIN_BLK_K) instead of per-row / per-tensor.
Outputs: q_int8, k_int8, q_scale, k_scale (float32), matching
SageAttention's input contract.
q_scale carries sm_scale (SageAttention folds sm_scale·log2e into the
Q scale); smooth_k (k − k_mean) can be kept as an optional flag.
V stays fp16/bf16 (see §5).

3.2 SageAttention side

Expose the existing low-level entry points as public API:

Triton path already accepts external int8 via forward(q, k, v, q_scale, k_scale, ...) — wrap it as e.g. sageattn_from_qk_int8(q8, k8, v, q_scale, k_scale, ...).
HIP-native path already has qk_int8_sv_bf16_attn_t(q8, k8, v, o, q_scale, k_scale, ...) — same wrapper.
The wrapper skips mean_seq + quant_qk_int8 entirely.

3.3 ComfyUI glue

A thin model-patch that, for quantized backbones, routes Q/K projection
outputs through the int8-output GEMM and feeds them to the SageAttention
wrapper. No changes to the attention kernel math.

4. Estimated benefit

Metric	Estimate	Basis
Attention time saving	3–14%, typical 6–10%	measured per-case quant-stage share above
iGPU Q/K bandwidth	~67% reduction in Q/K traffic (64 MB/step saved on Anima01)	bandwidth table above
Precision	single-rounding int8 vs double-rounding fp16→int8	analysis
Kernel launches	−2 (−3) per attention call	`mean_seq` + `quant_qk_int8` (+ possibly `v.to`)

Caveats:

The 3–14% is the upper bound (the stage that disappears). The int32→int8
block quantization fused into the GEMM epilogue costs a small fraction of a
standalone kernel, so net saving is close to the bound but slightly below.
Saving scales with seq length and head_dim: d128/long-seq (Anima-class)
benefit most; very short cross-attention (kv ≤ 1024, fp16 direct path) is
unaffected.

5. Backend compatibility analysis

5.1 Comfy Kitchen backends

Backend	int8_linear	Notes
`eager`	✅	pure-torch reference; block-quant epilogue is expressible in torch ops
`cuda`	✅	torch `_int_mm` / native kernels
`triton`	✅	Triton GEMM; scale semantics portable
`hip`	✅	WMMA GEMMs (RDNA3/3.5/4); RDNA2 falls back to triton/eager

All four already implement int8_linear (capability matrix in README). A
per-block quantize epilogue is a small, backend-local addition — the scale
layout contract is the only cross-backend interface needed.

5.2 SageAttention backends

Backend	HW	external int8 entry
Triton	CUDA (official focus: RTX 4090/3090), also runs on ROCm	`forward(..., q_scale, k_scale, ...)` exists
HIP native	RDNA3 (gfx11xx) fork (`sageattention-rdna3`)	`qk_int8_sv_bf16_attn_t` exists

5.3 Compatibility matrix for the proposal

Comfy Kitchen \ SageAttention	Triton (CUDA)	Triton (ROCm)	HIP native (RDNA3)
eager	✅ protocol works	✅	✅
cuda	✅	—	—
triton	✅	✅	✅
hip	✅ (via protocol)	✅	✅ direct

The key design choice: the integration contract is the data layout
(int8 Q/K + per-block float32 scales), not a kernel ABI.

Originally created by @LuXuxue on GitHub (Aug 9, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/101 ## 1. Summary When an int8-quantized model is served in ComfyUI (weights quantized by Comfy Kitchen), the Q/K/V projections emit **fp16** activations even though the GEMM itself ran in int8 (int32 accumulator). SageAttention then re-quantizes Q and K back to int8 with its own `mean_seq` + `quant_qk_int8` kernels. The fp16 round-trip is pure overhead: it costs two extra quantization passes (precision loss), two extra kernel launches and full-size fp16 tensor read/writes on a memory-bandwidth-limited iGPU. This request is to add a **direct int8 Q/K passthrough**: Comfy Kitchen's int8 GEMM optionally emits per-block-int8 Q/K + scale in the exact layout SageAttention already consumes, and SageAttention exposes the low-level "external int8 Q/K + scale" entry point it already has internally as a first-class API. The fp16 intermediate disappears entirely. ``` Current: int8 GEMM ──dequant──▶ fp16 Q/K ──mean+quant──▶ int8 Q/K + scale ──▶ int8 attn Proposed: int8 GEMM ──block-quant──▶ int8 Q/K + scale ───────────────────────▶ int8 attn ``` --- ## 2. Motivation ### 2.1 The fp16 round-trip is redundant Comfy Kitchen's `int8_linear` (all backends) works as: ``` x (fp16) → rowwise int8 quant → int8 GEMM (WMMA i8×i8→i32) → int32 result × (x_scale × w_scale) → fp16 output ``` SageAttention's int8 path (`qk_int8_pv_fp16`) then does, on the **fp16** activations: ``` mean_seq(k) → k_mean (smooth_k) quant_qk_int8 → per-block int8 Q/K + q_scale/k_scale qk_int8_sv_bf16_attn → int8 QK WMMA + fp16 PV ``` The Q/K int8 representation that SageAttention wants is exactly the representation the GEMM already produced internally (before dequantization). Nothing is gained by materializing fp16 in between; both precision and bandwidth are lost. ### 2.2 Measured cost of the quantization stage Measured on AMD Radeon 780M (gfx1103, 12 CU), ROCm 7.14, torch 2.12, SageAttention RDNA3 HIP-native backend, all int8-path self-attention cases in [benchmark_attn.py](https://github.com/user-attachments/files/30873562/benchmark_attn.py) (full pipeline vs. per-stage, same-process, warmed up): | Case | dtype | head_dim | seq | full (ms) | mean_seq | quant_qk_int8 | **Q/K quant total** | |-----------------|-------|----------|-----|-----------|----------|---------------|---------------------| | Anima01 | bf16 | 128 | 4096 | 18.17 | 1.4% | 12.4% | **13.8%** | | Anima03 | bf16 | 128 | 6144 | 38.35 | 1.0% | 8.5% | **9.5%** | | Anima05 | bf16 | 128 | 9216 | 82.68 | 0.7% | 5.8% | **6.5%** | | SDXL01 | fp16 | 64 | 4096 | 5.24 | 1.3% | 5.1% | **6.4%** | | SDXL07 | fp16 | 64 | 6144 | 11.54 | 0.9% | 3.3% | **4.2%** | | SDXL13 | fp16 | 64 | 9216 | 24.73 | 0.7% | 2.3% | **3.0%** | The quantization stage (`mean_seq` + `quant_qk_int8`) costs **3–14% of total attention time**, with the biggest share on head_dim=128 (Anima-class) cases (6.5–13.8%). The d128 `quant_qk_int8` kernel runs ~3.6× slower than its bandwidth roofline, i.e. there is no cleverness to keep here — the whole stage is best skipped. Additional bandwidth (per inference step, Anima01 shape: b=1, h=16, s=4096, d=128): | Tensor | current (fp16 round-trip) | proposed (int8 passthrough) | saved | |--------|---------------------------|------------------------------|-------| | Q write (GEMM out) | 16 MB fp16 | 8 MB int8 | 8 MB | | K write (GEMM out) | 16 MB fp16 | 8 MB int8 | 8 MB | | Q read (quant) | 16 MB | — | 16 MB | | K read (mean+quant) | 32 MB | — | 32 MB | | int8 Q/K write | 16 MB | 16 MB (same) | 0 | | **Total per step** | **96 MB** | **32 MB** | **64 MB** (~67%) | ### 2.3 Precision The current chain quantizes Q/K twice (fp16 intermediate is only 11-bit mantissa and truncates the int32 result; then int8 re-quantization). A direct int32→int8 block quantization is a **single** rounding step and is strictly less lossy. (Optionally fused with `sm_scale` and smooth_k semantics.) --- ## 3. Proposed change ### 3.1 Comfy Kitchen side Add an **"int8 output mode"** to the int8 GEMM path (or a separate `int8_linear_qk` entry): - int32 accumulator is quantized per **block** (Q: 32 rows, K: 16 rows — SageAttention's `MIN_BLK_Q` / `MIN_BLK_K`) instead of per-row / per-tensor. - Outputs: `q_int8`, `k_int8`, `q_scale`, `k_scale` (float32), matching SageAttention's input contract. - `q_scale` carries `sm_scale` (SageAttention folds `sm_scale·log2e` into the Q scale); smooth_k (`k − k_mean`) can be kept as an optional flag. - V stays fp16/bf16 (see §5). ### 3.2 SageAttention side Expose the existing low-level entry points as public API: - Triton path already accepts external int8 via `forward(q, k, v, q_scale, k_scale, ...)` — wrap it as e.g. `sageattn_from_qk_int8(q8, k8, v, q_scale, k_scale, ...)`. - HIP-native path already has `qk_int8_sv_bf16_attn_t(q8, k8, v, o, q_scale, k_scale, ...)` — same wrapper. - The wrapper skips `mean_seq` + `quant_qk_int8` entirely. ### 3.3 ComfyUI glue A thin model-patch that, for quantized backbones, routes Q/K projection outputs through the int8-output GEMM and feeds them to the SageAttention wrapper. No changes to the attention kernel math. --- ## 4. Estimated benefit | Metric | Estimate | Basis | |--------|----------|-------| | Attention time saving | **3–14%**, typical 6–10% | measured per-case quant-stage share above | | iGPU Q/K bandwidth | **~67% reduction** in Q/K traffic (64 MB/step saved on Anima01) | bandwidth table above | | Precision | single-rounding int8 vs double-rounding fp16→int8 | analysis | | Kernel launches | −2 (−3) per attention call | `mean_seq` + `quant_qk_int8` (+ possibly `v.to`) | Caveats: - The 3–14% is the **upper bound** (the stage that disappears). The int32→int8 block quantization fused into the GEMM epilogue costs a small fraction of a standalone kernel, so net saving is close to the bound but slightly below. - Saving scales with seq length and head_dim: d128/long-seq (Anima-class) benefit most; very short cross-attention (kv ≤ 1024, fp16 direct path) is unaffected. --- ## 5. Backend compatibility analysis ### 5.1 Comfy Kitchen backends | Backend | int8_linear | Notes | |---------|-------------|-------| | `eager` | ✅ | pure-torch reference; block-quant epilogue is expressible in torch ops | | `cuda` | ✅ | torch `_int_mm` / native kernels | | `triton` | ✅ | Triton GEMM; scale semantics portable | | `hip` | ✅ | WMMA GEMMs (RDNA3/3.5/4); RDNA2 falls back to triton/eager | All four already implement `int8_linear` (capability matrix in README). A **per-block quantize epilogue** is a small, backend-local addition — the scale layout contract is the only cross-backend interface needed. ### 5.2 SageAttention backends | Backend | HW | external int8 entry | |---------|----|---------------------| | Triton | CUDA (official focus: RTX 4090/3090), also runs on ROCm | `forward(..., q_scale, k_scale, ...)` exists | | HIP native | RDNA3 (gfx11xx) fork (`sageattention-rdna3`) | `qk_int8_sv_bf16_attn_t` exists | ### 5.3 Compatibility matrix for the proposal | Comfy Kitchen \ SageAttention | Triton (CUDA) | Triton (ROCm) | HIP native (RDNA3) | |-------------------------------|---------------|---------------|--------------------| | eager | ✅ protocol works | ✅ | ✅ | | cuda | ✅ | — | — | | triton | ✅ | ✅ | ✅ | | hip | ✅ (via protocol) | ✅ | ✅ direct | The key design choice: **the integration contract is the data layout** (int8 Q/K + per-block float32 scales), not a kernel ABI.

bsakaguchi closed this issue

2026-10-26 22:19:25 +00:00

bsakaguchi commented

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

Author

Owner

@0xDELUXA commented on GitHub (Aug 11, 2026):

Don’t open issues when the HIP side of comfy-kitchen’s int8 attention hasn’t even been merged yet.

@0xDELUXA commented on GitHub (Aug 11, 2026): Don’t open issues when the HIP side of comfy-kitchen’s int8 attention hasn’t even been merged yet.

bsakaguchi commented

2026-10-26 22:19:27 +00:00

Author

Owner

@0xDELUXA commented on GitHub (Aug 11, 2026):

Beyond timing, there is a structural problem with this as specified.

The diagram assumes int8 GEMM -> attention, but for every modern DiT the real path is int8 GEMM -> fp16 -> RMSNorm + RoPE -> attention. That is why comfy-kitchen ships nine rope variants (apply_rope, rms_rope, rms_rope_split_half, ...). RoPE rotates channel pairs, so it cannot run on int8 with block scales without dequantizing, which is the round trip this asks to remove. The fp16 intermediate is load
bearing, not redundant.

The benefit also lands on the wrong side of that: the largest wins listed are the d128 long sequence Anima cases, which are RoPE models. The cases where passthrough could work (UNet cross attention with no RoPE on Q/K) are the ones measured at 3-6%.

"The contract is the data layout" also does not hold. The CUDA and HIP backends already use different scale layouts on purpose: the HIP one is per query row and per 16 keys because it follows RDNA's WMMA fragment layout, which is not m16n8k16. One shared layout would force a backend onto foreign fragment geometry and give back what it is trying to win.

Note also that mean_seq is already gone: upstream removed smooth_k, and --use-ck-attention routes to comfy-kitchen's own kernel, which quantizes once internally.

@0xDELUXA commented on GitHub (Aug 11, 2026): Beyond timing, there is a structural problem with this as specified. The diagram assumes `int8 GEMM -> attention`, but for every modern DiT the real path is `int8 GEMM -> fp16 -> RMSNorm + RoPE -> attention`. That is why comfy-kitchen ships nine rope variants (`apply_rope, rms_rope, rms_rope_split_half`, ...). RoPE rotates channel pairs, so it cannot run on int8 with block scales without dequantizing, which is the round trip this asks to remove. The fp16 intermediate is load bearing, not redundant. The benefit also lands on the wrong side of that: the largest wins listed are the d128 long sequence Anima cases, which are RoPE models. The cases where passthrough could work (UNet cross attention with no RoPE on Q/K) are the ones measured at 3-6%. "The contract is the data layout" also does not hold. The CUDA and HIP backends already use different scale layouts on purpose: the HIP one is per query row and per 16 keys because it follows RDNA's WMMA fragment layout, which is not m16n8k16. One shared layout would force a backend onto foreign fragment geometry and give back what it is trying to win. Note also that `mean_seq` is already gone: upstream removed `smooth_k`, and `--use-ck-attention` routes to comfy-kitchen's own kernel, which quantizes once internally.

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#32

No description provided.

Rows
Columns