add OoTPatchedModuleFusedSDPA #2361
Open
+389
−6
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
Author
|
@yiliu30 @Wei-Lin-Intel @czhu15 |
| causal_res = self.fp8_fsdpa_fwd(q, causal_k, causal_v, causal_mask, dropout_p, scale, False, softmax_mode) | ||
| causal_out, causal_m, causal_linv = (gqa_output_reshape(x) if gqa else x for x in causal_res[:3]) | ||
| causal_m = causal_m.to(torch.float32) | ||
| causal_linv = causal_linv.to(torch.float32) * (128.0 if softmax_mode != "fp32" else 1.0) |
There was a problem hiding this comment.
Only fast mode requires * 128.0, for fp32 and None modes, scale is 1.0
yiliu30
approved these changes
Dec 16, 2025
Contributor
yiliu30
left a comment
yiliu30
left a comment
There was a problem hiding this comment.
LGTM
Please add usage in PR desc :)
Author
Done. |
yangulei
changed the title
|
So this solution only truncate the prefix, not padding? |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
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.
Motivation
Current implementation in vLLM will call FusedSDPA with a huge attention mask, which leads to bad performance. This PR introduced three implementations to get better performance.
Usage
Three environment variables are introduced to control the implementation:
PT_HPU_QKV_SLICE_SEQ_LEN_THLD:int, the threshold forkv_len(=q_len+prefix_len) to apply the implementations, defaults to4096.PT_HPU_QKV_SLICE_CHUNK_SIZE:int, chunk size for the slicing in the implementation, defaults toPT_HPU_QKV_SLICE_SEQ_LEN_THLD.PT_HPU_QKV_SLICE_IMPL:strwith choices in['split_kv', 'slice_causal', 'slice_qkv'], used to select the implementations, defaults toslice_qkv.Implementations
For a FusedSDPA with
q_len=11525andprefix_len=10752. The lengths will be padded and be truncated respectively toq_len=16384andprefix_len=8192before calling the FusedSDPA. The full attention mask is shown bellow, in which the Gray parts steads for the values to be masked out.Notation
The following images include rectangles with three colors:
rgb(255,0,0):is_causal=Falseandattn_mask is not Nonergb(255,255,0):is_causal=Trueandattn_mask=Nonergb(255,0,255):is_causal=Falseandattn_mask=NoneThe original implementation
The original implementation pass the full attention mask and set
is_causal=Falseandvalid_seq_len=None, which results in bad TPC/MME pipeline. The implementation could be visualized as the following image.The
SplitKVimplementationThis implementation call the FusedSPDA twice for the prefix part and causal part respectively. And do not pass
attn_maskfor the prefix part thus gives better performance.The
SliceCausalimplementationThis implementation further slice the causal part into smaller chunks as illustrated in the following image.
The
SliceQKVimplementationThis implementation further slice the prefix part as shown below.
