gpt-o3 / triton2e18c2

import torch
import triton
import triton.language as tl


@triton.jit
def _fused_add_rmsnorm_kernel(
    hidden_ptr,      # *bf16 [batch_size, 2048]
    residual_ptr,    # *bf16 [batch_size, 2048]
    weight_ptr,      # *bf16 [2048]
    output_ptr,      # *bf16 [batch_size, 2048]
    batch_size,      # int
    eps,             # float32
    BLOCK_SIZE: tl.constexpr,  # 2048
):
    row = tl.program_id(0)
    if row >= batch_size:
        return

    offs = tl.arange(0, BLOCK_SIZE)

    hidden_ptrs   = hidden_ptr   + row * BLOCK_SIZE + offs
    residual_ptrs = residual_ptr + row * BLOCK_SIZE + offs
    weight_ptrs   = weight_ptr   + offs
    out_ptrs      = output_ptr   + row * BLOCK_SIZE + offs

    hidden   = tl.load(hidden_ptrs).to(tl.float32)
    residual = tl.load(residual_ptrs).to(tl.float32)
    weight   = tl.load(weight_ptrs).to(tl.float32)

    x   = hidden + residual
    sq  = x * x
    ssq = tl.sum(sq, axis=0)
    inv_rms = tl.math.rsqrt(ssq / BLOCK_SIZE + eps)

    y = x * inv_rms * weight
    y = y.to(tl.bfloat16)

    tl.store(out_ptrs, y)


def run(hidden_states, residual, weight, eps=1e-6):
    # Basic validation
    if hidden_states.ndim != 2 or residual.ndim != 2:
        raise ValueError("hidden_states and residual must be 2-D tensors")
    if hidden_states.shape != residual.shape:
        raise ValueError("hidden_states and residual must have identical shapes")
    if hidden_states.shape[1] != 2048 or weight.numel() != 2048:
        raise ValueError("hidden_size must be 2048 for all tensors")
    if hidden_states.dtype != torch.bfloat16 or residual.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16:
        raise ValueError("All tensors must be of dtype torch.bfloat16")
    if hidden_states.device != residual.device or hidden_states.device.type != weight.device.type:
        raise ValueError("All input tensors must reside on the same device")

    batch_size = hidden_states.shape[0]
    src_device = hidden_states.device

    # Ensure data is on CUDA
    if src_device.type != "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available but is required for Triton kernels")
        device = torch.device("cuda")
        hidden_cuda   = hidden_states.to(device)
        residual_cuda = residual.to(device)
        weight_cuda   = weight.to(device)
    else:
        device = src_device
        hidden_cuda   = hidden_states.contiguous()
        residual_cuda = residual.contiguous()
        weight_cuda   = weight.contiguous()

    output_cuda = torch.empty_like(hidden_cuda)

    grid = (batch_size,)
    _fused_add_rmsnorm_kernel[grid](
        hidden_cuda,
        residual_cuda,
        weight_cuda,
        output_cuda,
        batch_size,
        eps,
        BLOCK_SIZE=2048,
        num_warps=8,
        num_stages=4,
    )

    return output_cuda if src_device.type == "cuda" else output_cuda.to(src_device)