claude-opus-4-1-20250805 / triton0a753b

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,
):
    """GEMM kernel optimized for B200 GPU with N=5120, K=2048."""
    # Get program IDs
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # Compute block offsets
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    # Create masks for boundary conditions
    mask_m = offs_m < M
    mask_n = offs_n < N
    
    # Initialize pointers to A and B blocks
    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 K-loop
    for k in range(0, K, BLOCK_SIZE_K):
        # Load A and B blocks with boundary checking
        mask_k = offs_k < K - k
        a_mask = mask_m[:, None] & mask_k[None, :]
        b_mask = mask_n[:, None] & mask_k[None, :]
        
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # Compute matrix multiplication for this block
        acc += tl.dot(a, tl.trans(b))
        
        # Advance pointers
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    
    # Convert accumulator to fp16 and store result
    c = acc.to(tl.float16)
    
    # Compute output pointer and mask
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm < M)[:, None] & (offs_cn < N)[None, :]
    
    tl.store(c_ptrs, c, mask=c_mask)


def run(*args, **kwargs):
    """Entry point function with complete device management."""
    # Handle both args and kwargs
    if len(args) >= 2:
        A, B = args[0], args[1]
    else:
        A = kwargs.get('A')
        B = kwargs.get('B')
    
    if A is None or B is None:
        raise ValueError("Missing required arguments A and B")
    
    # Store original device information
    orig_device_A = A.device
    orig_device_B = B.device
    
    # Check if CUDA is available for GPU operations
    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")
    
    # Move tensors to GPU if needed
    if torch.cuda.is_available():
        if not A.is_cuda:
            A = A.cuda()
        if not B.is_cuda:
            B = B.cuda()
    
    # Validate input shapes and dtypes
    assert A.dtype == torch.float16, f"Expected A to be float16, got {A.dtype}"
    assert B.dtype == torch.float16, f"Expected B to be float16, got {B.dtype}"
    assert A.shape[1] == 2048, f"Expected A.shape[1] == 2048, got {A.shape[1]}"
    assert B.shape[0] == 5120, f"Expected B.shape[0] == 5120, got {B.shape[0]}"
    assert B.shape[1] == 2048, f"Expected B.shape[1] == 2048, got {B.shape[1]}"
    
    M, K = A.shape
    N = B.shape[0]
    
    # Allocate output tensor
    C = torch.empty((M, N), dtype=torch.float16, device=A.device)
    
    # Define block sizes optimized for B200 GPU
    # B200 has large shared memory and high throughput
    BLOCK_SIZE_M = 128
    BLOCK_SIZE_N = 128
    BLOCK_SIZE_K = 64
    
    # Compute grid dimensions
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']),
        triton.cdiv(N, META['BLOCK_SIZE_N']),
    )
    
    # Launch kernel
    gemm_kernel[grid](
        A, B, C,
        M, N, K,
        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 orig_device_A.type == 'cpu':
        C = C.cpu()
    
    return C