Add attention_backend to let user choose #1456
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1456
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
There was a problem hiding this comment.
Thanks for making these changes! This looks good to me.
cc @Jack-Khuu
|
Looks like the ET pin needs bumping so I'll spin that up |
torchchat/generate.py
Outdated
| if self.builder_args.device == "cpu" and (self.builder_args.attention_backend == "efficient_attention" | ||
| or self.builder_args.attention_backend == "cudnn_attention"): | ||
| print(f"Warning: {self.builder_args.attention_backend} is not supported on CPU. Using math instead.") | ||
| self.builder_args.attention_backend = "math" |
There was a problem hiding this comment.
Can we bump this into the constructor of BuilderArgs?
| @@ -69,6 +69,7 @@ class BuilderArgs: | |||
| prefill_possible: bool = False | |||
| dynamic_shapes: bool = False | |||
| max_seq_length: Optional[int] = None | |||
| attention_backend: str = "math" | |||
There was a problem hiding this comment.
Thanks for adding the change!! Mind bumping the casting to the actual kernel into the body here?
That way we can do any sanity checks early and aren't passing around a string
There was a problem hiding this comment.
Have updated. Please review again!
|
@angelayi I recall you mentioning that there might be export issue with the other SDP backend? Do you know if this issue was addressed? Or should be gate the backends to non-export? |
|
Seems like yall had a chance to explore the export issue in the tagged issue Let's add a check in |
torchchat/generate.py
Outdated
| sdp_backend_dict = { | ||
| 'math': torch.nn.attention.SDPBackend.MATH, | ||
| 'flash_attention': torch.nn.attention.SDPBackend.FLASH_ATTENTION, | ||
| 'efficient_attention': torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, | ||
| 'cudnn_attention': torch.nn.attention.SDPBackend.CUDNN_ATTENTION, | ||
| } |
There was a problem hiding this comment.
Can we push this into the builder.py as well?
yanbing-j
force-pushed
the
yanbing/fix_1452
branch
from
d22fdb6 to
48f3c19
Compare
bump this into the constructor of BuilderArgs
This PR is submitted for #1452, to add an argument
attention_backendto let user choose the backend used in SDPA.For CPU, only
mathandflash_attentionare valid. For CUDA, all 4 backends are valid.