claude-opus-4-1 / tritonbf2710

import torch
import triton
import triton.language as tl

@triton.jit
def rmsnorm_h128_kernel(
    hidden_states_ptr,
    weight_ptr,
    output_ptr,
    batch_size,
    HIDDEN_SIZE: tl.constexpr,
    EPS: tl.constexpr,
):
    # Process one row per program
    row_idx = tl.program_id(0)
    
    if row_idx >= batch_size:
        return
    
    # Load the entire row (hidden_size = 128)
    row_start = row_idx * HIDDEN_SIZE
    offsets = tl.arange(0, HIDDEN_SIZE)
    
    # Load hidden states for this row
    hidden_states = tl.load(hidden_states_ptr + row_start + offsets).to(tl.float32)
    
    # Compute RMS normalization
    # Calculate mean of squares
    squared = hidden_states * hidden_states
    mean_squared = tl.sum(squared, axis=0) / HIDDEN_SIZE
    
    # Compute inverse RMS
    inv_rms = tl.rsqrt(mean_squared + EPS)
    
    # Load weights
    weights = tl.load(weight_ptr + offsets).to(tl.float32)
    
    # Apply normalization and scaling
    normalized = hidden_states * inv_rms
    output = normalized * weights
    
    # Store result
    tl.store(output_ptr + row_start + offsets, output.to(tl.bfloat16))

def run(hidden_states, weight):
    # Handle device management
    original_device = hidden_states.device
    
    # Move to GPU if needed
    if not hidden_states.is_cuda:
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available but GPU execution is required")
        hidden_states = hidden_states.cuda()
    
    if not weight.is_cuda:
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available but GPU execution is required")
        weight = weight.cuda()
    
    # Validate shapes and dtypes
    batch_size, hidden_size = hidden_states.shape
    assert hidden_size == 128, f"hidden_size must be 128, got {hidden_size}"
    assert weight.shape == (128,), f"weight shape must be [128], got {weight.shape}"
    
    # Convert to bfloat16 if needed
    if hidden_states.dtype != torch.bfloat16:
        hidden_states = hidden_states.to(torch.bfloat16)
    if weight.dtype != torch.bfloat16:
        weight = weight.to(torch.bfloat16)
    
    # Allocate output tensor
    output = torch.empty_like(hidden_states, device=hidden_states.device, dtype=torch.bfloat16)
    
    # Configure grid
    grid = (batch_size,)
    
    # Launch kernel
    rmsnorm_h128_kernel[grid](
        hidden_states,
        weight,
        output,
        batch_size,
        HIDDEN_SIZE=128,
        EPS=1e-6,
    )
    
    # Move result back to original device if needed
    if original_device != output.device:
        output = output.to(original_device)
    
    return output