Benchmark command fixes (#229)

danieldk · drbh · web-flow · commit 657683441a48 · 2026-01-24T08:36:25.000+01:00
* feat: upload benchmarks folder and subcommand cleanups

* feat: respect versions on upload and download

* fix: adjust for comat and mypy

* Simplify `kernels upload` changes to avoid blocking release

---------

Co-authored-by: David Holtz &lt;david.richard.holtz@gmail.com&gt;
diff --git a/kernels/src/kernels/benchmark.py b/kernels/src/kernels/benchmark.py
@@ -88,11 +88,11 @@ def verify_silu(self) -> torch.Tensor:
 
     seed: int | None = None  # Optional: seed for reproducibility
 
-    def __init__(self):
-        self.kernel = None
+    def __init__(self) -> None:
+        self.kernel: Any = None
         self.out: Any = None  # Output tensor, set by setup methods
 
-    def setup(self):
+    def setup(self) -> None:
         """Override to set up tensors as instance attributes."""
         pass
 
@@ -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_", "")
@@ -689,8 +689,8 @@ def submit_benchmark(
 
 def run_benchmark(
     repo_id: str,
-    # TODO: change default to 1 in the future
-    revision: str = "main",
+    branch: str | None,
+    version: int | None,
     iterations: int = 100,
     warmup: int = 10,
     upload: bool = False,
@@ -708,6 +708,27 @@ def run_benchmark(
     # Suppress progress bars for cleaner output (files are often cached)
     disable_progress_bars()
 
+    # Requires either branch or version or parses from repo_id
+    if branch is None and version is None:
+        if "@" not in repo_id:
+            print("Error: must specify either branch or version", file=sys.stderr)
+            sys.exit(1)
+
+        # Parse from repo_id
+        repo_id, rev = repo_id.split("@", 1)
+
+        if rev.startswith("v") and rev[1:].isdigit():
+            version = int(rev[1:])
+        elif rev.isdigit():
+            print("Error: version must be prefixed with 'v'", file=sys.stderr)
+            sys.exit(1)
+        else:
+            branch = rev
+
+    # Move version or branch into revision
+    revision = f"v{version}" if version is not None else branch
+    assert revision is not None  # Guaranteed by parsing logic above
+
     print(f"Downloading {repo_id}@{revision}...", file=sys.stderr)
     repo_path = Path(snapshot_download(repo_id=repo_id, revision=revision))
 
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
@@ -126,10 +126,10 @@ def main():
         help="Kernel repo ID (e.g., kernels-community/activation)",
     )
     benchmark_parser.add_argument(
-        "--revision",
-        type=str,
-        default="main",
-        help="Kernel revision (default: main)",
+        "--branch", type=str, help="Kernel branch to benchmark"
+    )
+    benchmark_parser.add_argument(
+        "--version", type=int, help="Kernel version to benchmark"
     )
     benchmark_parser.add_argument(
         "--output",
@@ -247,7 +247,8 @@ def run_benchmark(args):
 
     benchmark.run_benchmark(
         repo_id=args.repo_id,
-        revision=args.revision,
+        branch=args.branch,
+        version=args.version,
         iterations=args.iterations,
         warmup=args.warmup,
         output=args.output,
diff --git a/kernels/src/kernels/metadata.py b/kernels/src/kernels/metadata.py
@@ -2,12 +2,27 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from kernels.compat import tomllib
+
 
 @dataclass
 class Metadata:
     python_depends: list[str]
     version: int | None
 
+    @staticmethod
+    def load_from_build_toml(build_toml_path: Path) -> "Metadata":
+        if build_toml_path.exists():
+            with open(build_toml_path, "rb") as f:
+                data = tomllib.load(f)
+                version = data.get("general", {}).get("version", None)
+                return Metadata(
+                    version=version,
+                    python_depends=[],
+                )
+
+        return Metadata(version=None, python_depends=[])
+
     @staticmethod
     def load_from_variant(variant_path: Path) -> "Metadata":
         metadata_path = variant_path / "metadata.json"
diff --git a/kernels/src/kernels/upload.py b/kernels/src/kernels/upload.py
@@ -7,7 +7,11 @@
 
 
 def upload_kernels_dir(
-    kernel_dir: Path, *, repo_id: str, branch: str | None, private: bool
+    kernel_dir: Path,
+    *,
+    repo_id: str,
+    branch: str | None,
+    private: bool,
 ):
     kernel_dir = Path(kernel_dir).resolve()
 
@@ -54,6 +58,17 @@ def upload_kernels_dir(
         if build_variant.is_dir():
             delete_patterns.add(f"{build_variant.name}/**")
 
+    # in the case we have variants, upload to the same as the kernel_dir
+    upload_folder(
+        repo_id=repo_id,
+        folder_path=kernel_dir / "benchmarks",
+        revision=branch,
+        path_in_repo="benchmarks",
+        delete_patterns=["benchmark*.py"],
+        commit_message="Benchmarks uploaded using `kernels`.",
+        allow_patterns=["benchmark*.py"],
+    )
+
     upload_folder(
         repo_id=repo_id,
         folder_path=build_dir,
diff --git a/kernels/uv.lock b/kernels/uv.lock

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`	`]`