feat: Add CUDA SwiGLU, RMSNorm, and RoPE training kernels

CUDA kernels for three core training operations:

- SwiGLU forward+backward: h = silu(gate) * up, element-wise
- RMSNorm forward+backward: y = x * rsqrt(mean(x²) + eps) * w
  with Gemma variant (add_unit_offset: w + 1). One block per row,
  shared memory reduction. Forward stores rrms for backward.
- RoPE forward (backward reuses with -sin): in-place rotary
  position embedding. Supports arbitrary head_dim.

All kernels support fp16 and bf16 via C++ templates.
Includes extern "C" wrappers, torch.library op registrations,
backend dispatch, and torch.autograd.Function wrappers.

24 tests pass covering forward correctness against PyTorch
reference, backward gradient correctness, Gemma variant,
various sizes, and autograd integration.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/_ops.py

@@ -705,3 +705,80 @@ def _(
     )
     total_M = A_concat.shape[0]
     return torch.empty(total_M, N, device=A_concat.device, dtype=A_concat.dtype)
+
+
+# ============================================================================
+# Training Kernels: SwiGLU, RMSNorm, RoPE
+# ============================================================================
+
+# SwiGLU forward: h = silu(gate) * up
+torch.library.define(
+    "bitsandbytes::swiglu_forward",
+    "(Tensor gate, Tensor up) -> Tensor",
+)
+
+
+@register_fake("bitsandbytes::swiglu_forward")
+def _(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    torch._check(gate.shape == up.shape, lambda: "gate and up must have same shape")
+    return torch.empty_like(gate)
+
+
+# SwiGLU backward: (grad_gate, grad_up) from grad_h
+torch.library.define(
+    "bitsandbytes::swiglu_backward",
+    "(Tensor grad_h, Tensor gate, Tensor up) -> (Tensor, Tensor)",
+)
+
+
+@register_fake("bitsandbytes::swiglu_backward")
+def _(grad_h: torch.Tensor, gate: torch.Tensor, up: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.empty_like(gate), torch.empty_like(up)
+
+
+# RMSNorm forward: y = x * rsqrt(mean(x^2) + eps) * w, also returns rrms
+torch.library.define(
+    "bitsandbytes::rmsnorm_forward",
+    "(Tensor x, Tensor w, float eps, bool add_unit_offset) -> (Tensor, Tensor)",
+)
+
+
+@register_fake("bitsandbytes::rmsnorm_forward")
+def _(x: torch.Tensor, w: torch.Tensor, eps: float, add_unit_offset: bool) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(x.dim() == 2, lambda: "x must be 2D [rows, cols]")
+    rows = x.shape[0]
+    out = torch.empty_like(x)
+    rrms = torch.empty(rows, device=x.device, dtype=torch.float32)
+    return out, rrms
+
+
+# RMSNorm backward: (grad_x, grad_w) from grad_out
+torch.library.define(
+    "bitsandbytes::rmsnorm_backward",
+    "(Tensor grad_out, Tensor x, Tensor w, Tensor rrms, bool add_unit_offset) -> (Tensor, Tensor)",
+)
+
+
+@register_fake("bitsandbytes::rmsnorm_backward")
+def _(
+    grad_out: torch.Tensor,
+    x: torch.Tensor,
+    w: torch.Tensor,
+    rrms: torch.Tensor,
+    add_unit_offset: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    grad_x = torch.empty_like(x)
+    grad_w = torch.empty(x.shape[1], device=x.device, dtype=torch.float32)
+    return grad_x, grad_w
+
+
+# RoPE forward (in-place): applies rotary embeddings to Q (or Q+K)
+torch.library.define(
+    "bitsandbytes::rope_forward",
+    "(Tensor(a!) q, Tensor cos_cache, Tensor sin_cache, int n_heads) -> ()",
+)
+
+
+@register_fake("bitsandbytes::rope_forward")
+def _(q: torch.Tensor, cos_cache: torch.Tensor, sin_cache: torch.Tensor, n_heads: int) -> None:
+    pass

bitsandbytes/autograd/training_kernels.py

@@ -0,0 +1,147 @@
+"""torch.autograd.Function wrappers for CUDA training kernels.
+
+Wraps the low-level CUDA ops (SwiGLU, RMSNorm, RoPE) into autograd-aware
+functions that can be used directly in PyTorch training.
+"""
+
+import torch
+
+
+class SwiGLUFunction(torch.autograd.Function):
+    """SwiGLU activation: h = silu(gate) * up.
+
+    Forward:  h = (gate * sigmoid(gate)) * up
+    Backward: grad_gate = grad_h * up * sigmoid(gate) * (1 + gate * (1 - sigmoid(gate)))
+              grad_up   = grad_h * silu(gate)
+    """
+
+    @staticmethod
+    def forward(ctx, gate, up):
+        ctx.save_for_backward(gate, up)
+        return torch.ops.bitsandbytes.swiglu_forward(gate, up)
+
+    @staticmethod
+    def backward(ctx, grad_h):
+        gate, up = ctx.saved_tensors
+        grad_gate, grad_up = torch.ops.bitsandbytes.swiglu_backward(
+            grad_h.contiguous(), gate, up,
+        )
+        return grad_gate, grad_up
+
+
+def swiglu(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    """SwiGLU activation with autograd support.
+
+    Args:
+        gate: Gate tensor (any shape, fp16 or bf16).
+        up: Up tensor (same shape as gate).
+
+    Returns:
+        silu(gate) * up
+    """
+    return SwiGLUFunction.apply(gate, up)
+
+
+class RMSNormFunction(torch.autograd.Function):
+    """RMS normalization: y = x * rsqrt(mean(x^2) + eps) * w.
+
+    Supports Gemma variant with ``add_unit_offset=True`` (uses w + 1).
+    """
+
+    @staticmethod
+    def forward(ctx, x, w, eps=1e-6, add_unit_offset=False):
+        # Flatten to 2D for the CUDA kernel
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1]).contiguous()
+
+        out_2d, rrms = torch.ops.bitsandbytes.rmsnorm_forward(
+            x_2d, w, eps, add_unit_offset,
+        )
+
+        ctx.save_for_backward(x_2d, w, rrms)
+        ctx.add_unit_offset = add_unit_offset
+        ctx.orig_shape = orig_shape
+
+        return out_2d.reshape(orig_shape)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x_2d, w, rrms = ctx.saved_tensors
+        grad_out_2d = grad_out.reshape(x_2d.shape).contiguous()
+
+        grad_x_2d, grad_w = torch.ops.bitsandbytes.rmsnorm_backward(
+            grad_out_2d, x_2d, w, rrms, ctx.add_unit_offset,
+        )
+
+        grad_x = grad_x_2d.reshape(ctx.orig_shape)
+        return grad_x, grad_w.to(w.dtype), None, None
+
+
+def rmsnorm(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    eps: float = 1e-6,
+    add_unit_offset: bool = False,
+) -> torch.Tensor:
+    """RMS normalization with autograd support.
+
+    Args:
+        x: Input tensor (*, hidden_size), fp16 or bf16.
+        w: Weight tensor (hidden_size,).
+        eps: Epsilon for numerical stability.
+        add_unit_offset: If True, uses (w + 1) instead of w (Gemma convention).
+
+    Returns:
+        Normalized tensor of same shape as x.
+    """
+    return RMSNormFunction.apply(x, w, eps, add_unit_offset)
+
+
+class RoPEFunction(torch.autograd.Function):
+    """Rotary Position Embedding (in-place).
+
+    Forward:  q[..., :half] = q[..., :half] * cos - q[..., half:] * sin
+              q[..., half:] = q[..., half:] * cos + q[..., :half] * sin
+    Backward: same operation with sin negated.
+    """
+
+    @staticmethod
+    def forward(ctx, q, cos_cache, sin_cache, n_heads):
+        # q: [total_tokens, n_heads, head_dim]
+        ctx.save_for_backward(cos_cache, sin_cache)
+        ctx.n_heads = n_heads
+
+        q_out = q.clone()
+        torch.ops.bitsandbytes.rope_forward(q_out, cos_cache, sin_cache, n_heads)
+        return q_out
+
+    @staticmethod
+    def backward(ctx, grad_q):
+        cos_cache, sin_cache = ctx.saved_tensors
+
+        # Backward of RoPE is the same operation with sin negated
+        grad_q_out = grad_q.clone()
+        torch.ops.bitsandbytes.rope_forward(
+            grad_q_out, cos_cache, -sin_cache, ctx.n_heads,
+        )
+        return grad_q_out, None, None, None
+
+
+def rope(
+    q: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    n_heads: int,
+) -> torch.Tensor:
+    """Apply Rotary Position Embedding with autograd support.
+
+    Args:
+        q: Query tensor [total_tokens, n_heads, head_dim], fp16 or bf16.
+        cos_cache: Cosine cache [total_tokens, head_dim/2].
+        sin_cache: Sine cache [total_tokens, head_dim/2].
+        n_heads: Number of attention heads.
+
+    Returns:
+        Rotated query tensor (same shape).
+    """
+    return RoPEFunction.apply(q, cos_cache, sin_cache, n_heads)

bitsandbytes/backends/cuda/ops.py

@@ -1231,3 +1231,144 @@ def _(
         )
 
     return C_concat
+
+
+# ============================================================================
+# Training Kernels: SwiGLU, RMSNorm, RoPE
+# ============================================================================
+
+
+@register_kernel("bitsandbytes::swiglu_forward", "cuda")
+def _(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
+    torch._check(gate.shape == up.shape, lambda: "gate and up must have same shape")
+    torch._check(gate.is_contiguous(), lambda: "gate must be contiguous")
+    torch._check(up.is_contiguous(), lambda: "up must be contiguous")
+    torch._check(
+        gate.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"swiglu supports float16/bfloat16, got {gate.dtype}",
+    )
+
+    out = torch.empty_like(gate)
+    n = gate.numel()
+    dtype_suffix = "fp16" if gate.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(gate):
+        fn = getattr(lib, f"cswiglu_forward_{dtype_suffix}_c")
+        fn(get_ptr(gate), get_ptr(up), get_ptr(out), ct.c_int(n))
+
+    return out
+
+
+@register_kernel("bitsandbytes::swiglu_backward", "cuda")
+def _(grad_h: torch.Tensor, gate: torch.Tensor, up: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(grad_h.is_contiguous(), lambda: "grad_h must be contiguous")
+    torch._check(gate.is_contiguous(), lambda: "gate must be contiguous")
+    torch._check(up.is_contiguous(), lambda: "up must be contiguous")
+
+    grad_gate = torch.empty_like(gate)
+    grad_up = torch.empty_like(up)
+    n = gate.numel()
+    dtype_suffix = "fp16" if gate.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(gate):
+        fn = getattr(lib, f"cswiglu_backward_{dtype_suffix}_c")
+        fn(get_ptr(grad_h), get_ptr(gate), get_ptr(up), get_ptr(grad_gate), get_ptr(grad_up), ct.c_int(n))
+
+    return grad_gate, grad_up
+
+
+@register_kernel("bitsandbytes::rmsnorm_forward", "cuda")
+def _(
+    x: torch.Tensor,
+    w: torch.Tensor,
+    eps: float,
+    add_unit_offset: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(x.dim() == 2, lambda: "x must be 2D [rows, cols]")
+    torch._check(x.is_contiguous(), lambda: "x must be contiguous")
+    torch._check(
+        x.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"rmsnorm supports float16/bfloat16, got {x.dtype}",
+    )
+
+    rows, cols = x.shape
+    out = torch.empty_like(x)
+    rrms = torch.empty(rows, device=x.device, dtype=torch.float32)
+    dtype_suffix = "fp16" if x.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(x):
+        fn = getattr(lib, f"crmsnorm_forward_{dtype_suffix}_c")
+        fn(
+            get_ptr(x),
+            get_ptr(w),
+            get_ptr(out),
+            get_ptr(rrms),
+            ct.c_int(rows),
+            ct.c_int(cols),
+            ct.c_float(eps),
+            ct.c_bool(add_unit_offset),
+        )
+
+    return out, rrms
+
+
+@register_kernel("bitsandbytes::rmsnorm_backward", "cuda")
+def _(
+    grad_out: torch.Tensor,
+    x: torch.Tensor,
+    w: torch.Tensor,
+    rrms: torch.Tensor,
+    add_unit_offset: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(grad_out.is_contiguous(), lambda: "grad_out must be contiguous")
+    torch._check(x.is_contiguous(), lambda: "x must be contiguous")
+
+    rows, cols = x.shape
+    grad_x = torch.empty_like(x)
+    grad_w = torch.zeros(cols, device=x.device, dtype=torch.float32)
+    dtype_suffix = "fp16" if x.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(x):
+        fn = getattr(lib, f"crmsnorm_backward_{dtype_suffix}_c")
+        fn(
+            get_ptr(grad_out),
+            get_ptr(x),
+            get_ptr(w),
+            get_ptr(rrms),
+            get_ptr(grad_x),
+            get_ptr(grad_w),
+            ct.c_int(rows),
+            ct.c_int(cols),
+            ct.c_bool(add_unit_offset),
+        )
+
+    return grad_x, grad_w
+
+
+@register_kernel("bitsandbytes::rope_forward", "cuda")
+def _(
+    q: torch.Tensor,
+    cos_cache: torch.Tensor,
+    sin_cache: torch.Tensor,
+    n_heads: int,
+) -> None:
+    torch._check(q.is_contiguous(), lambda: "q must be contiguous")
+    torch._check(
+        q.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"rope supports float16/bfloat16, got {q.dtype}",
+    )
+
+    total_tokens = q.shape[0]
+    head_dim = q.shape[-1]
+    dtype_suffix = "fp16" if q.dtype == torch.float16 else "bf16"
+
+    with _cuda_device_of(q):
+        fn = getattr(lib, f"crope_forward_{dtype_suffix}_c")
+        fn(
+            get_ptr(q),
+            get_ptr(cos_cache),
+            get_ptr(sin_cache),
+            ct.c_int(total_tokens),
+            ct.c_int(n_heads),
+            ct.c_int(head_dim),
+        )