Current Behavior:
The torch.nn.functional.scaled_dot_product_attention function currently applies causal masking to the top left corner of the attention matrix. This well defined when seqlen_q equals seqlen_kv, but an alignment choice needs to made when they are not.
Current alignment choice
If seqlen_q != seqlen_kv and causal=True, the causal mask is aligned to the the top-left corner.
For example, if seqlen_q = 2 and seqlen_kv = 5, the causal mask (1 = keep, 0 = masked out) is:
If seqlen_q = 5 and seqlen_kv = 2, the causal mask is:
Proposal:
I propose changing the default behavior of the scaled_dot_product_attention function's causal masking to align with the bottom right corner of the attention matrix. This change would be backward incompatible for inputs when seqlen_q does not equal seqlen_kv, as it would shift the masking pattern.
For example, if seqlen_q = 2 and seqlen_kv = 5, the causal mask (1 = keep, 0 = masked out) is:
If seqlen_q = 5 and seqlen_kv = 2, the causal mask is:
If the row of the mask is all zero, the output will be zero.
Benefits:
This choice of mask is beneficial when performing autoregressive decoding for LLMs.
For example lets say that your prefill prompt consisted of 4 tokens and you had a max KV cache size of 8. Then at step 1 of decoding, your query sequence length would be size 1 (index position 5) and your KV input would be of size 5. The query's seqlen is 1 but it really represents the 5 token in the sequence. So the mask for this attention should be:
With the existing implementation of causal, the produced attention mask would be.
Drawbacks:
Backward compatibility: The proposed change would break backward compatibility for cases where seqlen_q doesn't equal seqlen_kv.
Existing code: Users who rely on the current masking alignment would need to update their code if they want to adopt the new behavior.
Implementation:
To implement this change, the following steps need to be taken:
- Modify the
scaled_dot_product_attention_math function to align the causal masking with the bottom right corner of the attention matrix.
- Change the default call
_scaled_dot_product_efficient_attention to
|
sdp::CustomMaskType custom_mask_type = is_causal |
to use the bottom right
- After updating to FlashAttentionV2 we would apply the changes from this commit: Dao-AILab/flash-attention@9e5e8bc
- Update the fused CPU implementation to obey the above semantics.
- Update the function's documentation and any relevant examples to reflect the new masking behavior.
Additional Context:
Both Xformers and FlashAttention have made this behavior the new default
cc @ezyang @gchanan @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch @erichan1
Current Behavior:
The
torch.nn.functional.scaled_dot_product_attentionfunction currently applies causal masking to the top left corner of the attention matrix. This well defined when seqlen_q equals seqlen_kv, but an alignment choice needs to made when they are not.Current alignment choice
If seqlen_q != seqlen_kv and causal=True, the causal mask is aligned to the the top-left corner.
For example, if seqlen_q = 2 and seqlen_kv = 5, the causal mask (1 = keep, 0 = masked out) is:
If seqlen_q = 5 and seqlen_kv = 2, the causal mask is:
Proposal:
I propose changing the default behavior of the scaled_dot_product_attention function's causal masking to align with the bottom right corner of the attention matrix. This change would be backward incompatible for inputs when seqlen_q does not equal seqlen_kv, as it would shift the masking pattern.
For example, if seqlen_q = 2 and seqlen_kv = 5, the causal mask (1 = keep, 0 = masked out) is:
If seqlen_q = 5 and seqlen_kv = 2, the causal mask is:
If the row of the mask is all zero, the output will be zero.
Benefits:
This choice of mask is beneficial when performing autoregressive decoding for LLMs.
For example lets say that your prefill prompt consisted of 4 tokens and you had a max KV cache size of 8. Then at step 1 of decoding, your query sequence length would be size 1 (index position 5) and your KV input would be of size 5. The query's seqlen is 1 but it really represents the 5 token in the sequence. So the mask for this attention should be:
With the existing implementation of causal, the produced attention mask would be.
Drawbacks:
Backward compatibility: The proposed change would break backward compatibility for cases where seqlen_q doesn't equal seqlen_kv.
Existing code: Users who rely on the current masking alignment would need to update their code if they want to adopt the new behavior.
Implementation:
To implement this change, the following steps need to be taken:
scaled_dot_product_attention_mathfunction to align the causal masking with the bottom right corner of the attention matrix._scaled_dot_product_efficient_attentiontopytorch/aten/src/ATen/native/transformers/cuda/attention.cu
Line 734 in 3488837
Additional Context:
Both Xformers and FlashAttention have made this behavior the new default
cc @ezyang @gchanan @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @bhosmer @cpuhrsch @erichan1