feat: align outputs and support backward method

compare_example.py

@@ -130,6 +130,10 @@ model.down_proj_bias.data = down_proj_bias
 model = model.cuda()
 model.eval()
 
+# Warmup
+for _ in range(5):
+    _ = model(hidden_states, router_indices, routing_weights)
+
 torch.cuda.synchronize()
 torch.cuda.reset_peak_memory_stats()
 start = time.perf_counter()
@@ -151,6 +155,19 @@ ref_output_reshaped = ref_output.view(kernel_output.shape)
 # Test yamoe_ref implementation
 expert_capacity = batch_seq * top_k // num_experts * 2  # Generous capacity
 
+# Warmup
+for _ in range(5):
+    _ = binned_experts_ref(
+        hidden_states,
+        router_indices,
+        routing_weights,
+        gate_up_proj,
+        gate_up_proj_bias,
+        down_proj,
+        down_proj_bias,
+        expert_capacity,
+    )
+
 torch.cuda.synchronize()
 torch.cuda.reset_peak_memory_stats()
 start = time.perf_counter()

gpt_oss_backward.py

@@ -0,0 +1,181 @@
+# /// script
+# requires-python = "==3.10"
+# dependencies = ["torch==2.7.0", "triton", "numpy", "kernels"]
+# [tool.uv.sources]
+# kernels = { git = "https://github.com/huggingface/kernels.git" }
+# ///
+
+import torch
+import sys
+import time
+from kernels import get_kernel, get_local_kernel
+from pathlib import Path
+
+load_method = 2  # 1: sym, 2: local, 3: hf
+
+if load_method == 1:
+    sys.path.insert(0, "./torch-ext")
+    import yamoe
+elif load_method == 2:
+    yamoe = get_local_kernel(Path("result"), "yamoe")
+elif load_method == 3:
+    yamoe = get_kernel("drbh/yamoe", revision="v0.1.0")
+
+torch.manual_seed(42)
+
+
+def zero_grads(m):
+    for p in m.parameters():
+        if p.grad is not None:
+            p.grad = None
+
+
+def benchmark_backward(model, x, tag: str, iters: int = 10, warmup: int = 10):
+    x_local = x.detach().clone().requires_grad_(True)
+    x_local.retain_grad()
+
+    # Warmup
+    for _ in range(warmup):
+        out = model(x_local)
+        out = out[0] if isinstance(out, tuple) else out
+        loss = out.mean()
+        zero_grads(model)
+        if x_local.grad is not None:
+            x_local.grad = None
+        loss.backward()
+
+    # Benchmark
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(iters):
+        out = model(x_local)
+        out = out[0] if isinstance(out, tuple) else out
+        loss = out.mean()
+        zero_grads(model)
+        if x_local.grad is not None:
+            x_local.grad = None
+        loss.backward()
+        torch.cuda.synchronize()
+    bwd_ms = (time.perf_counter() - start) * 1e3 / iters
+    peak_mem = torch.cuda.max_memory_allocated() / 1024**3  # Convert to GB
+
+    print(f"[{tag}] backward: {bwd_ms:.2f} ms | peak mem: {peak_mem:.2f} GB")
+    return bwd_ms
+
+
+def print_grad_norms(m, x, tag):
+    print(f"\n[{tag}] Gradient norms:")
+    xg = x.grad.norm().item() if x.grad is not None else 0.0
+    print(f"  input grad: {xg:.6f}")
+    for name, p in m.named_parameters():
+        if p.grad is None:
+            print(f"  {name}: None")
+        else:
+            print(f"  {name}: {p.grad.norm().item():.6f}")
+
+
+def main():
+    ref_moe_cls = yamoe.vendored.gpt_oss_mlp.GptOssMLP
+    new_moe_cls = yamoe.Yamoe
+
+    batch_size, seq_len, hidden_dim = 4, 1024, 2880
+    num_experts, top_k = 8, 2
+
+    config = type("Config", (), {})()
+    config.hidden_size = hidden_dim
+    config.intermediate_size = hidden_dim
+    config.num_local_experts = num_experts
+    config.num_experts_per_tok = top_k
+    ref_moe = ref_moe_cls(config)
+
+    print(ref_moe)
+
+    for p in ref_moe.parameters():
+        if p.dim() > 1:
+            torch.nn.init.xavier_uniform_(p)
+        else:
+            torch.nn.init.zeros_(p)
+
+    x = torch.randn(batch_size, seq_len, hidden_dim, device="cuda")
+    ref_moe = ref_moe.cuda()
+
+    # Test reference implementation backward
+    print("\nReference Implementation Backward")
+
+    # Small warmup
+    print(" Warming up...")
+    x_warmup = x.detach().requires_grad_(True)
+    for _ in range(3):
+        out = ref_moe(x_warmup)
+        out = out[0] if isinstance(out, tuple) else out
+        loss = out.mean()
+        zero_grads(ref_moe)
+        if x_warmup.grad is not None:
+            x_warmup.grad = None
+        loss.backward()
+    torch.cuda.synchronize()
+
+    # Run once to get gradient info
+    x_ref = x.detach().requires_grad_(True)
+    x_ref.retain_grad()
+    ref_output = ref_moe(x_ref)
+    out = ref_output[0] if isinstance(ref_output, tuple) else ref_output
+    print(f" Input shape: {x_ref.shape}")
+    print(f" Output shape: {out.shape}")
+    print(
+        f" Output mean: {out.mean():.6f}, std: {out.std():.6f}, norm: {out.norm():.6f}"
+    )
+
+    loss = out.mean()
+    zero_grads(ref_moe)
+    loss.backward()
+    print_grad_norms(ref_moe, x_ref, "reference")
+
+    benchmark_backward(ref_moe, x, tag="reference", warmup=10, iters=20)
+
+    # Switch to YAMOE-backed backward
+    print("\nYAMOE-backed Implementation Backward")
+    ref_moe.forward = new_moe_cls.forward.__get__(ref_moe)
+    ref_moe._routing_weights_buffer = None
+    ref_moe._batch_indices_buffer = None
+    ref_moe._last_batch_seq = None
+    ref_moe._last_num_experts = None
+    ref_moe.enable_router_grads = True
+    ref_moe.num_experts = num_experts
+    ref_moe.top_k = top_k
+
+    # Small warmup
+    print(" Warming up...")
+    x_warmup = x.detach().requires_grad_(True)
+    for _ in range(3):
+        out = ref_moe(x_warmup)
+        out = out[0] if isinstance(out, tuple) else out
+        loss = out.mean()
+        zero_grads(ref_moe)
+        if x_warmup.grad is not None:
+            x_warmup.grad = None
+        loss.backward()
+    torch.cuda.synchronize()
+
+    # Run once to get gradient info
+    x_cuda = x.detach().requires_grad_(True)
+    x_cuda.retain_grad()
+    cuda_output = ref_moe(x_cuda)
+    out = cuda_output[0] if isinstance(cuda_output, tuple) else cuda_output
+    print(f" Input shape: {x_cuda.shape}")
+    print(f" Output shape: {out.shape}")
+    print(
+        f" Output mean: {out.mean():.6f}, std: {out.std():.6f}, norm: {out.norm():.6f}"
+    )
+
+    loss = out.mean()
+    zero_grads(ref_moe)
+    loss.backward()
+    print_grad_norms(ref_moe, x_cuda, "yamoe-backed")
+
+    benchmark_backward(ref_moe, x, tag="yamoe-backed", warmup=10, iters=20)
+
+
+if __name__ == "__main__":
+    main()

