feat: upload benchmarks folder and subcommand cleanups

drbh · drbh · commit 7858dca0a09b · 2026-01-23T09:35:48.000-05:00
diff --git a/kernels/src/kernels/benchmark.py b/kernels/src/kernels/benchmark.py
@@ -452,16 +452,16 @@ def collect_machine_info() -> MachineInfo:
     )
 
 
-def get_kernel_sha_from_ops(kernel: Any) -> str:
-    ops_name = kernel.ops.__name__
-    # Format is torch.ops._<name>_<sha>, extract the last part after underscore
+def get_kernel_sha_from_build_name(kernel: Any) -> str:
+    ops_name = kernel.__name__
+    # Format is <name>_<sha>, extract the last part after underscore
     sha = ops_name.rsplit("_", 1)[-1]
     return sha
 
 
 def _synchronize() -> None:
     if torch.cuda.is_available():
-        _synchronize()
+        torch.cuda.synchronize()
     elif hasattr(torch, "xpu") and torch.xpu.is_available():
         torch.xpu.synchronize()
     elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
@@ -491,7 +491,7 @@ def run_benchmark_class(
     from kernels import get_kernel
 
     kernel = get_kernel(repo_id, revision=revision)
-    kernel_sha = get_kernel_sha_from_ops(kernel)
+    kernel_sha = get_kernel_sha_from_build_name(kernel)
 
     for method_name in benchmark_methods:
         workload_name = method_name.replace("benchmark_", "")
@@ -659,6 +659,8 @@ def run_benchmark_script(
         )
         for name, timing in results.items():
             all_results[f"{cls.__name__}.{name}"] = timing
+
+    print("test")
     return all_results, kernel_sha
 
 
diff --git a/kernels/src/kernels/benchmarks/__init__.py b/kernels/src/kernels/benchmarks/__init__.py
@@ -4,10 +4,13 @@
     FlashAttentionCausalBenchmark,
     FlashAttentionVarlenBenchmark,
 )
+from .layer_norm import LayerNormBenchmark, RMSNormBenchmark
 
 __all__ = [
     "FlashAttentionBenchmark",
     "FlashAttentionCausalBenchmark",
     "FlashAttentionVarlenBenchmark",
+    "LayerNormBenchmark",
+    "RMSNormBenchmark",
     "SiluAndMulBenchmark",
 ]
