claude-opus-4-1 / tritond18c66

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,
):
    # Program ID
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # Block starting positions
    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)
    
    # Initialize accumulator with float32 for better precision
    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):
        offs_k = k + tl.arange(0, BLOCK_SIZE_K)
        
        # Load A tile [BLOCK_SIZE_M, BLOCK_SIZE_K]
        a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
        a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        
        # Load B tile [BLOCK_SIZE_N, BLOCK_SIZE_K] - B is stored as [N, K]
        b_ptrs = b_ptr + (offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk)
        b_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # Perform matrix multiplication: A @ B.T
        # a is [BLOCK_SIZE_M, BLOCK_SIZE_K]
        # b is [BLOCK_SIZE_N, BLOCK_SIZE_K]
        # We need to compute a @ b.T which gives [BLOCK_SIZE_M, BLOCK_SIZE_N]
        acc = tl.dot(a, tl.trans(b), acc)
    
    # Store result
    c_ptrs = c_ptr + (offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn)
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, acc.to(tl.float16), mask=c_mask)

def run(A, B):
    # Input validation
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available. This kernel requires a GPU.")
    
    # Store original devices
    a_device = A.device
    b_device = B.device
    
    # Move to GPU if needed
    if A.device.type != 'cuda':
        A = A.cuda()
    if B.device.type != 'cuda':
        B = B.cuda()
    
    # Ensure correct dtypes
    if A.dtype != torch.float16:
        A = A.to(torch.float16)
    if B.dtype != torch.float16:
        B = B.to(torch.float16)
    
    # Get dimensions
    M = A.shape[0]
    N = 4096
    K = 4096
    
    # Validate shapes
    assert A.shape == (M, K), f"Expected A shape ({M}, {K}), got {A.shape}"
    assert B.shape == (N, K), f"Expected B shape ({N}, {K}), got {B.shape}"
    
    # Ensure contiguous memory layout
    A = A.contiguous()
    B = B.contiguous()
    
    # Allocate output
    C = torch.empty((M, N), dtype=torch.float16, device=A.device)
    
    # Block sizes optimized for B200
    BLOCK_SIZE_M = 128
    BLOCK_SIZE_N = 128
    BLOCK_SIZE_K = 32
    
    # Grid dimensions
    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, 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 a_device.type != 'cuda':
        C = C.cpu()
    
    return C