Kunshang/flash attn interface #11
Open
Conversation
This file contains hidden or 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
Signed-off-by: Kunshang Ji <[email protected]>
Signed-off-by: Kunshang Ji <[email protected]>
* add cutlass Signed-off-by: Kunshang Ji <[email protected]> * fix import Signed-off-by: Kunshang Ji <[email protected]> --------- Signed-off-by: Kunshang Ji <[email protected]>
This was referenced Aug 21, 2025
YizhouZ
force-pushed
the
kunshang/flash_attn_interface
branch
from
fb1b3ac to
0c93a3b
Compare
Signed-off-by: Yizhou Wang <[email protected]>
Signed-off-by: Yizhou Wang <[email protected]>
Signed-off-by: Yizhou Wang <[email protected]>
yma11
reviewed
Sep 8, 2025
| int head_size; | ||
| int max_blocks_per_seq; | ||
| int block_size; | ||
| bool is_causal; |
There was a problem hiding this comment.
please add a placeholder for sink support s_aux: Optional[torch.Tensor] = None,.
jikunshang
reviewed
Sep 12, 2025
| if(cuType == CutlassType::half) { | ||
| FMHAKernel<typename chunk_policy::ShapeQK, typename chunk_policy::ShapePV, | ||
| typename chunk_policy::ShapeOutPut, typename chunk_policy::SubgroupLayout, PipelineStages, | ||
| cutlass::half_t, XE_8x16x16_F32F16F16F32_TT>::dispatch(queue, args); |
There was a problem hiding this comment.
qq: do we support bf16?
| } else { | ||
| TORCH_INTERNAL_ASSERT( | ||
| false, | ||
| ""); |
There was a problem hiding this comment.
Please add some error log
| @@ -0,0 +1,127 @@ | |||
| /*************************************************************************************************** | |||
There was a problem hiding this comment.
maybe we can remove this file since the example is also removed.
| is_causal); | ||
|
|
||
| if(return_softmax) { | ||
| auto softmax_lse = torch::empty_like(out); |
There was a problem hiding this comment.
seems it will always return an empty tensor, please add some FIXME if not support for now.
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.
First PR for cutlass chunk_prefill