add sink and local for fla benchmark (#27)

* add sink and loca for fla benchmark

Author: sunjiweiswift

benchmark/bench_flash_attn.py

@@ -10,7 +10,9 @@ def flash_attn_baseline(
     k_cache,
     v_cache,
     causal,
+    window_size,
     softmax_scale,
+    softmax_sink,
     cache_seqlens,
     page_table,
     cu_seqlens_q,
@@ -22,6 +24,8 @@ def flash_attn_baseline(
         k_cache,
         v_cache,
         causal=causal,
+        softmax_sink=softmax_sink,
+        window_size=window_size,
         softmax_scale=softmax_scale,
         page_table=page_table,
         cache_seqlens=cache_seqlens,
@@ -34,20 +38,39 @@ def flash_attn_baseline(
 
 # Benchmark configurations
 causal = [True, False]
+local = [True, False]
+use_softmax_sink = [True, False]
 batch_size = [1, 16]
 q_seq_length_range = [1, 512, 1024]
 kv_seq_length_range = [512, 1024, 2048, 4096, 8192, 16384]
 page_size_range = [32, 64, 128]
 configs = list(
-    product(
-        causal, batch_size, q_seq_length_range, kv_seq_length_range, page_size_range
+    filter(
+        lambda cfg: not (cfg[0] and cfg[1]),
+        product(
+            causal,
+            local,
+            use_softmax_sink,
+            batch_size,
+            q_seq_length_range,
+            kv_seq_length_range,
+            page_size_range,
+        ),
     )
 )
 
 
 @triton.testing.perf_report(
     triton.testing.Benchmark(
-        x_names=["causal", "batch_size", "q_seq_length", "kv_seq_length", "page_size"],
+        x_names=[
+            "causal",
+            "local",
+            "use_softmax_sink",
+            "batch_size",
+            "q_seq_length",
+            "kv_seq_length",
+            "page_size",
+        ],
         x_vals=[list(c) for c in configs],
         line_arg="provider",
         line_vals=["flash_attn"],
@@ -58,7 +81,16 @@ def flash_attn_baseline(
         args={},
     )
 )
-def benchmark(causal, batch_size, q_seq_length, kv_seq_length, page_size, provider):
+def benchmark(
+    causal,
+    local,
+    use_softmax_sink,
+    batch_size,
+    q_seq_length,
+    kv_seq_length,
+    page_size,
+    provider,
+):
     dtype = torch.bfloat16
     device = torch.device("xpu")
 
@@ -93,6 +125,11 @@ def benchmark(causal, batch_size, q_seq_length, kv_seq_length, page_size, provid
         dtype=torch.int32,
     )
     max_seqlen_q = q_seq_length
+    window_size = (-1, -1) if not local else torch.randint(0, kv_seq_length, (2,))
+
+    softmax_sink = (
+        torch.randn(num_heads, device=device, dtype=dtype) if use_softmax_sink else None
+    )
 
     softmax_scale = 1.0 / (head_dim**0.5)
 
@@ -105,7 +142,9 @@ def benchmark(causal, batch_size, q_seq_length, kv_seq_length, page_size, provid
                 k_cache.clone(),
                 v_cache.clone(),
                 causal=causal,
+                window_size=window_size,
                 softmax_scale=softmax_scale,
+                softmax_sink=softmax_sink,
                 cache_seqlens=cache_seqlens,
                 page_table=page_table,
                 cu_seqlens_q=cu_seqlens_q,
@@ -119,3 +158,4 @@ def benchmark(causal, batch_size, q_seq_length, kv_seq_length, page_size, provid
 
 if __name__ == "__main__":
     benchmark.run(print_data=True)
+    print("Benchmark finished!")