-
Notifications
You must be signed in to change notification settings - Fork 270
Update/add to qr_ks_vs_whole_k_prefetch pipeline #3485
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
qianfengz
wants to merge
25
commits into
develop
Choose a base branch
from
whole_k_prefetch_n0loop
base: develop
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
25 commits
Select commit
Hold shift + click to select a range
5fada1c
Initial re-implementation of pipeline qr_ks_vs_whole_k_prefetch in lo…
qianfengz 98f9b4a
Add prefetching whole next iteration K path in the pipeline
qianfengz c32949b
Change in GetKVBlockGemm to let gemm1 to use WarpTile-16x16x16/32x32x…
qianfengz 25521a7
Switch the codes based on the iteration index (first/intermediate/last)
qianfengz 8b85919
Simplify the block_gemm codes
qianfengz 5722f8a
[Performance] Change __builtin_amdgcn_sched_barrier() in block_gemm
qianfengz 044f554
Refine the interleaving in the loop of Gemm0
qianfengz 2ea8d83
Using explicit vgpr-saved partition_index with store_tile(lds_window,…
qianfengz 12c8873
Separate kN0Sub from kK0 to be used for flexible tile tuning for whol…
qianfengz c3d3487
Load Q through Lds
qianfengz 409ec3b
Fix move_tile_window(k_dram_window, ..) step in the pipeline
qianfengz 370d386
Remove replicated codes in the pipeline
qianfengz d281c51
Adjust in GetNumPrefetchV()
qianfengz 384f470
Add support of loading QK tiles of hdim96 without padding to hdim128
qianfengz eb598a9
Add qr_ks_vs_whole_k_prefetch_trload pipeline
qianfengz 57abd10
Using is_using_trload_v to check the kUseTrLoad from pipeline
qianfengz 3f6d26e
Load Q directly from global memory to registers for BlockGemm
qianfengz 1ef76a6
Fix the static_assert expression in the pipeline
qianfengz db5c12d
Update to the non-whole-k-prefetch path in the whoke_k_prefetch pipeline
qianfengz 57cf989
Update to only pre-load one v_tile during Gemm0 loop
qianfengz b77fdbf
Move the loading of k_file for next iteration into the Gemm1 loop (no…
qianfengz e7e6ebc
Update to GetNumPrefetchV()
qianfengz 6c91b0c
Move the loading of k_tile for next iteration into the Gemm1 loop (no…
qianfengz 489e255
Update to GetNumPrefetchV() for kM0=64 path
qianfengz f5b4d5d
Update in whole_k_prefetch_trload pipeline to prefetch two k_tile for…
qianfengz File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,6 +8,52 @@ | |
|
|
||
| namespace ck_tile { | ||
|
|
||
| namespace detail { | ||
|
|
||
| template <typename DataType, index_t ElemPerThread> | ||
| CK_TILE_HOST_DEVICE static constexpr auto GetMaxVectorSize() | ||
| { | ||
| if constexpr(std::is_same_v<DataType, half_t> || std::is_same_v<DataType, bf16_t>) | ||
| { | ||
| // ToDo: need support in ck_tile for using buffer_load_dwordx3 | ||
| // if constexpr(ElemPerThread % 6 == 0) | ||
| // return 6; | ||
| if constexpr(ElemPerThread % 8 == 0) | ||
| return 8; | ||
| else if constexpr(ElemPerThread % 4 == 0) | ||
| return 4; | ||
| else if constexpr(ElemPerThread % 2 == 0) | ||
| return 2; | ||
| return 1; | ||
| } | ||
| else if constexpr(std::is_same_v<DataType, float>) | ||
| { | ||
| // ToDo: need support in ck_tile for using buffer_load_dwordx3 | ||
| // if constexpr(ElemPerThread % 3 == 0) | ||
| // return 3; | ||
| if constexpr(ElemPerThread % 4 == 0) | ||
| return 4; | ||
| else if constexpr(ElemPerThread % 2 == 0) | ||
| return 2; | ||
| return 1; | ||
| } | ||
| else | ||
| static_assert(false, "The data type is not supported!"); | ||
|
Comment on lines
+40
to
+41
|
||
| }; | ||
|
|
||
| template <typename DataType, | ||
| index_t kThreadBlockSize, | ||
| index_t kHigherDimSize, | ||
| index_t kLowerDimSize> | ||
| CK_TILE_HOST_DEVICE static constexpr auto GetDramTileAccessMaxVectorSize() | ||
| { | ||
| constexpr index_t ElemPerThread = (kHigherDimSize * kLowerDimSize) / kThreadBlockSize; | ||
|
|
||
| return GetMaxVectorSize<DataType, ElemPerThread>(); | ||
| } | ||
|
|
||
| }; // namespace detail | ||
|
|
||
| template <typename QDataType_, | ||
| typename KDataType_, | ||
| typename VDataType_, | ||
|
|
@@ -62,6 +108,33 @@ struct BlockFmhaPipelineProblem | |
| static constexpr bool kHasDropout = Traits::kHasDropout; | ||
| static constexpr auto QScaleEnum = Traits::QScaleEnum; | ||
| static constexpr index_t kBlockPerCu = Traits::kBlockPerCu; | ||
|
|
||
| CK_TILE_HOST_DEVICE static constexpr auto GetQDramTileAccessMaxVectorSize() | ||
| { | ||
| constexpr index_t kMPerBlock = BlockFmhaShape::kM0; | ||
| constexpr index_t kKPerBlock = BlockFmhaShape::kQKHeaddim; | ||
|
|
||
| return detail:: | ||
| GetDramTileAccessMaxVectorSize<QDataType, kBlockSize, kMPerBlock, kKPerBlock>(); | ||
| } | ||
|
|
||
| CK_TILE_HOST_DEVICE static constexpr auto GetKDramTileAccessMaxVectorSize() | ||
| { | ||
| constexpr index_t kNPerBlock = BlockFmhaShape::kN0; | ||
| constexpr index_t kKPerBlock = BlockFmhaShape::kK0; | ||
|
|
||
| return detail:: | ||
| GetDramTileAccessMaxVectorSize<KDataType, kBlockSize, kNPerBlock, kKPerBlock>(); | ||
| } | ||
|
|
||
| CK_TILE_HOST_DEVICE static constexpr auto GetVDramTileAccessMaxVectorSize() | ||
| { | ||
| constexpr index_t kNPerBlock = BlockFmhaShape::kN1; | ||
| constexpr index_t kKPerBlock = BlockFmhaShape::kK1; | ||
|
|
||
| return detail:: | ||
| GetDramTileAccessMaxVectorSize<VDataType, kBlockSize, kNPerBlock, kKPerBlock>(); | ||
| }; | ||
| }; | ||
|
|
||
| template <typename QDataType_, | ||
|
|
||
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Corrected spelling of 'detechting' to 'detecting'.