diff --git a/kernels/src/kernels/benchmarks/layer_norm.py b/kernels/src/kernels/benchmarks/layer_norm.py
@@ -0,0 +1,196 @@
+import torch
+
+from kernels.benchmark import Benchmark
+
+
+class RMSNormBenchmark(Benchmark):
+    seed: int = 42
+    eps: float = 1e-5
+
+    # Workload: small (B=2, S=128, D=768)
+    def setup_small(self):
+        B, S, D = 2, 128, 768
+        self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
+        self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
+        self.out = torch.empty_like(self.x)
+        self.B, self.S, self.D = B, S, D
+
+    def benchmark_small(self):
+        self.out = self.kernel.dropout_add_ln_fwd(
+            input=self.x.view(-1, self.D),
+            gamma=self.weight,
+            beta=None,
+            rowscale=None,
+            colscale=None,
+            x0_subset=None,
+            z_subset=None,
+            dropout_p=0.0,
+            epsilon=self.eps,
+            rowscale_const=1.0,
+            z_numrows=self.S,
+            gen=None,
+            residual_in_fp32=False,
+            is_rms_norm=True,
+        )[0].view(self.B, self.S, self.D)
+
+    def verify_small(self) -> torch.Tensor:
+        var = self.x.pow(2).mean(-1, keepdim=True)
+        return (self.x * torch.rsqrt(var + self.eps)) * self.weight
+
+    # Workload: medium (B=4, S=512, D=2048)
+    def setup_medium(self):
+        B, S, D = 4, 512, 2048
+        self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
+        self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
+        self.out = torch.empty_like(self.x)
+        self.B, self.S, self.D = B, S, D
+
+    def benchmark_medium(self):
+        self.out = self.kernel.dropout_add_ln_fwd(
+            input=self.x.view(-1, self.D),
+            gamma=self.weight,
+            beta=None,
+            rowscale=None,
+            colscale=None,
+            x0_subset=None,
+            z_subset=None,
+            dropout_p=0.0,
+            epsilon=self.eps,
+            rowscale_const=1.0,
+            z_numrows=self.S,
+            gen=None,
+            residual_in_fp32=False,
+            is_rms_norm=True,
+        )[0].view(self.B, self.S, self.D)
+
+    def verify_medium(self) -> torch.Tensor:
+        var = self.x.pow(2).mean(-1, keepdim=True)
+        return (self.x * torch.rsqrt(var + self.eps)) * self.weight
+
+    # Workload: large (B=8, S=1024, D=4096)
+    def setup_large(self):
+        B, S, D = 8, 1024, 4096
+        self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
+        self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
+        self.out = torch.empty_like(self.x)
+        self.B, self.S, self.D = B, S, D
+
+    def benchmark_large(self):
+        self.out = self.kernel.dropout_add_ln_fwd(
+            input=self.x.view(-1, self.D),
+            gamma=self.weight,
+            beta=None,
+            rowscale=None,
+            colscale=None,
+            x0_subset=None,
+            z_subset=None,
+            dropout_p=0.0,
+            epsilon=self.eps,
+            rowscale_const=1.0,
+            z_numrows=self.S,
+            gen=None,
+            residual_in_fp32=False,
+            is_rms_norm=True,
+        )[0].view(self.B, self.S, self.D)
+
+    def verify_large(self) -> torch.Tensor:
+        var = self.x.pow(2).mean(-1, keepdim=True)
+        return (self.x * torch.rsqrt(var + self.eps)) * self.weight
+
+
+class LayerNormBenchmark(Benchmark):
+    seed: int = 42
+    eps: float = 1e-5
+
+    # Workload: small (B=2, S=128, D=768)
+    def setup_small(self):
+        B, S, D = 2, 128, 768
+        self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
+        self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
+        self.out = torch.empty_like(self.x)
+        self.B, self.S, self.D = B, S, D
+
+    def benchmark_small(self):
+        self.out = self.kernel.dropout_add_ln_fwd(
+            input=self.x.view(-1, self.D),
+            gamma=self.weight,
+            beta=None,
+            rowscale=None,
+            colscale=None,
+            x0_subset=None,
+            z_subset=None,
+            dropout_p=0.0,
+            epsilon=self.eps,
+            rowscale_const=1.0,
+            z_numrows=self.S,
+            gen=None,
+            residual_in_fp32=False,
+            is_rms_norm=False,
+        )[0].view(self.B, self.S, self.D)
+
+    def verify_small(self) -> torch.Tensor:
+        return torch.nn.functional.layer_norm(
+            self.x, [self.D], self.weight, eps=self.eps
+        )
+
+    # Workload: medium (B=4, S=512, D=2048)
+    def setup_medium(self):
+        B, S, D = 4, 512, 2048
+        self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
+        self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
+        self.out = torch.empty_like(self.x)
+        self.B, self.S, self.D = B, S, D
+
+    def benchmark_medium(self):
+        self.out = self.kernel.dropout_add_ln_fwd(
+            input=self.x.view(-1, self.D),
+            gamma=self.weight,
+            beta=None,
+            rowscale=None,
+            colscale=None,
+            x0_subset=None,
+            z_subset=None,
+            dropout_p=0.0,
+            epsilon=self.eps,
+            rowscale_const=1.0,
+            z_numrows=self.S,
+            gen=None,
+            residual_in_fp32=False,
+            is_rms_norm=False,
+        )[0].view(self.B, self.S, self.D)
+
+    def verify_medium(self) -> torch.Tensor:
+        return torch.nn.functional.layer_norm(
+            self.x, [self.D], self.weight, eps=self.eps
+        )
+
+    # Workload: large (B=8, S=1024, D=4096)
+    def setup_large(self):
+        B, S, D = 8, 1024, 4096
+        self.x = torch.randn(B, S, D, device="cuda", dtype=torch.float16)
+        self.weight = torch.ones(D, device="cuda", dtype=torch.float16)
+        self.out = torch.empty_like(self.x)
+        self.B, self.S, self.D = B, S, D
+
+    def benchmark_large(self):
+        self.out = self.kernel.dropout_add_ln_fwd(
+            input=self.x.view(-1, self.D),
+            gamma=self.weight,
+            beta=None,
+            rowscale=None,
+            colscale=None,
+            x0_subset=None,
+            z_subset=None,
+            dropout_p=0.0,
+            epsilon=self.eps,
+            rowscale_const=1.0,
+            z_numrows=self.S,
+            gen=None,
+            residual_in_fp32=False,
+            is_rms_norm=False,
+        )[0].view(self.B, self.S, self.D)
+
+    def verify_large(self) -> torch.Tensor:
+        return torch.nn.functional.layer_norm(
+            self.x, [self.D], self.weight, eps=self.eps
+        )
diff --git a/kernels/src/kernels/cli.py b/kernels/src/kernels/cli.py
@@ -81,6 +81,18 @@ def main():
         action="store_true",
         help="If the repository should be private.",
     )
