Skip to content

FlashAttention FP8: Added support for asymmetric head dimensions, variable sequence lengths (varlen), and KV splitting - #1097

Merged
binding7012 merged 9 commits into
mainfrom
binding_fa_more_feature_support
Sep 8, 2026
Merged

binding7012 merged 9 commits into
mainfrom
binding_fa_more_feature_support

Conversation

@binding7012

@binding7012 binding7012 commented Sep 6, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

FlashAttention FP8: Added support for asymmetric head dimensions, variable sequence lengths (varlen), and KV splitting

Perf: MI355X, B=1, H=16, Hkv=1, causal, Skv = 2*Sq

image

@binding7012
binding7012 merged commit 826efd0 into main Sep 8, 2026
16 checks passed
@binding7012
binding7012 deleted the binding_fa_more_feature_support branch September 8, 2026 02:21
akii96 added a commit to akii96/FlyDSL that referenced this pull request Sep 10, 2026
Conflict in kernels/attention/flash_attn_interface.py.

Upstream ROCm#1097 (fp8 asymmetric head dims / varlen / split-K) hoisted the fp8
route out of the dense `else:` branch, which de-nested the dualwave-vs-generic
routing by one level. This branch had edited that same block at the old
nesting depth, so the whole region conflicted.

Resolved by keeping upstream's restructured control flow and re-applying the
two things this branch adds inside it:
  - causal_lpt=causal_lpt is not False on the dense dualwave build
  - the XCD-swizzle rationale comment

Verified: black + ruff clean; all 20 causal_lpt sites intact; numerics match
SDPA on gfx950 (max abs err <= 0.011) across dense/causal/generic routes and
for causal_lpt in {None, True, False}.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants