update frontend and mk_ck_lib #777
Conversation
facebook-github-bot
added
the
CLA Signed
|
@chenyang78 @ipiszy review this one too. |
|
|
||
| def forward(self, *args): | ||
| def forward(self, *args, seqlens=None): |
There was a problem hiding this comment.
What does seqlen represent? Where is seqlens used?
There was a problem hiding this comment.
It's used in rocm attention.
There was a problem hiding this comment.
I feel the API after this change is not very readable. This API is supposed to be backend-agnostic, and is shared between CUDA and ROCM. Can we avoid this change?
There was a problem hiding this comment.
rocm attention backend requires seqlen. Do you have any plan to add an API supports attention mask?
| @@ -395,7 +419,8 @@ def forward(self, *args): | |||
| else: | |||
| x = self.proj(attn_output) | |||
| x = self.proj_drop(x) | |||
| x = ops.reshape()(x, [batch, -1, self.dim]) | |||
| if not isinstance(batch, IntVar): | |||
| x = ops.reshape()(x, [batch, -1, self.dim]) | |||
There was a problem hiding this comment.
This looks weird. Why this change?
There was a problem hiding this comment.
ROCM attention support [batch * seqlen, nheads] tensor as inputs and output [batch * seqlen, nheads] tensor. The branch keeps original behavior.
|
I have reverted the attention |
|
@ipiszy has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
1 similar comment
|
@ipiszy has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
No description provided.