[FEATURE SUPPORT] Move scaling out of streaming loops, bias-initialized acc_s, and fix dQ double-scaling #203
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
Removes bias and scaling handling from the mask helper to reduce specialization paths and rely solely on masking behavior.
Rationalizes accumulator setup so bias kernels reuse shared-memory bias instead of clearing registers, trimming sync overhead. Simplifies mask application templates to drop unused bias handling, tightening specialization footprint.
Pre-scales query tiles before streaming to cut redundant softmax multiplications. Initializes accumulators from shared bias when active so mask paths can skip extra clears. Simplifies mask application by dropping per-iteration bias scaling logic.
Provides an inverse scaling factor so kernels can reuse precomputed softmax adjustments instead of recomputing them
Ensures the reciprocal scale is always populated so downstream kernels can undo the softmax amplification without branching.
Drops bias and scale handling from the masking helper so upstream code owns those adjustments, preventing duplicated math.
Waits for outstanding async loads and syncs threads so Q scaling never races ahead of shared-memory tiles.
Prevents applying the softmax factor twice in the backward preprocessing so downstream gradients stay correctly scaled.
Pre-scales the keys right after synchronization so later matmul steps reuse the scaled values and hide latency. Unifies the mask and bias hydration before streaming to keep accumulators coherent and drops the now redundant gradient scaling.
Removes the unused reverse scaling parameter from the forward configuration to avoid stale values when softcap toggles.
github-actions
bot
requested review from
Evanwu1125,
SNHuan,
Thanksyy,
ftgreat,
juliohsu,
wubingheng111 and
zacliu2023
github-actions
bot
assigned
Evanwu1125,
ftgreat,
juliohsu,
SNHuan,
Thanksyy,
wubingheng111 and
zacliu2023
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR refactors the flash attention implementation to move scaling operations earlier in the computation pipeline and removes bias handling from the masking logic. The key changes include:
- Moving Q and K scaling from the masking phase to before the streaming loops
- Simplifying
apply_maskby removing bias and scale_softmax parameters - Initializing attention score accumulators with bias values when present, instead of applying bias during masking
- Commenting out the final dQ scaling in backward passes
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| csrc/flash_dmattn/src/mask.h | Removed bias and scale_softmax parameters from apply_mask functions, simplified masking logic by removing redundant conditional branches |
| csrc/flash_dmattn/src/flash_fwd_kernel.h | Added Q scaling before streaming loops, moved bias application to accumulator initialization, updated apply_mask calls |
| csrc/flash_dmattn/src/flash_bwd_preprocess_kernel.h | Commented out final dQ scaling operation |
| csrc/flash_dmattn/src/flash_bwd_kernel.h | Added K scaling before streaming loops, moved bias application to accumulator initialization, commented out final dQ scaling, updated apply_mask calls |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1151
to
+1152
| for (int i = 0; i < size(tSrQ); ++i) { | ||
| tSsQ(i) = static_cast<Element>(tSsQ(i) * params.scale_softmax); |
There was a problem hiding this comment.
The loop iterates over size(tSrQ) but modifies tSsQ(i). This is inconsistent and likely incorrect. Based on the non-splitkv version (lines 429-439), there should be a conditional check for Kernel_traits::Is_Q_in_regs to determine whether to use tSrQ or tSsQ. Either add the conditional check or ensure both the loop bound and the tensor being modified are consistent.
Suggested change
| for (int i = 0; i < size(tSrQ); ++i) { | |
| tSsQ(i) = static_cast<Element>(tSsQ(i) * params.scale_softmax); | |
| if (Kernel_traits::Is_Q_in_regs) { | |
| for (int i = 0; i < size(tSrQ); ++i) { | |
| tSrQ(i) = static_cast<Element>(tSrQ(i) * params.scale_softmax); | |
| } | |
| } else { | |
| for (int i = 0; i < size(tSsQ); ++i) { | |
| tSsQ(i) = static_cast<Element>(tSsQ(i) * params.scale_softmax); | |
| } |
Comment on lines
+675
to
676
|
|
||
|
|
There was a problem hiding this comment.
Empty line with trailing whitespace. Remove the trailing whitespace for consistency.
Suggested change
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.
Summary
Motivation: reduce math inside hot loops for better perf, fix backward dQ correctness, and simplify masking semantics.
Design
Forward kernels:
acc_sfrom bias when present to avoid +bias in-loop.Backward kernels:
acc_dq *= params.scale_softmaxat dQ write-out (main kernel) and in seq-k pathconvert_dQso dQ is not double-scaled.Masking:
Alternatives considered: leaving scaling inside loops (more math per iteration), or keeping bias/scaling in the mask helper (more template paths, duplicated responsibilities). The chosen approach centralizes responsibilities and reduces inner-loop work.
Changes
Internal kernels:
flash_fwd_kernel.h.flash_bwd_kernel.handflash_bwd_preprocess_kernel.h.Mask API (internal-only):
apply_masktemplate signature to accept only mask (no bias/scale), and perform masking only inmask.h.Public Python API remains unchanged.
Implementation Notes
cp_async_waitand__syncthreads()before pre-scaling Q/K when needed to avoid races with async smem preloads.acc_sonce per tile; removes repeated +bias in the streaming loop.Tests
Functional equivalence:
convert_dQ.Benchmarks (provided):
Forward (before vs. after):
Backward (before vs. after):
Notes:
Docs
flash_fwd_kernel.hflash_bwd_kernel.hflash_bwd_preprocess_kernel.hmask.hChecklist