[DSv4][P5-5] Shared Expert MLP: strict CUDA + Triton kernels

benchmarks/benchmark_shared_expert_mlp.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""P5-5 (#64) shared_expert_mlp benchmark: torch-native vs Triton vs CUDA.
+
+torch-native is the non-deterministic cuBLAS/eager reference (speed ceiling);
+the Triton and CUDA rows are the strict ``oracle-fp32-serial-v1`` kernels this
+PR delivers. Alignment between the strict backends is asserted on every shape.
+
+    python benchmarks/benchmark_shared_expert_mlp.py [--tokens 16,256,2048]
+"""
+
+from __future__ import annotations
+
+import argparse
+import pathlib
+import sys
+
+import torch
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from rl_engine.moe.contract import SharedBatch, tensor_sha256  # noqa: E402
+
+
+def torch_native(batch: SharedBatch, dy: torch.Tensor):
+    """Eager BF16 reference (cuBLAS + fused silu): fast but not bit-stable."""
+    x = batch.x.detach().requires_grad_(True)
+    z = x @ batch.w_fc1.t()
+    ffn = z.shape[1] // 2
+    gate, up = z[:, :ffn], z[:, ffn:]
+    h = torch.nn.functional.silu(gate) * up
+    y = h @ batch.w_fc2.t()
+    y.backward(dy)
+    return y, x.grad
+
+
+def make_runner(provider):
+    def run(batch: SharedBatch, dy: torch.Tensor):
+        y, saved = provider.shared_expert_mlp_fwd(batch)
+        dx = provider.shared_expert_mlp_bwd(dy, batch, saved)
+        return y, dx
+
+    return run
+
+
+def time_ms(fn, *args, warmup: int = 3, iters: int = 10) -> float:
+    for _ in range(warmup):
+        fn(*args)
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn(*args)
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--hidden", type=int, default=4096)
+    parser.add_argument("--ffn", type=int, default=2048)
+    parser.add_argument("--tokens", default="16,256,2048")
+    parser.add_argument("--iters", type=int, default=10)
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        print("CUDA device required")
+        return 1
+
+    from rl_engine.moe.provider import resolve_provider
+
+    runners: dict[str, object] = {"torch-native": torch_native}
+    strict: dict[str, object] = {}
+    for label, spec in (
+        ("triton", "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider"),
+        ("cuda", "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider"),
+    ):
+        try:
+            runner = make_runner(resolve_provider(spec))
+            runners[label] = runner
+            strict[label] = runner
+        except NotImplementedError as exc:
+            print(f"[skip] {label}: {exc}")
+
+    device = torch.device("cuda")
+    gen = torch.Generator(device="cpu").manual_seed(2026)
+    header = f"{'T':>6} {'backend':>14} {'fwd+bwd ms':>12} {'vs native':>10}"
+    print(f"H={args.hidden} F={args.ffn} ({torch.cuda.get_device_name(0)})")
+    print(header)
+    for t in [int(v) for v in args.tokens.split(",")]:
+        x = torch.randn(t, args.hidden, generator=gen).to(torch.bfloat16).to(device)
+        w1 = (
+            (torch.randn(2 * args.ffn, args.hidden, generator=gen) / args.hidden**0.5)
+            .to(torch.bfloat16)
+            .to(device)
+        )
+        w2 = (
+            (torch.randn(args.hidden, args.ffn, generator=gen) / args.ffn**0.5)
+            .to(torch.bfloat16)
+            .to(device)
+        )
+        batch = SharedBatch(x=x, w_fc1=w1, w_fc2=w2)
+        dy = torch.randn(t, args.hidden, generator=gen).to(torch.bfloat16).to(device)
+
+        outputs = {}
+        base_ms = None
+        for label, fn in runners.items():
+            ms = time_ms(fn, batch, dy, iters=args.iters)
+            outputs[label] = fn(batch, dy)
+            if label == "torch-native":
+                base_ms = ms
+            rel = f"{ms / base_ms:8.2f}x" if base_ms else "        -"
+            print(f"{t:>6} {label:>14} {ms:12.3f} {rel:>10}")
+
+        strict_hashes = {
+            label: (tensor_sha256(outputs[label][0]), tensor_sha256(outputs[label][1]))
+            for label in strict
+        }
+        if len(strict_hashes) == 2 and len(set(strict_hashes.values())) != 1:
+            print(f"  !! strict backends diverged at T={t}: {strict_hashes}")
+            return 1
+        if strict_hashes:
+            print(f"  strict backends byte-equal: {len(strict_hashes)}/{len(strict_hashes)}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/cuda/moe/shared_expert_mlp.cu

@@ -0,0 +1,213 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// P5-5 (#64) Shared Expert MLP strict kernels: fc1 -> one-round SwiGLU -> fc2.
+//
+// Numeric profile ``oracle-fp32-serial-v1`` (see rl_engine/moe/oracle.py):
+// every accumulation is FP32, serial, ascending-k, with multiply and add
+// rounded separately (__fmul_rn / __fadd_rn; never contracted into FMA).
+// One thread owns one output element, so batch size and padding cannot
+// change a row's bytes (Axis-A bitwise invariance) and there is no
+// cross-thread floating-point reduction anywhere.
+//
+// The one-round SwiGLU core (FP32 math, single BF16 round on the output) is
+// shared with P5-2 (#63): the p_s / clamp variant extends the same device
+// functions in this translation unit rather than forking the math.
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_bf16.h>
+
+namespace {
+
+// out[m, n] = sum_{k ascending} fadd_rn(acc, fmul_rn(a[m, k], b(n, k)))
+// A is BF16 [M, K]; B is BF16 [N, K] (TRANS_B = false) or [K, N] (true).
+// Output stays FP32; the caller rounds to BF16 where the contract says so.
+template <bool TRANS_B>
+__global__ void p5_strict_gemm_kernel(
+    const __nv_bfloat16* __restrict__ a,
+    const __nv_bfloat16* __restrict__ b,
+    float* __restrict__ out,
+    const int64_t m_rows,
+    const int64_t n_cols,
+    const int64_t k_dim) {
+  const int64_t total = m_rows * n_cols;
+  const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
+  for (int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x; idx < total;
+       idx += stride) {
+    const int64_t m = idx / n_cols;
+    const int64_t n = idx - m * n_cols;
+    const __nv_bfloat16* a_row = a + m * k_dim;
+    float acc = 0.0f;
+    for (int64_t k = 0; k < k_dim; ++k) {
+      const float av = __bfloat162float(a_row[k]);
+      const float bv =
+          __bfloat162float(TRANS_B ? b[k * n_cols + n] : b[n * k_dim + k]);
+      acc = __fadd_rn(acc, __fmul_rn(av, bv));
+    }
+    out[idx] = acc;
+  }
+}
+
+// Matches torch.sigmoid on FP32 CUDA tensors: 1 / (1 + exp(-x)) with
+// IEEE div.rn and the accurate expf (no fast-math in this build).
+__device__ __forceinline__ float sigmoid_rn(float x) {
+  return 1.0f / (1.0f + expf(-x));
+}
+
+// One-round SwiGLU core, shared-expert mode (p_s = None: no clamp, no route
+// weight). z is the packed FP32 fc1 output [T, 2F] (gate columns then up);
+// h is the single BF16 round of SiLU(gate) * up.
+__global__ void p5_swiglu_shared_forward_kernel(
+    const float* __restrict__ z,
+    __nv_bfloat16* __restrict__ h,
+    const int64_t n,
+    const int64_t width) {
+  const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
+  for (int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x; idx < n;
+       idx += stride) {
+    const int64_t row = idx / width;
+    const int64_t col = idx - row * width;
+    const int64_t gate_index = row * (2 * width) + col;
+    const float g = z[gate_index];
+    const float u = z[gate_index + width];
+    const float sig = sigmoid_rn(g);
+    const float silu = __fmul_rn(g, sig);
+    h[idx] = __float2bfloat16(__fmul_rn(silu, u));
+  }
+}
+
+// Backward of the same graph (p_s = None): recomputes sig/silu from the saved
+// FP32 z with the identical instruction sequence, so the bits match the
+// forward. dz packs (dgate | dup), each rounded to BF16 exactly once at the
+// operator edge (mirrors the oracle's cat(...).to(bfloat16)).
+__global__ void p5_swiglu_shared_backward_kernel(
+    const __nv_bfloat16* __restrict__ dh,
+    const float* __restrict__ z,
+    __nv_bfloat16* __restrict__ dz,
+    const int64_t n,
+    const int64_t width) {
+  const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x;
+  for (int64_t idx = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x; idx < n;
+       idx += stride) {
+    const int64_t row = idx / width;
+    const int64_t col = idx - row * width;
+    const int64_t gate_index = row * (2 * width) + col;
+    const float g = z[gate_index];
+    const float u = z[gate_index + width];
+    const float dh32 = __bfloat162float(dh[idx]);
+    const float sig = sigmoid_rn(g);
+    const float silu = __fmul_rn(g, sig);
+    // dsilu = sig * (1 + g * (1 - sig)), each op rounded separately.
+    float t = __fsub_rn(1.0f, sig);
+    t = __fmul_rn(g, t);
+    t = __fadd_rn(1.0f, t);
+    const float dsilu = __fmul_rn(sig, t);
+    const float dgate = __fmul_rn(__fmul_rn(dh32, u), dsilu);
+    const float dup = __fmul_rn(dh32, silu);
+    dz[gate_index] = __float2bfloat16(dgate);
+    dz[gate_index + width] = __float2bfloat16(dup);
+  }
+}
+
+void launch_1d(int64_t n, int& threads, int64_t& blocks) {
+  threads = 256;
+  blocks = (n + threads - 1) / threads;
+  if (blocks == 0) {
+    blocks = 1;
+  }
+  if (blocks > 65535) {
+    blocks = 65535;  // grid-stride loops cover the rest
+  }
+}
+
+void check_cuda_2d(const torch::Tensor& t, at::ScalarType dtype, const char* name) {
+  TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(t.dim() == 2, name, " must be 2-D");
+  TORCH_CHECK(t.scalar_type() == dtype, name, " must be ", dtype, ", got ", t.scalar_type());
+}
+
+}  // namespace
+
+torch::Tensor p5_strict_gemm(torch::Tensor a, torch::Tensor b, bool trans_b) {
+  check_cuda_2d(a, at::kBFloat16, "a");
+  check_cuda_2d(b, at::kBFloat16, "b");
+  TORCH_CHECK(a.device() == b.device(), "a and b must be on the same CUDA device");
+  const int64_t m_rows = a.size(0);
+  const int64_t k_dim = a.size(1);
+  const int64_t n_cols = trans_b ? b.size(1) : b.size(0);
+  const int64_t bk = trans_b ? b.size(0) : b.size(1);
+  TORCH_CHECK(bk == k_dim, "K mismatch: a has K=", k_dim, ", b has K=", bk);
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
+  auto out = torch::empty({m_rows, n_cols}, a.options().dtype(at::kFloat));
+  const int64_t n = out.numel();
+  if (n == 0 || k_dim == 0) {
+    return n == 0 ? out : out.zero_();
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  const auto* a_ptr = reinterpret_cast<const __nv_bfloat16*>(a.data_ptr());
+  const auto* b_ptr = reinterpret_cast<const __nv_bfloat16*>(b.data_ptr());
+  if (trans_b) {
+    p5_strict_gemm_kernel<true><<<blocks, threads, 0, stream>>>(
+        a_ptr, b_ptr, out.data_ptr<float>(), m_rows, n_cols, k_dim);
+  } else {
+    p5_strict_gemm_kernel<false><<<blocks, threads, 0, stream>>>(
+        a_ptr, b_ptr, out.data_ptr<float>(), m_rows, n_cols, k_dim);
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return out;
+}
+
+torch::Tensor p5_swiglu_shared_forward(torch::Tensor z) {
+  check_cuda_2d(z, at::kFloat, "z");
+  TORCH_CHECK(z.size(1) % 2 == 0, "z width must be even (packed gate|up)");
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(z));
+  const int64_t width = z.size(1) / 2;
+  auto h = torch::empty({z.size(0), width}, z.options().dtype(at::kBFloat16));
+  const int64_t n = h.numel();
+  if (n == 0) {
+    return h;
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  p5_swiglu_shared_forward_kernel<<<blocks, threads, 0, stream>>>(
+      z.data_ptr<float>(), reinterpret_cast<__nv_bfloat16*>(h.data_ptr()), n, width);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return h;
+}
+
+torch::Tensor p5_swiglu_shared_backward(torch::Tensor dh, torch::Tensor z) {
+  check_cuda_2d(dh, at::kBFloat16, "dh");
+  check_cuda_2d(z, at::kFloat, "z");
+  TORCH_CHECK(dh.device() == z.device(), "dh and z must be on the same CUDA device");
+  TORCH_CHECK(z.size(1) % 2 == 0, "z width must be even (packed gate|up)");
+  TORCH_CHECK(
+      dh.size(0) == z.size(0) && dh.size(1) * 2 == z.size(1),
+      "dh shape must match the packed gate/up halves of z");
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(z));
+  auto dz = torch::empty_like(z, z.options().dtype(at::kBFloat16));
+  const int64_t n = dh.numel();
+  if (n == 0) {
+    return dz;
+  }
+  int threads = 0;
+  int64_t blocks = 0;
+  launch_1d(n, threads, blocks);
+  auto stream = at::cuda::getCurrentCUDAStream();
+  p5_swiglu_shared_backward_kernel<<<blocks, threads, 0, stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(dh.data_ptr()),
+      z.data_ptr<float>(),
+      reinterpret_cast<__nv_bfloat16*>(dz.data_ptr()),
+      n,
+      dh.size(1));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return dz;
+}

csrc/ops.cpp

@@ -131,6 +131,11 @@ std::vector<torch::Tensor> swiglu_packed_backward_cuda(
     torch::Tensor dy,
     torch::Tensor gate_up);
 
+// P5-5 (#64) Shared Expert MLP strict kernels (oracle-fp32-serial-v1)
+torch::Tensor p5_strict_gemm(torch::Tensor a, torch::Tensor b, bool trans_b);
+torch::Tensor p5_swiglu_shared_forward(torch::Tensor z);
+torch::Tensor p5_swiglu_shared_backward(torch::Tensor dh, torch::Tensor z);
+
 // RMSNorm Declarations & Wrappers
 
 void rmsnorm_forward_cuda(
@@ -503,6 +508,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("swiglu_packed_backward", &swiglu_packed_backward,
           "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]");
 
+    // P5-5 (#64) Shared Expert MLP strict kernels (oracle-fp32-serial-v1)
+    m.def("p5_strict_gemm", &p5_strict_gemm,
+          "Strict BF16-in/FP32-out GEMM, serial ascending-k, mul-then-add");
+    m.def("p5_swiglu_shared_forward", &p5_swiglu_shared_forward,
+          "One-round SwiGLU forward, shared-expert mode (p_s = None)");
+    m.def("p5_swiglu_shared_backward", &p5_swiglu_shared_backward,
+          "One-round SwiGLU backward, shared-expert mode (p_s = None)");
+
     // Deterministic standard-softmax attention (issue #147)
     m.def(
         "deterministic_attention_forward",