claude-opus-4-1 / triton48d048

import torch
import triton
import triton.language as tl
import math

@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,
    GROUP_M: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)
    
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_bn[:, None] * stride_bn + offs_k[None, :] * stride_bk)

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        mask_k = (k * BLOCK_K + offs_k) < K
        
        a = tl.load(a_ptrs, mask=(offs_am[:, None] < M) & mask_k[None, :], other=0.0)
        b = tl.load(b_ptrs, mask=(offs_bn[:, None] < N) & mask_k[None, :], other=0.0)
        
        accumulator = tl.dot(a, tl.trans(b), accumulator)
        
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    
    c = accumulator.to(tl.float16)
    tl.store(c_ptrs, c, mask=c_mask)

def run(A, B):
    # Handle device management
    original_device_A = A.device
    original_device_B = B.device
    
    if not A.is_cuda:
        if torch.cuda.is_available():
            A = A.cuda()
        else:
            raise RuntimeError("CUDA is not available for GPU tensor operations")
    
    if not B.is_cuda:
        if torch.cuda.is_available():
            B = B.cuda()
        else:
            raise RuntimeError("CUDA is not available for GPU tensor operations")
    
    # Get dimensions
    M, K_A = A.shape
    N, K_B = B.shape
    
    assert K_A == K_B, f"Dimension mismatch: A has K={K_A}, B has K={K_B}"
    K = K_A
    
    # Ensure inputs are float16
    A = A.to(torch.float16)
    B = B.to(torch.float16)
    
    # Allocate output
    C = torch.empty((M, N), device=A.device, dtype=torch.float16)
    
    # Configure kernel parameters for B200
    BLOCK_M = 128
    BLOCK_N = 128
    BLOCK_K = 32
    GROUP_M = 8
    
    # Calculate grid
    num_blocks = triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N)
    
    # Launch kernel
    gemm_kernel[(num_blocks,)](
        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_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        GROUP_M=GROUP_M,
    )
    
    # Move result back to original device
    if original_device_A.type == 'cpu' and original_device_B.type == 'cpu':
        C = C.cpu()
    elif original_device_A != C.device:
        C = C.to(original_device_A)
    
    return C