claude-opus-4-1-20250805 / tritonfaa4d0

import torch
import triton
import triton.language as tl


@triton.jit
def gemm_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,
):
    """
    Compute C = A @ B.T where B is stored as [N, K]
    """
    # Get program ids
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # Create block pointers
    rm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    rn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    rk = tl.arange(0, BLOCK_SIZE_K)
    
    # Create pointers to A and B
    A = a_ptr + rm[:, None] * stride_am + rk[None, :] * stride_ak
    B = b_ptr + rn[:, None] * stride_bn + rk[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 tiles with boundary checks
        a_mask = (rm[:, None] < M) & ((k + rk[None, :]) < K)
        b_mask = (rn[:, None] < N) & ((k + rk[None, :]) < K)
        
        a = tl.load(A, mask=a_mask, other=0.0)
        b = tl.load(B, mask=b_mask, other=0.0)
        
        # Compute dot product (note: B is already transposed in memory layout)
        acc += tl.dot(a, tl.trans(b), allow_tf32=True)
        
        # Advance pointers
        A += BLOCK_SIZE_K * stride_ak
        B += BLOCK_SIZE_K * stride_bk
    
    # Convert accumulator to float16
    c = acc.to(tl.float16)
    
    # Store output with boundary checks
    C = c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
    c_mask = (rm[:, None] < M) & (rn[None, :] < N)
    tl.store(C, c, mask=c_mask)


def gemm_n2048_k4096(A, B):
    """
    Optimized GEMM for N=2048, K=4096 on B200 GPU
    A: [M, K] float16
    B: [N, K] float16
    Returns: C = A @ B.T [M, N] float16
    """
    # Get dimensions
    M, K_a = A.shape
    N, K_b = B.shape
    
    assert K_a == K_b == 4096, f"K dimension mismatch or not 4096: A.shape={A.shape}, B.shape={B.shape}"
    assert N == 2048, f"N dimension must be 2048, got {N}"
    assert A.dtype == torch.float16 and B.dtype == torch.float16, "Inputs must be float16"
    
    # Allocate output
    C = torch.empty((M, N), device=A.device, dtype=torch.float16)
    
    # Block sizes optimized for B200
    # B200 has large shared memory and high compute throughput
    BLOCK_SIZE_M = 128
    BLOCK_SIZE_N = 128
    BLOCK_SIZE_K = 64
    
    # Calculate grid dimensions
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']),
        triton.cdiv(N, META['BLOCK_SIZE_N']),
        1
    )
    
    # Launch kernel
    gemm_kernel[grid](
        A, B, C,
        M, N, K_a,
        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,
    )
    
    return C


def run(*args, **kwargs):
    """
    Entry point function with complete device management
    Handles both positional and keyword arguments
    """
    # Handle both positional and keyword arguments
    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 2 positional arguments (A, B) or keyword arguments 'A' and 'B'")
    
    # Check CUDA availability
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available. This kernel requires a GPU.")
    
    # Store original devices
    original_device_A = A.device
    original_device_B = B.device
    
    # Move tensors to GPU if needed
    if A.device.type == 'cpu':
        A = A.cuda()
    elif A.device.type != 'cuda':
        raise ValueError(f"Unsupported device type for tensor A: {A.device.type}")
    
    if B.device.type == 'cpu':
        B = B.cuda()
    elif B.device.type != 'cuda':
        raise ValueError(f"Unsupported device type for tensor B: {B.device.type}")
    
    # Ensure correct dtypes
    if A.dtype != torch.float16:
        A = A.to(torch.float16)
    if B.dtype != torch.float16:
        B = B.to(torch.float16)
    
    # Validate shapes
    if len(A.shape) != 2 or len(B.shape) != 2:
        raise ValueError(f"Expected 2D tensors, got A.shape={A.shape}, B.shape={B.shape}")
    
    if A.shape[1] != 4096:
        raise ValueError(f"Expected A to have K=4096, got {A.shape[1]}")
    
    if B.shape[0] != 2048 or B.shape[1] != 4096:
        raise ValueError(f"Expected B to have shape [2048, 4096], got {B.shape}")
    
    # Run the kernel
    C = gemm_n2048_k4096(A, B)
    
    # Move result back to original device (prioritize A's device)
    if original_device_A.type == 'cpu':
        C = C.cpu()
    
    return C