gpt-5 / triton14f133

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


@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=8, num_stages=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_warps=8, num_stages=4),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=8, num_stages=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_warps=8, num_stages=4),
    ],
    key=['M'],
)
@triton.jit
def _gemm_n4096_k4096_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,
):
    tl.static_assert(BLOCK_K % 16 == 0, "BLOCK_K must be a multiple of 16 for tensor cores")
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

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

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

    M_mask = offs_m[:, None] < M
    N_mask = offs_n[None, :] < N

    for k0 in range(0, K, BLOCK_K):
        offs_k = k0 + tl.arange(0, BLOCK_K)

        # Pointers
        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)

        a = tl.load(a_ptrs, mask=M_mask & (offs_k[None, :] < K), other=0.0)
        b = tl.load(b_ptrs, mask=(offs_n[:, None] < N) & (offs_k[None, :] < K), other=0.0)

        acc += tl.dot(a, tl.trans(b))

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


def run(A, B, **kwargs):
    """
    Compute C = A @ B.T where:
    - A: [M, 4096] float16
    - B: [4096, 4096] float16
    Returns C: [M, 4096] float16
    """
    if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor):
        raise TypeError("A and B must be torch.Tensor")

    if A.ndim != 2 or B.ndim != 2:
        raise ValueError(f"Expected 2D tensors, got A.ndim={A.ndim}, B.ndim={B.ndim}")

    M, K_a = A.shape
    N_b, K_b = B.shape

    if K_a != 4096 or K_b != 4096 or N_b != 4096:
        raise ValueError(f"Invalid shapes: A is {A.shape}, B is {B.shape}; expected A: [M,4096], B: [4096,4096]")

    # Dtype checks/conversions
    if A.dtype != torch.float16:
        A = A.to(torch.float16)
    if B.dtype != torch.float16:
        B = B.to(torch.float16)

    # Device management
    orig_dev_A = A.device
    orig_dev_B = B.device
    any_cuda_input = (A.is_cuda or B.is_cuda)

    if not torch.cuda.is_available():
        if any_cuda_input:
            raise RuntimeError("CUDA is not available but one or more inputs are CUDA tensors.")
        # Triton requires CUDA; no CPU fallback provided
        raise RuntimeError("CUDA is required to run this Triton kernel, but no CUDA device is available.")

    # Choose compute device
    if A.is_cuda:
        compute_device = A.device
    elif B.is_cuda:
        compute_device = B.device
    else:
        compute_device = torch.device(f"cuda:{torch.cuda.current_device()}")

    # Move to compute device and make contiguous for optimal access
    A_dev = A.to(device=compute_device, non_blocking=True).contiguous()
    B_dev = B.to(device=compute_device, non_blocking=True).contiguous()

    # Allocate output on compute device
    N = 4096
    K = 4096
    C_dev = torch.empty((M, N), dtype=torch.float16, device=compute_device)

    # Kernel launch parameters
    def grid(meta):
        return (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N']))

    # Call kernel
    _gemm_n4096_k4096_kernel[grid](
        A_dev, B_dev, C_dev,
        M, N, K,
        A_dev.stride(0), A_dev.stride(1),
        B_dev.stride(0), B_dev.stride(1),
        C_dev.stride(0), C_dev.stride(1),
    )

    # Decide output device: preserve original locations; if both were CPU, return CPU; otherwise prefer A's device if CUDA, else B's
    if orig_dev_A.type == 'cpu' and orig_dev_B.type == 'cpu':
        out_device = torch.device('cpu')
    elif orig_dev_A.type == 'cuda':
        out_device = orig_dev_A
    elif orig_dev_B.type == 'cuda':
        out_device = orig_dev_B
    else:
        out_device = torch.device('cpu')

    C_out = C_dev if C_dev.device == out_device else C_dev.to(out_device, non_blocking=True)
    return C_out