gpt-5 / triton1bd4a7

import torch
import triton
import triton.language as tl


@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 128}, num_warps=8, num_stages=5),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_warps=8, num_stages=5),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_warps=8, num_stages=5),
        triton.Config({'BLOCK_M': 64,  'BLOCK_N': 256, 'BLOCK_K': 128}, num_warps=4, num_stages=5),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_warps=4, num_stages=5),
        triton.Config({'BLOCK_M': 256, 'BLOCK_N': 256, 'BLOCK_K': 64},  num_warps=16, num_stages=4),
    ],
    key=['M']
)
@triton.jit
def _gemm_mk_kn_to_mn_kernel(
    A_ptr, B_ptr, C_ptr,
    M, N, K,
    stride_am, stride_ak,   # A: [M, K]
    stride_bn, stride_bk,   # B: [N, K] but loaded as [K, N] via strides
    stride_cm, stride_cn,   # C: [M, N]
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    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)
    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_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

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

    k = 0
    while k < 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=(k + offs_k[:, None] < K) & (offs_n[None, :] < N),
            other=0.0
        )
        acc += tl.dot(a, b)

        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
        k += BLOCK_K

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


def run(A, B):
    if not isinstance(A, torch.Tensor) or not isinstance(B, torch.Tensor):
        raise TypeError("Inputs A and B must be torch.Tensor")

    if A.ndim != 2 or B.ndim != 2:
        raise ValueError("A and B must be 2D tensors")

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

    # Constants from specification
    REQUIRED_N = 4096
    REQUIRED_K = 14336

    if K_a != REQUIRED_K or K_b != REQUIRED_K:
        raise ValueError(f"K must be {REQUIRED_K}. Got A.shape[1]={K_a}, B.shape[1]={K_b}")
    if N_b != REQUIRED_N:
        raise ValueError(f"N must be {REQUIRED_N}. Got B.shape[0]={N_b}")

    if A.dtype != torch.float16 or B.dtype != torch.float16:
        raise TypeError("A and B must be of dtype torch.float16")

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required to run this Triton kernel, but torch.cuda.is_available() is False.")

    # Pick a CUDA device
    if A.is_cuda:
        cuda_dev = A.device
    elif B.is_cuda:
        cuda_dev = B.device
    else:
        cuda_dev = torch.device('cuda')

    # Preserve original devices without modifying inputs
    dev_A_orig = A.device
    dev_B_orig = B.device

    # Move to chosen CUDA device if needed
    A_gpu = A.to(device=cuda_dev, non_blocking=True) if A.device != cuda_dev else A
    B_gpu = B.to(device=cuda_dev, non_blocking=True) if B.device != cuda_dev else B

    # Shapes
    M = A_gpu.shape[0]
    N = B_gpu.shape[0]
    K = A_gpu.shape[1]

    # Allocate output on GPU
    C_gpu = torch.empty((M, N), dtype=torch.float16, device=cuda_dev)

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

    # Launch kernel
    _gemm_mk_kn_to_mn_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),
    )

    # Move result back to the original device of A
    return C_gpu.to(dev_A_orig, non_blocking=True)