Add more input sweep in Triton Bench

Differential Revision: D82252277

Pull Request resolved: #441

Author: sryap

tritonbench/operators/blackwell_attentions/generate_inputs.py

@@ -78,3 +78,16 @@ def fa3_paper_inputs(dtype, device) -> Generator:
         yield _generated_qkv_inputs(
             shape=(BATCH, H, H, N_CTX, N_CTX, D_HEAD), dtype=dtype, device=device
         )
+
+
+def sweep_inputs(dtype, device) -> Generator:
+    D = 128
+    batch_sizes = [2**i for i in range(6)]
+    num_heads = [5, 8, 16, 24]
+    seqlen = [512 * (2**i) for i in range(6)]
+    for B in batch_sizes:
+        for H in num_heads:
+            for S in seqlen:
+                yield _generated_qkv_inputs(
+                    shape=(B, H, H, S, S, D), dtype=dtype, device=device
+                )

tritonbench/operators/blackwell_attentions/operator.py

@@ -87,17 +87,14 @@
     register_x_val,
 )
 
-from .generate_inputs import customized_inputs, fa3_paper_inputs
+from .generate_inputs import customized_inputs, fa3_paper_inputs, sweep_inputs
 
 HAS_CUDA_124 = (
     torch.cuda.is_available() and torch.version.cuda and torch.version.cuda >= "12.4"
 )
 
 IS_B200 = is_cuda() and "B200" in get_nvidia_gpu_model()
 
-CUSTOMIZED_SHAPES = "CUSTOMIZED_SHAPES"
-FA3_PAPER_SHAPES = "FA3_PAPER_SHAPES"
-
 
 def parse_op_args(args: List[str]):
     parser = argparse.ArgumentParser()
@@ -134,8 +131,8 @@ def parse_op_args(args: List[str]):
     parser.add_argument(
         "--input-types",
         type=str,
-        default=CUSTOMIZED_SHAPES,
-        choices=[CUSTOMIZED_SHAPES, FA3_PAPER_SHAPES],
+        default="CUSTOMIZED_SHAPES",
+        choices=["CUSTOMIZED_SHAPES", "FA3_PAPER_SHAPES", "SWEEP_SHAPES"],
         help="specify input types",
     )
     return parser.parse_args(args)
@@ -514,7 +511,11 @@ def get_input_iter(self) -> Generator:
                 dtype=self.dtype,
                 device=self.device,
             )
-
+        elif self.input_types == "SWEEP_SHAPES":
+            return sweep_inputs(
+                dtype=self.dtype,
+                device=self.device,
+            )
         else:
             raise AssertionError(f"Unknown input type {self.input_types}")