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 #177] int8_attention: masked path costs ~1.8x unmasked, and a request for optional return_lse #55

New issue

Open

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

bsakaguchi commented

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

Owner

Originally created by @shootthesound on GitHub (Sep 14, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/177

We run banded attention over long MiniMax H3 videos: the sequence is a packed [text | refs | audio | video] layout of up to ~266k rows, and each block of queries attends only to the keys inside a time window around it. Windows overlap, and the edge of each window carries a raised-cosine fade — a per-key additive bias, 0.0 across the core and falling to log(1e-4) across the skirt. So every call passes an attn_mask of shape [1, 1, 1, K]: one value per key, constant across heads and queries.

The measurement

4096 queries × 33,600 keys, 56 heads of 128, bf16 in, RTX 5090. Two runs, both shown:

mask	time	vs unmasked
none	8.64 / 8.45 ms	—
bf16 `[1,1,1,K]`	16.08 / 16.08 ms	1.86× / 1.90×
fp32 `[1,1,1,K]`	15.28 / 15.62 ms	1.77× / 1.85×
fp16 `[1,1,1,K]`	15.89 / 16.19 ms	1.84× / 1.92×
bf16 `[1,1,Q,K]` (expanded)	15.93 / 16.32 ms	1.84× / 1.93×
bool `[1,1,1,K]`	14.15 / 14.48 ms	1.64× / 1.71×

For contrast, the same bias costs F.scaled_dot_product_attention 1.05–1.21× on the same shapes. So masked int8 is still comfortably faster than masked SDPA — it just gives back about half of its advantage.

Two things we noticed, offered as observations rather than a diagnosis, since we have not looked inside the kernel:

It does not look bandwidth-bound. A 4-byte fp32 mask costs the same as a 2-byte bf16 one, and bool (1 byte) is only ~10% cheaper rather than several times. Differences between the float dtypes are within run-to-run noise.
A broadcast mask costs the same as a fully expanded one. [1,1,1,K] and [1,1,Q,K] are indistinguishable, so the stride-0 view does not appear to be exploited today.

Question 1

Is ~1.8× the expected cost of the masked path, or is there a cheaper route we are missing for the common "one additive bias per key, shared across heads and queries" case? That shape seems likely to recur — sliding-window and banded attention, ALiBi-style position biases, and soft-edged local attention all produce it.

Question 2: would you consider `return_lse` on `int8_attention`?

The reason is specific. The natural way to avoid the masked path is to split each window in two — the core with no mask, the skirt with one — and merge the partial softmaxes. Measured on the same shapes, that would be worth 1.38–1.41×:

core  20160 keys, unmasked   5.19 ms
skirt 13440 keys, masked     5.86 ms
                             ----
two calls                   11.05 ms   vs 15.20 ms for one masked call

But merging two softmaxes exactly requires each call's log-sum-exp, and int8_attention returns only the output. The online-softmax normaliser already exists inside the kernel; an optional flag to write it out would make this exact rather than approximate.

We would rather not approximate it — this path is quality-sensitive — so for now we take the masked cost. sageattn exposes return_lse for comparison, though it takes no attn_mask, so neither library can do the split alone today.

Happy to test any branch on a 5090 against a real workload and report back.

Repro

import math, time, torch, comfy_kitchen

DEV, HEADS, HD, Q, K = "cuda", 56, 128, 4096, 33600
SCALE = HD ** -0.5
q = torch.randn(1, HEADS, Q, HD, device=DEV, dtype=torch.bfloat16)
k = torch.randn(1, HEADS, K, HD, device=DEV, dtype=torch.bfloat16)
v = torch.randn(1, HEADS, K, HD, device=DEV, dtype=torch.bfloat16)

t = torch.linspace(0, 1, K, device=DEV)
w = torch.where(t < 0.6, torch.ones_like(t),
                0.5 * (1.0 + torch.cos(math.pi * ((t - 0.6) / 0.4).clamp(max=1.0))))
log_w = torch.log(w.clamp_min(1e-4))

def bench(fn, n=10):
    for _ in range(3):
        fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0) / n * 1000.0

base = bench(lambda: comfy_kitchen.int8_attention(q, k, v, scale=SCALE))
print("no mask %.2f ms" % base)
for name, m in (("bf16 [1,1,1,K]", log_w.to(torch.bfloat16).view(1, 1, 1, -1)),
                ("fp32 [1,1,1,K]", log_w.float().view(1, 1, 1, -1)),
                ("bf16 [1,1,Q,K]", log_w.to(torch.bfloat16).view(1, 1, 1, -1).expand(1, 1, Q, K)),
                ("bool [1,1,1,K]", (w > 1e-3).view(1, 1, 1, -1))):
    ms = bench(lambda m=m: comfy_kitchen.int8_attention(q, k, v, scale=SCALE, attn_mask=m))
    print("%-16s %.2f ms  %.2fx" % (name, ms, ms / base))

Environment

comfy-kitchen 0.2.31 · RTX 5090, driver 616.86 · torch 2.9.1+cu130 · Python 3.10 · Windows 11

