int8_attention: masked path costs ~1.8x unmasked, and a request for optional return_lse

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

Start with the supplied Python repro and the comfy_kitchen.int8_attention entry point, confirming the masked and unmasked timings on the stated shapes. Read the implementation that handles attn_mask and output normalization. Done means the requested behavior is implemented and the repro verifies the intended performance and optional log-sum-exp results.

Written by the indexing model from the issue text.