gpt_oss_match.py

@@ -0,0 +1,124 @@
+# /// script
+# requires-python = "==3.10"
+# dependencies = ["torch==2.7.0", "triton", "numpy", "kernels"]
+# [tool.uv.sources]
+# kernels = { git = "https://github.com/huggingface/kernels.git" }
+# ///
+
+import torch
+import sys
+import time
+from kernels import get_kernel, get_local_kernel
+from pathlib import Path
+
+load_method = 2  # 1: sym, 2: local, 3: hf
+
+if load_method == 1:
+    sys.path.insert(0, "./torch-ext")
+    import yamoe
+elif load_method == 2:
+    yamoe = get_local_kernel(Path("result"), "yamoe")
+elif load_method == 3:
+    yamoe = get_kernel("drbh/yamoe", revision="v0.1.0")
+
+torch.manual_seed(42)
+
+
+def benchmark_forward(model, x, tag: str, iters: int = 10, warmup: int = 10):
+    x_local = x.detach().clone().requires_grad_(False)
+
+    for _ in range(warmup):
+        out = model(x_local)
+        out = out[0] if isinstance(out, tuple) else out
+
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(iters):
+        out = model(x_local)
+        out = out[0] if isinstance(out, tuple) else out
+        torch.cuda.synchronize()
+    fwd_ms = (time.perf_counter() - start) * 1e3 / iters
+    peak_mem = torch.cuda.max_memory_allocated() / 1024**3  # Convert to GB
+
+    print(f"[{tag}] fwd: {fwd_ms:.2f} ms | peak mem: {peak_mem:.2f} GB")
+    return fwd_ms
+
+
+def main():
+    ref_moe_cls = yamoe.vendored.gpt_oss_mlp.GptOssMLP
+    new_moe_cls = yamoe.Yamoe
+
+    batch_size, seq_len, hidden_dim = 4, 1024, 2880
+    num_experts, top_k = 8, 2
+
+    config = type("Config", (), {})()
+    config.hidden_size = hidden_dim
+    config.intermediate_size = hidden_dim
+    config.num_local_experts = num_experts
+    config.num_experts_per_tok = top_k
+    ref_moe = ref_moe_cls(config)
+
+    print(ref_moe)
+
+    for p in ref_moe.parameters():
+        if p.dim() > 1:
+            torch.nn.init.xavier_uniform_(p)
+        else:
+            torch.nn.init.zeros_(p)
+
+    x = torch.randn(batch_size, seq_len, hidden_dim, device="cuda")
+    ref_moe = ref_moe.cuda()
+    ref_moe = ref_moe.eval()
+
+    # Test reference implementation
+    print("\nReference Implementation")
+
+    # Small warmup
+    print(" Warming up...")
+    for _ in range(3):
+        _ = ref_moe(x)
+    torch.cuda.synchronize()
+
+    x_ref = x.detach().requires_grad_(False)
+    ref_output = ref_moe(x_ref)
+    out = ref_output[0] if isinstance(ref_output, tuple) else ref_output
+    print(f" Input shape: {x_ref.shape}")
+    print(f" Output shape: {out.shape}")
+    print(
+        f" Output mean: {out.mean():.6f}, std: {out.std():.6f}, norm: {out.norm():.6f}"
+    )
+
+    benchmark_forward(ref_moe, x, tag="reference", warmup=10, iters=20)
+
+    # Switch to YAMOE-backed forward
+    print("\nYAMOE-backed Implementation")
+    ref_moe.forward = new_moe_cls.forward.__get__(ref_moe)
+    ref_moe._routing_weights_buffer = None
+    ref_moe._batch_indices_buffer = None
+    ref_moe._last_batch_seq = None
+    ref_moe._last_num_experts = None
+    ref_moe.enable_router_grads = False
+    ref_moe.num_experts = num_experts
+    ref_moe.top_k = top_k
+
+    # Small warmup
+    print(" Warming up...")
+    for _ in range(3):
+        _ = ref_moe(x)
+    torch.cuda.synchronize()
+
+    x_cuda = x.detach().requires_grad_(False)
+    cuda_output = ref_moe(x_cuda)
+    out = cuda_output[0] if isinstance(cuda_output, tuple) else cuda_output
+    print(f" Input shape: {x_cuda.shape}")
+    print(f" Output shape: {out.shape}")
+    print(
+        f" Output mean: {out.mean():.6f}, std: {out.std():.6f}, norm: {out.norm():.6f}"
+    )
+
+    benchmark_forward(ref_moe, x, tag="yamoe-backed", warmup=10, iters=20)
+
+
+if __name__ == "__main__":
+    main()

