Remove packed tensor functions and improve dynamic mask attention #142
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 FlashDMAttnQKVPackedFunc, FlashDMAttnVarlenQKVPackedFunc, FlashDMAttnKVPackedFunc, and FlashDMAttnVarlenKVPackedFunc classes along with their corresponding helper functions to simplify the interface. Adds sequence length padding to ensure compatibility with hardware requirements by padding sequences to multiples of 128 tokens and properly handling the unpadding in backward passes. Fixes mask and bias tensor dimension alignment to use number of key heads instead of query heads for better multi-query and grouped-query attention support.
Adds comprehensive docstring explaining function parameters and return values Separates query and key length variables for better clarity and accuracy Removes redundant kwargs to prevent duplicate parameter passing Updates implementation parameter to use explicit string value instead of config reference Adds support for keep_window_size parameter from module configuration
Removes complex padding/unpadding logic and varlen functions in favor of a streamlined approach that handles attention masking directly through bias manipulation. Adds support for keep_window_size parameter to enable top-k attention window selection when key length exceeds the specified window size. Updates function signature to include key_length parameter and removes unnecessary helper functions for index manipulation and tensor reshaping.
LoserCheems
requested review from
Evanwu1125,
SNHuan,
Thanksyy,
Copilot,
ftgreat and
wubingheng111
and removed request for
Copilot
LoserCheems
assigned
LoserCheems,
Copilot,
Evanwu1125,
ftgreat,
SNHuan,
Thanksyy and
wubingheng111
Contributor
There was a problem hiding this comment.
Pull Request Overview
This PR removes packed tensor functions and improves the dynamic mask attention implementation. The changes simplify the interface by eliminating QKV and KV packed function variants, while enhancing the core attention function with new features like window-based attention and improved parameter handling.
- Removed multiple packed tensor function classes and their corresponding public API functions
- Enhanced dynamic mask attention with new parameters like
keep_window_sizeand improved mask handling - Added sequence padding for hardware compatibility requirements (128-token alignment)
Reviewed Changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| flash_dmattn/integrations/modeling_flash_dynamic_mask_attention_utils.py | Removed packed tensor utility functions and improved dynamic mask attention implementation with new windowing features |
| flash_dmattn/integrations/flash_dynamic_mask_attention.py | Enhanced wrapper function with better documentation and support for new parameters |
| flash_dmattn/flash_dmattn_interface.py | Removed packed tensor function classes and added sequence padding for hardware alignment |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| min_dtype = torch.finfo(dtype).min | ||
| batch_size, _, num_kv_heads, _ = key_states.shape | ||
|
|
||
| if not all(k in globals() for k in ("_flash_fn")): |
There was a problem hiding this comment.
The all() function with a single-element tuple is unnecessary and inefficient. This should be simplified to a direct membership check.
Suggested change
| if not all(k in globals() for k in ("_flash_fn")): | |
| if "_flash_fn" not in globals(): |
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.
Eliminate packed tensor functions to simplify the interface and add sequence padding for hardware compatibility. Enhance the dynamic mask attention function with clearer parameters and improved implementation, while streamlining the overall logic and supporting new configuration options.