diff --git a/benchmarks/benchmark_shared_expert_mlp.py b/benchmarks/benchmark_shared_expert_mlp.py new file mode 100644 index 00000000..6ed904de --- /dev/null +++ b/benchmarks/benchmark_shared_expert_mlp.py @@ -0,0 +1,136 @@ +#!/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, is_strict in ( + ("triton", "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider", True), + ("cuda", "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider", True), + ("triton-det", "rl_engine.moe.backends.shared_expert:TritonDetSharedExpertProvider", False), + ("cuda-det", "rl_engine.moe.backends.shared_expert:CudaDetSharedExpertProvider", False), + ): + try: + runner = make_runner(resolve_provider(spec)) + runners[label] = runner + if is_strict: + 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()) diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 1b535d2a..a907534e 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -546,6 +546,26 @@ torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b) { return gemm_dispatch(a, b); } +// FP32-output variants: same kernels and reduction order, but the FP32 +// accumulator is stored without the final BF16 round. Used by operators whose +// contract keeps an intermediate in FP32 (e.g. P5-5 fc1 output, dX). +torch::Tensor det_gemm_fwd_out_fp32(torch::Tensor a, torch::Tensor b) { + check_in(a, "A"); check_in(b, "B"); + a = a.contiguous(); b = b.contiguous(); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "det_gemm_fwd_out_fp32: expect 2D [M,K]@[K,N]"); + TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd_out_fp32: K mismatch"); + return gemm_dispatch(a, b, RhsLayout::kKN, OutputLayout::kMN, /*output_fp32=*/true); +} + +torch::Tensor det_gemm_fwd_rhs_transposed_out_fp32(torch::Tensor a, torch::Tensor bt) { + check_in(a, "A"); check_in(bt, "Bt"); + a = a.contiguous(); bt = bt.contiguous(); + TORCH_CHECK(a.dim() == 2 && bt.dim() == 2, + "det_gemm_fwd_rhs_transposed_out_fp32: expect A[M,K] and Bt[N,K]"); + TORCH_CHECK(bt.size(1) == a.size(1), "det_gemm_fwd_rhs_transposed_out_fp32: K mismatch"); + return gemm_dispatch(a, bt, RhsLayout::kNK, OutputLayout::kMN, /*output_fp32=*/true); +} + torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b) { check_in(dc, "dC"); check_in(b, "B"); dc = dc.contiguous(); b = b.contiguous(); diff --git a/csrc/cuda/moe/shared_expert_mlp.cu b/csrc/cuda/moe/shared_expert_mlp.cu new file mode 100644 index 00000000..a1cd7b1b --- /dev/null +++ b/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 +#include +#include +#include +#include + +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 +__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(blockDim.x) * gridDim.x; + for (int64_t idx = blockIdx.x * static_cast(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(blockDim.x) * gridDim.x; + for (int64_t idx = blockIdx.x * static_cast(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(blockDim.x) * gridDim.x; + for (int64_t idx = blockIdx.x * static_cast(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(a.data_ptr()); + const auto* b_ptr = reinterpret_cast(b.data_ptr()); + if (trans_b) { + p5_strict_gemm_kernel<<>>( + a_ptr, b_ptr, out.data_ptr(), m_rows, n_cols, k_dim); + } else { + p5_strict_gemm_kernel<<>>( + a_ptr, b_ptr, out.data_ptr(), 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<<>>( + z.data_ptr(), 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<<>>( + reinterpret_cast(dh.data_ptr()), + z.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(dz.data_ptr()), + n, + dh.size(1)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dz; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..fb8f8872 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -115,6 +115,8 @@ void deterministic_collective_all_gather_fused( bool det_gemm_sm90_compiled(); torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_fwd_rhs_transposed(torch::Tensor a, torch::Tensor bt); +torch::Tensor det_gemm_fwd_out_fp32(torch::Tensor a, torch::Tensor b); +torch::Tensor det_gemm_fwd_rhs_transposed_out_fp32(torch::Tensor a, torch::Tensor bt); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc); torch::Tensor det_gemm_db_transposed(torch::Tensor a, torch::Tensor dc); @@ -131,6 +133,11 @@ std::vector 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( @@ -476,6 +483,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "det_gemm_fwd_rhs_transposed", &det_gemm_fwd_rhs_transposed, "Batch-invariant deterministic GEMM with physical Bt[N,K] (C=A@Bt^T)"); + m.def( + "det_gemm_fwd_out_fp32", + &det_gemm_fwd_out_fp32, + "det_gemm_fwd storing the FP32 accumulator (no final BF16 round)"); + m.def( + "det_gemm_fwd_rhs_transposed_out_fp32", + &det_gemm_fwd_rhs_transposed_out_fp32, + "det_gemm_fwd_rhs_transposed storing the FP32 accumulator"); m.def("det_gemm_da", &det_gemm_da, "Batch-invariant deterministic GEMM backward dA (dC@B^T)"); m.def("det_gemm_db", &det_gemm_db, "Batch-invariant deterministic GEMM backward dB (A^T@dC)"); m.def( @@ -503,6 +518,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", diff --git a/docs/operators/shared-expert-mlp.md b/docs/operators/shared-expert-mlp.md new file mode 100644 index 00000000..e06086e5 --- /dev/null +++ b/docs/operators/shared-expert-mlp.md @@ -0,0 +1,50 @@ +# Shared Expert MLP (P5-5, issue #64) + +Shared expert for the DSv4 MoE block: every valid token runs +`fc1 -> one-round SwiGLU -> fc2` once. BF16 frozen weights, backward returns +only `dX` (FP32 accumulator dtype); the shared output stays independent of the +routed path (the combine belongs to P6). + +## Fixed math (`oracle-fp32-serial-v1`) + +``` +z = x @ w_fc1.T # BF16 operands, FP32 serial ascending-k, mul-then-add +h = BF16(SiLU(gate) * up) # FP32 math, single round; no clamp, no p_s +y = BF16(h @ w_fc2.T) # FP32 accumulate, one round +dX = FP32(dz @ w_fc1) # dh, dz round BF16 at operator edges +``` + +Strict kernels reproduce the oracle byte-for-byte on the same device: one lane +owns one output element and reduces serially in ascending k with +`__fmul_rn`/`__fadd_rn` (CUDA) or uncontracted IEEE fp32 arith (Triton), so +there is no cross-lane floating-point reduction and results are +batch/padding invariant. + +Sigmoid is transcendental, so its bits follow the libm implementation: nvcc +`expf` (used by both `torch.sigmoid` and the CUDA kernel) bit-matches, while +`tl.exp` and libdevice `__nv_expf` do not (~10-45% of values differ by 1 ulp). +The Triton path therefore sources the sigmoid tensor from `torch.sigmoid` and +fuses the remaining SwiGLU math in the kernel. + +## Backends + +| backend | entry point | +| --- | --- | +| CUDA | `rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider` (`csrc/cuda/moe/shared_expert_mlp.cu`) | +| Triton | `rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider` (`rl_engine/kernels/ops/triton/moe/shared_expert.py`) | + +The one-round SwiGLU core runs in shared mode (`p_s = None`, no clamp) and is +the reuse point for P5-2 (#63), which extends the same core with clamp, +route weight, and `dp_s`. + +## Acceptance + +```bash +python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider --device cuda +python scripts/check_p5.py --provider rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider --device cuda +pytest tests/test_shared_expert_mlp.py +python benchmarks/benchmark_shared_expert_mlp.py +``` + +Fail-closed: non-CUDA input, a missing extension/triton install, or a foreign +numeric profile raises instead of falling back to the oracle. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..fff3f125 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -194,6 +194,11 @@ def det_gemm_fwd_rhs_transposed( a: torch.Tensor, bt: torch.Tensor, ) -> torch.Tensor: ... +def det_gemm_fwd_out_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ... +def det_gemm_fwd_rhs_transposed_out_fp32( + a: torch.Tensor, + bt: torch.Tensor, +) -> torch.Tensor: ... def det_gemm_da(dc: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ... def det_gemm_db(a: torch.Tensor, dc: torch.Tensor) -> torch.Tensor: ... def det_gemm_db_transposed(a: torch.Tensor, dc: torch.Tensor) -> torch.Tensor: ... @@ -205,6 +210,9 @@ def swiglu_backward( gate: torch.Tensor, up: torch.Tensor, ) -> list[torch.Tensor]: ... +def p5_strict_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: ... +def p5_swiglu_shared_forward(z: torch.Tensor) -> torch.Tensor: ... +def p5_swiglu_shared_backward(dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: ... def rmsnorm_forward( x: torch.Tensor, weight: torch.Tensor, diff --git a/rl_engine/kernels/ops/triton/moe/__init__.py b/rl_engine/kernels/ops/triton/moe/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/rl_engine/kernels/ops/triton/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py new file mode 100644 index 00000000..211f2022 --- /dev/null +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) Shared Expert MLP strict Triton kernels. + +Numeric profile ``oracle-fp32-serial-v1`` (see ``rl_engine/moe/oracle.py``): +FP32 accumulation, serial ascending-k, multiply and add rounded separately. +Each output element is owned by one lane and reduced serially, so there is no +cross-lane floating-point reduction and results are batch/padding invariant. + +The one-round SwiGLU core (FP32 math, single BF16 round on the output) is the +shared-mode (``p_s = None``, no clamp) variant shared with P5-2 (#63). + +Sigmoid is transcendental and its bits depend on the libm implementation: +neither ``tl.exp`` (exp2-based) nor libdevice ``__nv_expf`` bit-matches the +nvcc ``expf`` inside ``torch.sigmoid``. The Triton path therefore sources the +sigmoid tensor from ``torch.sigmoid`` (same-device oracle parity, per the P5 +transcendental rule) and fuses all remaining SwiGLU math in the kernel. + +All mul/add/sub go through libdevice ``*_rn`` (the Triton spelling of CUDA's +``__fmul_rn``/``__fadd_rn``/``__fsub_rn``): the compiler is allowed to +contract a plain ``a * b + c`` into an FMA, which changes the rounding (seen +as 1-ulp dgate drift at T=256), and the ``_rn`` intrinsics forbid that. +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + import triton.language.extra.libdevice as tld + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover - exercised on non-GPU installs + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _strict_gemm_kernel( + A, + B, + C, + K, + N, + stride_bn, + stride_bk, + BLOCK_N: tl.constexpr, + ): + # C[m, n] = sum_{k ascending} A[m, k] * B(n, k); FP32 accumulator, + # one lane per output element, mul and add rounded separately + # (fp32 arith in Triton is IEEE by default: no FMA contraction). + m = tl.program_id(0) + pn = tl.program_id(1) + offs_n = pn * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_n < N + acc = tl.zeros([BLOCK_N], dtype=tl.float32) + a_row = A + m * K + b_cols = B + offs_n * stride_bn + for k in range(0, K): + a = tl.load(a_row + k).to(tl.float32) + b = tl.load(b_cols + k * stride_bk, mask=mask_n, other=0.0).to(tl.float32) + acc = tld.add_rn(acc, tld.mul_rn(a, b)) + tl.store(C + m * N + offs_n, acc, mask=mask_n) + + @triton.jit + def _det_dot_gemm_kernel( + A, + B, + C, + M, + N, + K, + stride_bn, + stride_bk, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + # Performance GEMM (profile p5-triton-dot-v1): tensor-core tl.dot with + # FP32 accumulators, fixed constexpr tiles, ascending-k tile order and + # NO split-K -> deterministic and batch-invariant (a row's product uses + # only its own A row plus B; masked padding rows are zero). The + # reduction order inside a tile differs from the strict serial path, + # so this is NOT byte-equal to oracle-fp32-serial-v1. + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + for k0 in range(0, K, BLOCK_K): + offs_k = k0 + tl.arange(0, BLOCK_K) + a = tl.load( + A + offs_m[:, None] * K + offs_k[None, :], + mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), + other=0.0, + ) + b = tl.load( + B + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn, + mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), + other=0.0, + ) + acc = tl.dot(a, b, acc) + c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(C + offs_m[:, None] * N + offs_n[None, :], acc, mask=c_mask) + + @triton.jit + def _swiglu_shared_fwd_kernel( + Z, + SIG, + H, + n_elem, + width, + BLOCK: tl.constexpr, + ): + # One-round SwiGLU, shared mode: h = BF16(SiLU(gate) * up), FP32 math, + # gate = z[:, :F], up = z[:, F:] packed in one [T, 2F] FP32 tensor. + # SIG is torch.sigmoid(gate) (see module docstring). + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elem + row = offs // width + col = offs - row * width + gate_index = row * (2 * width) + col + g = tl.load(Z + gate_index, mask=mask, other=0.0) + u = tl.load(Z + gate_index + width, mask=mask, other=0.0) + sig = tl.load(SIG + offs, mask=mask, other=0.0) + silu = tld.mul_rn(g, sig) + h = tld.mul_rn(silu, u).to(tl.bfloat16) + tl.store(H + offs, h, mask=mask) + + @triton.jit + def _swiglu_shared_bwd_kernel( + DH, + Z, + SIG, + DZ, + n_elem, + width, + BLOCK: tl.constexpr, + ): + # dgate = ((dh * u) * dsilu); dup = dh * silu; both round to BF16 once + # at the operator edge (mirrors the oracle's cat(...).to(bfloat16)). + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elem + row = offs // width + col = offs - row * width + gate_index = row * (2 * width) + col + g = tl.load(Z + gate_index, mask=mask, other=0.0) + u = tl.load(Z + gate_index + width, mask=mask, other=0.0) + dh = tl.load(DH + offs, mask=mask, other=0.0).to(tl.float32) + sig = tl.load(SIG + offs, mask=mask, other=0.0) + silu = tld.mul_rn(g, sig) + # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. + t = tld.sub_rn(1.0, sig) + t = tld.mul_rn(g, t) + t = tld.add_rn(1.0, t) + dsilu = tld.mul_rn(sig, t) + dgate = tld.mul_rn(tld.mul_rn(dh, u), dsilu) + dup = tld.mul_rn(dh, silu) + tl.store(DZ + gate_index, dgate.to(tl.bfloat16), mask=mask) + tl.store(DZ + gate_index + width, dup.to(tl.bfloat16), mask=mask) + + +def _check_cuda_2d(t: torch.Tensor, dtype: torch.dtype, name: str) -> None: + if not t.is_cuda: + raise NotImplementedError(f"{name} must be a CUDA tensor for the Triton backend") + if not t.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if t.dim() != 2: + raise ValueError(f"{name} must be 2-D") + if t.dtype != dtype: + raise TypeError(f"{name} must be {dtype}, got {t.dtype}") + + +def strict_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + """``a @ b.T`` (or ``a @ b`` when ``trans_b``): BF16 in, FP32 out, strict.""" + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(a, torch.bfloat16, "a") + _check_cuda_2d(b, torch.bfloat16, "b") + m, k = a.shape + if trans_b: + bk, n = b.shape + stride_bn, stride_bk = 1, n + else: + n, bk = b.shape + stride_bn, stride_bk = k, 1 + if bk != k: + raise ValueError(f"K mismatch: a has K={k}, b has K={bk}") + out = torch.empty(m, n, dtype=torch.float32, device=a.device) + if out.numel() == 0: + return out + if k == 0: + return out.zero_() + block_n = min(triton.next_power_of_2(n), 256) + grid = (m, triton.cdiv(n, block_n)) + _strict_gemm_kernel[grid](a, b, out, k, n, stride_bn, stride_bk, BLOCK_N=block_n) + return out + + +def det_dot_gemm(a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + """Performance GEMM (p5-triton-dot-v1): tl.dot tiles, FP32 out, no split-K. + + Same signature and round positions as :func:`strict_gemm`; only the + reduction order inside a tile differs (tensor-core MMA vs serial). + """ + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(a, torch.bfloat16, "a") + _check_cuda_2d(b, torch.bfloat16, "b") + m, k = a.shape + if trans_b: + bk, n = b.shape + stride_bn, stride_bk = 1, n + else: + n, bk = b.shape + stride_bn, stride_bk = k, 1 + if bk != k: + raise ValueError(f"K mismatch: a has K={k}, b has K={bk}") + out = torch.empty(m, n, dtype=torch.float32, device=a.device) + if out.numel() == 0: + return out + if k == 0: + return out.zero_() + # Fixed tiles (no autotune): tuning by shape would change the reduction + # tree with batch size and break invariance. + block_m, block_n, block_k = 64, 64, 32 + grid = (triton.cdiv(m, block_m), triton.cdiv(n, block_n)) + _det_dot_gemm_kernel[grid]( + a, + b, + out, + m, + n, + k, + stride_bn, + stride_bk, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + num_warps=4, + num_stages=3, + ) + return out + + +def swiglu_shared_fwd(z: torch.Tensor) -> torch.Tensor: + """One-round SwiGLU forward, shared mode: FP32 [T, 2F] -> BF16 [T, F].""" + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(z, torch.float32, "z") + if z.shape[1] % 2 != 0: + raise ValueError("z width must be even (packed gate|up)") + width = z.shape[1] // 2 + h = torch.empty(z.shape[0], width, dtype=torch.bfloat16, device=z.device) + n_elem = h.numel() + if n_elem == 0: + return h + sig = torch.sigmoid(z[:, :width]).contiguous() + block = 1024 + grid = (triton.cdiv(n_elem, block),) + _swiglu_shared_fwd_kernel[grid](z, sig, h, n_elem, width, BLOCK=block) + return h + + +def swiglu_shared_bwd(dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + """One-round SwiGLU backward, shared mode: returns packed BF16 dz [T, 2F].""" + if not TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + _check_cuda_2d(dh, torch.bfloat16, "dh") + _check_cuda_2d(z, torch.float32, "z") + if z.shape[1] % 2 != 0: + raise ValueError("z width must be even (packed gate|up)") + if dh.shape[0] != z.shape[0] or dh.shape[1] * 2 != z.shape[1]: + raise ValueError("dh shape must match the packed gate/up halves of z") + dz = torch.empty_like(z, dtype=torch.bfloat16) + n_elem = dh.numel() + if n_elem == 0: + return dz + width = dh.shape[1] + sig = torch.sigmoid(z[:, :width]).contiguous() + block = 1024 + grid = (triton.cdiv(n_elem, block),) + _swiglu_shared_bwd_kernel[grid](dh, z, sig, dz, n_elem, width, BLOCK=block) + return dz diff --git a/rl_engine/moe/backends/__init__.py b/rl_engine/moe/backends/__init__.py new file mode 100644 index 00000000..6eec9430 --- /dev/null +++ b/rl_engine/moe/backends/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5 kernel backends. Each sub-issue registers its providers here.""" diff --git a/rl_engine/moe/backends/shared_expert.py b/rl_engine/moe/backends/shared_expert.py new file mode 100644 index 00000000..ba156c28 --- /dev/null +++ b/rl_engine/moe/backends/shared_expert.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) Shared Expert MLP providers (CUDA and Triton strict backends). + +Both backends implement the frozen math ``fc1 -> one-round SwiGLU -> fc2`` +under the ``oracle-fp32-serial-v1`` numeric profile and are byte-equal to the +FP32 oracle running on the same device. Only the two shared-expert methods are +overridden; every other operator stays on the oracle per the S0 start kit, so +the full acceptance command runs unchanged: + + python scripts/check_p5.py \ + --provider rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider \ + --device cuda + +Fail-closed: unsupported input (non-CUDA device, missing extension/triton, +schema violations) raises instead of falling back to another implementation. +The shared output is produced from ``SharedBatch`` alone -- no route weight, +no routed combine (that boundary belongs to P6). +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.moe.contract import ORACLE_PROFILE, SharedBatch +from rl_engine.moe.provider import ReferenceProvider + + +class _StrictSharedExpertProvider(ReferenceProvider): + """Common composite: strict GEMMs + one-round SwiGLU, dX only (frozen base).""" + + name = "shared-expert-strict" + numeric_profile = ORACLE_PROFILE + # Performance profiles keep the contract's round positions but change the + # in-GEMM reduction order, so they are not byte-equal to the oracle. + # Selecting such a provider by name is the explicit opt-in; it is never a + # silent substitute for the strict path (fail-closed rule, P5-6). + strict_profile = True + + # Backend hooks ------------------------------------------------------- + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + raise NotImplementedError + + def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + # Provider surface ---------------------------------------------------- + def capabilities(self) -> dict[str, Any]: + return { + "backend": self.name, + "operators": ["shared_expert_mlp_fwd", "shared_expert_mlp_bwd"], + "geometry": ["one-row", "packed"], + "devices": ["cuda"], + } + + def provenance(self) -> dict[str, Any]: + return { + "requested_backend": self.name, + "actual_backend": self.name, + "numeric_profile": self.numeric_profile, + "torch_version": torch.__version__, + # Changing any of these changes the addition order (P5-5 s4). + "split_k": 1, + "reduction": "serial-ascending-k", + "rounding": "mul-then-add, no FMA", + "workspace": "none", + } + + def _check_batch(self, batch: SharedBatch) -> None: + batch.validate() + if self.strict_profile and batch.numeric_profile != ORACLE_PROFILE: + raise NotImplementedError( + f"{self.name} only implements {ORACLE_PROFILE!r}, " + f"got {batch.numeric_profile!r} (fail-closed, no fallback)" + ) + if not batch.x.is_cuda: + raise NotImplementedError( + f"{self.name} requires CUDA tensors, got device {batch.x.device} " + "(fail-closed, no fallback)" + ) + + def shared_expert_mlp_fwd(self, batch: SharedBatch) -> tuple[torch.Tensor, dict[str, Any]]: + self._check_batch(batch) + x = batch.x.contiguous() + w_fc1 = batch.w_fc1.contiguous() + w_fc2 = batch.w_fc2.contiguous() + z = self._gemm(x, w_fc1, False) # [T, 2F] FP32, kept for backward + h_bf16 = self._swiglu_fwd(z) # [T, F] BF16, the one round + y = self._gemm(h_bf16, w_fc2, False).to(torch.bfloat16) # [T, H] + saved: dict[str, Any] = {"z32": z, "h_bf16": h_bf16} + return y, saved + + def shared_expert_mlp_bwd( + self, dy: torch.Tensor, batch: SharedBatch, saved: dict[str, Any] + ) -> torch.Tensor: + self._check_batch(batch) + z = saved["z32"] + dy_bf16 = dy.to(torch.bfloat16).contiguous() + # dh = BF16(dY @ W2), dz = swiglu_bwd, dX = dz @ W1 (FP32 accumulator). + dh = self._gemm(dy_bf16, batch.w_fc2.contiguous(), True).to(torch.bfloat16) + dz = self._swiglu_bwd(dh, z) + dx = self._gemm(dz, batch.w_fc1.contiguous(), True) + return dx + + +class CudaSharedExpertProvider(_StrictSharedExpertProvider): + """CUDA backend: csrc/cuda/moe/shared_expert_mlp.cu via rl_engine._C.""" + + name = "shared-expert-cuda" + + def __init__(self) -> None: + try: + from rl_engine import _C + except ImportError as exc: # fail-closed: no oracle fallback + raise NotImplementedError( + "rl_engine._C is not built; install with RL_KERNEL_REQUIRE_EXT=1" + ) from exc + for symbol in ("p5_strict_gemm", "p5_swiglu_shared_forward", "p5_swiglu_shared_backward"): + if not hasattr(_C, symbol): + raise NotImplementedError(f"rl_engine._C lacks {symbol}; rebuild the extension") + self._ext = _C + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + return self._ext.p5_strict_gemm(a, b, trans_b) + + def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: + return self._ext.p5_swiglu_shared_forward(z) + + def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return self._ext.p5_swiglu_shared_backward(dh, z) + + +class TritonSharedExpertProvider(_StrictSharedExpertProvider): + """Triton backend: rl_engine/kernels/ops/triton/moe/shared_expert.py.""" + + name = "shared-expert-triton" + + def __init__(self) -> None: + from rl_engine.kernels.ops.triton.moe import shared_expert as tk + + if not tk.TRITON_AVAILABLE: + raise NotImplementedError("triton is not installed (fail-closed, no fallback)") + self._tk = tk + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + return self._tk.strict_gemm(a, b, trans_b) + + def _swiglu_fwd(self, z: torch.Tensor) -> torch.Tensor: + return self._tk.swiglu_shared_fwd(z) + + def _swiglu_bwd(self, dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: + return self._tk.swiglu_shared_bwd(dh, z) + + +class CudaDetSharedExpertProvider(CudaSharedExpertProvider): + """Performance CUDA backend: det_gemm (csrc/cuda/gemm/det_gemm_kernel.cu). + + Deterministic and batch-invariant (fixed K order, no split-K; TMA+mma.sync + on SM90+, scalar K-tree fallback elsewhere). Round positions match the + P5-5 contract (fc1 out and dX stay FP32; y/dh round once to BF16), but the + in-GEMM reduction order differs from the oracle, so outputs are close, not + byte-equal. SwiGLU stays on the strict CUDA core. + """ + + name = "shared-expert-cuda-det" + numeric_profile = "p5-det-gemm-v1" + strict_profile = False + # det_gemm rounds each BK=32 partial to BF16 and merges the K dimension + # with a BF16 mid-split tree (its TP-equivalence design), so its deviation + # from the FP32-serial oracle is BF16-tree-sized, not FP32-sized. + oracle_tolerance = {"rtol": 1e-1, "atol": 6e-2} + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + if trans_b: # b is the logical [K, N] operand + return self._ext.det_gemm_fwd_out_fp32(a, b) + return self._ext.det_gemm_fwd_rhs_transposed_out_fp32(a, b) + + def provenance(self) -> dict[str, Any]: + info = super().provenance() + info.update( + { + "split_k": 1, + "reduction": "fixed-k-tile-tree (det_gemm)", + "rounding": "FP32 accumulate, FMA/MMA inside tiles", + "sm90_tensor_core": bool(getattr(self._ext, "det_gemm_sm90_compiled")()), + } + ) + return info + + +class TritonDetSharedExpertProvider(TritonSharedExpertProvider): + """Performance Triton backend: tl.dot tiles with fixed geometry. + + Same guarantees and caveats as :class:`CudaDetSharedExpertProvider`, with + profile ``p5-triton-dot-v1`` (tile reduction order differs per backend). + """ + + name = "shared-expert-triton-det" + numeric_profile = "p5-triton-dot-v1" + strict_profile = False + # Full-FP32 accumulators (only the contract's BF16 rounds), so deviation + # from the oracle is reduction-order noise only. + oracle_tolerance = {"rtol": 2e-2, "atol": 2e-2} + + def _gemm(self, a: torch.Tensor, b: torch.Tensor, trans_b: bool) -> torch.Tensor: + return self._tk.det_dot_gemm(a, b, trans_b) + + def provenance(self) -> dict[str, Any]: + info = super().provenance() + info.update( + { + "split_k": 1, + "reduction": "tl.dot 64x64x32 tiles, ascending-k", + "rounding": "FP32 accumulate, MMA inside tiles", + } + ) + return info diff --git a/setup.py b/setup.py index 79f882d9..6d8bbb4b 100644 --- a/setup.py +++ b/setup.py @@ -1,313 +1,314 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import warnings +from pathlib import Path + +from setuptools import find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/moe/shared_expert_mlp.cu", + "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + return {"build_ext": BuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/test_shared_expert_mlp.py b/tests/test_shared_expert_mlp.py new file mode 100644 index 00000000..549f5ac9 --- /dev/null +++ b/tests/test_shared_expert_mlp.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""P5-5 (#64) shared_expert_mlp: bit-wise alignment against the FP32 oracle. + +Every comparison is byte-equality (sha256 over raw little-endian bytes) with +the oracle executed on the same device, per the P5 start-kit acceptance rules. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.moe import fixtures, oracle +from rl_engine.moe.contract import SharedBatch, tensor_sha256 + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +PROVIDER_SPECS = { + "cuda": "rl_engine.moe.backends.shared_expert:CudaSharedExpertProvider", + "triton": "rl_engine.moe.backends.shared_expert:TritonSharedExpertProvider", +} + +# Performance profiles: deterministic and batch-invariant, close to (but not +# byte-equal with) the oracle -- the in-GEMM reduction order differs. +PERF_PROVIDER_SPECS = { + "cuda-det": "rl_engine.moe.backends.shared_expert:CudaDetSharedExpertProvider", + "triton-det": "rl_engine.moe.backends.shared_expert:TritonDetSharedExpertProvider", +} + + +def _resolve_or_skip(label: str, spec: str): + from rl_engine.moe.provider import resolve_provider + + try: + return resolve_provider(spec) + except NotImplementedError as exc: + pytest.skip(f"{label} backend unavailable: {exc}") + + +@pytest.fixture(params=sorted(PROVIDER_SPECS)) +def provider(request): + return _resolve_or_skip(request.param, PROVIDER_SPECS[request.param]) + + +@pytest.fixture(params=sorted(PERF_PROVIDER_SPECS)) +def perf_provider(request): + return _resolve_or_skip(request.param, PERF_PROVIDER_SPECS[request.param]) + + +def _random_batch(t: int, hidden: int, ffn: int, seed: int = 2026): + gen = torch.Generator(device="cpu").manual_seed(seed) + batch = SharedBatch( + x=torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda(), + w_fc1=(torch.randn(2 * ffn, hidden, generator=gen) / hidden**0.5).to(torch.bfloat16).cuda(), + w_fc2=(torch.randn(hidden, ffn, generator=gen) / ffn**0.5).to(torch.bfloat16).cuda(), + ) + dy = torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda() + return batch, dy + + +def _run_oracle(batch: SharedBatch, dy: torch.Tensor): + y, saved = oracle.shared_expert_mlp_fwd(batch) + dx = oracle.shared_expert_mlp_bwd(dy, batch, saved) + return y, dx + + +def _run_provider(provider, 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 + + +@requires_cuda +@pytest.mark.parametrize("case", sorted(fixtures.SHARED_CASES)) +def test_shared_cases_byte_equal(provider, case): + batch = fixtures.make_shared_batch(case).to("cuda") + y_gold, saved_gold = oracle.shared_expert_mlp_fwd(batch) + dy = fixtures.make_grad_output(case, tuple(y_gold.shape)).to("cuda") + dx_gold = oracle.shared_expert_mlp_bwd(dy, batch, saved_gold) + y, dx = _run_provider(provider, batch, dy) + assert y.dtype == torch.bfloat16 and dx.dtype == torch.float32 + assert tensor_sha256(y) == tensor_sha256(y_gold) + assert tensor_sha256(dx) == tensor_sha256(dx_gold) + + +@requires_cuda +def test_batch_padding_invariance(provider): + """fwd(x)[t] must equal fwd(x[t:t+1]) byte-for-byte (Axis-A invariance).""" + batch = fixtures.make_shared_batch("shared_t16").to("cuda") + y_full, _ = provider.shared_expert_mlp_fwd(batch) + for t in range(batch.x.shape[0]): + row_batch = SharedBatch( + x=batch.x[t : t + 1].contiguous(), + w_fc1=batch.w_fc1, + w_fc2=batch.w_fc2, + ) + y_row, _ = provider.shared_expert_mlp_fwd(row_batch) + assert tensor_sha256(y_row) == tensor_sha256(y_full[t : t + 1]), f"row {t} diverged" + + +@requires_cuda +def test_frozen_weights_no_dw(provider): + """Backward returns only dX; the shared base weights stay frozen.""" + batch = fixtures.make_shared_batch("shared_t16").to("cuda") + w1_before = tensor_sha256(batch.w_fc1) + w2_before = tensor_sha256(batch.w_fc2) + y, saved = provider.shared_expert_mlp_fwd(batch) + dy = fixtures.make_grad_output("shared_t16", tuple(y.shape)).to("cuda") + dx = provider.shared_expert_mlp_bwd(dy, batch, saved) + assert dx.shape == batch.x.shape and dx.dtype == torch.float32 + assert batch.w_fc1.grad is None and batch.w_fc2.grad is None + assert not batch.w_fc1.requires_grad and not batch.w_fc2.requires_grad + assert tensor_sha256(batch.w_fc1) == w1_before + assert tensor_sha256(batch.w_fc2) == w2_before + + +@requires_cuda +def test_shared_output_independent_of_routed(provider): + """Shared output is not premixed with the routed path (boundary fixture). + + Running the full routed pipeline (any p_s, any expert batch) between two + shared calls must not change a single byte of the shared output, and the + shared output must equal the standalone oracle result (no p_s applied). + """ + shared = fixtures.make_shared_batch("shared_t16").to("cuda") + y_gold, _ = oracle.shared_expert_mlp_fwd(shared) + y_before, _ = provider.shared_expert_mlp_fwd(shared) + + routed = fixtures.make_expert_batch("base_plus_lora").to("cuda") + y_routed, saved_routed = oracle.routed_expert_forward(routed, ops=provider) + dy_routed = fixtures.make_grad_output("base_plus_lora", tuple(y_routed.shape)).to("cuda") + oracle.routed_expert_backward(routed, saved_routed, dy_routed, ops=provider) + + y_after, _ = provider.shared_expert_mlp_fwd(shared) + assert tensor_sha256(y_before) == tensor_sha256(y_gold) + assert tensor_sha256(y_after) == tensor_sha256(y_gold) + assert y_after.data_ptr() != shared.x.data_ptr() + + +@requires_cuda +@pytest.mark.parametrize("shape", [(16, 128, 64), (256, 1024, 512)]) +def test_cuda_triton_byte_equal(shape): + """The two backends agree with each other bit-for-bit. + + The larger shape samples enough values to expose rare transcendental + 1-ulp divergences that survive the BF16 round (caught once at T=256). + """ + from rl_engine.moe.provider import resolve_provider + + providers = [] + for spec in PROVIDER_SPECS.values(): + try: + providers.append(resolve_provider(spec)) + except NotImplementedError as exc: + pytest.skip(f"backend unavailable: {exc}") + t, hidden, ffn = shape + gen = torch.Generator(device="cpu").manual_seed(hash(shape) % (2**31)) + batch = SharedBatch( + x=torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda(), + w_fc1=(torch.randn(2 * ffn, hidden, generator=gen) / hidden**0.5).to(torch.bfloat16).cuda(), + w_fc2=(torch.randn(hidden, ffn, generator=gen) / ffn**0.5).to(torch.bfloat16).cuda(), + ) + dy = torch.randn(t, hidden, generator=gen).to(torch.bfloat16).cuda() + results = [_run_provider(p, batch, dy) for p in providers] + (y_a, dx_a), (y_b, dx_b) = results + assert tensor_sha256(y_a) == tensor_sha256(y_b) + assert tensor_sha256(dx_a) == tensor_sha256(dx_b) + + +@requires_cuda +def test_fail_closed_on_cpu_input(provider): + batch = fixtures.make_shared_batch("shared_t1") # stays on CPU + with pytest.raises(NotImplementedError): + provider.shared_expert_mlp_fwd(batch) + + +@requires_cuda +def test_perf_backend_deterministic(perf_provider): + """Two runs of the performance profile are byte-identical.""" + batch, dy = _random_batch(256, 1024, 512) + runs = [_run_provider(perf_provider, batch, dy) for _ in range(2)] + (y_a, dx_a), (y_b, dx_b) = runs + assert tensor_sha256(y_a) == tensor_sha256(y_b) + assert tensor_sha256(dx_a) == tensor_sha256(dx_b) + + +@requires_cuda +def test_perf_backend_batch_invariance(perf_provider): + """fwd(x)[t] == fwd(x[t:t+1]) byte-for-byte also on the performance path.""" + batch, _ = _random_batch(16, 256, 128) + y_full, _ = perf_provider.shared_expert_mlp_fwd(batch) + for t in range(batch.x.shape[0]): + row_batch = SharedBatch( + x=batch.x[t : t + 1].contiguous(), + w_fc1=batch.w_fc1, + w_fc2=batch.w_fc2, + ) + y_row, _ = perf_provider.shared_expert_mlp_fwd(row_batch) + assert tensor_sha256(y_row) == tensor_sha256(y_full[t : t + 1]), f"row {t} diverged" + + +@requires_cuda +def test_perf_backend_close_to_oracle(perf_provider): + """Same round positions, different reduction order: close, not byte-equal. + + Each provider declares its own tolerance: det_gemm merges the K dimension + with a BF16 mid-split tree (BF16-tree-sized deviation), while the tl.dot + path keeps FP32 accumulators (reduction-order noise only). + """ + tol = perf_provider.oracle_tolerance + batch, dy = _random_batch(64, 512, 256) + y_gold, saved_gold = oracle.shared_expert_mlp_fwd(batch) + dx_gold = oracle.shared_expert_mlp_bwd(dy, batch, saved_gold) + y, dx = _run_provider(perf_provider, batch, dy) + assert y.dtype == torch.bfloat16 and dx.dtype == torch.float32 + torch.testing.assert_close(y.float(), y_gold.float(), **tol) + torch.testing.assert_close(dx, dx_gold, **tol)