torch-ext/yamoe/layers.py

@@ -1,95 +1,190 @@
 import torch
+import time
 from ._ops import ops
 
 
-class Yamoe(torch.nn.Module):
-    """Yamoe MoE layer with routing and expert computation"""
+class _ExpertsFn(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        hidden_flat: torch.Tensor,
+        router_indices: torch.Tensor,
+        routing_weights_dense: torch.Tensor,
+        gate_up_proj: torch.Tensor,
+        gate_up_proj_bias: torch.Tensor,
+        down_proj: torch.Tensor,
+        down_proj_bias: torch.Tensor,
+        expert_capacity: int,
+        num_experts: int,
+        top_k: int,
+        enable_router_grads: bool,
+    ):
+        out = ops.experts(
+            hidden_flat,
+            router_indices,
+            routing_weights_dense,
+            gate_up_proj,
+            gate_up_proj_bias,
+            down_proj,
+            down_proj_bias,
+            expert_capacity,
+            num_experts,
+            top_k,
+        )
+        ctx.expert_capacity = expert_capacity
+        ctx.num_experts = num_experts
+        ctx.top_k = top_k
+        ctx.enable_router_grads = bool(enable_router_grads)
+        ctx.save_for_backward(
+            hidden_flat,
+            router_indices,
+            routing_weights_dense,
+            gate_up_proj,
+            gate_up_proj_bias,
+            down_proj,
+            down_proj_bias,
+        )
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        (
+            hidden_flat,
+            router_indices,
+            routing_weights_dense,
+            gate_up_proj,
+            gate_up_proj_bias,
+            down_proj,
+            down_proj_bias,
+        ) = ctx.saved_tensors
+
+        (
+            grad_hidden_flat,
+            grad_routing_weights,
+            grad_gate_up_proj,
+            grad_gate_up_proj_bias,
+            grad_down_proj,
+            grad_down_proj_bias,
+        ) = ops.experts_backward(
+            grad_output,
+            hidden_flat,
+            router_indices,
+            routing_weights_dense,
+            gate_up_proj,
+            gate_up_proj_bias,
+            down_proj,
+            down_proj_bias,
+            ctx.expert_capacity,
+            ctx.num_experts,
+            ctx.top_k,
+        )
 
