webgpu: support MultiHeadAttention operator #22144
Draft
+681
−40
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
Closed
fs-eire
reviewed
Sep 19, 2024
|
|
||
| class TransferBSDToBNSHProgram final : public Program<TransferBSDToBNSHProgram> { | ||
| public: | ||
| TransferBSDToBNSHProgram(const std::string& kernel_name, bool has_bias) : Program{kernel_name}, has_bias_(has_bias) {} |
There was a problem hiding this comment.
Suggested change
| TransferBSDToBNSHProgram(const std::string& kernel_name, bool has_bias) : Program{kernel_name}, has_bias_(has_bias) {} | |
| TransferBSDToBNSHProgram(bool has_bias) : Program{"TransferBSDToBNSH"}, has_bias_(has_bias) {} |
fs-eire
reviewed
Sep 19, 2024
| shader.AddOutput("present_key", ShaderVariable::UseUniform); | ||
| } | ||
|
|
||
| shader.AppendImplementation("const TILE_SIZE = ", tile_size_, "u;\n") |
There was a problem hiding this comment.
it seems that cache keys for the programs are not set correctly.
for example, here tile_size_ is used as a part of the shader source code, but it is not set in the cache key. Use program.CacheHint() to set the cache key.
There was a problem hiding this comment.
TILE_SIZE can also be declared in overridable constants.
fs-eire
reviewed
Sep 19, 2024
| shader.AppendImplementation("var<workgroup> thread_max: array<f32, ", work_group_size_, ">;\n") | ||
| .AppendImplementation("var<workgroup> thread_sum: array<f32, ", work_group_size_, ">;\n"); | ||
|
|
||
| std::string f32_str = components_ == 4 ? "vec4<f32>" : (components_ == 2 ? "vec2<f32>" : "f32"); |
There was a problem hiding this comment.
can use x_value_t for the value type of x.
There was a problem hiding this comment.
No matter x's type is f32 or f16, the program only uses f32 to define max and sum values.
|
I didn't see any call to set program cache key. This may be correct (if necessary information is already in uniform). need to confirm. |
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.
Description
Motivation and Context