+    # by default dont include benchmarks, but enable with flag or
+    # only upload benchmarks with separate flag
+    upload_parser.add_argument(
+        "--benchmarks",
+        action="store_true",
+        help="If set, upload both benchmarks and build variants (default).",
+    )
+    upload_parser.add_argument(
+        "--benchmarks-only",
+        action="store_true",
+        help="If set, only upload the benchmarks directory.",
+    )
     upload_parser.set_defaults(func=upload_kernels)
 
     lock_parser = subparsers.add_parser("lock", help="Lock kernel revisions")
@@ -202,6 +214,28 @@ def upload_kernels(args):
     # Resolve `kernel_dir` to be uploaded.
     kernel_dir = Path(args.kernel_dir).resolve()
 
+    repo_id = create_repo(
+        repo_id=args.repo_id, private=args.private, exist_ok=True
+    ).repo_id
+
+    if args.branch is not None:
+        create_branch(repo_id=repo_id, branch=args.branch, exist_ok=True)
+
+    # benchmarks directory upload (doesn't require build variants)
+    if args.benchmarks or args.benchmarks_only:
+        upload_folder(
+            repo_id=repo_id,
+            folder_path=kernel_dir / "benchmarks",
+            revision=args.branch,
+            path_in_repo="benchmarks",
+            delete_patterns=["benchmark*.py"],
+            commit_message="Benchmarks uploaded using `kernels`.",
+            allow_patterns=["benchmark*.py"],
+        )
+
+    if args.benchmarks_only:
+        return  # Exit if only benchmarks are to be uploaded
+
     build_dir = None
     for candidate in [kernel_dir / "build", kernel_dir]:
         variants = [
@@ -217,13 +251,6 @@ def upload_kernels(args):
             f"Couldn't find any build variants in: {kernel_dir.absolute()} or {(kernel_dir / 'build').absolute()}"
         )
 
-    repo_id = create_repo(
-        repo_id=args.repo_id, private=args.private, exist_ok=True
-    ).repo_id
-
-    if args.branch is not None:
-        create_branch(repo_id=repo_id, branch=args.branch, exist_ok=True)
-
     delete_patterns: set[str] = set()
     for build_variant in build_dir.iterdir():
         if build_variant.is_dir():

Original file line number	Diff line number	Diff line change
`@@ -4,10 +4,13 @@`
`4`	`4`	`FlashAttentionCausalBenchmark,`
`5`	`5`	`FlashAttentionVarlenBenchmark,`
`6`	`6`	`)`
	`7`	`+from .layer_norm import LayerNormBenchmark, RMSNormBenchmark`
`7`	`8`
`8`	`9`	`__all__ = [`
`9`	`10`	`"FlashAttentionBenchmark",`
`10`	`11`	`"FlashAttentionCausalBenchmark",`
`11`	`12`	`"FlashAttentionVarlenBenchmark",`
	`13`	`+ "LayerNormBenchmark",`
	`14`	`+ "RMSNormBenchmark",`
`12`	`15`	`"SiluAndMulBenchmark",`
`13`	`16`	`]`