Merge branch 'FlagOpen:master' into fill(cpp-wrapper)

Author: AdvancedCompiler

.gitignore

@@ -9,6 +9,9 @@ bin-release/
 __pycache__/
 .vscode/
 
+# Python venv directory
+.venv
+
 # Distribution files
 build/
 *.egg-info/

benchmark/core_shapes.yaml

@@ -4,7 +4,7 @@ outer:
     - [1024, 1024]
     - [4096, 4096]
     - [8192, 8192]
-    - [10240, 10240]  #from perf
+    - [10240, 10240] #from perf
 
 single_dim_shapes: &single_dim_shapes
   shapes:
@@ -23,31 +23,31 @@ randperm:
 
 gelu backward:
   shapes:
-    - [1048576,]  # 1024 * 1024
+    - [1048576] # 1024 * 1024
     - [64, 64]
     - [4096, 4096]
     - [64, 512, 512]
     - [256, 1024, 1024]
 
 softmax backward:
   shapes:
-    - [1048576,]  # 1024 * 1024
+    - [1048576] # 1024 * 1024
     - [64, 64]
     - [4096, 4096]
     - [64, 512, 512]
     - [256, 1024, 1024]
 
 diag:
   shapes:
-    - [64,]
-    - [1024,]
+    - [64]
+    - [1024]
     - [64, 64]
     - [256, 1024]
     - [1024, 1024]
 
 quantile:
   shapes:
-    - [1048576,]  # 1024 * 1024
+    - [1048576] # 1024 * 1024
     - [64, 64]
     - [4096, 256]
     - [64, 512, 128]
@@ -60,7 +60,7 @@ BlasBenchmark:
     - [16, 1024, 1024, 1024]
     - [16, 2048, 2048, 2048]
     - [16, 4096, 4096, 4096]
-  shape_desc: "B, M, N, K"   # shapes are defined as (B, M, N, K)
+  shape_desc: "B, M, N, K" # shapes are defined as (B, M, N, K)
 
 # NORM shapes can be either 3D or 4D:
 # - 3D shapes are represented as [batch_size, channels, hidden_size]
@@ -75,7 +75,6 @@ NormBenchmark:
     - [20, 6, 65536] # from perf
   shape_desc: "N, C, *"
 
-
 TensorSelectBenchmark:
   shapes:
     - [64, 64]
@@ -130,8 +129,8 @@ GenericBenchmarkExcluse1D:
 
 GenericBenchmarkExcluse3D:
   shapes:
-    - [1048576,]     # 1024 * 1024
-    - [1073741824,]  # 1024 * 1024 * 1024
+    - [1048576] # 1024 * 1024
+    - [1073741824] # 1024 * 1024 * 1024
     - [1024, 512]
     - [1024, 1024]
     - [4096, 4096]
@@ -147,15 +146,15 @@ GenericBenchmark2DOnly:
 
 UnaryReductionBenchmark:
   shapes:
-    - [1048576,]  # 1024 * 1024
+    - [1048576] # 1024 * 1024
     - [64, 64]
     - [4096, 4096]
     - [64, 512, 512]
     - [1024, 1024, 1024]
 
 UnaryPointwiseBenchmark:
   shapes:
-    - [1073741824,]  # 1024 * 1024 * 1024
+    - [1073741824] # 1024 * 1024 * 1024
     - [64, 64]
     - [4096, 4096]
     - [64, 512, 512]
@@ -165,7 +164,7 @@ UnaryPointwiseBenchmark:
 # the default shapes info
 Benchmark:
   shapes:
-    - [1073741824,]  # 1024 * 1024 * 1024
+    - [1073741824] # 1024 * 1024 * 1024
     - [64, 64]
     - [4096, 4096]
     - [64, 512, 512]
@@ -174,16 +173,16 @@ Benchmark:
 # default Ni Ci Hi WI  Co Hk Wk stride padding groups
 ConvBenchmark:
   shapes:
