From 16d3bfc3f8f4dab35d112a6e8752bf7ccb286540 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 03:26:48 -0700 Subject: [PATCH 1/4] [DSv4][P5-5] Shared Expert MLP: strict CUDA + Triton kernels (#64) Implements shared_expert_mlp_fwd/bwd per the P5-S0 contract: every valid token runs fc1 -> one-round SwiGLU -> fc2 on BF16 frozen weights, backward returns dX only (FP32 accumulator), and the shared output stays independent of the routed path. Both backends reproduce the FP32 oracle's numeric profile oracle-fp32-serial-v1 byte-for-byte on the same device: one lane owns one output element and reduces serially in ascending k, multiply and add rounded separately (__fmul_rn/__fadd_rn on CUDA, uncontracted IEEE fp32 in Triton), sigmoid computed as 1/(1+expf(-x)) to match torch.sigmoid on FP32 CUDA tensors. No cross-lane floating-point reduction exists anywhere, so results are batch/padding invariant by construction (fwd(x)[t] == fwd(x[t:t+1]) byte-equal). The one-round SwiGLU core runs in shared mode (p_s=None, no clamp, per S0 decision D6) and is the reuse point for P5-2 (#63). Providers subclass ReferenceProvider and override only the two shared-expert methods, so the full acceptance command runs unchanged; unsupported input (non-CUDA device, missing extension or triton, foreign numeric profile) raises instead of falling back (fail-closed). Provenance records split_k=1 / serial-ascending-k / no-FMA per the P5-5 provenance requirement. Acceptance: 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 --- benchmarks/benchmark_shared_expert_mlp.py | 133 ++++ csrc/cuda/moe/shared_expert_mlp.cu | 213 ++++++ csrc/ops.cpp | 13 + docs/operators/shared-expert-mlp.md | 44 ++ rl_engine/_C.pyi | 3 + rl_engine/kernels/ops/triton/moe/__init__.py | 2 + .../kernels/ops/triton/moe/shared_expert.py | 185 ++++++ rl_engine/moe/backends/__init__.py | 3 + rl_engine/moe/backends/shared_expert.py | 152 +++++ setup.py | 625 +++++++++--------- tests/test_shared_expert_mlp.py | 137 ++++ 11 files changed, 1198 insertions(+), 312 deletions(-) create mode 100644 benchmarks/benchmark_shared_expert_mlp.py create mode 100644 csrc/cuda/moe/shared_expert_mlp.cu create mode 100644 docs/operators/shared-expert-mlp.md create mode 100644 rl_engine/kernels/ops/triton/moe/__init__.py create mode 100644 rl_engine/kernels/ops/triton/moe/shared_expert.py create mode 100644 rl_engine/moe/backends/__init__.py create mode 100644 rl_engine/moe/backends/shared_expert.py create mode 100644 tests/test_shared_expert_mlp.py diff --git a/benchmarks/benchmark_shared_expert_mlp.py b/benchmarks/benchmark_shared_expert_mlp.py new file mode 100644 index 00000000..8ef6fe5a --- /dev/null +++ b/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()) 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..33bb88bd 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -131,6 +131,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( @@ -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", diff --git a/docs/operators/shared-expert-mlp.md b/docs/operators/shared-expert-mlp.md new file mode 100644 index 00000000..c4ac1809 --- /dev/null +++ b/docs/operators/shared-expert-mlp.md @@ -0,0 +1,44 @@ +# 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. + +## 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..8f7879fb 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -205,6 +205,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..125d8a83 --- /dev/null +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -0,0 +1,185 @@ +# 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). +""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + 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) + prod = a * b + acc = acc + prod + tl.store(C + m * N + offs_n, acc, mask=mask_n) + + @triton.jit + def _swiglu_shared_fwd_kernel( + Z, + 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. + 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 = 1.0 / (1.0 + tl.exp(-g)) + silu = g * sig + h = (silu * u).to(tl.bfloat16) + tl.store(H + offs, h, mask=mask) + + @triton.jit + def _swiglu_shared_bwd_kernel( + DH, + Z, + 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 = 1.0 / (1.0 + tl.exp(-g)) + silu = g * sig + # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. + t = 1.0 - sig + t = g * t + t = 1.0 + t + dsilu = sig * t + dgate = (dh * u) * dsilu + dup = 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 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 + block = 1024 + grid = (triton.cdiv(n_elem, block),) + _swiglu_shared_fwd_kernel[grid](z, 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 + block = 1024 + grid = (triton.cdiv(n_elem, block),) + _swiglu_shared_bwd_kernel[grid](dh, z, dz, n_elem, dh.shape[1], 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..ef11df76 --- /dev/null +++ b/rl_engine/moe/backends/shared_expert.py @@ -0,0 +1,152 @@ +# 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 + + # 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 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) 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..6590ceea --- /dev/null +++ b/tests/test_shared_expert_mlp.py @@ -0,0 +1,137 @@ +# 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", +} + + +@pytest.fixture(params=sorted(PROVIDER_SPECS)) +def provider(request): + from rl_engine.moe.provider import resolve_provider + + try: + return resolve_provider(PROVIDER_SPECS[request.param]) + except NotImplementedError as exc: + pytest.skip(f"{request.param} backend unavailable: {exc}") + + +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 +def test_cuda_triton_byte_equal(): + """The two backends agree with each other bit-for-bit.""" + 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}") + batch = fixtures.make_shared_batch("shared_t16").to("cuda") + dy = fixtures.make_grad_output("shared_t16", (batch.x.shape[0], batch.x.shape[1])).to("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) From 74087d3c3cc4b7d6ab5f92acc0f39c097b1e2eda Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 04:02:02 -0700 Subject: [PATCH 2/4] fix(p5-5): use libdevice exp in the Triton SwiGLU core tl.exp is the fast exp2-based path and does not bit-match torch.sigmoid; libdevice __nv_expf does (0/4M mismatches on the device probe). --- rl_engine/kernels/ops/triton/moe/shared_expert.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index 125d8a83..95b415bc 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -18,6 +18,7 @@ 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 @@ -72,7 +73,8 @@ def _swiglu_shared_fwd_kernel( 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 = 1.0 / (1.0 + tl.exp(-g)) + # libdevice exp (__nv_expf) bit-matches torch.sigmoid; tl.exp does not. + sig = 1.0 / (1.0 + tld.exp(-g)) silu = g * sig h = (silu * u).to(tl.bfloat16) tl.store(H + offs, h, mask=mask) @@ -97,7 +99,7 @@ def _swiglu_shared_bwd_kernel( 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 = 1.0 / (1.0 + tl.exp(-g)) + sig = 1.0 / (1.0 + tld.exp(-g)) silu = g * sig # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. t = 1.0 - sig From c73959c4f260da8984f3d614e76d38265950dfcc Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 09:28:09 -0700 Subject: [PATCH 3/4] fix(p5-5): source Triton sigmoid from torch.sigmoid Neither tl.exp (exp2-based) nor libdevice __nv_expf bit-matches the nvcc expf inside torch.sigmoid (~45% / ~10% of fp32 values differ by 1 ulp); the tiny fixtures passed only because the BF16 round absorbed the difference, and the T=256 benchmark cross-check caught the divergence. The Triton path now takes torch.sigmoid(gate) as a kernel input and fuses the remaining SwiGLU math; a (256, 1024, 512) cross-backend byte-equality test locks the regression in. --- docs/operators/shared-expert-mlp.md | 6 +++++ .../kernels/ops/triton/moe/shared_expert.py | 22 ++++++++++++++----- tests/test_shared_expert_mlp.py | 19 ++++++++++++---- 3 files changed, 37 insertions(+), 10 deletions(-) diff --git a/docs/operators/shared-expert-mlp.md b/docs/operators/shared-expert-mlp.md index c4ac1809..e06086e5 100644 --- a/docs/operators/shared-expert-mlp.md +++ b/docs/operators/shared-expert-mlp.md @@ -20,6 +20,12 @@ owns one output element and reduces serially in ascending k with 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 | diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index 95b415bc..f510f239 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -9,6 +9,12 @@ 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. """ from __future__ import annotations @@ -18,7 +24,6 @@ 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 @@ -58,6 +63,7 @@ def _strict_gemm_kernel( @triton.jit def _swiglu_shared_fwd_kernel( Z, + SIG, H, n_elem, width, @@ -65,6 +71,7 @@ def _swiglu_shared_fwd_kernel( ): # 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 @@ -73,8 +80,7 @@ def _swiglu_shared_fwd_kernel( 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) - # libdevice exp (__nv_expf) bit-matches torch.sigmoid; tl.exp does not. - sig = 1.0 / (1.0 + tld.exp(-g)) + sig = tl.load(SIG + offs, mask=mask, other=0.0) silu = g * sig h = (silu * u).to(tl.bfloat16) tl.store(H + offs, h, mask=mask) @@ -83,6 +89,7 @@ def _swiglu_shared_fwd_kernel( def _swiglu_shared_bwd_kernel( DH, Z, + SIG, DZ, n_elem, width, @@ -99,7 +106,7 @@ def _swiglu_shared_bwd_kernel( 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 = 1.0 / (1.0 + tld.exp(-g)) + sig = tl.load(SIG + offs, mask=mask, other=0.0) silu = g * sig # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. t = 1.0 - sig @@ -161,9 +168,10 @@ def swiglu_shared_fwd(z: torch.Tensor) -> torch.Tensor: 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, h, n_elem, width, BLOCK=block) + _swiglu_shared_fwd_kernel[grid](z, sig, h, n_elem, width, BLOCK=block) return h @@ -181,7 +189,9 @@ def swiglu_shared_bwd(dh: torch.Tensor, z: torch.Tensor) -> torch.Tensor: 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, dz, n_elem, dh.shape[1], BLOCK=block) + _swiglu_shared_bwd_kernel[grid](dh, z, sig, dz, n_elem, width, BLOCK=block) return dz diff --git a/tests/test_shared_expert_mlp.py b/tests/test_shared_expert_mlp.py index 6590ceea..a114d711 100644 --- a/tests/test_shared_expert_mlp.py +++ b/tests/test_shared_expert_mlp.py @@ -112,8 +112,13 @@ def test_shared_output_independent_of_routed(provider): @requires_cuda -def test_cuda_triton_byte_equal(): - """The two backends agree with each other bit-for-bit.""" +@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 = [] @@ -122,8 +127,14 @@ def test_cuda_triton_byte_equal(): providers.append(resolve_provider(spec)) except NotImplementedError as exc: pytest.skip(f"backend unavailable: {exc}") - batch = fixtures.make_shared_batch("shared_t16").to("cuda") - dy = fixtures.make_grad_output("shared_t16", (batch.x.shape[0], batch.x.shape[1])).to("cuda") + 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) From 20e122b6fcdfc3ed5dc4bf4ddfa5a2e2fdaee7f3 Mon Sep 17 00:00:00 2001 From: Yizheng Jiao Date: Fri, 4 Sep 2026 09:58:41 -0700 Subject: [PATCH 4/4] fix(p5-5): forbid FMA contraction in the Triton kernels via libdevice _rn ops The compiler may contract a * b + c into an FMA; at T=256 that rounded dsilu = sig * (1 + g * (1 - sig)) differently on 2/262144 dgate elements (1 ulp after the BF16 round). All mul/add/sub in the Triton strict GEMM and SwiGLU kernels now go through libdevice add_rn/mul_rn/sub_rn, the exact Triton spelling of the CUDA kernel's __fadd_rn/__fmul_rn/__fsub_rn. --- .../kernels/ops/triton/moe/shared_expert.py | 27 +++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/rl_engine/kernels/ops/triton/moe/shared_expert.py b/rl_engine/kernels/ops/triton/moe/shared_expert.py index f510f239..d65a10f4 100644 --- a/rl_engine/kernels/ops/triton/moe/shared_expert.py +++ b/rl_engine/kernels/ops/triton/moe/shared_expert.py @@ -15,6 +15,11 @@ 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 @@ -24,6 +29,7 @@ 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 @@ -56,8 +62,7 @@ def _strict_gemm_kernel( 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) - prod = a * b - acc = acc + prod + acc = tld.add_rn(acc, tld.mul_rn(a, b)) tl.store(C + m * N + offs_n, acc, mask=mask_n) @triton.jit @@ -81,8 +86,8 @@ def _swiglu_shared_fwd_kernel( 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 = g * sig - h = (silu * u).to(tl.bfloat16) + silu = tld.mul_rn(g, sig) + h = tld.mul_rn(silu, u).to(tl.bfloat16) tl.store(H + offs, h, mask=mask) @triton.jit @@ -107,14 +112,14 @@ def _swiglu_shared_bwd_kernel( 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 = g * sig + silu = tld.mul_rn(g, sig) # dsilu = sig * (1 + g * (1 - sig)), each op rounded separately. - t = 1.0 - sig - t = g * t - t = 1.0 + t - dsilu = sig * t - dgate = (dh * u) * dsilu - dup = dh * silu + 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)