[FEATURE SUPPORT] Variable-Length Attention with Padding-Free Execution

Author: LoserCheems

.github/ISSUE_TEMPLATE/performance_issue.yml

@@ -1,5 +1,5 @@
 name: Performance issue
-description: Report performance problems or optimisation opportunities
+description: Report performance problems or optimization opportunities
 title: "[PERFORMANCE] "
 labels:
   - performance

.github/PULL_REQUEST_TEMPLATE/performance_optimization.yml

@@ -7,7 +7,7 @@ body:
   - type: markdown
     attributes:
       value: |
-        Document the optimisation, methodology, and results so reviewers can validate gains and correctness.
+        Document the optimization, methodology, and results so reviewers can validate gains and correctness.
   - type: textarea
     id: summary
     attributes:

README.md

@@ -195,7 +195,7 @@ output = flash_dmattn_func(
     attn_mask=attention_mask,
     attn_bias=attention_bias,
     is_causal=True,
-    scale=1.0/math.sqrt(head_dim),
+    softmax_scale=1.0/math.sqrt(head_dim),
 )
 
 print(f"Output shape: {output.shape}")  # [1, 256, 2, 64]
@@ -216,7 +216,7 @@ output = flash_dmattn_func(
     attn_mask=attention_mask,
     attn_bias=attention_bias,
     is_causal=True,
-    scale=1.0/math.sqrt(head_dim)
+    softmax_scale=1.0/math.sqrt(head_dim)
 )
 
 # Backward pass

README_zh.md

@@ -195,7 +195,7 @@ output = flash_dmattn_func(
     attn_mask=attention_mask,
     attn_bias=attention_bias,
     is_causal=True,
-    scale=1.0/math.sqrt(head_dim),
+    softmax_scale=1.0/math.sqrt(head_dim),
 )
 
 print(f"输出形状: {output.shape}")  # [1, 256, 2, 64]
@@ -216,7 +216,7 @@ output = flash_dmattn_func(
     attn_mask=attention_mask,
     attn_bias=attention_bias,
     is_causal=True,
-    scale=1.0/math.sqrt(head_dim)
+    softmax_scale=1.0/math.sqrt(head_dim)
 )
 
 # 反向传播

benchmarks/backward_equivalence.py

