triton_kernels.py

from __future__ import annotations

import torch
import triton
import triton.language as tl


def _settings(n_cols: int) -> tuple[int, int]:
    block_size = triton.next_power_of_2(n_cols)
    if block_size > 65_536:
        raise ValueError(f"Unsupported SwiGLU width: {n_cols}")
    num_warps = 4 if block_size < 2_048 else 8
    return block_size, num_warps


@triton.jit
def _packed_swiglu_forward_kernel(
    packed_ptr,
    output_ptr,
    packed_stride: tl.constexpr,
    output_stride: tl.constexpr,
    n_cols: tl.constexpr,
    block_size: tl.constexpr,
):
    row = tl.program_id(0).to(tl.int64)
    offsets = tl.arange(0, block_size)
    mask = offsets < n_cols
    packed_row = packed_ptr + row * packed_stride
    gate = tl.load(packed_row + offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(packed_row + n_cols + offsets, mask=mask, other=0.0)
    activated = (gate * tl.sigmoid(gate)).cast(up.dtype) * up
    tl.store(output_ptr + row * output_stride + offsets, activated, mask=mask)


@triton.jit
def _packed_swiglu_backward_kernel(
    grad_output_ptr,
    packed_ptr,
    grad_packed_ptr,
    grad_output_stride: tl.constexpr,
    packed_stride: tl.constexpr,
    n_cols: tl.constexpr,
    block_size: tl.constexpr,
):
    row = tl.program_id(0).to(tl.int64)
    offsets = tl.arange(0, block_size)
    mask = offsets < n_cols
    packed_row = packed_ptr + row * packed_stride
    grad_packed_row = grad_packed_ptr + row * packed_stride
    grad_output = tl.load(
        grad_output_ptr + row * grad_output_stride + offsets,
        mask=mask,
        other=0.0,
    )
    gate = tl.load(packed_row + offsets, mask=mask, other=0.0).to(tl.float32)
    up = tl.load(packed_row + n_cols + offsets, mask=mask, other=0.0)
    sigmoid_gate = tl.sigmoid(gate)
    silu_gate = gate * sigmoid_gate
    grad_up = grad_output * silu_gate
    grad_gate = grad_output * up * (sigmoid_gate + silu_gate * (1.0 - sigmoid_gate))
    tl.store(grad_packed_row + offsets, grad_gate, mask=mask)
    tl.store(grad_packed_row + n_cols + offsets, grad_up, mask=mask)


class PackedSwiGLUFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, packed: torch.Tensor) -> torch.Tensor:
        if not packed.is_cuda:
            raise ValueError("PackedSwiGLUFunction requires a CUDA tensor")
        if packed.dtype not in (torch.float16, torch.bfloat16, torch.float32):
            raise TypeError(f"Unsupported PackedSwiGLU dtype: {packed.dtype}")
        if packed.shape[-1] % 2:
            raise ValueError("Packed gate/up dimension must be even")
        packed = packed.contiguous()
        n_cols = packed.shape[-1] // 2
        packed_2d = packed.view(-1, 2 * n_cols)
        output = torch.empty((packed_2d.shape[0], n_cols), dtype=packed.dtype, device=packed.device)
        block_size, num_warps = _settings(n_cols)
        _packed_swiglu_forward_kernel[(packed_2d.shape[0],)](
            packed_2d,
            output,
            packed_2d.stride(0),
            output.stride(0),
            n_cols=n_cols,
            block_size=block_size,
            num_warps=num_warps,
        )
        ctx.save_for_backward(packed_2d)
        ctx.original_shape = packed.shape
        ctx.n_cols = n_cols
        ctx.block_size = block_size
        ctx.num_warps = num_warps
        return output.view(*packed.shape[:-1], n_cols)

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
        (packed_2d,) = ctx.saved_tensors
        grad_output_2d = grad_output.contiguous().view(-1, ctx.n_cols)
        grad_packed = torch.empty_like(packed_2d)
        _packed_swiglu_backward_kernel[(packed_2d.shape[0],)](
            grad_output_2d,
            packed_2d,
            grad_packed,
            grad_output_2d.stride(0),
            packed_2d.stride(0),
            n_cols=ctx.n_cols,
            block_size=ctx.block_size,
            num_warps=ctx.num_warps,
        )
        return (grad_packed.view(ctx.original_shape),)