feat(op): support varlen npu flash attention #209
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
SolenoidWGT
force-pushed
the
feat/add_npu_fa
branch
from
a257888
to
14f0106
Compare
SolenoidWGT
force-pushed
the
feat/add_npu_fa
branch
from
82758c5
to
dbc6869
Compare
mwiacx
reviewed
May 24, 2024
| @@ -472,11 +604,13 @@ def _qkv_without_cu_seqlens(self, qkv, softmax_scale=None, causal=None, key_padd | |||
| return _torch_fixedlen_qkvpacked_attn(qkv, self.dropout, softmax_scale, causal, key_padding_mask) | |||
|
|
|||
| @forward.register(conditions=(str(QKVPackType.KVPACKED), str(CuSeqlenType.WithOut))) | |||
| def _q_kv_without_cu_seqlens(self, q, kv, softmax_scale=None, causal=None, key_padding_mask=None): | |||
| def _q_kv_without_cu_seqlens( | |||
| self, q, kv, softmax_scale=None, causal=None, key_padding_mask=None, use_flash_attn=True | |||
There was a problem hiding this comment.
不用把use_flash_attn作为attention forward的参数，这意味着model需要感知底层是否为fa的算子，我们不希望为model的开发者引入复杂度，这个东西作为算子选择系统的自动逻辑或者配置的一部分就好。
mwiacx
reviewed
May 24, 2024
| @@ -341,7 +341,7 @@ def output_hook(self, module: Embedding1D, args: Any, output: Tuple[Any]) -> Tup | |||
| """ | |||
| _emb_dim = 2 # [bsz, seqlen, emb_dim] | |||
|
|
|||
| return gather_forward_split_backward(output, self._parallel_mode, dim=_emb_dim) | |||
| return gather_forward_split_backward(output, self._parallel_mode, dim=_emb_dim), DUMMY_HANDLE_CONST | |||
There was a problem hiding this comment.
这里为啥要加DUMMY_HANDLE_CONST哇，他是作为register_forward_hook来用的，下同
mwiacx
reviewed
May 24, 2024
| @@ -78,18 +77,15 @@ def _nyi_attn(func_name, *args, **kwargs): # pylint: disable=W0613 | |||
|
|
|||
|
|
|||
| def _flash_float32_compatibility_wrapper(input_idxs: Tuple, flash_func: Callable, *args, **kwargs): | |||
| if gpc.config.model.dtype is torch.float32: | |||
|
|
||
|
|
||
| def _npu_varlen_qkvpacked_attn( | ||
| qkv: torch.Tensor, cu_seqlens, max_seqlen, dropout_p, softmax_scale=None, causal=False # pylint: disable=W0613 | ||
| def __npu_varlen_qkvsplited_attn( |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Motivation
支持
torch_npu的 var flash attention.Modification
移除了在pack数据下，给q和k进行padding和unpadding下的操作
BC-breaking (Optional)
Does the modification introduce changes that break the backward compatibility of the downstream repositories?
If so, please describe how it breaks the compatibility and how the downstream projects should modify their code to keep compatibility with this PR.
Use cases (Optional)
If this PR introduces a new feature, it is better to list some use cases here and update the documentation.
Checklist
Before PR:
After PR: