claude-opus-4-1 / triton91c9a3

import torch
import triton
import triton.language as tl
import math

@triton.jit
def rmsnorm_kernel(
    hidden_states_ptr,
    weight_ptr,
    output_ptr,
    hidden_size,
    eps,
    BLOCK_SIZE: tl.constexpr,
):
    # Get the batch index
    batch_idx = tl.program_id(0)
    
    # Initialize accumulator for variance calculation
    acc = tl.zeros([1], dtype=tl.float32)
    
    # Compute variance in multiple passes if needed
    num_iters = tl.cdiv(hidden_size, BLOCK_SIZE)
    
    for iter in range(num_iters):
        # Load block of hidden states
        offsets = iter * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        mask = offsets < hidden_size
        
        hidden_states_offset = batch_idx * hidden_size + offsets
        x = tl.load(hidden_states_ptr + hidden_states_offset, mask=mask, other=0.0).to(tl.float32)
        
        # Accumulate squared values
        acc += tl.sum(x * x, axis=0)
    
    # Compute inverse RMS
    mean_sq = acc / hidden_size
    inv_rms = tl.rsqrt(mean_sq + eps)
    
    # Apply normalization and weight in a second pass
    for iter in range(num_iters):
        offsets = iter * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        mask = offsets < hidden_size
        
        hidden_states_offset = batch_idx * hidden_size + offsets
        x = tl.load(hidden_states_ptr + hidden_states_offset, mask=mask, other=0.0).to(tl.float32)
        
        # Load weights
        w = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
        
        # Apply RMSNorm
        y = (x * inv_rms) * w
        
        # Store output
        output_offset = batch_idx * hidden_size + offsets
        tl.store(output_ptr + output_offset, y.to(tl.bfloat16), mask=mask)


@triton.jit
def rmsnorm_kernel_optimized(
    hidden_states_ptr,
    weight_ptr,
    output_ptr,
    hidden_size,
    eps,
    BLOCK_SIZE: tl.constexpr,
):
    # Get the batch index
    batch_idx = tl.program_id(0)
    
    # For B200, we can use larger block sizes and more efficient memory access
    # Process the entire row in chunks
    acc = 0.0
    
    # First pass: compute variance
    for offset in range(0, hidden_size, BLOCK_SIZE):
        cols = offset + tl.arange(0, BLOCK_SIZE)
        mask = cols < hidden_size
        
        hidden_states_offset = batch_idx * hidden_size + cols
        x = tl.load(hidden_states_ptr + hidden_states_offset, mask=mask, other=0.0).to(tl.float32)
        
        # Accumulate squared values
        acc += tl.sum(x * x, axis=0)
    
    # Compute inverse RMS
    mean_sq = acc / hidden_size
    inv_rms = tl.rsqrt(mean_sq + eps)
    
    # Second pass: apply normalization and weight
    for offset in range(0, hidden_size, BLOCK_SIZE):
        cols = offset + tl.arange(0, BLOCK_SIZE)
        mask = cols < hidden_size
        
        hidden_states_offset = batch_idx * hidden_size + cols
        x = tl.load(hidden_states_ptr + hidden_states_offset, mask=mask, other=0.0).to(tl.float32)
        
        # Load weights
        w = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
        
        # Apply RMSNorm
        y = (x * inv_rms) * w
        
        # Store output
        output_offset = batch_idx * hidden_size + cols
        tl.store(output_ptr + output_offset, y.to(tl.bfloat16), mask=mask)


def run(hidden_states, weight):
    # Handle device management
    original_device = hidden_states.device
    
    if not torch.cuda.is_available() and (hidden_states.is_cuda or weight.is_cuda):
        raise RuntimeError("CUDA is not available but GPU tensors were provided")
    
    # Move tensors to GPU if needed
    if torch.cuda.is_available():
        if not hidden_states.is_cuda:
            hidden_states = hidden_states.cuda()
        if not weight.is_cuda:
            weight = weight.cuda()
    else:
        raise RuntimeError("CUDA is required for Triton kernels")
    
    batch_size, hidden_size = hidden_states.shape
    
    # Check constants
    assert hidden_size == 7168, f"hidden_size must be 7168, got {hidden_size}"
    
    # Allocate output tensor
    output = torch.empty_like(hidden_states, device=hidden_states.device)
    
    # Constants
    EPS = 1e-6
    BLOCK_SIZE = 1024  # Optimized for B200 GPU
    
    # Launch kernel
    grid = (batch_size,)
    
    # Use optimized kernel for B200
    rmsnorm_kernel_optimized[grid](
        hidden_states,
        weight,
        output,
        hidden_size,
        EPS,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    
    # Move result back to original device if needed
    if original_device != output.device:
        output = output.to(original_device)
    
    return output