Merge branch 'master' into masked_select

Author: tongxin

.github/workflows/gems-cpp-extension.yaml

@@ -3,6 +3,17 @@ name: gems-cpp-extension
 on:
   push:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
   pull_request:
     branches: [ "master" ]
     paths:

.github/workflows/gems-test-on-hopper.yaml

@@ -3,8 +3,30 @@ name: flag-gems-test-on-hopper
 on:
   push:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
   pull_request:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
 
 jobs:
   op-test-on-hopper:

.github/workflows/gems-test-on-metax.yaml

@@ -3,8 +3,30 @@ name: flag-gems-test-on-metax
 on:
   push:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
   pull_request:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
 
 jobs:
   op-test-on-metax:

.github/workflows/python-coverage.yaml

@@ -3,8 +3,30 @@ name: python-coverage
 on:
   push:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
   pull_request:
     branches: [ "master" ]
+    paths:
+      - '**.py'
+      - '**.cpp'
+      - '**.cu'
+      - '**.h'
+      - '**.hpp'
+      - '**.cc'
+      - 'CMakeLists.txt'
+      - '**/*.cmake'
+      - '**/*.mk'
+      - 'Makefile'
 
 jobs:
   blas-op-test:

benchmark/core_shapes.yaml

@@ -187,11 +187,11 @@ ConvBenchmark:
 
 AttentionBenchmark:
   shapes:
-  - [4, 8, 512, 128]
-  - [4, 8, 1024, 128]
-  - [4, 8, 2048, 128]
-  - [4, 8, 3072, 128]
-  - [4, 8, 4096, 128]
+  - [4, 32, 1024, 64]
+  - [4, 32, 1024, 128]
+  - [4, 32, 2048, 128]
+  - [4, 32, 4096, 128]
+  - [4, 32, 8192, 128]
 
 KronBenchmark:
   shapes:

benchmark/test_attention_perf.py

@@ -22,17 +22,33 @@ def set_more_shapes(self):
     flag_gems.device == "musa" or vendor_name == "hygon", reason="RuntimeError"
 )
 @pytest.mark.attention
-def test_perf_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):
     def scaled_dot_product_attention_kwargs(shape, dtype, device):
         query = torch.randn(shape, device=device, dtype=dtype)
         key = torch.randn(shape, device=device, dtype=dtype)
         value = torch.randn(shape, device=device, dtype=dtype)
-        yield query, key, value, None, 0.0, True
+        yield query, key, value, dropout_p, is_causal
+
+    def sdpa_flash(query, key, value, dropout_p=dropout_p, is_causal=is_causal):
+        from torch.nn.attention import SDPBackend, sdpa_kernel
+
+        with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
+            torch.nn.functional.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=None,
+                dropout_p=dropout_p,
+                is_causal=is_causal,
+            )
 
     bench = AttentionBenchmark(
         op_name="scaled_dot_product_attention",
         input_fn=scaled_dot_product_attention_kwargs,
-        torch_op=torch.nn.functional.scaled_dot_product_attention,
+        # torch_op=torch.nn.functional.scaled_dot_product_attention,
+        torch_op=sdpa_flash,
         dtypes=[
             torch.float16,
             torch.bfloat16,

benchmark/test_special_perf.py

@@ -118,7 +118,7 @@ def unique_input_fn(shape, dtype, device):
 def test_perf_sort():
     class SortBenchmark(GenericBenchmark2DOnly):
         def set_more_shapes(self):
-            return [(1024, 1), (1024, 512)]
+            return [(1024, 1), (1024, 512), (16, 128 * 1024), (8, 256 * 1024)]
 
     def sort_input_fn(shape, dtype, device):
         inp = generate_tensor_input(shape, dtype, device)

benchmark/test_tensor_concat_perf.py

@@ -197,7 +197,6 @@ def repeat_interleave_tensor_input_fn(shape, dtype, device):
 def test_tensor_repeat_benchmark(op_name, torch_op, input_fn, dtypes):
     if vendor_name == "kunlunxin" and op_name in [
         "repeat_interleave_self_tensor",
-        "repeat_interleave_tensor",
     ]:
         pytest.skip("RUNTIME TODOFIX")
     bench = TensorRepeatBenchmark(

benchmark/test_unary_pointwise_perf.py

@@ -12,7 +12,7 @@
     FLOAT_DTYPES,
     INT_DTYPES,
 )
-from .performance_utils import Benchmark, generate_tensor_input, vendor_name
+from .performance_utils import Benchmark, generate_tensor_input
 
 
 class UnaryPointwiseBenchmark(Benchmark):
@@ -94,8 +94,6 @@ def get_tflops(self, op, *args, **kwargs):
     ],
 )
 def test_general_unary_pointwise_perf(op_name, torch_op, dtypes):
-    if vendor_name == "kunlunxin" and op_name == "elu":
-        pytest.skip("RUNTIME TODOFIX")
     bench = UnaryPointwiseBenchmark(op_name=op_name, torch_op=torch_op, dtypes=dtypes)
     bench.run()
 

ctests/CMakeLists.txt

@@ -16,3 +16,13 @@ add_executable(test_triton_norm test_triton_norm.cpp)
 target_link_libraries(test_triton_norm
   PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
 add_test(NAME test_triton_norm COMMAND test_triton_norm)
+
+add_executable(test_triton_rope test_triton_rope.cpp)
+target_link_libraries(test_triton_rope
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_rope COMMAND test_triton_rope)
+
+add_executable(test_triton_bmm test_triton_bmm.cpp)
+target_link_libraries(test_triton_bmm
+  PRIVATE Torch::Torch operators GTest::gtest GTest::gtest_main)
+add_test(NAME test_triton_bmm COMMAND test_triton_bmm)