@@ -266,7 +266,7 @@ def dynamic_mask_attention_cuda(
         attn_mask=attn_mask,                                        # mask: [batch, num_kv_heads, query_len, key_len]
         attn_bias=attn_bias,                                        # bias: [batch, num_kv_heads, query_len, key_len]
         is_causal=is_causal,                                        # causal masking
-        scale=scaling,                                              # scaling factor
+        softmax_scale=scaling,                                              # scaling factor
         softcap=0.0,
         deterministic=False,
         return_attn_probs=False
@@ -351,7 +351,7 @@ def dynamic_mask_attention_triton(
         attn_mask=attn_mask,                                        # mask: [batch, num_heads, seqlen_q, seqlen_k]
         attn_bias=attn_bias,                                        # bias: [batch, num_heads, seqlen_q, seqlen_k]
         is_causal=is_causal,                                        # causal masking
-        scale=scaling                                               # scaling factor
+        softmax_scale=scaling                                               # scaling factor
     )
 
     # Backward pass
@@ -424,7 +424,7 @@ def dynamic_mask_attention_flex(
         attn_mask=attn_mask,                        # attn_mask: [batch, num_heads, query_len, key_len]
         attn_bias=attn_bias,                        # attn_bias: [batch, num_heads, query_len, key_len]
         is_causal=is_causal,                        # is_causal: whether to apply causal masking
-        scale=scaling                               # scaling factor
+        softmax_scale=scaling                               # scaling factor
     )
 
     # Backward pass

benchmarks/backward_performance.py

@@ -183,7 +183,7 @@ def scaled_dot_product_attention_backward(
             key_states,                      # [batch, num_kv_heads, key_len, head_dim]
             value_states,                    # [batch, num_kv_heads, key_len, head_dim]
             attn_mask=causal_mask,
-            scale=scaling,
+            softmax_scale=scaling,
             # is_causal=is_causal if query_len == key_len else False,
             enable_gqa=True
         )
@@ -262,7 +262,7 @@ def dynamic_mask_attention_backward_cuda(
             attn_mask=attn_mask,                                        # mask: [batch, num_kv_heads, query_len, key_len]
             attn_bias=attn_bias,                                        # bias: [batch, num_kv_heads, query_len, key_len]
             is_causal=is_causal,                                        # causal masking
-            scale=scaling,                                              # scaling factor
+            softmax_scale=scaling,                                              # scaling factor
             softcap=0.0,
             deterministic=False,
             return_attn_probs=False
@@ -351,7 +351,7 @@ def dynamic_mask_attention_backward_triton(
             attn_mask=attn_mask,                                        # mask: [batch, num_heads, seqlen_q, seqlen_k]
             attn_bias=attn_bias,                                        # bias: [batch, num_heads, seqlen_q, seqlen_k]
             is_causal=is_causal,                                        # causal masking
-            scale=scaling                                               # scaling factor
+            softmax_scale=scaling                                               # scaling factor
         )
 
         torch.cuda.synchronize()
@@ -433,7 +433,7 @@ def dynamic_mask_attention_backward_flex(
             attn_mask=attn_mask,                        # attn_mask: [batch, num_heads, query_len, key_len]
             attn_bias=attn_bias,                        # attn_bias: [batch, num_heads, query_len, key_len]
             is_causal=is_causal,                        # is_causal: whether to apply causal masking
-            scale=scaling                               # scaling factor
+            softmax_scale=scaling                               # scaling factor
         )
 
         torch.cuda.synchronize()

benchmarks/forward_equivalence.py

@@ -253,7 +253,7 @@ def dynamic_mask_attention_cuda(
         attn_mask=attn_mask,        # [batch, num_kv_heads, query_len, key_len]
         attn_bias=attn_bias,        # [batch, num_kv_heads, query_len, key_len]
         is_causal=is_causal,
-        scale=scaling,
+        softmax_scale=scaling,
         softcap=0.0,
         deterministic=True,
         return_attn_probs=return_softmax
@@ -329,7 +329,7 @@ def dynamic_mask_attention_triton(
         attn_mask=attn_mask,        # mask: [batch, num_heads, seqlen_q, seqlen_k]
         attn_bias=attn_bias,        # bias: [batch, num_heads, seqlen_q, seqlen_k]
         is_causal=is_causal,        # causal masking
-        scale=scaling               # scaling factor
+        softmax_scale=scaling               # scaling factor
     )
 
     return attn_outputs  # [batch, query_len, num_heads, head_dim]
@@ -398,7 +398,7 @@ def dynamic_mask_attention_flex(
         attn_mask=attn_mask,                        # attn_mask: [batch, num_heads, query_len, key_len]
         attn_bias=attn_bias,                        # attn_bias: [batch, num_heads, query_len, key_len]
         is_causal=is_causal,                        # is_causal: whether to apply causal masking
-        scale=scaling                               # scaling factor
+        softmax_scale=scaling                               # scaling factor
     )
 
     return attn_outputs  # [batch, query_len, num_heads, head_dim]

benchmarks/forward_performance.py

@@ -186,7 +186,7 @@ def scaled_dot_product_attention_cuda(
             key_states,
             value_states,
             attn_mask=causal_mask,
-            scale=scaling,
+            softmax_scale=scaling,
             # is_causal=is_causal if query_len == key_len else False,
             enable_gqa=True
         )
@@ -262,7 +262,7 @@ def dynamic_mask_attention_cuda(
             attn_mask=attn_mask,        # [batch, num_kv_heads, query_len, key_len]
             attn_bias=attn_bias,        # [batch, num_kv_heads, query_len, key_len]
             is_causal=is_causal,
-            scale=scaling,
+            softmax_scale=scaling,
             softcap=0.0,
             deterministic=False,
             return_attn_probs=return_softmax
@@ -348,7 +348,7 @@ def dynamic_mask_attention_triton(
             attn_mask=attn_mask,        # mask: [batch, num_heads, seqlen_q, seqlen_k]
             attn_bias=attn_bias,        # bias: [batch, num_heads, seqlen_q, seqlen_k]
             is_causal=is_causal,        # causal masking
-            scale=scaling               # scaling factor
+            softmax_scale=scaling               # scaling factor
         )
 
         torch.cuda.synchronize()
@@ -427,7 +427,7 @@ def dynamic_mask_attention_flex(
             attn_mask=attn_mask,                        # attn_mask: [batch, num_heads, query_len, key_len]
             attn_bias=attn_bias,                        # attn_bias: [batch, num_heads, query_len, key_len]
             is_causal=is_causal,                        # is_causal: whether to apply causal masking
-            scale=scaling                               # scaling factor
+            softmax_scale=scaling                               # scaling factor
         )
 
         torch.cuda.synchronize()