Add bias gradient computation support in backward kernel #69
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
Introduces dedicated offset calculation and tensor configuration for bias gradient computation. Adds row_offset_dbias calculation using dbias-specific stride parameters and creates gdBias tensor with proper memory layout. Reorganizes tensor declarations with improved formatting and adds shared memory tensors for mask, bias, and bias gradient operations to support the enhanced backward pass functionality.
LoserCheems
requested review from
Evanwu1125,
SNHuan,
Copilot and
wubingheng111
and removed request for
Copilot
Contributor
There was a problem hiding this comment.
Pull Request Overview
Introduce bias gradient computation in the backward kernel by adding dedicated offset calculations and tensor configurations, while also refactoring tensor construction formatting.
- Compute and apply
row_offset_dbiasfor bias gradients - Declare
gdBiastensor and extend shared memory layout with bias-related buffers - Reformat multi-line
make_tensorcalls for improved readability
csrc/src/flash_bwd_kernel.h
Outdated
| typename Kernel_traits::SmemLayoutQdO{}); | ||
| Tensor sQt = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutQdOtransposed{}); | ||
| Tensor sQtNoSwizzle = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutQdOtransposedNoSwizzle{}); | ||
| // Golobal memory tensor configuration |
There was a problem hiding this comment.
Correct the spelling of 'Golobal' to 'Global' in this comment.
Suggested change
| // Golobal memory tensor configuration | |
| // Global memory tensor configuration |
Corrects spelling error in comment describing global memory tensor configuration across multiple kernel files.
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.
Introduce bias gradient computation in the backward kernel, including dedicated offset calculations and tensor configurations. Enhance tensor declarations and memory layouts to support improved backward pass functionality.