SIGN IN SIGN UP

replace enable_gpa with repeat_interleave for krea2 (#14523)

* Krea 2: repeat key/value heads instead of `enable_gqa`

Krea 2 always attends with a text padding mask, and no fused SDPA kernel takes
a mask together with mismatched query/key head counts — flash rejects the mask,
the memory-efficient kernel rejects the mismatch. Attention therefore fell back
to the math backend, which materializes the full
[batch_size, num_heads, seq_len, seq_len] score matrix with no error or warning.

Repeating the key/value heads in the processor computes the same thing and keeps
the memory-efficient kernel eligible: at 1024x1024 (48/12 heads, 4608 tokens) one
attention call goes from 9.02 GiB / 26.7 ms to 0.16 GiB / 4.1 ms.

It also unpins the model from the native backend, since cuDNN, flash, FA3, sage
and the hub kernels all raise on `enable_gqa`.

Fixes #14518

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* docs: how to choose between `enable_gqa` and repeating key/value heads

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* docs: recommend a choice rather than just reporting one

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* docs: add enable_gqa reference model, reword performance note

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Y
YiYi Xu committed
11a82a15fe473ed974ff35111dd629b05fb1b3ed
Parent: 425b113
Committed by GitHub <noreply@github.com> on 8/18/2026, 8:15:09 PM