Hstu attention n0loop fused unroll pr #2896
Open
+8,340
−15
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
…uilding succeeded)
…oLocal to save vgprs for non-local situations
…st-K/last-V or last-K/first-V
…n when both causal=true and local=true
qianfengz
requested review from
illsilin,
carlushuang,
aosewski,
poyenc,
geyyer,
bartekxk,
andriy-ca,
afagaj,
asleepzzz,
ThomasNing,
coderfeli,
aska-0096,
cgmillette,
shumway,
vidyasagar-amd,
ddembeckAMD,
a team and
tenpercent
as code owners
spolifroni-amd
requested changes
Sep 22, 2025
| @@ -0,0 +1,64 @@ | |||
| # HSTU attention operator | |||
|
|
|||
| HSTU-attention operator is an operator which takes tensor `q: [batches, seqlen, nhead, hdim_qk]`, `k: [batches, seqlen, nhead, hdim_qk`, | |||
There was a problem hiding this comment.
Suggested change
| HSTU-attention operator is an operator which takes tensor `q: [batches, seqlen, nhead, hdim_qk]`, `k: [batches, seqlen, nhead, hdim_qk`, | |
| The HSTU-attention operator is an operator which takes as input three tensor `q: [batches, seqlen, nhead, hdim_qk]`, `k: [batches, seqlen, nhead, hdim_qk`, |
| # HSTU attention operator | ||
|
|
||
| HSTU-attention operator is an operator which takes tensor `q: [batches, seqlen, nhead, hdim_qk]`, `k: [batches, seqlen, nhead, hdim_qk`, | ||
| `v: [batches, seqlen, nhead, hdim_v]` and some parameters for defining the functional masking as inputs, and do the following: |
There was a problem hiding this comment.
Are the parameters for definiing functional masking, or is the operator for functional masking? It's not clear from the sentence.
But if that's what it is, then this can be changed to:
Suggested change
| `v: [batches, seqlen, nhead, hdim_v]` and some parameters for defining the functional masking as inputs, and do the following: | |
| `v: [batches, seqlen, nhead, hdim_v]`, as well as parameters that define functional masking to do the following: | |
| `` |
| HSTU-attention operator is an operator which takes tensor `q: [batches, seqlen, nhead, hdim_qk]`, `k: [batches, seqlen, nhead, hdim_qk`, | ||
| `v: [batches, seqlen, nhead, hdim_v]` and some parameters for defining the functional masking as inputs, and do the following: | ||
|
|
||
| * Multiply `q: [batches, seqlen, nhead, hdim_qk]` with `k: [batches, seqlen, nhead, hdim_k]` to get temporary tensor `s: [batches, nhead, seqlen, seqlen]` |
There was a problem hiding this comment.
Suggested change
| * Multiply `q: [batches, seqlen, nhead, hdim_qk]` with `k: [batches, seqlen, nhead, hdim_k]` to get temporary tensor `s: [batches, nhead, seqlen, seqlen]` | |
| * Multiply `q: [batches, seqlen, nhead, hdim_qk]` with `k: [batches, seqlen, nhead, hdim_k]` to get the intermediate tensor `s: [batches, nhead, seqlen, seqlen]` |
| `v: [batches, seqlen, nhead, hdim_v]` and some parameters for defining the functional masking as inputs, and do the following: | ||
|
|
||
| * Multiply `q: [batches, seqlen, nhead, hdim_qk]` with `k: [batches, seqlen, nhead, hdim_k]` to get temporary tensor `s: [batches, nhead, seqlen, seqlen]` | ||
| * Update `s` by filtering its values according to a special functional mask, which includes the logics of lower-triangular and diagonal window causal mask |
There was a problem hiding this comment.
Suggested change
| * Update `s` by filtering its values according to a special functional mask, which includes the logics of lower-triangular and diagonal window causal mask | |
| * Update `s` by filtering it with a functional mask that includes a lower-triangular mask, a diagonal window causal mask, and |
|
|
||
| * Multiply `q: [batches, seqlen, nhead, hdim_qk]` with `k: [batches, seqlen, nhead, hdim_k]` to get temporary tensor `s: [batches, nhead, seqlen, seqlen]` | ||
| * Update `s` by filtering its values according to a special functional mask, which includes the logics of lower-triangular and diagonal window causal mask | ||
| as well assequence mask |
There was a problem hiding this comment.
Suggested change
| as well assequence mask | |
| a sequence mask. |
| * Update `s` by filtering its values according to a special functional mask, which includes the logics of lower-triangular and diagonal window causal mask | ||
| as well assequence mask | ||
| * Do element-wise SiLu on the `lower seqlen` dimension of `s` to get temporary tensor `p: [batches, nhead, seqlen, seqlen]` | ||
| * Multiply `p : [batches, nhead, seqlen, seqlen]` with `v: [batches, seqlen, nhead, hdim_v]` to get final output `o: [batches, seqlen_q, nhead, headsz_v]` |
There was a problem hiding this comment.
Suggested change
| * Multiply `p : [batches, nhead, seqlen, seqlen]` with `v: [batches, seqlen, nhead, hdim_v]` to get final output `o: [batches, seqlen_q, nhead, headsz_v]` | |
| * Multiply `p : [batches, nhead, seqlen, seqlen]` with `v: [batches, seqlen, nhead, hdim_v]` to get the final tensor `o: [batches, seqlen_q, nhead, headsz_v]` |
| as well assequence mask | ||
| * Do element-wise SiLu on the `lower seqlen` dimension of `s` to get temporary tensor `p: [batches, nhead, seqlen, seqlen]` | ||
| * Multiply `p : [batches, nhead, seqlen, seqlen]` with `v: [batches, seqlen, nhead, hdim_v]` to get final output `o: [batches, seqlen_q, nhead, headsz_v]` | ||
| * Jagged inputs are also supported, where each batch has separate seqlen defined by the `sequence_offsets[]` |
There was a problem hiding this comment.
Suggested change
| * Jagged inputs are also supported, where each batch has separate seqlen defined by the `sequence_offsets[]` | |
| Jagged inputs are also supported, where each batch has separate seqlen defined by the `sequence_offsets[]` |
This isn't a thing that the operator does, so it shouldn't be in the same bullet list
|
|
||
| ## implementation | ||
|
|
||
| The operator is implemented using a fused kernel in the example: |
There was a problem hiding this comment.
Suggested change
| The operator is implemented using a fused kernel in the example: | |
| The operator is implemented using a fused kernel: |
|
|
||
| The operator is implemented using a fused kernel in the example: | ||
|
|
||
| * Tensor S and Tensor P only exist in VGPRs as per-workgroup tiles, no global memory access is needed |
There was a problem hiding this comment.
This doesn't need to be a bullet point. You can combine it with the sentence above.
| #> . example/ck_tile/07_hstu_attention/test_hstu_attention.sh | ||
| ``` | ||
|
|
||
| Check the example file `example_hstu_attention.cpp` for an understanding of the command-line arguments. Which is like the following: |
There was a problem hiding this comment.
Suggested change
| Check the example file `example_hstu_attention.cpp` for an understanding of the command-line arguments. Which is like the following: | |
| Check the example file `example_hstu_attention.cpp` for more information about the command-line arguments. |
To be honest, I'd rather see the explanations here, or at least have the code snippet commented. It doesn't need to be everything but some of it.
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.
This PR brings an implementation of HSTU attention on ck_tile. HSTU attention is very different from the
fmhaimplemented in ck_tile, for details, please refer to the hstu paperThe implementation is well verified on MI300 for both functionalities and targeted performance, but it does not make any optimization for MI350.
To build
#> cd build; ../scripts/cmake-ck-dev.sh .. gfx942; make -j 128 tile_example_hstu_attentionTo verify
#> . examples/ck_tile/23_hstu_attention/scripts/test_hstu_attention.shThe codes of HSTU are all located under the folder examples/ck_tile/23_hstu_attention, but this PR also made some tiny change to the core ck_tile codes under
include/ck_tile/core/tensor