Benchmarks uploaded using `kernels`.

benchmarks/benchmark.py

@@ -0,0 +1,233 @@
+import torch
+import torch.nn.functional as F
+from collections import namedtuple
+
+from kernels.benchmark import Benchmark
+
+
+def moe_mlp_reference(
+    x: torch.Tensor,
+    router_weight: torch.Tensor,
+    router_bias: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    gate_up_proj_bias: torch.Tensor,
+    down_proj: torch.Tensor,
+    down_proj_bias: torch.Tensor,
+    top_k: int = 4,
+    alpha: float = 1.702,
+    limit: float = 7.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    in_shape = x.shape
+    num_experts = router_weight.shape[0]
+    hidden_size = x.shape[-1]
+
+    # Flatten to (num_tokens, hidden_size)
+    hidden_states = x.view(-1, hidden_size)
+    num_tokens = hidden_states.shape[0]
+
+    # Router: compute logits and get top-k experts
+    logits = F.linear(hidden_states, router_weight, router_bias)
+    expert_weights, router_indices = torch.topk(logits, top_k, dim=-1)
+    routing_weights = F.softmax(expert_weights, dim=-1)
+
+    # Initialize output
+    next_states = torch.zeros_like(hidden_states)
+
+    # Create expert mask using one_hot
+    with torch.no_grad():
+        expert_mask = F.one_hot(router_indices, num_classes=num_experts)
+        expert_mask = expert_mask.permute(2, 1, 0)  # (num_experts, top_k, num_tokens)
+        # Find which experts are hit
+        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+    # Process each expert that has tokens
+    for expert_idx in expert_hit:
+        expert_idx = expert_idx[0]
+        with torch.no_grad():
+            top_k_idx, token_idx = torch.where(expert_mask[expert_idx])
+
+        current_state = hidden_states[token_idx]
+
+        # Up projection
+        gate_up = (
+            current_state @ gate_up_proj[expert_idx] + gate_up_proj_bias[expert_idx]
+        )
+
+        # Split into gate and up
+        gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+
+        # Clamp
+        gate = gate.clamp(min=None, max=limit)
+        up = up.clamp(min=-limit, max=limit)
+
+        # SwiGLU-like activation
+        glu = gate * torch.sigmoid(gate * alpha)
+        gated_output = (up + 1) * glu
+
+        # Down projection
+        out = gated_output @ down_proj[expert_idx] + down_proj_bias[expert_idx]
+
+        # Get the routing weight for this expert at the correct top_k position
+        weights_for_expert = routing_weights[token_idx, top_k_idx]
+        weighted_output = out * weights_for_expert[:, None]
+        next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype))
+
+    return next_states.view(in_shape), routing_weights
+
+
+class MegaBlocksMoeBenchmark(Benchmark):
+    seed: int = 42
+
+    def setup(self):
+        # Config matching readme_example.py
+        ne, hs, isz = 128, 1152, 3072
+        batch, seq = 8, 1
+
+        # Router
+        self.router_weight = torch.randn(
+            ne, hs, device=self.device, dtype=torch.float32
+        )
+        torch.nn.init.kaiming_uniform_(self.router_weight)
+        self.router_bias = torch.zeros(ne, device=self.device, dtype=torch.float32)
+
+        # Expert weights
+        self.gate_up_proj = (
+            torch.randn(ne, hs, isz, device=self.device, dtype=torch.float32) * 0.02
+        )
+        self.gate_up_proj_bias = torch.zeros(
+            ne, isz, device=self.device, dtype=torch.float32
+        )
+        self.down_proj = (
+            torch.randn(ne, isz // 2, hs, device=self.device, dtype=torch.float32)
+            * 0.02
+        )
+        self.down_proj_bias = torch.zeros(
+            ne, hs, device=self.device, dtype=torch.float32
+        )
+
+        # Input
+        self.x = (
+            torch.randn(seq, batch, hs, device=self.device, dtype=torch.float32) * 0.1
+        )
+
+        # Setup the model
+        self.model = self.kernel.layers.MegaBlocksMoeMLP()
+        self.model.router = torch.nn.Linear(hs, ne, device=self.device)
+        self.model.router.weight.data = self.router_weight.clone()
+        self.model.router.bias.data = self.router_bias.clone()
+
+        Experts = namedtuple(
+            "Experts",
+            [
+                "gate_up_proj",
+                "gate_up_proj_bias",
+                "down_proj",
+                "down_proj_bias",
+                "hidden_size",
+                "num_experts",
+            ],
+        )
+        self.model.experts = Experts(
+            gate_up_proj=torch.nn.Parameter(self.gate_up_proj.clone()),
+            gate_up_proj_bias=torch.nn.Parameter(self.gate_up_proj_bias.clone()),
+            down_proj=torch.nn.Parameter(self.down_proj.clone()),
+            down_proj_bias=torch.nn.Parameter(self.down_proj_bias.clone()),
+            hidden_size=hs,
+            num_experts=ne,
+        )
+
+        self.out = torch.empty(seq, batch, hs, device=self.device, dtype=torch.float32)
+
+    def benchmark_base(self):
+        self.out, self.expert_weights = self.model(self.x)
+
+    def verify_base(self) -> torch.Tensor:
+        ref_out, _ = moe_mlp_reference(
+            self.x,
+            self.router_weight,
+            self.router_bias,
+            self.gate_up_proj,
+            self.gate_up_proj_bias,
+            self.down_proj,
+            self.down_proj_bias,
+            top_k=4,
+        )
+        return ref_out
+
+    def setup_large(self):
+        # Larger config with more tokens
+        ne, hs, isz = 128, 1152, 3072
+        batch, seq = 32, 16
+
+        # Router
+        self.router_weight = torch.randn(
+            ne, hs, device=self.device, dtype=torch.float32
+        )
+        torch.nn.init.kaiming_uniform_(self.router_weight)
+        self.router_bias = torch.zeros(ne, device=self.device, dtype=torch.float32)
+
+        # Expert weights
+        self.gate_up_proj = (
+            torch.randn(ne, hs, isz, device=self.device, dtype=torch.float32) * 0.02
+        )
+        self.gate_up_proj_bias = torch.zeros(
+            ne, isz, device=self.device, dtype=torch.float32
+        )
+        self.down_proj = (
+            torch.randn(ne, isz // 2, hs, device=self.device, dtype=torch.float32)
+            * 0.02
+        )
+        self.down_proj_bias = torch.zeros(
+            ne, hs, device=self.device, dtype=torch.float32
+        )
+
+        # Input
+        self.x = (
+            torch.randn(seq, batch, hs, device=self.device, dtype=torch.float32) * 0.1
+        )
+
+        # Setup the model
+        self.model = self.kernel.layers.MegaBlocksMoeMLP()
+        self.model.router = torch.nn.Linear(hs, ne, device=self.device)
+        self.model.router.weight.data = self.router_weight.clone()
+        self.model.router.bias.data = self.router_bias.clone()
+
+        Experts = namedtuple(
+            "Experts",
+            [
+                "gate_up_proj",
+                "gate_up_proj_bias",
+                "down_proj",
+                "down_proj_bias",
+                "hidden_size",
+                "num_experts",
+                "capacity_factor",
+            ],
+        )
+        self.model.experts = Experts(
+            gate_up_proj=torch.nn.Parameter(self.gate_up_proj.clone()),
+            gate_up_proj_bias=torch.nn.Parameter(self.gate_up_proj_bias.clone()),
+            down_proj=torch.nn.Parameter(self.down_proj.clone()),
+            down_proj_bias=torch.nn.Parameter(self.down_proj_bias.clone()),
+            hidden_size=hs,
+            num_experts=ne,
+            capacity_factor=4.0,  # Higher capacity to avoid token dropping
+        )
+
+        self.out = torch.empty(seq, batch, hs, device=self.device, dtype=torch.float32)
+
+    def benchmark_large(self):
+        self.out, self.expert_weights = self.model(self.x)
+
+    def verify_large(self) -> torch.Tensor:
+        ref_out, _ = moe_mlp_reference(
+            self.x,
+            self.router_weight,
+            self.router_bias,
+            self.gate_up_proj,
+            self.gate_up_proj_bias,
+            self.down_proj,
+            self.down_proj_bias,
+            top_k=4,
+        )
+        return ref_out