-  - [32, 64, 128, 128, 32, 3, 3, 1, 1, 1]
-  - [32, 64, 128, 128, 32, 3, 3, 1, 2, 1]
-  - [104, 16, 32, 32, 32, 4, 4, 1, 0, 1]
-  - [64, 32, 18, 18, 32, 5, 5, 2, 1, 1]
-  - [64, 32, 210, 210, 64, 5, 5, 2, 1, 1]
-  - [32, 64, 210, 210, 16, 5, 5, 2, 1, 1]
-  - [16, 32, 12, 12, 24, 3, 3, 2, 1, 1]
-  - [16, 32, 24, 24, 24, 3, 3, 1, 1, 2]
-  - [16, 32, 24, 24, 24, 3, 3, 2, 2, 2]
-  - [16, 32, 24, 24, 24, 3, 3, 1, 2, 2]
+    - [32, 64, 128, 128, 32, 3, 3, 1, 1, 1]
+    - [32, 64, 128, 128, 32, 3, 3, 1, 2, 1]
+    - [104, 16, 32, 32, 32, 4, 4, 1, 0, 1]
+    - [64, 32, 18, 18, 32, 5, 5, 2, 1, 1]
+    - [64, 32, 210, 210, 64, 5, 5, 2, 1, 1]
+    - [32, 64, 210, 210, 16, 5, 5, 2, 1, 1]
+    - [16, 32, 12, 12, 24, 3, 3, 2, 1, 1]
+    - [16, 32, 24, 24, 24, 3, 3, 1, 1, 2]
+    - [16, 32, 24, 24, 24, 3, 3, 2, 2, 2]
+    - [16, 32, 24, 24, 24, 3, 3, 1, 2, 2]
 
 AttentionBenchmark:
   shapes:
@@ -195,40 +194,48 @@ AttentionBenchmark:
 
 KronBenchmark:
   shapes:
-  - [16,16]
-  - [64,64]
-  - [128,128]
-  - [256,256]
-  - [4, 8, 16, 32]
+    - [16, 16]
+    - [64, 64]
+    - [128, 128]
+    - [256, 256]
+    - [4, 8, 16, 32]
 
 IndexPutAccFalseBenchmark:
   shapes:
-  - [[268435456,], [[65536,],], [65536,]]
-  - [[32, 32], [[8,], [2, 8]], [8,]]
-  - [[1024, 1024], [[4, 64],], [1024,]]
-  - [[512, 512, 512], [[2, 128], [128,], [128,]], [128,]]
-  - [[512, 512, 512], [[2, 128],], [512,]]
+    - [[268435456], [[65536]], [65536]]
+    - [[32, 32], [[8], [2, 8]], [8]]
+    - [[1024, 1024], [[4, 64]], [1024]]
+    - [[512, 512, 512], [[2, 128], [128], [128]], [128]]
+    - [[512, 512, 512], [[2, 128]], [512]]
 
 IndexPutAccTrueBenchmark:
   shapes:
-  - [[268435456,], [[65536,],], [65536,]]
-  - [[32, 32], [[8,], [8,]], [8,]]
-  - [[1024, 1024], [[64,], [64,]], [64,]]
-  - [[512, 512, 512], [[128,], [128,], [128,]], [128,]]
-  - [[512, 512, 512], [[2, 128], [2, 128], [2, 128]], [2, 128]]
+    - [[268435456], [[65536]], [65536]]
+    - [[32, 32], [[8], [8]], [8]]
+    - [[1024, 1024], [[64], [64]], [64]]
+    - [[512, 512, 512], [[128], [128], [128]], [128]]
+    - [[512, 512, 512], [[2, 128], [2, 128], [2, 128]], [2, 128]]
 
 glu:
   shapes:
