Add Paged Attention Op for CUDA SM80 support - #24595
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| # Some rows might be completely masked out so we fill them with zero instead of NaN | ||
| if window_size[0] >= 0 or window_size[1] >= 0: | ||
| attention = attention.masked_fill(torch.all(local_mask, dim=-1, keepdim=True), 0.0) |
Check failure
Code scanning / CodeQL
Potentially uninitialized local variable
Show autofix suggestion
Hide autofix suggestion
Copilot Autofix
AI about 1 year ago
To fix the issue, we need to ensure that local_mask is always initialized before it is used. The best approach is to initialize local_mask to a default value (e.g., None) at the start of the function or relevant block. Then, before using local_mask on line 410, we should check whether it has been initialized (i.e., not None) and only use it if it is valid. This ensures that the code does not attempt to use an uninitialized variable.
| @@ -391,2 +391,3 @@ | ||
| scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")) | ||
| local_mask = None | ||
| if window_size[0] >= 0 or window_size[1] >= 0: | ||
| @@ -408,3 +409,3 @@ | ||
| # Some rows might be completely masked out so we fill them with zero instead of NaN | ||
| if window_size[0] >= 0 or window_size[1] >= 0: | ||
| if local_mask is not None and (window_size[0] >= 0 or window_size[1] >= 0): | ||
| attention = attention.masked_fill(torch.all(local_mask, dim=-1, keepdim=True), 0.0) |
| slot_mappings = get_slot_mappings(config, block_table, total_seqlens, cum_seqlens) | ||
|
|
||
| # Set window size for local / causal | ||
| window_size = (-1, -1) |
Check warning
Code scanning / CodeQL
Variable defined multiple times
Show autofix suggestion
Hide autofix suggestion
Copilot Autofix
AI about 1 year ago
To fix the issue, we will remove the unnecessary assignment to window_size on line 548. This will ensure that the variable is only defined when it is actually needed, either on line 552 or 555, depending on the branch of the if statement. No additional changes are required, as the logic of the code remains intact.
| @@ -547,3 +547,2 @@ | ||
| # Set window size for local / causal | ||
| window_size = (-1, -1) | ||
| left_window_size = -1 |
|
|
||
| # Set window size for local / causal | ||
| window_size = (-1, -1) | ||
| left_window_size = -1 |
Check warning
Code scanning / CodeQL
Variable defined multiple times
Show autofix suggestion
Hide autofix suggestion
Copilot Autofix
AI about 1 year ago
To fix the issue, we should remove the redundant assignment to left_window_size on line 549. This will eliminate unnecessary code and improve clarity without altering the functionality of the program. The subsequent if-else block already ensures that left_window_size is assigned an appropriate value based on the condition.
| @@ -548,3 +548,2 @@ | ||
| window_size = (-1, -1) | ||
| left_window_size = -1 | ||
| if config.local: |
tianleiwu
left a comment
tianleiwu
left a comment
There was a problem hiding this comment.
What design change needed if we want to support FP8 or FP4 paged attention in the future?
New kernel necessary |
4 participants
Description
Adds Paged Attention Op which enables of Paged KV Cache. Inputs to this op are unpadded (packed / varlen) so Cumulative Sequence Lengths are a required input.
Motivation and Context
Adding this op to ONNXRuntime is necessary to allow the GenAI team to enable a continuous batching server API.