Feature Request: Direct int8 Q/K passthrough between Comfy Kitchen and SageAttention

[LuXuxue]

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.

[0xDELUXA]

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

[0xDELUXA]

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.