perf(krea2): keep the K/V heads unexpanded where a fused kernel takes them

[Pfannkuchensack]

Summary

Performance / simplification. Krea2MemoryEfficientAttnProcessor expands Krea-2's 12 K/V heads to 48 with repeat_interleave before every SDPA call, so that enable_gqa is not needed. The module docstring gives the reason: the fused kernels did not support grouped-query attention, and the fallback is the math backend, which materialises the full [heads, seq, seq] score matrix at ~9 GB.

That premise has become build-dependent. cuDNN serves grouped-query attention, mask included. Measured on the real shape ([1, 48q/12kv, 4608, 128], bf16), RTX 4090:

Backend	GQA, unmasked	GQA, masked
cudnn	3.64 ms	3.53 ms
efficient	refuses GQA	refuses GQA
flash	absent in this build	refuses the additive mask anyway
math	46.98 ms / 9.3 GB	58.44 ms / 9.3 GB

So the expansion is now decided per call by asking the dispatcher, and skipped where a permitted kernel takes the shape as it is.

What it is worth

Live allocation at the SDPA call, and time for the whole attention block, on an RTX 4090:

seq	expanded (today)	grouped	live saved	time
4608	216.0 MB, 13.09 ms	135.5 MB, 12.44 ms	80.5 MB	1.05×
9216	432.0 MB, 36.53 ms	271.0 MB, 33.10 ms	161.0 MB	1.10×

repeat_interleave on K and V allocates +108 MB at 4608 tokens, every call, and the copy is work as well as memory.

What this does not do, stated because it is easy to read the other way: it does not lower the peak of the attention block. That peak is set after SDPA by the output path — transpose, flatten, gate, to_out — at a moment when the K/V tensors are already free, and it is unchanged at 460.8 MB (4608) and 921.7 MB (9216). An isolated-kernel measurement on a 30-series card shows 270 MB → 197 MB, which is the same 80 MB but reads like a ceiling reduction; it is not one. This is less allocator traffic and a few percent of time.

Why the check is conservative

Guessing wrong in the optimistic direction does not cost speed, it costs ~9 GB — the exact failure this processor exists to prevent. Two things follow:

• Only backends this call is permitted to reach are consulted. Asking whether cuDNN serves the shape while the resolved list excludes it — which INVOKE_KREA2_SDPA_BACKEND=efficient does — answers for a kernel the dispatcher cannot use, and the call then fails outright with No available kernel. That was a real bug in the first version of this change, caught by an existing CUDA test rather than by reasoning.
• Anything unexpected answers no. A non-CUDA tensor, or a torch build whose SDPAParams signature differs, falls back to the expansion — correct everywhere, merely costs memory. ROCm lands here too for masked blocks: no cuDNN, and flash refuses the mask.

The answer depends only on the call shape, dtype, device, mask presence and permitted backends — all fixed for a given block within a generation — so it is cached per processor and queried once per kind.

Related Issues / Discussions

Stacked on #172 (feat/krea2-cudnn-attention), which this PR's base branch is set to. It depends on that PR's ranked backend list: with the old permissive list the memory-efficient kernel would be reached first and refuse the unexpanded shape.

The measurements that prompted this are in the local .ideas/krea2-sdpa-shape-probe.py output from a 30-series card, which showed cudnn gqa=True at 19.02 ms / 197 MB against 20.15 ms / 279 MB for the expanded path.

QA Instructions

Measured on RTX 4090, torch 2.7.1+cu128, Windows; the shape probe that prompted it ran on a Linux 30-series card.

• tests/backend/krea2/test_grouped_query_attention.py — 9 tests: the decision answers no on non-CUDA tensors, no when the permitted list excludes the capable kernels, and no when SDPAParams raises; the answer is cached once per call kind and re-asked when the mask state changes; the CPU path still matches the stock Krea2AttnProcessor exactly; and on CUDA the two paths agree numerically.
• The CUDA comparison asserts which path each run took — it records the K/V head count reaching SDPA and fails if the forced run did not expand, skipping only if no fused kernel on the device serves the grouped shape. Without that it would pass just as happily if both runs expanded, which is the thing it exists to rule out. On the 4090 it passes rather than skips, so the grouped path really ran.
• tests/backend/krea2 tests/app/util tests/app/invocations — 1119 passed, 1 skipped, 6 xfailed. pytest --collect-only clean. ruff check / format --check clean. No openapi.json change.

To verify by hand: generate with Krea-2 and compare against main; the images should be indistinguishable (bf16 kernels differ in accumulation order, so the tests bound correlation at > 0.9999 and max difference at 0.05 rather than demanding equality). Then set INVOKE_KREA2_SDPA_BACKEND=efficient and generate again — that path must still work, because it falls back to the expansion.

Not run: a generation on ROCm. The masked blocks there have no fused kernel that takes the grouped shape, so they should fall back to the expansion — which is today's behaviour, and the tests cover the decision rather than the platform.

Merge Plan

Merge after #172. No DB schema, no redux slice, no API change, no dependency change.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable) — 9 tests, including a CUDA comparison that asserts which path ran
• ❗Changes to a redux slice have a corresponding migration — n/a, no redux changes
• Documentation added / updated (if applicable) — the module docstring stated the old premise as fact; it now says when it holds
• Updated What's New copy (if doing a release after this PR) — probably not worth a line on its own

🤖 Generated with Claude Code