diff --git a/benchmarks/benchmark_swiglu.py b/benchmarks/benchmark_swiglu.py new file mode 100644 index 00000000..7ecb88d0 --- /dev/null +++ b/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() diff --git a/ci/run_gpu_ci.sh b/ci/run_gpu_ci.sh index 5a757464..03de8847 100644 --- a/ci/run_gpu_ci.sh +++ b/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() { diff --git a/csrc/cuda/activation/swiglu_sm90.cu b/csrc/cuda/activation/swiglu_sm90.cu new file mode 100644 index 00000000..9ac3d6e6 --- /dev/null +++ b/csrc/cuda/activation/swiglu_sm90.cu @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include + +#include +#include +#include +#include +#include + +#include +#include + +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(blockDim.x) * gridDim.x * 2; + for (int64_t index = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) * 2; + index < numel; index += stride) { + if (index + 1 < numel) { + const auto gate2 = + reinterpret_cast(gate)[index / 2]; + const auto up2 = reinterpret_cast(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(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( + std::min((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<<>>( + reinterpret_cast(gate.data_ptr()), + reinterpret_cast(up.data_ptr()), + reinterpret_cast<__nv_bfloat16 *>(output.data_ptr()), gate.numel()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index dc03ab58..58c0c7e4 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -73,6 +73,11 @@ torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::optional 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"); diff --git a/docs/operators/activation.md b/docs/operators/activation.md index 6487d4a8..d06e2ee4 100644 --- a/docs/operators/activation.md +++ b/docs/operators/activation.md @@ -1,13 +1,12 @@ # SiLU / SwiGLU Activation -The activation operators are the element-wise core of the Qwen3/Llama gated MLP. They are -**WS1 ground-truth references** (issue #108): pure-PyTorch, fp32-accumulating definitions of -the "correct answer" that downstream fused CUDA/Triton MLP kernels are validated against. +The activation operators are the element-wise core of the Qwen3/Llama gated MLP. The +pure-PyTorch implementations are **WS1 ground-truth references** (issue #108): +fp32-accumulating definitions of the correct answer for optimized backends. - **SiLU** (`NativeSiLUOp`): `silu(x) = x * sigmoid(x)` — the `hidden_act="silu"` gate. -- **SwiGLU** (`NativeSwiGLUOp`): `swiglu(gate, up) = silu(gate) * up` — the gated MLP middle - stage. `gate` / `up` are the `gate_proj` / `up_proj` outputs (already at the intermediate - width); the following `down_proj` is a plain Matmul and is **not** part of this operator. +- **SwiGLU** (`NativeSwiGLUOp`): `swiglu(gate, up) = silu(gate) * up` — the gated MLP + middle stage. The following `down_proj` is a separate operator. ```text hidden --gate_proj--> gate --\ @@ -15,57 +14,65 @@ hidden --gate_proj--> gate --\ hidden --up_proj----> up ----/ ``` -## Entry Point +## Entry point + ```python from rl_engine.kernels.registry import kernel_registry silu = kernel_registry.get_op("silu") swiglu = kernel_registry.get_op("swiglu") -# SiLU: single element-wise activation -y = silu(x) # [..., N] -> [..., N] - -# SwiGLU: gated activation (gate and up must share shape) -h = swiglu(gate, up) # [..., I], [..., I] -> [..., I] +y = silu(x) +h = swiglu(gate, up) ``` -Both ops expose the WS1 dual-path contract: +The native reference ops expose the WS1 dual-path contract: -- `forward(...)` — computes in fp32, casts back to the input dtype (Axis-B accuracy - candidate / dtype-behavior path). -- `forward_fp32(...)` — computes and returns fp32 (the ground-truth golden path). +- `forward(...)` computes in fp32 and casts back to the input dtype. +- `forward_fp32(...)` computes and returns fp32 ground truth. -## Backends +## Fused Qwen3 forward contract -| Backend | Wrapper | Native symbol | Status | -| --- | --- | --- | --- | -| PyTorch fallback | `NativeSiLUOp` / `NativeSwiGLUOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused MLP kernels validate against this reference. | +The optimized operator is the local activation boundary in issue #239's Qwen3-8B TP=2 +pipeline: -## Tensor Contract +```text +gate_local [M_local, 6144] --\ + SiLU(gate) * up --> hidden_local [M_local, 6144] +up_local [M_local, 6144] --/ +``` -| Argument | Shape | Dtype | Requirements | +| Tensor | Shape | Dtype | Layout/device | | --- | --- | --- | --- | -| `x` (SiLU) | `[..., N]` | float (fp16/bf16/fp32) | Any shape; last dim arbitrary (Qwen3-8B `I=12288`). | -| `gate` (SwiGLU) | `[..., I]` | float | `gate_proj` output. | -| `up` (SwiGLU) | `[..., I]` | float | `up_proj` output; **must share `gate`'s shape**. | -| output | same as input | `forward`: input dtype · `forward_fp32`: float32 | Same shape as input. | +| `gate_local` | `[..., I/TP]` | BF16 | CUDA; arbitrary input strides accepted | +| `up_local` | same as `gate_local` | BF16 | same shape and device as `gate_local` | +| `hidden_local` | same as inputs | BF16 | contiguous CUDA output | + +Both optimized backends compute every coordinate in FP32 and round once when storing BF16: + +```text +sigmoid_gate = 1 / (1 + exp(-float(gate))) +hidden = float(gate) * sigmoid_gate * float(up) +``` -Element-wise and shape-agnostic: the Qwen3-8B intermediate dim `I=12288` is just one valid -last-dim size, not a hard requirement. Pure functions — no randomness, no in-place -mutation, device/dtype follow the inputs. +The fused activation has no collective, reduction, random state, in-place mutation, or Down +GEMM. It runs on the caller's current stream and returns a tensor on the inputs' device. -## Dispatch Behavior +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| CUDA SM90 | `SwiGLUSM90Op` | `swiglu_forward_sm90` | BF16x2 fixed mapping with scalar odd tail | +| Triton | `TritonSwiGLUOp` | None | Fixed block size; no autotune | +| PyTorch | `NativeSiLUOp` / `NativeSwiGLUOp` | None | FP32 ground-truth reference and fallback | -`kernel_registry.get_op("silu" | "swiglu")` resolves through the `OpBackend` priority map. -On `cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_SILU` / `PYTORCH_NATIVE_SWIGLU`), so every device dispatches to the -fp32 reference. When fused kernels land, they are prepended to the priority list and the -native op becomes the fallback. +On CUDA, registry dispatch prefers the compiled CUDA implementation on SM90, then Triton, +then the PyTorch reference. Constructing a backend class directly provides explicit backend +selection for validation. -## Accuracy +## Accuracy and invariance -Reference semantics (`forward_fp32`, fp32 accumulation): +Reference semantics are: ```python # SiLU @@ -76,37 +83,47 @@ gate_f = gate.float() out = gate_f * torch.sigmoid(gate_f) * up.float() ``` -- **Ground truth**: `forward_fp32` always accumulates in and returns fp32. -- **Dtype path**: `forward` runs the same fp32 math, then casts back to the input dtype; - it is bitwise-equal to `forward_fp32(x).to(dtype)`. -- **Axis A — batch invariance**: element-wise and row-independent, so a row's output is - bitwise-identical regardless of batch size or padding (`torch.equal`, `atol=0`). -- **Axis B — tolerance**: as `elementwise` ops, low-precision tolerance follows the - `elementwise` row of the WS1 numerical contract. +- The native dtype path is bitwise equal to its fp32 formula cast to the input dtype. +- Every backend must be bitwise batch/chunk/padding invariant: a coordinate is independent + of unrelated rows. +- Optimized output is currently checked against the independent fp32 oracle with issue + #108's elementwise threshold. +- CUDA/Triton cross-backend bitwise equality is not currently part of the implementation + contract because their exponential implementations may differ. This must be confirmed + with the integration owner before the contract is frozen. -## Performance Notes +## Build and validation -Reference operators — no fused kernel or benchmark yet. Downstream fused MLP kernels carry -their own benchmarks and are measured against this reference for correctness. +Build the CUDA backend with: -## Tests +```bash +KERNEL_ALIGN_ACTIVATION_SM90=1 pip install --no-build-isolation -e . +``` + +Validate the Qwen3-8B TP-local width on an H100: ```bash -python -m pytest tests/test_swiglu.py -v +python scripts/check_operator.py --op swiglu --candidate cuda-sm90 \ + --device cuda --dtype bf16 --batch 1 --seq 4096 --intermediate-dim 6144 --arch-key sm90 +python scripts/check_operator.py --op swiglu --candidate triton \ + --device cuda --dtype bf16 --batch 1 --seq 4096 --intermediate-dim 6144 --arch-key sm90 +python -m pytest tests/test_swiglu.py tests/test_swiglu_forward_backends.py -v ``` -Covers: correctness vs an independent fp32 formula, dtype paths, Axis-A batch invariance -(slice + padding), input purity, gradient flow, the SwiGLU shape guard, and registry -dispatch. +The forward benchmark defaults to `[M_local, 6144]`: -## Implementation Files +```bash +python benchmarks/benchmark_swiglu.py --rows 4096 --width 6144 +``` + +## Implementation files - `rl_engine/kernels/ops/pytorch/activation/swiglu.py` -- `rl_engine/kernels/registry.py` +- `rl_engine/kernels/ops/cuda/activation/swiglu.py` +- `rl_engine/kernels/ops/triton/activation/swiglu.py` +- `csrc/cuda/activation/swiglu_sm90.cu` - `tests/test_swiglu.py` +- `tests/test_swiglu_forward_backends.py` -## Known Limitations - -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). -- SwiGLU requires `gate` and `up` to share shape (raises `ValueError` otherwise); no - broadcasting. +Backward for the optimized CUDA and Triton backends is intentionally outside this +forward-stage implementation. diff --git a/envs.py b/envs.py index 34aded7c..9f2cbacd 100644 --- a/envs.py +++ b/envs.py @@ -26,3 +26,4 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO" KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" +KERNEL_ALIGN_ACTIVATION_SM90 = "KERNEL_ALIGN_ACTIVATION_SM90" diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index babbdc81..86659fa7 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -9,6 +9,7 @@ def batch_invariant_logp_sm90( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... +def swiglu_forward_sm90(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: ... def fused_linear_logp_sm90( hidden: torch.Tensor, weight: torch.Tensor, diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 835ee0e4..cbba0bcb 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -57,8 +57,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "batch_invariant_logp": f"{batch}x{seq}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", - "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", - "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", + "silu": f"{batch}x{seq}x{_intermediate_dim(args)}", + "swiglu": f"{batch}x{seq}x{_intermediate_dim(args)}", "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", @@ -193,7 +193,7 @@ def _make_silu_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) return { - "x": _floating_tensor((batch, seq, DEFAULT_INTERMEDIATE), args, dtype, device, 0), + "x": _floating_tensor((batch, seq, _intermediate_dim(args)), args, dtype, device, 0), } @@ -201,9 +201,10 @@ def _make_swiglu_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: batch, seq = _batch_seq(args) + intermediate_dim = _intermediate_dim(args) return { - "gate": _floating_tensor((batch, seq, DEFAULT_INTERMEDIATE), args, dtype, device, 0), - "up": _floating_tensor((batch, seq, DEFAULT_INTERMEDIATE), args, dtype, device, 1), + "gate": _floating_tensor((batch, seq, intermediate_dim), args, dtype, device, 0), + "up": _floating_tensor((batch, seq, intermediate_dim), args, dtype, device, 1), } @@ -302,6 +303,10 @@ def _normalized_dim(args: argparse.Namespace) -> int: return _arg_int(args, "normalized_dim", DEFAULT_HIDDEN) +def _intermediate_dim(args: argparse.Namespace) -> int: + return _arg_int(args, "intermediate_dim", DEFAULT_INTERMEDIATE) + + def _matmul_k(args: argparse.Namespace) -> int: return _arg_int(args, "k_dim", DEFAULT_HIDDEN) diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 09304845..badf7757 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -122,6 +122,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("a", "b"), ), + "swiglu": OperatorSpec( + name="swiglu", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUSM90Op", + "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp", + }, + ), "rope": OperatorSpec( name="rope", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/activation/__init__.py b/rl_engine/kernels/ops/cuda/activation/__init__.py new file mode 100644 index 00000000..a5f4b851 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/activation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .swiglu import SwiGLUSM90Op + +__all__ = ["SwiGLUSM90Op"] diff --git a/rl_engine/kernels/ops/cuda/activation/swiglu.py b/rl_engine/kernels/ops/cuda/activation/swiglu.py new file mode 100644 index 00000000..8f684c86 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/activation/swiglu.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic BF16 ``SiLU(gate) * up`` forward kernel for Hopper.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch import Tensor + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.utils.logger import logger + + +def _validate_inputs(gate: Tensor, up: Tensor) -> None: + if gate.device.type != "cuda" or up.device.type != "cuda": + raise RuntimeError("gate and up must be CUDA tensors") + if gate.dtype is not torch.bfloat16 or up.dtype is not torch.bfloat16: + raise TypeError("gate and up must have dtype torch.bfloat16") + if gate.shape != up.shape: + raise ValueError( + f"gate and up must share shape, got {tuple(gate.shape)} vs {tuple(up.shape)}" + ) + if gate.device != up.device: + raise ValueError(f"gate and up must share device, got {gate.device} vs {up.device}") + + +def _as_aligned_contiguous(tensor: Tensor) -> Tensor: + tensor = tensor.contiguous() + # A contiguous view may still begin at an odd BF16 storage offset. + return tensor if tensor.data_ptr() % 4 == 0 else tensor.clone() + + +class SwiGLUSM90Op(nn.Module): + """``gate, up -> SiLU(gate) * up`` with BF16 I/O and FP32 element math.""" + + op_class = "elementwise" + + def __init__(self) -> None: + super().__init__() + if not _EXT_AVAILABLE or not hasattr(_C, "swiglu_forward_sm90"): + raise RuntimeError( + "SM90 SwiGLU is not compiled into rl_engine._C. Rebuild with " + "'KERNEL_ALIGN_ACTIVATION_SM90=1 pip install --no-build-isolation -e .'." + ) + logger.info("Successfully linked to the SM90 SwiGLU forward kernel.") + + def forward(self, gate: Tensor, up: Tensor) -> Tensor: + _validate_inputs(gate, up) + return _C.swiglu_forward_sm90(_as_aligned_contiguous(gate), _as_aligned_contiguous(up)) diff --git a/rl_engine/kernels/ops/triton/activation/__init__.py b/rl_engine/kernels/ops/triton/activation/__init__.py new file mode 100644 index 00000000..a6b4585b --- /dev/null +++ b/rl_engine/kernels/ops/triton/activation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .swiglu import TritonSwiGLUOp + +__all__ = ["TritonSwiGLUOp"] diff --git a/rl_engine/kernels/ops/triton/activation/swiglu.py b/rl_engine/kernels/ops/triton/activation/swiglu.py new file mode 100644 index 00000000..5a37508f --- /dev/null +++ b/rl_engine/kernels/ops/triton/activation/swiglu.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic BF16 ``SiLU(gate) * up`` forward kernel, Triton backend.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import triton +import triton.language as tl +from torch import Tensor + +from rl_engine.utils.logger import logger + +_BLOCK_SIZE = 256 + + +@triton.jit +def _swiglu_forward_kernel( + gate_ptr, + up_ptr, + output_ptr, + numel, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + gate = tl.load(gate_ptr + offsets, mask=mask).to(tl.float32) + up = tl.load(up_ptr + offsets, mask=mask).to(tl.float32) + sigmoid_gate = 1.0 / (1.0 + tl.exp(-gate)) + output = (gate * sigmoid_gate) * up + tl.store(output_ptr + offsets, output, mask=mask) + + +def _validate_inputs(gate: Tensor, up: Tensor) -> None: + if gate.device.type != "cuda" or up.device.type != "cuda": + raise RuntimeError("gate and up must be CUDA tensors") + if gate.dtype is not torch.bfloat16 or up.dtype is not torch.bfloat16: + raise TypeError("gate and up must have dtype torch.bfloat16") + if gate.shape != up.shape: + raise ValueError( + f"gate and up must share shape, got {tuple(gate.shape)} vs {tuple(up.shape)}" + ) + if gate.device != up.device: + raise ValueError(f"gate and up must share device, got {gate.device} vs {up.device}") + + +class TritonSwiGLUOp(nn.Module): + """``gate, up -> SiLU(gate) * up`` with BF16 I/O and FP32 element math.""" + + op_class = "elementwise" + + def __init__(self) -> None: + super().__init__() + logger.info("TritonSwiGLUOp ready (fixed elementwise schedule, no autotune).") + + def forward(self, gate: Tensor, up: Tensor) -> Tensor: + _validate_inputs(gate, up) + gate = gate.contiguous() + up = up.contiguous() + output = torch.empty_like(gate) + if gate.numel() == 0: + return output + grid = (triton.cdiv(gate.numel(), _BLOCK_SIZE),) + _swiglu_forward_kernel[grid]( + gate, + up, + output, + gate.numel(), + BLOCK_SIZE=_BLOCK_SIZE, + num_warps=4, + ) + return output diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index fb2feb6f..2fc796b3 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -86,6 +86,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" + CUDA_SWIGLU_SM90 = "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUSM90Op" + TRITON_SWIGLU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp" PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" @@ -214,7 +216,11 @@ def __init__(self): "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], - "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "swiglu": [ + OpBackend.CUDA_SWIGLU_SM90, + OpBackend.TRITON_SWIGLU, + OpBackend.PYTORCH_NATIVE_SWIGLU, + ], # Default dispatch logic for new operators "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rope": [ @@ -358,6 +364,7 @@ def _adjust_priority_for_hardware(self): lm_head_list = self._priority_map["cuda"]["lm_head"] if OpBackend.CUDA_SM90_LM_HEAD not in lm_head_list: lm_head_list.insert(0, OpBackend.CUDA_SM90_LM_HEAD) + except Exception as e: logger.warning(f"Failed to probe device capability: {e}") @@ -369,6 +376,9 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: + if not self._backend_supports_device(backend, device): + continue + if backend.name in self._instance_cache: return self._instance_cache[backend.name] @@ -389,6 +399,22 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + @staticmethod + def _backend_supports_device(backend: OpBackend, device: torch.device | str | None) -> bool: + """Return whether a hardware-specific backend can serve the requested device.""" + if backend is not OpBackend.CUDA_SWIGLU_SM90: + return True + + try: + requested = torch.device(device) if device is not None else None + if requested is not None and requested.type != "cuda": + return False + capability = torch.cuda.get_device_capability(requested) + return capability[0] == 9 + except Exception as e: + logger.warning(f"Failed to probe requested device for {backend.name}: {e}") + return False + def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: if device_ctx.is_rocm: diff --git a/scripts/check_operator.py b/scripts/check_operator.py index 2cc33682..a22a100b 100644 --- a/scripts/check_operator.py +++ b/scripts/check_operator.py @@ -84,6 +84,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--constant-value", type=float, default=0.25) parser.add_argument("--token-value", type=int, default=0) parser.add_argument("--normalized-dim", type=int, default=4096) + parser.add_argument("--intermediate-dim", type=int, default=12288) parser.add_argument("--k-dim", type=int, default=4096) parser.add_argument("--n-dim", type=int, default=4096) parser.add_argument("--theta", type=float, default=1.0e6) diff --git a/scripts/ci_smoke.py b/scripts/ci_smoke.py index 0cbe1cda..f17b9d76 100644 --- a/scripts/ci_smoke.py +++ b/scripts/ci_smoke.py @@ -11,9 +11,11 @@ succeeds (``dlopen`` does not check arch) but the first kernel launch raises ``cudaErrorNoKernelImageForDevice`` once the stream is synchronized. -Uses only ``fused_logp`` - the op registered unconditionally in ``csrc/ops.cpp`` - -so it does not require ``KERNEL_ALIGN_FORCE_SM90=1`` / a Hopper build. +The baseline launch uses ``fused_logp``, which is registered unconditionally. When an SM90 +activation build is requested on Hopper, the smoke check additionally requires and launches +the fused SwiGLU forward symbol. """ +import os import sys import torch @@ -70,6 +72,35 @@ def main() -> int: ) return 1 + activation_requested = ( + os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + or os.environ.get("KERNEL_ALIGN_ACTIVATION_SM90") == "1" + ) + if cc[0] == 9 and activation_requested: + activation_symbols = ("swiglu_forward_sm90",) + missing = [symbol for symbol in activation_symbols if not hasattr(_C, symbol)] + if missing: + print( + "[smoke] FATAL: SM90 activation build is missing symbols: " + ", ".join(missing), + file=sys.stderr, + ) + return 1 + try: + gate = torch.randn(3, 17, device="cuda", dtype=torch.bfloat16) + up = torch.randn_like(gate) + activation_out = _C.swiglu_forward_sm90(gate, up) + torch.cuda.synchronize() + except Exception as exc: + print( + "[smoke] FATAL: SM90 SwiGLU forward launch failed: " f"{type(exc).__name__}: {exc}", + file=sys.stderr, + ) + return 1 + expected_shape = tuple(gate.shape) + if tuple(activation_out.shape) != expected_shape: + print("[smoke] FATAL: SM90 SwiGLU returned an unexpected shape", file=sys.stderr) + return 1 + print(f"[smoke] OK: rl_engine._C built and fused_logp ran on sm_{cc[0]}{cc[1]}.") return 0 diff --git a/setup.py b/setup.py index 288c23e6..62628ab9 100644 --- a/setup.py +++ b/setup.py @@ -177,6 +177,19 @@ def get_extensions(): nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + # Qwen3 SiLU/SwiGLU is independently buildable, following det_gemm's + # fail-closed pattern. This avoids forcing unrelated TMA/logp sources + # into an activation-only build. KERNEL_ALIGN_FORCE_SM90 remains an + # umbrella switch for users that intentionally build every SM90 op. + enable_activation_sm90 = enable_sm90 or envs.env_flag(envs.KERNEL_ALIGN_ACTIVATION_SM90) + activation_sm90_source = "csrc/cuda/activation/swiglu_sm90.cu" + if enable_activation_sm90 and not os.path.exists(activation_sm90_source): + raise FileNotFoundError(f"SM90 activation source is missing: {activation_sm90_source}") + if enable_activation_sm90: + cuda_sources.append(activation_sm90_source) + nvcc_flags.append("-DRL_KERNEL_ENABLE_ACTIVATION_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_ACTIVATION_SM90") + extensions.append( CUDAExtension( name="rl_engine._C", diff --git a/tests/test_kernel_registry.py b/tests/test_kernel_registry.py index 66a09297..03297800 100644 --- a/tests/test_kernel_registry.py +++ b/tests/test_kernel_registry.py @@ -119,3 +119,32 @@ class FakeExtension: assert OpBackend.CUDA_SM90_EMBEDDING not in registry._priority_map["cuda"]["embedding"] assert OpBackend.CUDA_SM90_LM_HEAD not in registry._priority_map["cuda"]["lm_head"] + + +def test_swiglu_dispatch_checks_requested_device_before_using_cached_backend(monkeypatch): + class FakeCudaSwiGLU: + pass + + class FakeTritonSwiGLU: + pass + + classes = { + OpBackend.CUDA_SWIGLU_SM90: FakeCudaSwiGLU, + OpBackend.TRITON_SWIGLU: FakeTritonSwiGLU, + } + capabilities = {"cuda:0": (9, 0), "cuda:1": (8, 6)} + + registry = KernelRegistry() + registry._priority_map["cuda"]["swiglu"] = [ + OpBackend.CUDA_SWIGLU_SM90, + OpBackend.TRITON_SWIGLU, + ] + monkeypatch.setattr(registry, "_load_backend", lambda backend: classes[backend]) + monkeypatch.setattr( + registry_module.torch.cuda, + "get_device_capability", + lambda device=None: capabilities[str(device)], + ) + + assert isinstance(registry.get_op("swiglu", device="cuda:0"), FakeCudaSwiGLU) + assert isinstance(registry.get_op("swiglu", device="cuda:1"), FakeTritonSwiGLU) diff --git a/tests/test_swiglu.py b/tests/test_swiglu.py index bd9ff32c..1db4c29a 100644 --- a/tests/test_swiglu.py +++ b/tests/test_swiglu.py @@ -165,6 +165,9 @@ def test_swiglu_backward_batch_invariance_slice(): assert torch.equal(up_slice.grad, grad_up_full_sliced) -def test_registry_dispatches_native_activation_ops(): +def test_registry_dispatches_native_silu(): assert isinstance(kernel_registry.get_op("silu"), NativeSiLUOp) - assert isinstance(kernel_registry.get_op("swiglu"), NativeSwiGLUOp) + + +def test_registry_dispatches_native_swiglu_on_cpu(): + assert isinstance(kernel_registry.get_op("swiglu", device="cpu"), NativeSwiGLUOp) diff --git a/tests/test_swiglu_forward_backends.py b/tests/test_swiglu_forward_backends.py new file mode 100644 index 00000000..5014e800 --- /dev/null +++ b/tests/test_swiglu_forward_backends.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.activation.swiglu import NativeSwiGLUOp + +try: + from rl_engine.kernels.ops.triton.activation.swiglu import TritonSwiGLUOp + + _HAS_TRITON = True +except ImportError: + TritonSwiGLUOp = None + _HAS_TRITON = False + +try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + from rl_engine.kernels.ops.cuda.activation.swiglu import SwiGLUSM90Op + + _HAS_CUDA_SM90_OP = _EXT_AVAILABLE and hasattr(_C, "swiglu_forward_sm90") +except ImportError: + SwiGLUSM90Op = None + _HAS_CUDA_SM90_OP = False + +_IS_SM90 = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9 +_HAS_CUDA_BF16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported() +requires_cuda_sm90 = pytest.mark.skipif( + not (_IS_SM90 and _HAS_CUDA_SM90_OP), + reason="requires Hopper and the compiled SM90 SwiGLU forward kernel", +) +requires_triton_cuda = pytest.mark.skipif( + not (_HAS_CUDA_BF16 and _HAS_TRITON), reason="requires CUDA with BF16 support and Triton" +) + +_ELEMENTWISE_ATOL = 2e-2 +_ELEMENTWISE_RTOL = 1.6e-2 +_TP_LOCAL_INTERMEDIATE = 6144 + + +def _inputs(shape, seed=239): + generator = torch.Generator(device="cuda").manual_seed(seed) + gate = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator) + up = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator) + return gate, up + + +def _assert_matches_reference(op, gate, up): + output = op(gate, up) + reference = NativeSwiGLUOp().forward_fp32(gate, up) + assert output.shape == gate.shape + assert output.dtype is torch.bfloat16 + assert output.is_contiguous() + torch.testing.assert_close( + output.float(), reference, atol=_ELEMENTWISE_ATOL, rtol=_ELEMENTWISE_RTOL + ) + + +@requires_cuda_sm90 +@pytest.mark.parametrize("shape", [(3, 257), (2, _TP_LOCAL_INTERMEDIATE)]) +def test_cuda_sm90_forward_matches_fp32_reference(shape): + _assert_matches_reference(SwiGLUSM90Op(), *_inputs(shape)) + + +@requires_triton_cuda +@pytest.mark.parametrize("shape", [(3, 257), (2, _TP_LOCAL_INTERMEDIATE)]) +def test_triton_forward_matches_fp32_reference(shape): + _assert_matches_reference(TritonSwiGLUOp(), *_inputs(shape)) + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +def test_forward_is_batch_and_padding_invariant(backend): + if backend == "cuda" and not (_IS_SM90 and _HAS_CUDA_SM90_OP): + pytest.skip("requires Hopper and the compiled SM90 SwiGLU forward kernel") + if backend == "triton" and not (_HAS_CUDA_BF16 and _HAS_TRITON): + pytest.skip("requires CUDA with BF16 support and Triton") + + op = SwiGLUSM90Op() if backend == "cuda" else TritonSwiGLUOp() + gate, up = _inputs((8, 257), seed=240) + gate_before, up_before = gate.clone(), up.clone() + full = op(gate, up) + assert torch.equal(op(gate[3:5], up[3:5]), full[3:5]) + + pad_gate, pad_up = _inputs((4, 257), seed=241) + padded = op(torch.cat((gate, pad_gate)), torch.cat((up, pad_up))) + assert torch.equal(padded[:8], full) + assert torch.equal(gate, gate_before) + assert torch.equal(up, up_before) + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +def test_forward_contract_guards(backend): + if backend == "cuda" and not (_IS_SM90 and _HAS_CUDA_SM90_OP): + pytest.skip("requires Hopper and the compiled SM90 SwiGLU forward kernel") + if backend == "triton" and not (_HAS_CUDA_BF16 and _HAS_TRITON): + pytest.skip("requires CUDA with BF16 support and Triton") + + op = SwiGLUSM90Op() if backend == "cuda" else TritonSwiGLUOp() + gate, up = _inputs((2, 8), seed=242) + with pytest.raises(ValueError, match="share shape"): + op(gate, up[:, :-1]) + with pytest.raises(TypeError, match="bfloat16"): + op(gate.float(), up.float()) + + empty = torch.empty((0, _TP_LOCAL_INTERMEDIATE), device="cuda", dtype=torch.bfloat16) + assert op(empty, empty).shape == empty.shape + + noncontiguous_gate = gate.t() + noncontiguous_up = up.t() + assert not noncontiguous_gate.is_contiguous() + _assert_matches_reference(op, noncontiguous_gate, noncontiguous_up) + + +@requires_cuda_sm90 +def test_cuda_and_triton_agree_within_elementwise_contract(): + if not _HAS_TRITON: + pytest.skip("requires Triton") + gate, up = _inputs((4, 513), seed=243) + cuda_output = SwiGLUSM90Op()(gate, up) + triton_output = TritonSwiGLUOp()(gate, up) + torch.testing.assert_close( + cuda_output.float(), + triton_output.float(), + atol=_ELEMENTWISE_ATOL, + rtol=_ELEMENTWISE_RTOL, + )