gpt-5 / tritonbfd137

import torch
import triton
import triton.language as tl

EPS = 1e-6
HIDDEN_SIZE = 128


@triton.jit
def rmsnorm_h128_kernel(
    X_ptr, W_ptr, Y_ptr,
    stride_x_bs, stride_x_h,
    stride_w,
    stride_y_bs, stride_y_h,
    B: tl.int32,
    H: tl.constexpr,
    EPS: tl.constexpr,
):
    row = tl.program_id(0)
    offs = tl.arange(0, H)

    # Guards
    row_in_range = row < B
    vec_mask = row_in_range & (offs < H)

    x_ptrs = X_ptr + row * stride_x_bs + offs * stride_x_h
    w_ptrs = W_ptr + offs * stride_w
    y_ptrs = Y_ptr + row * stride_y_bs + offs * stride_y_h

    # Load inputs in bf16 then cast to f32
    x_bf16 = tl.load(x_ptrs, mask=vec_mask, other=0.0)
    x = x_bf16.to(tl.float32)

    # Compute RMS
    ss = tl.sum(x * x, axis=0)
    mean = ss * (1.0 / H)
    inv_rms = tl.rsqrt(mean + EPS)

    # Load weight
    w_bf16 = tl.load(w_ptrs, mask=offs < H, other=0.0)
    w = w_bf16.to(tl.float32)

    # Apply normalization and scale
    y = x * inv_rms * w
    y_bf16 = y.to(tl.bfloat16)

    # Store result
    tl.store(y_ptrs, y_bf16, mask=vec_mask)


def run(hidden_states, weight):
    if not isinstance(hidden_states, torch.Tensor) or not isinstance(weight, torch.Tensor):
        raise TypeError("hidden_states and weight must be torch.Tensor")

    if hidden_states.ndim != 2:
        raise ValueError(f"hidden_states must be 2D [batch_size, hidden_size], got shape {tuple(hidden_states.shape)}")
    if weight.ndim != 1:
        raise ValueError(f"weight must be 1D [hidden_size], got shape {tuple(weight.shape)}")

    batch_size, hidden_size = hidden_states.shape
    if hidden_size != HIDDEN_SIZE:
        raise ValueError(f"hidden_size must be {HIDDEN_SIZE}, got {hidden_size}")
    if weight.shape[0] != HIDDEN_SIZE:
        raise ValueError(f"weight size must be [{HIDDEN_SIZE}], got {tuple(weight.shape)}")

    if hidden_states.dtype != torch.bfloat16:
        raise TypeError(f"hidden_states dtype must be torch.bfloat16, got {hidden_states.dtype}")
    if weight.dtype != torch.bfloat16:
        raise TypeError(f"weight dtype must be torch.bfloat16, got {weight.dtype}")

    # Handle empty batch early
    if batch_size == 0:
        return hidden_states.new_empty((0, HIDDEN_SIZE))

    cuda_available = torch.cuda.is_available()
    orig_device = hidden_states.device

    # Device management
    if hidden_states.is_cuda and weight.is_cuda:
        if hidden_states.device != weight.device:
            target_device = hidden_states.device
            hidden_cuda = hidden_states
            weight_cuda = weight.to(target_device, non_blocking=True)
        else:
            target_device = hidden_states.device
            hidden_cuda = hidden_states
            weight_cuda = weight
    elif hidden_states.is_cuda and not weight.is_cuda:
        if not cuda_available:
            raise RuntimeError("CUDA is not available but hidden_states is on CUDA.")
        target_device = hidden_states.device
        hidden_cuda = hidden_states
        weight_cuda = weight.to(target_device, non_blocking=True)
    elif not hidden_states.is_cuda and weight.is_cuda:
        if not cuda_available:
            raise RuntimeError("CUDA is not available but weight is on CUDA.")
        target_device = weight.device
        hidden_cuda = hidden_states.to(target_device, non_blocking=True)
        weight_cuda = weight
    else:
        if not cuda_available:
            raise RuntimeError("CUDA is not available; Triton kernel requires CUDA.")
        target_device = torch.device("cuda")
        hidden_cuda = hidden_states.to(target_device, non_blocking=True)
        weight_cuda = weight.to(target_device, non_blocking=True)

    out_cuda = torch.empty((batch_size, HIDDEN_SIZE), dtype=torch.bfloat16, device=target_device)

    grid = (batch_size,)
    rmsnorm_h128_kernel[grid](
        hidden_cuda, weight_cuda, out_cuda,
        hidden_cuda.stride(0), hidden_cuda.stride(1),
        weight_cuda.stride(0),
        out_cuda.stride(0), out_cuda.stride(1),
        batch_size,
        H=HIDDEN_SIZE,
        EPS=EPS,
        num_warps=4,
        num_stages=2,
    )

    if out_cuda.device != orig_device:
        out = out_cuda.to(orig_device, non_blocking=True)
    else:
        out = out_cuda
    return out