[BUG FIX] Unify masking utilities and improve performance #209
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
Replaces the ad‑hoc mask builder with a shared utility to centralize top‑k/causal masking across CUDA, Triton, and Flex paths. Passes explicit tensor shapes and dtype min to ensure correctness, and adds block‑sparse support (block size 64) for dynamic CUDA. Handles optional masks safely when repeating KV for GQA. Updates benchmark bias to broadcast over queries (B,H,1,K) to reduce memory and match masking expectations. Improves consistency, reduces duplication, and prepares for extensible masking strategies.
Removes an unused symbol to clean up imports and silence linter warnings. Reduces clutter and avoids confusion with unreferenced utilities. Introduces no functional changes.
Replaces the local mask builder with a centralized utility to standardize top-k/causal masking across Python, CUDA, Triton, and Flex paths. Passes explicit batch/query/key sizes and dtype min, repeats masks only when present, and skips masked_fill when unneeded. Reduces duplication, improves consistency and maintainability, and streamlines GQA handling.
Replaces ad-hoc mask construction with a shared mask utility to unify top‑k/causal masking across Python, CUDA, Triton, and Flex paths. Reduces duplication and allows safely skipping masking when not required. Fixes gradient reporting in the Flex path by returning grads w.r.t. the original input tensors. Also clarifies shape handling and guards masked fill, improving robustness.
Replaces duplicated mask logic with a shared utility to standardize top‑k/causal masking and dtype handling across CUDA, Triton, and Flex backends. Aligns attention bias to a broadcastable per‑query shape to cut memory and simplify kernel expectations. Removes redundant KV/mask/bias repetition in the Triton path, repeats conditionally for Flex, and makes GQA fan‑out explicit for correctness and performance.
Updates forward/backward equivalence benchmarks to create attention bias with a singleton query dimension so it broadcasts across queries. Aligns shapes with kernel expectations during cached decoding, reduces memory footprint, and prevents shape mismatches across CUDA, Triton, and Flex paths.
github-actions
bot
assigned
Evanwu1125,
ftgreat,
SNHuan,
Thanksyy,
wubingheng111 and
zacliu2023
Copilot finished reviewing on behalf of
LoserCheems
November 13, 2025 04:43
Contributor
There was a problem hiding this comment.
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
Comments suppressed due to low confidence (1)
benchmarks/backward_equivalence.py:268
- Missing repeat_kv calls for Triton implementation. The Triton function expects key_states, value_states, attn_mask, and attn_bias to have
num_headsdimension (notnum_kv_heads). The following code should be added before the transpose operations:
# Repeat KV for multi-head attention (GQA support)
key_states = repeat_kv(key_states, num_queries_per_kv)
value_states = repeat_kv(value_states, num_queries_per_kv)
attn_mask = repeat_kv(attn_mask, num_queries_per_kv)
attn_bias = repeat_kv(attn_bias_leaf, num_queries_per_kv)This is consistent with the forward implementation in forward_equivalence.py lines 242-245.
# Ensure correct data types and memory layout for Triton function
query_states = query_states.transpose(1, 2) # [batch, query_len, num_heads, head_dim]
key_states = key_states.transpose(1, 2) # [batch, key_len, num_heads, head_dim]
value_states = value_states.transpose(1, 2) # [batch, key_len, num_heads, head_dim]
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
280
to
283
| # Ensure correct data types and memory layout for Triton function | ||
| query_states = query_states.transpose(1, 2).contiguous() # [batch, query_len, num_heads, head_dim] | ||
| key_states = key_states.transpose(1, 2).contiguous() # [batch, key_len, num_heads, head_dim] | ||
| value_states = value_states.transpose(1, 2).contiguous() # [batch, key_len, num_heads, head_dim] |
There was a problem hiding this comment.
Missing repeat_kv calls for Triton implementation. The Triton function expects key_states, value_states, attn_mask, and attn_bias to have num_heads dimension (not num_kv_heads). The following code should be added before the transpose operations:
# Repeat KV for multi-head attention (GQA support)
key_states = repeat_kv(key_states, num_queries_per_kv)
value_states = repeat_kv(value_states, num_queries_per_kv)
attn_mask = repeat_kv(attn_mask, num_queries_per_kv)
attn_bias = repeat_kv(attn_bias, num_queries_per_kv)This is consistent with the forward implementation in forward_performance.py lines 273-276 and forward_equivalence.py lines 242-245.
| batch_size, num_heads, query_len, _ = query_states.shape | ||
| _, num_kv_heads, key_len, _ = key_states.shape | ||
|
|
||
| num_queries_per_kv = num_heads // num_kv_heads |
There was a problem hiding this comment.
Variable num_queries_per_kv is not used.
Suggested change
| num_queries_per_kv = num_heads // num_kv_heads |
Comment on lines
+187
to
+188
| num_queries_per_kv = num_heads // num_kv_heads | ||
|
|
There was a problem hiding this comment.
Variable num_queries_per_kv is not used.
Suggested change
| num_queries_per_kv = num_heads // num_kv_heads |
Comment on lines
+161
to
+162
| num_queries_per_kv = num_heads // num_kv_heads | ||
|
|
There was a problem hiding this comment.
Variable num_queries_per_kv is not used.
Suggested change
| num_queries_per_kv = num_heads // num_kv_heads |
Comment on lines
+183
to
+184
| num_queries_per_kv = num_heads // num_kv_heads | ||
|
|
There was a problem hiding this comment.
Variable num_queries_per_kv is not used.
Suggested change
| num_queries_per_kv = num_heads // num_kv_heads |
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
Root Cause
Changes
Reproduction
Tests
Compatibility
Checklist