claude-opus-4-1-20250805 / tritona20c42

import torch
import triton
import triton.language as tl


@triton.jit
def gemm_n128_k2048_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
):
    """Optimized GEMM kernel for N=128, K=2048 configuration."""
    # Program ID
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # Block indices
    block_start_m = pid_m * BLOCK_SIZE_M
    block_start_n = pid_n * BLOCK_SIZE_N
    
    # Thread block offsets
    offs_m = block_start_m + tl.arange(0, BLOCK_SIZE_M)
    offs_n = block_start_n + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    # Pointers to first blocks of A and B
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk)
    
    # Initialize accumulator
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # Main loop over K dimension
    for k in range(0, K, BLOCK_SIZE_K):
        # Load blocks from A and B with boundary checks
        mask_m = offs_m < M
        mask_n = offs_n < N
        mask_k = (k + offs_k) < K
        
        a_block = tl.load(a_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0)
        b_block = tl.load(b_ptrs, mask=mask_n[:, None] & mask_k[None, :], other=0.0)
        
        # Compute dot product for this K block
        # B is transposed in memory access pattern
        acc += tl.dot(a_block, tl.trans(b_block))
        
        # Advance pointers to next K block
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    
    # Store result with boundary check
    c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc.to(tl.float16), mask=mask)


def run(*args, **kwargs):
    """Entry point function that handles device management and kernel execution."""
    # Handle both args and kwargs
    if len(args) == 2:
        A, B = args
    elif 'A' in kwargs and 'B' in kwargs:
        A = kwargs['A']
        B = kwargs['B']
    else:
        raise ValueError("Expected either (A, B) as positional args or as keyword args")
    
    # Check input shapes and dtypes
    assert A.ndim == 2 and B.ndim == 2, "Input tensors must be 2D"
    M, K_a = A.shape
    N, K_b = B.shape
    assert K_a == 2048 and K_b == 2048, f"Expected K=2048, got K_a={K_a}, K_b={K_b}"
    assert N == 128, f"Expected N=128, got N={N}"
    
    # Store original devices
    device_a = A.device
    device_b = B.device
    
    # Move to GPU if needed
    if not torch.cuda.is_available():
        if A.is_cuda or B.is_cuda:
            raise RuntimeError("CUDA is not available but GPU tensors were provided")
        raise RuntimeError("CUDA is not available for GPU computation")
    
    # Move CPU tensors to GPU
    if not A.is_cuda:
        A = A.cuda()
    if not B.is_cuda:
        B = B.cuda()
    
    # Ensure correct dtype
    if A.dtype != torch.float16:
        A = A.to(torch.float16)
    if B.dtype != torch.float16:
        B = B.to(torch.float16)
    
    # Ensure tensors are on the same device
    if A.device != B.device:
        B = B.to(A.device)
    
    # Allocate output tensor
    C = torch.empty((M, N), dtype=torch.float16, device=A.device)
    
    # Configure kernel parameters optimized for B200
    # B200 has large shared memory and high compute throughput
    BLOCK_SIZE_M = 128
    BLOCK_SIZE_N = 64  # N=128, so we use 2 blocks
    BLOCK_SIZE_K = 64  # Tile K dimension for better cache usage
    
    # Compute grid dimensions
    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
    
    # Launch kernel
    gemm_n128_k2048_kernel[grid](
        A, B, C,
        M, N, 2048,
        A.stride(0), A.stride(1),
        B.stride(0), B.stride(1),
        C.stride(0), C.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
    )
    
    # Move result back to original device if needed
    if device_a.type == 'cpu':
        C = C.cpu()
    elif device_a != C.device:
        C = C.to(device_a)
    
    return C