huggingface / huggingface/pytorch-image-models

[FEATURE] [RFC] Support for interchangable attention backends

Open
#2,607 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
37.1k
Forks
5.2k
Avg merge
1d 11h
Merged PRs (30d)
37

Description

**Is your feature request related to a problem? Please describe.**
Currently, many models rely on a standard multi-head self-attention operator. Timm currently allows the user choose between 2 versions, an eager pytorch implementation and a fused implementation provided by PyTorch (`torch.nn.functional.scaled_dot_product_attention`), along with the 3 implementations available through PyTorch (FA2, memory-efficient attention, eager). This can be restrictive (better implementation available elsewhere, upstream issues that prevent PT SDPA from working correctly) or leave performance on the table (FA3 and other newer implementations). Adding more supported backends to timm for the user to choose from (and eventually allowing the user to register their own) will alleviate this restriction. Overall, the current way eager vs sdpa is handled is also somewhat hacky imo.

**Describe the solution you'd like**
My thoughts are to create a registry for backends, similar to how models are managed. Supported backends should be attempted to be imported (flash_attn, xformers, others) and registered on success. The user should also have access to this, if they want to provide some other implementation with the same call signature. I'm not sure if this is the best approach.

**Describe alternatives you've considered**
Alternatives would be to modify/monkeypatch the model implementation to call another attention implementation. Not sure of how necessary this is, since I'm not sure of the performance advantages of FA3/others over PT sdpa for vision models. Part of the reason other libraries keep an attention impl registry seems to be that language has much more variation in attention compared to vision.

Contributor guide

Open the contributing guide

Research direction

Start by examining timm's current eager-versus-SDPA attention handling, then compare the proposed registry approach with the existing model-management registry. Done should include a decided design for optional backend discovery, user registration, and selecting among available attention implementations.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.