flash-attention

mirror of https://github.com/Dao-AILab/flash-attention.git synced 2026-10-12 02:01:54 +00:00

Files

T

History

Driss Guessous 94e22c9066 [CuTe, SM100] Key MLA forward compilation on the KV head count (#2978 )

The MLA forward kernels take nheads_kv as a compile-time parameter, but the
forward compile key only recorded the per-KV-head query count. Shapes with
equal GQA ratios and different KV head counts, such as 128/2 and 64/1, reused
one kernel specialized for the wrong packed layout and scheduler head count:
reusing a kernel built for fewer KV heads never writes the extra head, and the
reverse launches a phantom head whose rows alias the next token.

stack-info: PR: https://github.com/Dao-AILab/flash-attention/pull/2978, branch: drisspg/stack/78

2026-10-06 21:28:37 -07:00

cute

[CuTe, SM100] Key MLA forward compilation on the KV head count (#2978 )

2026-10-06 21:28:37 -07:00

layers

[Rotary] Don't need to wrap in custom_op, just need wrap_triton

2025-04-25 17:14:20 -04:00

losses

[CrossEntropy] Support precomputed LSE

2024-09-08 09:24:43 -07:00

models

fix: add weights_only=True to all torch.load call sites (#2622 )