Originally created by @shootthesound on GitHub (Sep 14, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/177 We run banded attention over long MiniMax H3 videos: the sequence is a packed `[text | refs | audio | video]` layout of up to ~266k rows, and each block of queries attends only to the keys inside a time window around it. Windows overlap, and the edge of each window carries a **raised-cosine fade** — a per-key additive bias, `0.0` across the core and falling to `log(1e-4)` across the skirt. So every call passes an `attn_mask` of shape `[1, 1, 1, K]`: one value per key, constant across heads and queries. ## The measurement 4096 queries × 33,600 keys, 56 heads of 128, bf16 in, RTX 5090. Two runs, both shown: | mask | time | vs unmasked | |---|---|---| | none | 8.64 / 8.45 ms | — | | bf16 `[1,1,1,K]` | 16.08 / 16.08 ms | 1.86× / 1.90× | | fp32 `[1,1,1,K]` | 15.28 / 15.62 ms | 1.77× / 1.85× | | fp16 `[1,1,1,K]` | 15.89 / 16.19 ms | 1.84× / 1.92× | | bf16 `[1,1,Q,K]` (expanded) | 15.93 / 16.32 ms | 1.84× / 1.93× | | bool `[1,1,1,K]` | 14.15 / 14.48 ms | 1.64× / 1.71× | For contrast, the same bias costs `F.scaled_dot_product_attention` **1.05–1.21×** on the same shapes. So masked int8 is still comfortably faster than masked SDPA — it just gives back about half of its advantage. Two things we noticed, offered as observations rather than a diagnosis, since we have not looked inside the kernel: - **It does not look bandwidth-bound.** A 4-byte fp32 mask costs the same as a 2-byte bf16 one, and bool (1 byte) is only ~10% cheaper rather than several times. Differences between the float dtypes are within run-to-run noise. - **A broadcast mask costs the same as a fully expanded one.** `[1,1,1,K]` and `[1,1,Q,K]` are indistinguishable, so the stride-0 view does not appear to be exploited today. ## Question 1 Is ~1.8× the expected cost of the masked path, or is there a cheaper route we are missing for the common "one additive bias per key, shared across heads and queries" case? That shape seems likely to recur — sliding-window and banded attention, ALiBi-style position biases, and soft-edged local attention all produce it. ## Question 2: would you consider `return_lse` on `int8_attention`? The reason is specific. The natural way to avoid the masked path is to split each window in two — the core with no mask, the skirt with one — and merge the partial softmaxes. Measured on the same shapes, that would be worth **1.38–1.41×**: ``` core 20160 keys, unmasked 5.19 ms skirt 13440 keys, masked 5.86 ms ---- two calls 11.05 ms vs 15.20 ms for one masked call ``` But merging two softmaxes exactly requires each call's log-sum-exp, and `int8_attention` returns only the output. The online-softmax normaliser already exists inside the kernel; an optional flag to write it out would make this exact rather than approximate. We would rather not approximate it — this path is quality-sensitive — so for now we take the masked cost. `sageattn` exposes `return_lse` for comparison, though it takes no `attn_mask`, so neither library can do the split alone today. Happy to test any branch on a 5090 against a real workload and report back. ## Repro ```python import math, time, torch, comfy_kitchen DEV, HEADS, HD, Q, K = "cuda", 56, 128, 4096, 33600 SCALE = HD ** -0.5 q = torch.randn(1, HEADS, Q, HD, device=DEV, dtype=torch.bfloat16) k = torch.randn(1, HEADS, K, HD, device=DEV, dtype=torch.bfloat16) v = torch.randn(1, HEADS, K, HD, device=DEV, dtype=torch.bfloat16) t = torch.linspace(0, 1, K, device=DEV) w = torch.where(t < 0.6, torch.ones_like(t), 0.5 * (1.0 + torch.cos(math.pi * ((t - 0.6) / 0.4).clamp(max=1.0)))) log_w = torch.log(w.clamp_min(1e-4)) def bench(fn, n=10): for _ in range(3): fn() torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(n): fn() torch.cuda.synchronize() return (time.perf_counter() - t0) / n * 1000.0 base = bench(lambda: comfy_kitchen.int8_attention(q, k, v, scale=SCALE)) print("no mask %.2f ms" % base) for name, m in (("bf16 [1,1,1,K]", log_w.to(torch.bfloat16).view(1, 1, 1, -1)), ("fp32 [1,1,1,K]", log_w.float().view(1, 1, 1, -1)), ("bf16 [1,1,Q,K]", log_w.to(torch.bfloat16).view(1, 1, 1, -1).expand(1, 1, Q, K)), ("bool [1,1,1,K]", (w > 1e-3).view(1, 1, 1, -1))): ms = bench(lambda m=m: comfy_kitchen.int8_attention(q, k, v, scale=SCALE, attn_mask=m)) print("%-16s %.2f ms %.2fx" % (name, ms, ms / base)) ``` ## Environment `comfy-kitchen` 0.2.31 · RTX 5090, driver 616.86 · torch 2.9.1+cu130 · Python 3.10 · Windows 11

bsakaguchi referenced this issue

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

[GH-ISSUE #196] INT8 ConvRot stalls on model re-entry on Windows/NVIDIA — 100% GPU at ~80W, high VRAM, sometimes OOM in int8_linear #58

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#55

No description provided.

Rows
Columns