Backward Kernel: Dynamic Mask Skip, Race Fixes, and 5–6× Speedup vs SDPA (Fixes #132) #145
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
Reduces redundant mask copying and clearing operations by moving mask checks earlier in the loop and eliminating unnecessary register allocations for fully masked blocks. Consolidates async fence calls to reduce synchronization overhead and removes redundant clear operations on accumulator fragments when blocks are skipped due to masking.
Adds early mask checking to determine if any threads are active before performing expensive computations. This optimization prevents unnecessary work when entire blocks are masked out, improving performance by: - Moving mask evaluation earlier in the computation pipeline - Conditionally executing bias loading and gemm operations - Tracking active state across iterations to avoid redundant work - Reducing memory transfers and computation overhead for masked regions The change maintains correctness while significantly reducing wasted cycles in attention backward pass kernels when dealing with padded sequences.
Tracks whether bias parameter was provided during forward pass and only returns bias gradient during backward pass when bias was originally given. Prevents unnecessary computation and memory allocation for bias gradients when no bias is used in the attention mechanism.
LoserCheems
requested review from
Evanwu1125,
SNHuan,
Thanksyy,
ftgreat,
juliohsu and
wubingheng111
LoserCheems
assigned
ftgreat,
Thanksyy,
juliohsu,
SNHuan,
Evanwu1125,
wubingheng111,
LoserCheems and
Copilot
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…lash-dmattn into support-integration
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR implements a comprehensive optimization of the Flash Dynamic Mask Attention backward kernel, focusing on dynamic mask skip functionality and performance improvements. The changes eliminate race conditions, add early termination for fully masked blocks, and achieve 4.9×–6.2× speedups over PyTorch SDPA backward.
- Dynamic mask skip with block-level OR reduction to avoid computation on fully masked tiles
- Race condition fixes through proper shared memory aliasing and synchronization barrier optimization
- Performance optimization by removing redundant barriers and adding conditional execution paths
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| flash_dmattn/flash_dmattn_interface.py | Adds conditional dbias return logic to avoid unnecessary gradient computation when bias is None |
| csrc/src/flash_bwd_kernel.h | Implements dynamic mask skip, fixes race conditions, and optimizes synchronization in backward CUDA kernel |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
Comment on lines
+423
to
+425
| return_dbias = True | ||
| if bias is None: | ||
| return_dbias = False |
There was a problem hiding this comment.
The logic can be simplified to a single line: return_dbias = bias is not None. This makes the intent clearer and reduces code verbosity.
Suggested change
| return_dbias = True | |
| if bias is None: | |
| return_dbias = False | |
| return_dbias = bias is not None |
csrc/src/flash_bwd_kernel.h
Outdated
Comment on lines
565
to
571
| Tensor tSsMask_copy_view = smem_thr_copy_PdS.retile_S(tSsMask); | ||
| bool any_active_local = false; | ||
| bool any_active_local_next = false; // to be updated later for next iteration | ||
| #pragma unroll | ||
| for (int i = 0; i < size(tSsMask_copy_view); ++i) { any_active_local |= (tSsMask_copy_view(i) != Element(0)); } | ||
| bool any_active = __syncthreads_or(any_active_local); | ||
| bool any_active_next = any_active; // to be updated later for next iteration |
There was a problem hiding this comment.
[nitpick] The variable any_active_next is initialized to any_active but then immediately reassigned later. Consider initializing it to false or adding a comment explaining why it needs this initial value.
Suggested change
| Tensor tSsMask_copy_view = smem_thr_copy_PdS.retile_S(tSsMask); | |
| bool any_active_local = false; | |
| bool any_active_local_next = false; // to be updated later for next iteration | |
| #pragma unroll | |
| for (int i = 0; i < size(tSsMask_copy_view); ++i) { any_active_local |= (tSsMask_copy_view(i) != Element(0)); } | |
| bool any_active = __syncthreads_or(any_active_local); | |
| bool any_active_next = any_active; // to be updated later for next iteration | |
| bool any_active_next = false; // to be updated later for next iteration |
Comment on lines
+773
to
+774
| // Tensor dS_reshaped = make_tensor(dS.data(), acc_dp.layout()); | ||
| // Convert dS from fp32 to fp16 |
There was a problem hiding this comment.
[nitpick] The commented-out line above this shows the original approach using dS.data(). Consider removing the commented line or adding a comment explaining why the direct acc_dp conversion approach was chosen for clarity.
Suggested change
| // Tensor dS_reshaped = make_tensor(dS.data(), acc_dp.layout()); | |
| // Convert dS from fp32 to fp16 | |
| // Previous approach: Tensor dS_reshaped = make_tensor(dS.data(), acc_dp.layout()); | |
| // We now convert acc_dp directly to the target type for efficiency and correctness, | |
| // avoiding potential issues with pointer aliasing and layout mismatches. |
Initializes any_active_next to false instead of any_active to prevent potential issues with unintended carry-over of active state between iterations in the kernel loops. Changes affect both forward and backward kernel implementations to ensure consistent behavior across the codebase.
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
This PR resolves issue #132 by completing the dynamic-mask aware backward kernel refactor for Flash Dynamic Mask Attention.
Key goals achieved:
dS).__syncthreads()while preserving required barriers around alias reuse and cooperative GEMMs.any_active/any_active_nextare CTA-uniform (__syncthreads_or) so that shared loads / stores are never partially executed.Type of Change
Related Issues
Changes Made
Code Changes
Backward CUDA kernel (
compute_dq_dk_dv_1colblock):any_active/any_active_next.acc_dpbuffer directly (removed cross-scope dependency).__syncthreads_orcp_async_wait<1>barrier in V-in-regs pathany_active).Other:
dSto avoid compile error.Documentation
Testing
No Python API surface change. No dependency changes.
Testing
Performed:
CUDA_LAUNCH_BLOCKING=1and cuda-memcheck (spot tests)Planned / Deferred:
Test Configuration
.cpython-312build)Performance Impact
Benchmark Protocol
Command:
Mode: Local window attention sweeping window size (W) while keeping Q=K=16384, B=1, Hq=2, Hkv=1, D=32, causal (C).
Results (Averaged over 3 runs)
Observations:
Before (Legacy)
No earlier native dynamic mask backward existed in-tree (baseline = PyTorch SDPA). If you want a “Before” column for internal review later, we can add raw times from pre-synchronization-slimming commit.
Breaking Changes
None.
Checklist
CUDA-specific
Additional Notes
Future work (tracked or to be filed):
Safety rationale for removed barriers:
__syncthreads()was immediately followed by a__syncthreads_or, or only guarded accesses to per-thread exclusive SMEM slices postcp_async_wait.