-    can_torch_compile: bool = True
+        # Return grad for dense routing; autograd handles scatter->softmax->linear
+        grad_routing_weights_dense = (
+            grad_routing_weights if ctx.enable_router_grads else None
+        )
+
+        # Return gradients per input (None for non-differentiable)
+        return (
+            grad_hidden_flat,  # hidden_flat
+            None,  # router_indices
+            grad_routing_weights_dense,  # routing_weights_dense
+            grad_gate_up_proj,
+            grad_gate_up_proj_bias,
+            grad_down_proj,
+            grad_down_proj_bias,
+            None,  # expert_capacity
+            None,  # num_experts
+            None,  # top_k
+            None,  # enable_router_grads
+        )
 
-    def __init__(self):
-        super().__init__()
-        # Pre-allocate buffers to avoid repeated allocations
-        self._routing_weights_buffer = None
-        self._batch_indices_buffer = None
-        self._last_batch_seq = None
-        self._last_num_experts = None
+
+class Yamoe(torch.nn.Module):
+    can_torch_compile: bool = False
+
+    _routing_weights_buffer: torch.Tensor = None
+    _batch_indices_buffer: torch.Tensor = None
+    _last_batch_seq: int = None
+    _last_num_experts: int = None
+    enable_router_grads: bool = True
 
     def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Initialize runtime attrs when this forward is bound onto a different module (e.g., GptOssMLP)
+        if not hasattr(self, "enable_router_grads"):
+            # self.enable_router_grads = True
+            self.enable_router_grads = False
+        if not hasattr(self, "_routing_weights_buffer"):
+            self._routing_weights_buffer = None
+            self._batch_indices_buffer = None
+            self._last_batch_seq = None
+            self._last_num_experts = None
+            self._timing_enabled = False
+            self._timing_stats = {}
+
         batch_size, seq_len, hidden_dim = hidden_states.shape
         batch_seq = batch_size * seq_len
 
         num_experts = getattr(self, "num_experts", 128)
         top_k = getattr(self, "top_k", 4)
 
+        # Enable timing if requested
+        timing = getattr(self, "_timing_enabled", False)
+
+        if timing:
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+
         # Route tokens to experts
         x_flat = hidden_states.view(-1, hidden_dim)
         logits = torch.nn.functional.linear(
             x_flat, self.router.weight, self.router.bias
         )
 
+        if timing:
+            torch.cuda.synchronize()
+            t1 = time.perf_counter()
+            self._timing_stats["router"] = (t1 - t0) * 1000
+
         # Compute top-k
         if top_k == 1:
-            routing_weights, router_indices = logits.max(dim=-1, keepdim=True)
+            routing_logits_topk, router_indices = logits.max(dim=-1, keepdim=True)
         else:
-            routing_weights, router_indices = torch.topk(logits, top_k, dim=-1)
+            routing_logits_topk, router_indices = torch.topk(logits, top_k, dim=-1)
 
-        routing_weights = routing_weights.softmax(dim=-1)
+        # Match reference path exactly: use F.softmax with explicit dtype
+        routing_weights_topk = torch.nn.functional.softmax(
+            routing_logits_topk, dim=-1, dtype=routing_logits_topk.dtype
+        )
 
         # Create router scores
