Flash Attention for Neuron #939
Conversation
apoorvtintin
requested review from
ruomingp,
markblee and
a team
as code owners
There was a problem hiding this comment.
Maybe wait until this PR is checked in. From what i can tell, your PR also has the remat bug not fixed. #942 (review)
|
|
||
|
|
||
| def _mha_forward(query, key, value, bias, causal, softmax_scale, dropout_rate): | ||
| # Get the batch size, sequence lengths, number of heads, and hidden dimension |
There was a problem hiding this comment.
Nit: end comments with . (here and everywhere)
apoorvtintin
force-pushed
the
mainline_upstream_fa
branch
3 times, most recently
from
f40c4cc to
8a92182
Compare
apoorvtintin
force-pushed
the
mainline_upstream_fa
branch
from
8a92182 to
73a2808
Compare
| key: Tensor, | ||
| value: Tensor, | ||
| bias: Tensor, | ||
| causal: bool = False, |
There was a problem hiding this comment.
Can we support segment ID? Or a more general masking fn (with optimized handling) is even better.
There was a problem hiding this comment.
If not, I am fine with leaving a TODO here, but it is a hard blocker for enabling it for our internal training.
There was a problem hiding this comment.
Can we do segment IDs in a separate PR? That involves non-trivial work and needs some time.
|
Thanks for all the reviews @ruomingp @kelvin-zou. I resolved all the comments, please let me know if any more changes are needed. |
This PR adds support for flash attention kernel for Neuron implemented through Neuron Kernel Interface (NKI).
The flash attention kernel works with TRN1 and TRN2.
This PR is a newer version of #883 from a different fork. All comments from the previous PR are addressed in this one. It has dropout support.
Dropout and Segment ID support in the flash attention kernel is in progress and will be available at a later date.