Cleaner attention interfaces #35342
Cleaner attention interfaces #35342
Conversation
|
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. |
| # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions | ||
| # Reference: https://github.com/pytorch/pytorch/issues/112577. | ||
| if _needs_contiguous_inputs: | ||
| query = query.contiguous() | ||
| key = key.contiguous() | ||
| value = value.contiguous() |
There was a problem hiding this comment.
we could also just always apply contiguous!
There was a problem hiding this comment.
Agree, calling contiguous doesn't affect performance much either way so the comment clarification suffices imo
There was a problem hiding this comment.
Just as a heads up, maybe it would be nice to doc what attn type supports what feature, e.g. sdpa and fa cannot output attn weights or use head mask while flex attn cannot use dropout.
Which brings me to the point to ask if these things have been considered (flex attn dropout not supported for example).
| # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions | ||
| # Reference: https://github.com/pytorch/pytorch/issues/112577. | ||
| if _needs_contiguous_inputs: | ||
| query = query.contiguous() | ||
| key = key.contiguous() | ||
| value = value.contiguous() |
There was a problem hiding this comment.
Agree, calling contiguous doesn't affect performance much either way so the comment clarification suffices imo
|
LGTM |
What does this PR do?
Improve the attention interfaces with comments/dtype fix