claude-opus-4-1-20250805 / triton9c959c

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_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
):
    # Program ID and grid dimensions
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # Compute block boundaries
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    
    # Initialize accumulator
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    
    # Compute pointers to first blocks of A and B
    a_base = a_ptr + rm[:, None] * stride_am
    b_base = b_ptr + rn[:, None] * stride_bn
    
    # Main loop over K dimension
    for k in range(0, K, BLOCK_K):
        rk = k + tl.arange(0, BLOCK_K)
        
        # Load A block with masking
        a_mask = (rm[:, None] < M) & (rk[None, :] < K)
        a = tl.load(a_base + rk[None, :] * stride_ak, mask=a_mask, other=0.0)
        
        # Load B block with masking
        b_mask = (rn[:, None] < N) & (rk[None, :] < K)
        b = tl.load(b_base + rk[None, :] * stride_bk, mask=b_mask, other=0.0)
        
        # Accumulate dot product
        acc += tl.dot(a, tl.trans(b), allow_tf32=True)
    
    # Write result with masking
    c_mask = (rm[:, None] < M) & (rn[None, :] < N)
    c = c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
    tl.store(c, acc.to(tl.float16), mask=c_mask)

def run(*args, **kwargs):
    """Entry point function for GEMM operation."""
    # Handle both positional and keyword arguments
    if len(args) == 2:
        A, B = args
    elif len(args) == 0 and 'A' in kwargs and 'B' in kwargs:
        A = kwargs['A']
        B = kwargs['B']
    else:
        raise ValueError("Expected exactly 2 arguments (A, B)")
    
    # Store original device
    original_device_A = A.device
    original_device_B = B.device
    
    # Move to GPU if needed
    if A.device.type == 'cpu':
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available but CPU tensors were provided")
        A = A.cuda()
    elif A.device.type != 'cuda':
        raise ValueError(f"Unsupported device type: {A.device.type}")
    
    if B.device.type == 'cpu':
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA is not available but CPU tensors were provided")
        B = B.cuda()
    elif B.device.type != 'cuda':
        raise ValueError(f"Unsupported device type: {B.device.type}")
    
    # 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.dim() == 2, f"Expected A to be 2D, got {A.dim()}D"
    assert B.dim() == 2, f"Expected B to be 2D, got {B.dim()}D"
    
    M, K_A = A.shape
    N, K_B = B.shape
    
    assert K_A == 7168, f"Expected K dimension of A to be 7168, got {K_A}"
    assert K_B == 7168, f"Expected K dimension of B to be 7168, got {K_B}"
    assert N == 256, f"Expected N dimension of B to be 256, got {N}"
    
    # Allocate output tensor on GPU
    C = torch.empty((M, N), dtype=torch.float16, device=A.device)
    
    # Configure block sizes optimized for B200
    # B200 has high memory bandwidth and compute capability
    BLOCK_M = 128
    BLOCK_N = 128
    BLOCK_K = 64
    
    # Calculate grid dimensions
    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    
    # Launch kernel
    gemm_kernel[grid](
        a_ptr=A,
        b_ptr=B,
        c_ptr=C,
        M=M,
        N=N,
        K=7168,
        stride_am=A.stride(0),
        stride_ak=A.stride(1),
        stride_bn=B.stride(0),
        stride_bk=B.stride(1),
        stride_cm=C.stride(0),
        stride_cn=C.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
    )
    
    # Move result back to original device if needed
    if original_device_A.type == 'cpu':
        C = C.cpu()
    
    return C