[WS2][kernels] Deterministic Qwen3 SwiGLU forward (CUDA SM90 + Triton)

benchmarks/benchmark_swiglu.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""Benchmark Qwen3 TP-local SwiGLU forward backends on Hopper."""
+
+from __future__ import annotations
+
+import argparse
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+
+from rl_engine.kernels.ops.cuda.activation.swiglu import SwiGLUSM90Op
+from rl_engine.kernels.ops.triton.activation.swiglu import TritonSwiGLUOp
+
+
+def _bench(fn: Callable[[], torch.Tensor], warmup: int, iterations: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) / iterations
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--rows", type=int, default=4096, help="M_local token rows")
+    parser.add_argument(
+        "--width", type=int, default=6144, help="Qwen3-8B TP=2 local intermediate width"
+    )
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iterations", type=int, default=100)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("benchmark_swiglu.py requires CUDA")
+    major, minor = torch.cuda.get_device_capability()
+    if major != 9:
+        raise RuntimeError(f"benchmark_swiglu.py requires Hopper SM90, got sm_{major}{minor}")
+
+    generator = torch.Generator(device="cuda").manual_seed(239)
+    shape = (args.rows, args.width)
+    gate = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator)
+    up = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator)
+    cuda_op = SwiGLUSM90Op()
+    triton_op = TritonSwiGLUOp()
+
+    timings = {
+        "PyTorch": _bench(
+            lambda: F.silu(gate.float()).mul(up.float()).bfloat16(),
+            args.warmup,
+            args.iterations,
+        ),
+        "CUDA SM90": _bench(lambda: cuda_op(gate, up), args.warmup, args.iterations),
+        "Triton": _bench(lambda: triton_op(gate, up), args.warmup, args.iterations),
+    }
+
+    print(f"device={torch.cuda.get_device_name()} shape={shape} dtype=bf16")
+    print("backend        latency_ms")
+    for name, latency in timings.items():
+        print(f"{name:<14} {latency:>10.4f}")
+
+
+if __name__ == "__main__":
+    main()

ci/run_gpu_ci.sh

@@ -15,6 +15,7 @@ TARGET_SM="${TARGET_SM:-}"
 
 # Forwarded to the remote build; setup.py compiles the Hopper (sm90) kernels only when "1".
 KERNEL_ALIGN_FORCE_SM90="${KERNEL_ALIGN_FORCE_SM90:-}"
+KERNEL_ALIGN_ACTIVATION_SM90="${KERNEL_ALIGN_ACTIVATION_SM90:-}"
 
 CI_IMAGE="${CI_IMAGE:-runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04}"
 DISK_GB=40
@@ -137,6 +138,7 @@ echo "[remote] Using interpreter: $PY"
 export FORCE_CUDA=1
 export MAX_JOBS=8
 export KERNEL_ALIGN_FORCE_SM90="'"${KERNEL_ALIGN_FORCE_SM90}"'"
+export KERNEL_ALIGN_ACTIVATION_SM90="'"${KERNEL_ALIGN_ACTIVATION_SM90}"'"
 
 # normalize_sm: compact (90) or dotted (9.0) compute cap -> torch dotted form, keeping +PTX.
 normalize_sm() {

csrc/cuda/activation/swiglu_sm90.cu

@@ -0,0 +1,98 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+
+#include <torch/extension.h>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+#include <cstdint>
+
+namespace {
+
+constexpr int kThreads = 256;
+constexpr int64_t kMaxBlocks = 65535;
+
+__device__ __forceinline__ float swiglu_fp32(float gate, float up) {
+  const float sigmoid_gate = 1.0f / (1.0f + expf(-gate));
+  return (gate * sigmoid_gate) * up;
+}
+
+__global__ void swiglu_forward_kernel(const __nv_bfloat16 *__restrict__ gate,
+                                      const __nv_bfloat16 *__restrict__ up,
+                                      __nv_bfloat16 *__restrict__ output,
+                                      int64_t numel) {
+  const int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x * 2;
+  for (int64_t index =
+           (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) * 2;
+       index < numel; index += stride) {
+    if (index + 1 < numel) {
+      const auto gate2 =
+          reinterpret_cast<const __nv_bfloat162 *>(gate)[index / 2];
+      const auto up2 = reinterpret_cast<const __nv_bfloat162 *>(up)[index / 2];
+      reinterpret_cast<__nv_bfloat162 *>(output)[index / 2] =
+          __floats2bfloat162_rn(
+              swiglu_fp32(__low2float(gate2), __low2float(up2)),
+              swiglu_fp32(__high2float(gate2), __high2float(up2)));
+    } else {
+      output[index] = __float2bfloat16_rn(swiglu_fp32(
+          __bfloat162float(gate[index]), __bfloat162float(up[index])));
+    }
+  }
+}
+
+void check_bf16_tensor(const torch::Tensor &tensor, const char *name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.scalar_type() == torch::kBFloat16, name,
+              " must have dtype torch.bfloat16");
+  TORCH_CHECK(reinterpret_cast<uintptr_t>(tensor.data_ptr()) %
+                      alignof(__nv_bfloat162) ==
+                  0,
+              name, " must be 4-byte aligned for bfloat162 access");
+}
+
+void check_sm90(const torch::Tensor &tensor) {
+  c10::cuda::CUDAGuard device_guard(tensor.device());
+  const auto *properties = at::cuda::getCurrentDeviceProperties();
+  TORCH_CHECK(properties->major == 9,
+              "SwiGLU SM90 kernel requires Hopper compute capability 9.x, got "
+              "sm_",
+              properties->major, properties->minor);
+}
+
+int launch_blocks(int64_t numel) {
+  const int64_t work_items = (numel + 1) / 2;
+  return static_cast<int>(
+      std::min<int64_t>((work_items + kThreads - 1) / kThreads, kMaxBlocks));
+}
+
+} // namespace
+
+torch::Tensor swiglu_forward_sm90(torch::Tensor gate, torch::Tensor up) {
+  check_bf16_tensor(gate, "gate");
+  check_bf16_tensor(up, "up");
+  TORCH_CHECK(gate.sizes() == up.sizes(),
+              "gate and up must have the same shape");
+  TORCH_CHECK(gate.device() == up.device(),
+              "gate and up must be on the same device");
+  check_sm90(gate);
+
+  auto output = torch::empty_like(gate);
+  if (gate.numel() == 0) {
+    return output;
+  }
+
+  c10::cuda::CUDAGuard device_guard(gate.device());
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  swiglu_forward_kernel<<<launch_blocks(gate.numel()), kThreads, 0, stream>>>(
+      reinterpret_cast<const __nv_bfloat16 *>(gate.data_ptr()),
+      reinterpret_cast<const __nv_bfloat16 *>(up.data_ptr()),
+      reinterpret_cast<__nv_bfloat16 *>(output.data_ptr()), gate.numel());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}

csrc/ops.cpp

@@ -73,6 +73,11 @@ torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden,
                                         torch::optional<torch::Tensor> bias);
 #endif
 
+#if defined(__CUDACC__) || defined(RL_KERNEL_ENABLE_ACTIVATION_SM90)
+// Qwen3 fused activation boundary for SM90; BF16 I/O with FP32 element math.
+torch::Tensor swiglu_forward_sm90(torch::Tensor gate, torch::Tensor up);
+#endif
+
 #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)
 torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output);
 torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids);
@@ -311,6 +316,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           "Single-card SM90 batch-invariant LM-head forward with fp32 output");
 #endif
 
+#if defined(__CUDACC__) || defined(RL_KERNEL_ENABLE_ACTIVATION_SM90)
+    // Qwen3 fused SiLU + Multiply activation boundary, SM90.
+    m.def("swiglu_forward_sm90", &swiglu_forward_sm90,
+          "Qwen3 fused SiLU(gate) * up forward (BF16 I/O, FP32 element math), SM90");
+#endif
+
 #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA)
     m.def("fused_logp_forward_out", &fused_logp_forward_out, "Fused logp out");
     m.def("fused_logp_forward_fp32", &fused_logp_forward_fp32, "Fused logp fp32");