-  - [4, 8, 512, 128]
-  - [4, 8, 1024, 128]
-  - [4, 8, 2048, 128]
-  - [4, 8, 3072, 128]
-  - [4, 8, 4096, 128]
+    - [4, 8, 512, 128]
+    - [4, 8, 1024, 128]
+    - [4, 8, 2048, 128]
+    - [4, 8, 3072, 128]
+    - [4, 8, 4096, 128]
 
 index:
   shapes:
-  - [[268435456,], [[65536,],]]
-  - [[32, 32], [[8,], [2, 8]]]
-  - [[1024, 1024], [[64,], [64,]]]
-  - [[512, 512, 512], [[2, 128],]]
-  - [[64, 64, 64], [[2, 8], [2, 8]]]
+    - [[268435456], [[65536]]]
+    - [[32, 32], [[8], [2, 8]]]
+    - [[1024, 1024], [[64], [64]]]
+    - [[512, 512, 512], [[2, 128]]]
+    - [[64, 64, 64], [[2, 8], [2, 8]]]
+
+FlashMLABenchmark:
+  shapes:
+    - [1024]
+    - [2048]
+    - [4096]
+    - [8192]
+    - [16384]

benchmark/performance_utils.py

@@ -258,7 +258,11 @@ def get_latency(self, op, *args, **kwargs):
         if self.is_backward:
             out = fn()
             dout = torch.randn_like(out)
-            fn = lambda: out.backward(dout, retain_graph=True)
+            # fn = lambda: out.backward(dout, retain_graph=True)
+            xs = list(filter(lambda x: torch.is_tensor(x) and x.requires_grad, args))
+            fn = lambda: torch.autograd.grad(
+                (out,), xs, grad_outputs=(dout,), retain_graph=True
+            )
         if Config.cpu_mode:
             for i in range(Config.warm_up):
                 fn()
@@ -280,6 +284,7 @@ def get_latency(self, op, *args, **kwargs):
                 warmup=Config.warm_up,
                 rep=Config.repetition,
                 return_mode="median",
+                grad_to_none=xs if self.is_backward else None,
             )
         # average latency in ms
         return latency

benchmark/test_attention_perf.py

@@ -1,5 +1,8 @@
+import math
+
 import pytest
 import torch
+import triton
 
 import flag_gems
 
@@ -21,7 +24,7 @@ def set_more_shapes(self):
 @pytest.mark.skipif(
     flag_gems.device == "musa" or vendor_name == "hygon", reason="RuntimeError"
 )
-@pytest.mark.attention
+@pytest.mark.scaled_dot_product_attention
 @pytest.mark.parametrize("dropout_p", [0.0, 0.25])
 @pytest.mark.parametrize("is_causal", [True, False])
 def test_perf_scaled_dot_product_attention(dropout_p, is_causal):
@@ -55,3 +58,112 @@ def sdpa_flash(query, key, value, dropout_p=dropout_p, is_causal=is_causal):
         ],
     )
     bench.run()
