gpt-o3 / tritonfc01fa

import math
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,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    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)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (k + offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_n[None, :] < N) & (k + offs_k[:, None] < K), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    acc = acc.to(tl.float16)
    c_ptrs = C_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))


def _ensure_device(t):
    if t.is_cuda:
        return t
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA device required but not available.")
    return t.cuda()


def run(A, B):
    """
    Optimized GEMM: C = A @ B.T
    A: [M, 4096] float16
    B: [6144, 4096] float16
    Returns C: [M, 6144] float16
    """
    if A.dtype != torch.float16 or B.dtype != torch.float16:
        raise TypeError("Only float16 tensors are supported.")
    if A.shape[1] != 4096 or B.shape != (6144, 4096):
        raise ValueError("Incorrect tensor shapes for A or B.")

    orig_device_A = A.device
    orig_device_B = B.device

    A_gpu = _ensure_device(A.contiguous())
    B_gpu = _ensure_device(B.contiguous())

    M = A_gpu.shape[0]
    N = 6144
    K = 4096

    C_gpu = torch.empty((M, N), dtype=torch.float16, device=A_gpu.device)

    BLOCK_M = 128
    BLOCK_N = 128
    BLOCK_K = 32
    num_warps = 8
    num_ctas = 1  # per program_id—Triton handles grid dims below

    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))

    _gemm_kernel[grid](
        A_gpu, B_gpu, C_gpu,
        M, N, K,
        A_gpu.stride(0), A_gpu.stride(1),
        B_gpu.stride(0), B_gpu.stride(1),
        C_gpu.stride(0), C_gpu.stride(1),
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        num_warps=num_warps,
        num_ctas=num_ctas,
    )

    if not orig_device_A.type == 'cuda':
        C_out = C_gpu.cpu()
    else:
        C_out = C_gpu

    return C_out


if __name__ == "__main__":
    # Quick correctness test
    M_test = 256
    A_test = torch.randn((M_test, 4096), dtype=torch.float16)
    B_test = torch.randn((6144, 4096), dtype=torch.float16)
    C_ref = torch.matmul(A_test.cuda(), B_test.cuda().T).cpu()
    C_triton = run(A_test, B_test)
    assert torch.allclose(C_ref, C_triton, atol=1e-2, rtol=1e-2)
    print("Triton GEMM passed the correctness test.")