choose specific shapes for softmax / embedding / exp

Author: yf225

benchmarks/run.py

@@ -55,6 +55,34 @@
         "tritonbench.operators.softmax.operator",
         "examples.softmax",
         "softmax",
+        {
+            "only_shapes": [
+                [4096, 6912],
+                [4096, 6976],
+                [4096, 7040],
+                [4096, 7104],
+                [4096, 7168],
+                [4096, 7232],
+                [4096, 7296],
+                [4096, 7360],
+                [4096, 7424],
+                [4096, 7488],
+                [4096, 7552],
+                [4096, 7616],
+                [4096, 7680],
+                [4096, 7744],
+                [4096, 7808],
+                [4096, 7872],
+                [4096, 7936],
+                [4096, 8000],
+                [4096, 8064],
+                [4096, 8128],
+                [4096, 8192],
+                [4096, 8256],
+                [4096, 8320],
+                [4096, 8384],
+            ]
+        },
     ),
     "cross_entropy": (
         "tritonbench.operators.cross_entropy.operator",
@@ -78,11 +106,35 @@
         "tritonbench.operators.embedding.operator",
         "examples.embedding",
         "embedding_tritonbench",
+        {
+            "only_shapes": [
+                (8, 2048, 4096, 16384),
+                (8, 2048, 4096, 32768),
+                (8, 2048, 4096, 65536),
+                (8, 2048, 4096, 131072),
+            ]
+        },
     ),
     "vector_exp": (
         "tritonbench.operators.vector_exp.operator",
         "examples.exp",
         "exp_tritonbench",
+        {
+            "only_shapes": [
+                65536,
+                131072,
+                262144,
+                524288,
+                1048576,
+                2097152,
+                4194304,
+                8388608,
+                16777216,
+                33554432,
+                67108864,
+                134217728,
+            ]
+        },
     ),
     # "fp8_gemm": (
     #     "tritonbench.operators.fp8_gemm.fp8_gemm",