+
+
+class FlashMLABenchmark(GenericBenchmark):
+    """
+    benchmark for flash_mla
+    """
+
+    def set_more_shapes(self):
+        # self.shapes is a list of tuples, each containing three elements:
+        # (batch, num_heads, seq_len, head_size).
+        return None
+
+
+@pytest.mark.skipif(vendor_name == "kunlunxin", reason="RESULT TODOFIX")
+@pytest.mark.skipif(
+    flag_gems.device == "musa" or vendor_name == "hygon", reason="RuntimeError"
+)
+@pytest.mark.flash_mla
+def test_perf_flash_mla():
+    def flash_mla_kwargs(shape, dtype, device):
+        seqlen = shape[0]
+        b = 128
+        s_q = 1
+        h_q = 128
+        h_kv = 1
+        d = 576
+        dv = 512
+        causal = True
+        block_size = 64
+        cache_seqlens = torch.tensor(
+            [seqlen + 2 * i for i in range(b)], dtype=torch.int32, device=device
+        )
+        max_seqlen = cache_seqlens.max().item()
+        max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256
+
+        q = torch.randn([b, s_q, h_q, d], dtype=dtype, device=device)
+        block_table = torch.arange(
+            b * max_seqlen_pad // block_size, dtype=torch.int32, device=device
+        ).view(b, max_seqlen_pad // block_size)
+        blocked_k = torch.randn(
+            [block_table.numel(), block_size, h_kv, d], dtype=dtype, device=device
+        )
+        yield q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal
+
+    def scaled_dot_product_attention(query, key, value, h_q, h_kv, is_causal=False):
+        query = query.float()
+        key = key.float()
+        value = value.float()
+        key = key.repeat_interleave(h_q // h_kv, dim=0)
+        value = value.repeat_interleave(h_q // h_kv, dim=0)
+        attn_weight = query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))
+        if is_causal:
+            s_q = query.shape[-2]
+            s_k = key.shape[-2]
+            attn_bias = torch.zeros(s_q, s_k, dtype=query.dtype, device=query.device)
+            temp_mask = torch.ones(
+                s_q, s_k, dtype=torch.bool, device=query.device
+            ).tril(diagonal=s_k - s_q)
+            attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+            attn_bias.to(query.dtype)
+            attn_weight += attn_bias
+        lse = attn_weight.logsumexp(dim=-1)
+        attn_weight = torch.softmax(attn_weight, dim=-1, dtype=torch.float32)
+        return attn_weight @ value, lse
+
+    def ref_mla(
+        q,
+        block_table,
+        blocked_k,
+        max_seqlen_pad,
+        block_size,
+        b,
+        s_q,
+        cache_seqlens,
+        h_q,
+        h_kv,
+        d,
+        dv,
+        causal,
+    ):
+        device = q.device
+        blocked_v = blocked_k[..., :dv]
+        out = torch.empty(b, s_q, h_q, dv, dtype=torch.float32, device=device)
+        lse = torch.empty(b, h_q, s_q, dtype=torch.float32, device=device)
+        for i in range(b):
+            begin = i * max_seqlen_pad
+            end = begin + cache_seqlens[i]
+            O, LSE = scaled_dot_product_attention(
+                q[i].transpose(0, 1),
+                blocked_k.view(-1, h_kv, d)[begin:end].transpose(0, 1),
+                blocked_v.view(-1, h_kv, dv)[begin:end].transpose(0, 1),
+                h_q=h_q,
+                h_kv=h_kv,
+                is_causal=causal,
+            )
+            out[i] = O.transpose(0, 1)
+            lse[i] = LSE
+        return out, lse
+
+    bench = FlashMLABenchmark(
+        op_name="flash_mla",
+        input_fn=flash_mla_kwargs,
+        torch_op=ref_mla,
+        dtypes=[
+            torch.bfloat16,
+        ],
+    )
+    bench.set_gems(flag_gems.flash_mla)
+    bench.run()

benchmark/test_binary_pointwise_perf.py

@@ -49,7 +49,6 @@ def get_tflops(self, op, *args, **kwargs):
             ("mul", torch.mul, FLOAT_DTYPES),
             ("sub", torch.sub, FLOAT_DTYPES),
             ("pow", torch.pow, FLOAT_DTYPES),
-            ("rsub", torch.rsub, FLOAT_DTYPES),
             *(
                 [
                     ("polar", torch.polar, [torch.float32]),
@@ -75,7 +74,6 @@ def get_tflops(self, op, *args, **kwargs):
             # Bitwise operations
             ("bitwise_and", torch.bitwise_and, INT_DTYPES + BOOL_DTYPES),
             ("bitwise_or", torch.bitwise_or, INT_DTYPES + BOOL_DTYPES),
-            ("or_", torch.bitwise_or, INT_DTYPES + BOOL_DTYPES),
             # Numerical Checks
             ("isclose", torch.isclose, FLOAT_DTYPES + INT_DTYPES),
             ("allclose", torch.allclose, FLOAT_DTYPES + INT_DTYPES),