-
Notifications
You must be signed in to change notification settings - Fork 23k
CUDA: XOR swizzle flash attn K,V smem fp16 tiles #25635
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
Merged
JohannesGaessler
merged 11 commits into
ggml-org:master
from
ynankani:ynankani/swizzle_flash_attn
Aug 31, 2026
+194
−26
Merged
Changes from all commits
Commits
Show all changes
11 commits
Select commit
Hold shift + click to select a range
e13512a
CUDA: XOR swizzle flash attn K,V smem fp16 tiles
ynankani 8cb5912
Fix use 64bit generic pointer instead of 32bit shared pointer
ynankani 73e5d1e
fix shared memory race in FA on DGX Spark
ynankani 50b0d08
Handle corener case
ynankani 24da808
Add swizzle test cases and gate sync for swizzled path only
ynankani bc56f68
gate CUDA PTX
ynankani c9a90f7
offset calculation specific for swizzle branch
ynankani 2a57183
Reafctor code
ynankani de01a0c
Refactor FA swizzle ldmatrix if/else into helpers (K row/col, V offset)
ynankani 8ccc1da
rebase and update test case args
ynankani 088d425
Allow swizzle for non-pow2 shapes, for which nbatch_2%32==0
ynankani 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,126 @@ | ||
| #pragma once | ||
|
|
||
| #include "common.cuh" | ||
| #include "mma.cuh" | ||
|
|
||
| // XOR swizzle for K/V SMEM tiles to avoid bank conflicts without row padding (Turing+ only). | ||
| // Stride must be a multiple of 32 half2 columns, otherwise we keep +4 row padding. | ||
|
|
||
| namespace ggml_cuda_fattn_smem_swizzle { | ||
|
|
||
| static __host__ __device__ constexpr bool bank_aligned(const int nbatch_2) { | ||
| return nbatch_2 >= 32 && nbatch_2 % 32 == 0; | ||
| } | ||
|
|
||
| static __device__ constexpr bool enabled(const int nbatch_2) { | ||
| #if defined(TURING_MMA_AVAILABLE) | ||
| return bank_aligned(nbatch_2); | ||
| #else | ||
| GGML_UNUSED(nbatch_2); | ||
| return false; | ||
| #endif // defined(TURING_MMA_AVAILABLE) | ||
| } | ||
|
|
||
| static __host__ bool enabled(const int nbatch_2, const int cc) { | ||
| #ifdef GGML_USE_HIP | ||
| GGML_UNUSED(nbatch_2); | ||
| GGML_UNUSED(cc); | ||
| return false; | ||
| #else | ||
| return turing_mma_available(cc) && bank_aligned(nbatch_2); | ||
| #endif // GGML_USE_HIP | ||
| } | ||
|
|
||
| static __device__ constexpr int tile_stride(const int nbatch_2) { | ||
| return enabled(nbatch_2) ? nbatch_2 : nbatch_2 + 4; | ||
| } | ||
|
|
||
| static __host__ int tile_stride(const int nbatch_2, const int cc) { | ||
| return enabled(nbatch_2, cc) ? nbatch_2 : nbatch_2 + 4; | ||
| } | ||
|
|
||
| // Swizzled byte offset for tile element (row, col_h2), same map used for writes and reads. | ||
| template<int stride_h2> | ||
| static __device__ __forceinline__ int bytes_rc(const int row, const int col_h2) { | ||
| static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32"); | ||
| return ((row * stride_h2 + col_h2) * (int) sizeof(half2)) ^ ((row & 7) << 4); | ||
| } | ||
|
|
||
| // ldmatrix.x4 via 64-bit generic pointer. | ||
| static __device__ __forceinline__ void ldmatrix_x4(int * xi, const half2 * addr) { | ||
| #if defined(TURING_MMA_AVAILABLE) | ||
| asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" | ||
| : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) | ||
| : "l"(addr)); | ||
| #else | ||
| GGML_UNUSED_VARS(xi, addr); | ||
| NO_DEVICE_CODE; | ||
| #endif // defined(TURING_MMA_AVAILABLE) | ||
| } | ||
|
|
||
| static __device__ __forceinline__ void ldmatrix_x4_trans(int * xi, const half2 * addr) { | ||
| #if defined(TURING_MMA_AVAILABLE) | ||
| asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" | ||
| : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) | ||
| : "l"(addr)); | ||
| #else | ||
| GGML_UNUSED_VARS(xi, addr); | ||
| NO_DEVICE_CODE; | ||
| #endif // defined(TURING_MMA_AVAILABLE) | ||
| } | ||
|
|
||
| // Per-lane swizzled address for one tile<16, 8, half2> ldmatrix: 16 rows, 4 half2 columns per lane. | ||
| template<int stride_h2> | ||
| static __device__ __forceinline__ const half2 * lane_addr( | ||
| const half2 * tile_base, const int base_row, const int base_col_h2, const int I, const int J) { | ||
| static_assert(bank_aligned(stride_h2), "swizzled tile needs a stride that is a multiple of 32"); | ||
| const int lane_row = threadIdx.x % I; | ||
| const int lane_col = (threadIdx.x / I) * (J / 2); | ||
| uint32_t byte_off = (uint32_t) ((base_row + lane_row)*stride_h2 + base_col_h2 + lane_col) * (uint32_t) sizeof(half2); | ||
| byte_off ^= (uint32_t) (((base_row + lane_row) & 7) << 4); | ||
| return (const half2 *) ((const char *) tile_base + byte_off); | ||
| } | ||
|
|
||
| template<int stride_h2, bool swz, typename TileT> | ||
| static __device__ __forceinline__ void load_ldmatrix( | ||
| TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) { | ||
| if constexpr (swz) { | ||
| static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>, | ||
| "the swizzled layout is only supported for tile<16, 8, half2>"); | ||
| ldmatrix_x4((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J)); | ||
| } else { | ||
| ggml_cuda_mma::load_ldmatrix(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2); | ||
| } | ||
| } | ||
|
|
||
| template<int stride_h2, bool swz, typename TileT> | ||
| static __device__ __forceinline__ void load_ldmatrix(TileT & t, const half2 * tile_base, const int off_h2) { | ||
| if constexpr (swz) { | ||
| load_ldmatrix<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2); | ||
| } else { | ||
| ggml_cuda_mma::load_ldmatrix(t, tile_base + off_h2, stride_h2); | ||
| } | ||
| } | ||
|
|
||
| template<int stride_h2, bool swz, typename TileT> | ||
| static __device__ __forceinline__ void load_ldmatrix_trans( | ||
| TileT & t, const half2 * tile_base, const int base_row, const int base_col_h2) { | ||
| if constexpr (swz) { | ||
| static_assert(std::is_same_v<TileT, ggml_cuda_mma::tile<16, 8, half2>>, | ||
| "the swizzled layout is only supported for tile<16, 8, half2>"); | ||
| ldmatrix_x4_trans((int *) t.x, lane_addr<stride_h2>(tile_base, base_row, base_col_h2, TileT::I, TileT::J)); | ||
| } else { | ||
| ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + base_row*stride_h2 + base_col_h2, stride_h2); | ||
| } | ||
| } | ||
|
|
||
| template<int stride_h2, bool swz, typename TileT> | ||
| static __device__ __forceinline__ void load_ldmatrix_trans(TileT & t, const half2 * tile_base, const int off_h2) { | ||
| if constexpr (swz) { | ||
| load_ldmatrix_trans<stride_h2, swz>(t, tile_base, off_h2 / stride_h2, off_h2 % stride_h2); | ||
| } else { | ||
| ggml_cuda_mma::load_ldmatrix_trans(t, tile_base + off_h2, stride_h2); | ||
| } | ||
| } | ||
|
|
||
| } // namespace ggml_cuda_fattn_smem_swizzle |
Oops, something went wrong.
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.
What is the perf impact of this change in isolation?
Uh oh!
There was an error while loading. Please reload this page.
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.
Collected ncu traces on different config swizzle+occupancy_update , swizzle_only, master, master+occupancy_update to isolate the impact of the change. Main idea for the change was to increase occupancy and was seeing some regression
Ncu values for Qwen3.6_35B collected on RTX6000-PRO
Decode:
Prefill:
on DGX Spark:
Decode:
Prefill:
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.
At least on DGX Spark, it seems (128,2) is worse than (64,4).
Also, bank conflicts are high with this change during prefill. Is this expected?
Uh oh!
There was an error while loading. Please reload this page.
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.
The change is made only for decode config, for prefill it is most probably noise.
it is for op_ld and op_st those values are negligible in end to end run. we should be looking at op_ldgsts(cp.async) which is the swizzle change touching for K/V fill(There is no excessive wavefront on K read and V read). Somehow in ncu i don't see that value.
Collected wavefront data on DGX Spark with CTK 13.3(JFYI i was not able to find the ldgsts bank conflict field so had to look into SASS for which took help from Claude to export the SASS and count the "L1 Wavefronts Shared, L1 Wavefronts Shared Ideal, and L1 Wavefronts Shared Excessive") and below are the observations :
Decode:
Prefill:
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.
According to these numbers, please also change the config for
ncols == 16and apply the same changes for the configs withDKQ == 192 && DV == 128.