-        router_scores = (
-            torch.zeros_like(logits)
-            .scatter_(1, router_indices, routing_weights)
-            .transpose(0, 1)
+        router_scores = torch.zeros_like(logits).scatter_(
+            1, router_indices, routing_weights_topk
         )
 
-        # Convert routing_weights to sparse format [batch_seq, num_experts]
-        # Reuse buffer if possible to reduce allocations
-        if (
-            self._routing_weights_buffer is None
-            or self._last_batch_seq != batch_seq
-            or self._last_num_experts != num_experts
-            or self._routing_weights_buffer.device != routing_weights.device
-        ):
-            self._routing_weights_buffer = torch.zeros(
-                batch_seq,
-                num_experts,
-                device=routing_weights.device,
-                dtype=routing_weights.dtype,
-            )
-            self._batch_indices_buffer = (
-                torch.arange(batch_seq, device=routing_weights.device)
-                .unsqueeze(1)
-                .expand(-1, top_k)
-            )
-            self._last_batch_seq = batch_seq
-            self._last_num_experts = num_experts
-        else:
-            self._routing_weights_buffer.zero_()
+        if timing:
+            torch.cuda.synchronize()
+            t2 = time.perf_counter()
+            self._timing_stats["topk_softmax"] = (t2 - t1) * 1000
 
-        # Fill sparse routing weights
-        flat_indices = router_indices.view(batch_seq, top_k)
-        flat_weights = routing_weights.view(batch_seq, top_k)
-        self._routing_weights_buffer[self._batch_indices_buffer, flat_indices] = (
-            flat_weights
-        )
-
-        # FIX: Use the correct expert projections
-        gate_up = self.experts.gate_up_proj[:, :, : hidden_dim * top_k].contiguous()
-        gate_up_bias = self.experts.gate_up_proj_bias[
-            :, : hidden_dim * top_k
-        ].contiguous()
+        dense_routing = router_scores  # [B*S, E]
 
+        # Kernel expects shapes: [E, H, 2H] and [E, H, H]
+        gate_up = self.experts.gate_up_proj[:, :, : 2 * hidden_dim].contiguous()
+        gate_up_bias = self.experts.gate_up_proj_bias[:, : 2 * hidden_dim].contiguous()
         down_proj = self.experts.down_proj[:, :hidden_dim, :].contiguous()
-
         expert_capacity = batch_seq * top_k // num_experts * 2
 
-        with torch.no_grad():
-            # Compute expert output
-            output = ops.experts(
+        if timing:
+            torch.cuda.synchronize()
+            t3 = time.perf_counter()
+
+        # Compute expert output with custom backward
+        if self.enable_router_grads:
+            output = _ExpertsFn.apply(
                 hidden_states.view(-1, hidden_dim),
                 router_indices,
-                self._routing_weights_buffer,
+                dense_routing,
                 gate_up,
                 gate_up_bias,
                 down_proj,
@@ -97,8 +192,39 @@ class Yamoe(torch.nn.Module):
                 expert_capacity,
                 num_experts,
                 top_k,
+                self.enable_router_grads,
             )
+        else:
+            with torch.no_grad():
+                output = ops.experts(
+                    hidden_states.view(-1, hidden_dim),
+                    router_indices,
+                    dense_routing,
+                    gate_up,
+                    gate_up_bias,
+                    down_proj,
+                    self.experts.down_proj_bias,
+                    expert_capacity,
+                    num_experts,
+                    top_k,
+                )
+
+        if timing:
+            torch.cuda.synchronize()
+            t4 = time.perf_counter()
+            self._timing_stats["experts_kernel"] = (t4 - t3) * 1000
 
         # Reshape output back to [B, S, H]
         output = output.view(batch_size, seq_len, hidden_dim)
+
+        if timing:
+            torch.cuda.synchronize()
+            t5 = time.perf_counter()
+            self._timing_stats["total"] = (t5 - t0) * 1000
+            print(f"\n[Yamoe.forward timing in ms]")
+            print(f"  Router linear: {self._timing_stats['router']:.3f}")
+            print(f"  TopK + Softmax: {self._timing_stats['topk_softmax']:.3f}")
+            print(f"  Experts kernel: {self._timing_stats['experts_kernel']:.3f}")
+            print(f"  Total: {self._timing_stats['total']:.3f}")
+
         return output, router_scores