Make attention parameters optional with defaults and simplify API documentation #95
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
Improves flexibility by making attn_mask, attn_bias, is_causal, and scale parameters optional with sensible defaults. Creates default attention mask and bias tensors when not provided, sets causal attention to true by default, and calculates scale from head dimension when not specified. Adds proper null checking before tensor slicing operations to prevent errors when optional parameters are None.
Streamlines parameter descriptions and removes verbose explanations to improve readability. Updates code examples to use the simplified high-level interface consistently across all sections. Clarifies that the auto function returns a callable rather than direct output, reducing potential confusion for new users.
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
This PR enhances the Flash Dynamic Mask Attention API by making attention parameters optional with sensible defaults and improving documentation clarity. The main purpose is to simplify the API while maintaining backward compatibility and improving developer experience.
- Made core attention parameters (attn_mask, attn_bias, is_causal) optional with sensible defaults
- Added automatic scale calculation based on head dimension when not provided
- Significantly simplified and reorganized API documentation for better readability
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| flash_dmattn/flash_dmattn_flex.py | Added optional parameters with default tensor creation and automatic scale calculation |
| docs/api_reference.md | Comprehensive documentation rewrite with simplified examples and clearer formatting |
| if attn_mask is not None: | ||
| attn_mask = attn_mask[:, :, :, : key.shape[-2]] | ||
| else: | ||
| attn_mask = torch.ones((batch, nheads, seqlen_q, seqlen_k), device=query.device, dtype=query.dtype) |
There was a problem hiding this comment.
Creating large default tensors with torch.ones can be memory-intensive for long sequences. Consider using a more memory-efficient approach or lazy evaluation for default masks.
Suggested change
| attn_mask = torch.ones((batch, nheads, seqlen_q, seqlen_k), device=query.device, dtype=query.dtype) | |
| attn_mask = None # Avoid allocating a large dense tensor; treat None as "no mask" |
Comment on lines
+29
to
37
| else: | ||
| attn_bias = torch.zeros((batch, nheads, seqlen_q, seqlen_k), device=query.device, dtype=query.dtype) | ||
| if is_causal is None: | ||
| is_causal = True | ||
| if scale is None: | ||
| scale = 1.0 / math.sqrt(dhead) | ||
|
|
||
| def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): | ||
| score = score + attn_bias[batch_idx][head_idx][q_idx][kv_idx] |
There was a problem hiding this comment.
Creating large default tensors with torch.zeros can be memory-intensive for long sequences. Consider using a more memory-efficient approach or lazy evaluation for default biases.
Suggested change
| else: | |
| attn_bias = torch.zeros((batch, nheads, seqlen_q, seqlen_k), device=query.device, dtype=query.dtype) | |
| if is_causal is None: | |
| is_causal = True | |
| if scale is None: | |
| scale = 1.0 / math.sqrt(dhead) | |
| def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): | |
| score = score + attn_bias[batch_idx][head_idx][q_idx][kv_idx] | |
| # else: leave attn_bias as None to avoid allocating a large zero tensor | |
| if is_causal is None: | |
| is_causal = True | |
| if scale is None: | |
| scale = 1.0 / math.sqrt(dhead) | |
| def score_mod(score, batch_idx, head_idx, q_idx, kv_idx): | |
| if attn_bias is not None: | |
| score = score + attn_bias[batch_idx][head_idx][q_idx][kv_idx] |
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.
Enhance flexibility by making attention parameters optional with sensible defaults, including default tensors for masks and biases. Improve API documentation for better readability and clarity, ensuring consistent use of the simplified interface.