Skip to content

fix: prefer flash attn2 public api - #15000

Open
drbh wants to merge 3 commits into
mainfrom
prefer-public-attention-apis
Open

drbh wants to merge 3 commits into
mainfrom
prefer-public-attention-apis

Conversation

@drbh

@drbh drbh commented Oct 9, 2026

Copy link
Copy Markdown

this pr update the attention dispatch to avoid using private flash attention apis that have been removed in all hub flash attention versions >1

@github-actions github-actions Bot added models size/M PR with diff < 200 LOC labels Oct 9, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@github-actions github-actions Bot added size/L PR with diff > 200 LOC and removed size/M PR with diff < 200 LOC labels Oct 9, 2026
grad_query, grad_key, grad_value = torch.empty_like(query), torch.empty_like(key), torch.empty_like(value)

_ = wrapped_backward_fn(
grad_out, query, key, value, out = (

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the grad_* tensors would need to be contiguous as well right?

@github-actions github-actions Bot added size/M PR with diff < 200 LOC and removed size/L PR with diff > 200 LOC labels Oct 9, 2026
@drbh

drbh commented Oct 9, 2026

Copy link
Copy Markdown
Author

note this pr is a work in progress - and may require some small changes upstream to the flash attn2 to improve the exposure of public apis.

will add updates here as this work progresses

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models size/M PR with diff < 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants