Merge branch 'master' into fix/benchmark-bugs

Signed-off-by: HuangYiQun <63465798+huangyiqun@users.noreply.github.com>

Author: huangyiqun

.github/CODEOWNERS

@@ -3,14 +3,14 @@
 # @global-owner1 and @global-owner2 will be requested for
 # review when someone opens a pull request.
 
-/src/flag_gems/*      @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
+/src/flag_gems/*      @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit
 /src/flag_gems/runtime/backend/*
 
-/tests/*             @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
-/modules_tests/*     @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
-/benchmark/*         @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
+/tests/*             @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit
+/modules_tests/*     @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit
+/benchmark/*         @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit
 
 
-/lib/*                @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
-/ctests/*             @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
-/include/*            @kiddyjinjin   @0x45f   @zhangpeiyang1  @huangyiqun
+/lib/*                @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit
+/ctests/*             @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit
+/include/*            @0x45f   @huangyiqun @douxetpur @bin913 @w1120029931-bit

benchmark/conftest.py

@@ -137,14 +137,19 @@ def pytest_addoption(parser):
         help="Specify the shape file name for benchmarks. If not specified, a default shape list will be used.",
     )
 
-    parser.addoption(
-        "--record",
-        action="store",
-        default="none",
-        required=False,
-        choices=["none", "log"],
-        help="Benchmark info recorded in log files or not",
-    )
+    try:
+        parser.addoption(
+            "--record",
+            action="store",
+            default="none",
+            required=False,
+            choices=["none", "log"],
+            help="Benchmark info recorded in log files or not",
+        )
+    except ValueError:
+        # Mixed test+benchmark pytest runs may already register --record in
+        # tests/conftest.py. Reuse the existing option in that case.
+        pass
 
     parser.addoption(
         "--parallel",
@@ -158,11 +163,16 @@ def pytest_addoption(parser):
         ),
     )
 
-    parser.addoption(
-        "--collect-marks",
-        action="store_true",
-        help="Collect the tests with marker information without executing them",
-    )
+    try:
+        parser.addoption(
+            "--collect-marks",
+            action="store_true",
+            help="Collect the tests with marker information without executing them",
+        )
+    except ValueError:
+        # Mixed test+benchmark pytest runs may already register this option in
+        # tests/conftest.py. Reuse the existing option in that case.
+        pass
 
 
 def pytest_configure(config):

benchmark/core_shapes.yaml

@@ -63,6 +63,14 @@ log10_out:
     - [1024, 65536]
   shape_desc: "(B), M, N"
 
+gcd:
+  shapes:
+    - [65536]
+    - [1024, 1024]
+    - [256, 4096]
+    - [64, 256, 256]
+  shape_desc: "(B), M, N"
+
 softmax_backward:
   shapes:
     - [1048576] # 1024 * 1024
@@ -111,6 +119,19 @@ MvAndOuterBenchmark:
     - [8192, 8192]
     - [10240, 10240]  #from perf
 
+mm_self_transpose:
+  shapes:
+    - [257, 96]
+    - [384, 384]
+    - [1536, 320]
+    - [3072, 768]
+    - [1024, 1024]
+    - [4096, 1024]
+    - [5333, 71]
+    - [8192, 2048]
+    - [10240, 4096]
+  shape_desc: "M, K"
+
 # NORM shapes can be either 3D or 4D:
 # - 3D shapes are represented as [batch_size, channels, hidden_size]
 # - 4D shapes are represented as [batch_size, channels, height, width]

benchmark/test_attention_perf.py

@@ -24,6 +24,76 @@ def set_more_shapes(self):
         return None
 
 
+#
+# sparse_attention shape layout:
+# (batch, seq_len, kv_len, topk, heads, dim)
+#
+SPARSE_ATTENTION_SHAPES = [
+    (16, 1, 136, 136, 8, 512),
+    (16, 1, 392, 385, 8, 512),
+    (16, 1, 392, 386, 8, 512),
+    (16, 1, 392, 387, 8, 512),
+    (32, 1, 392, 388, 8, 512),
+    (32, 1, 392, 389, 8, 512),
+    (32, 1, 392, 390, 8, 512),
+    (32, 1, 392, 391, 8, 512),
+    (64, 1, 136, 136, 8, 512),
+    (64, 1, 392, 385, 8, 512),
+    (64, 1, 392, 388, 8, 512),
+    (64, 1, 392, 389, 8, 512),
+]
+
+
+def torch_sparse_attention(q, kv, attn_sink, topk_idxs, softmax_scale):
+    batch, seq_len, heads, dim = q.shape
+    topk = topk_idxs.shape[-1]
+
+    kv_expanded = kv[:, None, :, :].expand(batch, seq_len, -1, dim)
+    idx_expanded = topk_idxs[:, :, :, None].expand(batch, seq_len, topk, dim).long()
+    gathered_kv = torch.gather(kv_expanded, 2, idx_expanded)
+
+    scores = (
+        torch.einsum("bmhd,bmtd->bmht", q.float(), gathered_kv.float()) * softmax_scale
+    )
+    sink = attn_sink[None, None, :, None].expand(batch, seq_len, heads, 1)
+    attn = torch.softmax(torch.cat([scores, sink], dim=-1), dim=-1)
+
+    out = torch.einsum("bmht,bmtd->bmhd", attn[:, :, :, :-1], gathered_kv.float())
+    return out.to(q.dtype)
+
+
+class SparseAttentionBenchmark(Benchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = SPARSE_ATTENTION_SHAPES[:]
+        self.shape_desc = "B, M, KV_LEN, TOPK, H, D"
+
+    def set_more_shapes(self):
+        return None
+
+    def get_input_iter(self, cur_dtype):
+        for seed, (batch, seq_len, kv_len, topk, heads, dim) in enumerate(self.shapes):
+            torch.manual_seed(2026 + seed)
+            q = torch.randn(
+                (batch, seq_len, heads, dim),
+                dtype=cur_dtype,
+                device=self.device,
+            )
+            kv = torch.randn(
+                (batch, kv_len, dim),
+                dtype=cur_dtype,
+                device=self.device,
+            )
+            attn_sink = torch.zeros((heads,), dtype=torch.float32, device=self.device)
+            topk_idxs = torch.randint(
+                0,
+                kv_len,
+                (batch, seq_len, topk),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            yield q, kv, attn_sink, topk_idxs, 1.0 / math.sqrt(dim)
+
+
 def torch_flash_attention_forward(
     q, k, v, scale, is_causal, dropout_p=0.0, return_debug_mask=False, **extra_kwargs
 ):
@@ -315,6 +385,18 @@ def sdpa_flash(
         del os.environ["TRITON_HIP_USE_NEW_STREAM_PIPELINE"]
 
 
+@pytest.mark.skipif(flag_gems.device == "cpu", reason="Unsupported in CPU mode")
+@pytest.mark.sparse_attention
+def test_perf_sparse_attention():
+    bench = SparseAttentionBenchmark(
+        op_name="sparse_attention",
+        torch_op=torch_sparse_attention,
+        dtypes=[torch.bfloat16],
+    )
+    bench.set_gems(flag_gems.sparse_attn_triton)
+    bench.run()
+
+
 class FlashMLABenchmark(GenericBenchmark):
     """
     benchmark for flash_mla

benchmark/test_binary_pointwise_perf.py

@@ -75,6 +75,7 @@ def get_tflops(self, op, *args, **kwargs):
             ("bitwise_and", torch.bitwise_and, INT_DTYPES + BOOL_DTYPES),
             ("bitwise_or", torch.bitwise_or, INT_DTYPES + BOOL_DTYPES),
             ("div", torch.div, FLOAT_DTYPES + COMPLEX_DTYPES),
+            ("gcd", torch.gcd, INT_DTYPES),
             ("dunder_or", lambda a, b: a | b, INT_DTYPES + BOOL_DTYPES),
             ("eq", torch.eq, FLOAT_DTYPES),
             ("equal", torch.equal, FLOAT_DTYPES),

benchmark/test_blas_perf.py

@@ -135,6 +135,21 @@ def get_tflops(self, op, *args, **kwargs):
         return total_flops
 
 
+class MmSelfTransposeBenchmark(GenericBenchmark2DOnly):
+    """
+    Benchmark for the mm(a, a.t()) fast path.
+    """
+
+    DEFAULT_METRICS = DEFAULT_METRICS[:] + ["tflops"]
+
+    def set_more_shapes(self):
+        return None
+
+    def get_tflops(self, op, *args, **kwargs):
+        m, k = args[0].shape
+        return 2 * m * m * k
+
+
 def addmm_input_fn(b, m, n, k, cur_dtype, device, b_column_major):
     inp1 = torch.randn([m, k], dtype=cur_dtype, device=device)
     bias = torch.randn([m, n], dtype=cur_dtype, device=device)
@@ -184,6 +199,16 @@ def mm_input_fn(b, m, n, k, cur_dtype, device, b_column_major):
         yield inp1, inp2
 
 
+def torch_mm_self_transpose(inp):
+    return torch.mm(inp, inp.t())
+
+
+def mm_self_transpose_input_fn(shape, cur_dtype, device):
+    m, k = shape
+    inp = torch.randn([k, m], dtype=cur_dtype, device=device).t()
+    yield inp,
+
+
 def group_mm_input_fn(groups, N, K, cur_dtype, device):
     assert cur_dtype == torch.bfloat16
     import random
@@ -439,6 +464,17 @@ def test_mv_and_outer_benchmark(op_name, torch_op, input_fn):
     bench.run()
 
 
+@pytest.mark.mm
+def test_mm_self_transpose_benchmark():
+    bench = MmSelfTransposeBenchmark(
+        input_fn=mm_self_transpose_input_fn,
+        op_name="mm_self_transpose",
+        torch_op=torch_mm_self_transpose,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
 class AddmvBenchmark(GenericBenchmark2DOnly):
     """
     Benchmark for addmv