submission 512775

# submission_matmul_v2.py
# Matrix multiplication C = A @ B for float16 matrices.
# Interface: custom_kernel((A, B, C)) -> C
#   A: (M, K) float16,  B: (K, N) float16,  C: (M, N) float16 output
#
# Set CUBLAS_WORKSPACE_CONFIG so cuBLAS can run in deterministic mode when
# check_implementation calls the reference (a @ b in float16 under
# DeterministicContext).  Must be set before any cuBLAS handle is created.

import os
os.environ.setdefault('CUBLAS_WORKSPACE_CONFIG', ':4096:8')

import torch
import triton
import triton.language as tl
from task import input_t, output_t


@triton.autotune(
    configs=[
        triton.Config({'BM': 128, 'BN': 256, 'BK': 64,  'GROUP': 8}, num_warps=8,  num_stages=4),
        triton.Config({'BM': 64,  'BN': 256, 'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=4),
        triton.Config({'BM': 128, 'BN': 128, 'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=4),
        triton.Config({'BM': 128, 'BN': 64,  'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=4),
        triton.Config({'BM': 64,  'BN': 128, 'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=4),
        triton.Config({'BM': 128, 'BN': 32,  'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=4),
        triton.Config({'BM': 64,  'BN': 32,  'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=2),
        triton.Config({'BM': 32,  'BN': 64,  'BK': 32,  'GROUP': 8}, num_warps=4,  num_stages=2),
        triton.Config({'BM': 128, 'BN': 256, 'BK': 128, 'GROUP': 8}, num_warps=8,  num_stages=3),
        triton.Config({'BM': 256, 'BN': 128, 'BK': 128, 'GROUP': 8}, num_warps=8,  num_stages=3),
        triton.Config({'BM': 256, 'BN': 64,  'BK': 128, 'GROUP': 8}, num_warps=8,  num_stages=3),
        triton.Config({'BM': 64,  'BN': 256, 'BK': 128, 'GROUP': 8}, num_warps=8,  num_stages=3),
        triton.Config({'BM': 64,  'BN': 64,  'BK': 64,  'GROUP': 8}, num_warps=4,  num_stages=3),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M: int, N: int, K: int,
    stride_am: int, stride_ak: int,
    stride_bk: int, stride_bn: int,
    stride_cm: int, stride_cn: int,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr, GROUP: tl.constexpr,
):
    # Grouped tile ordering for better L2 cache reuse (from Triton tutorial).
    # group_size_m is clamped so the last group is handled correctly when
    # num_m is not a multiple of GROUP.
    pid   = tl.program_id(0)
    num_m = tl.cdiv(M, BM)
    num_n = tl.cdiv(N, BN)

    num_in_group = GROUP * num_n
    group_id     = pid // num_in_group
    first_m      = group_id * GROUP
    group_size_m = tl.minimum(num_m - first_m, GROUP)   # ← key fix

    pid_in_group = pid % num_in_group
    m_id = first_m + (pid_in_group % group_size_m)
    n_id = pid_in_group // group_size_m

    offs_m = (m_id * BM + tl.arange(0, BM)) % M
    offs_n = (n_id * BN + tl.arange(0, BN)) % N
    offs_k = tl.arange(0, BK)

    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((BM, BN), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BK)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BK, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BK, other=0.0)
        acc = tl.dot(a, b, acc)
        a_ptrs += BK * stride_ak
        b_ptrs += BK * stride_bk

    c_m    = m_id * BM + tl.arange(0, BM)
    c_n    = n_id * BN + tl.arange(0, BN)
    c_mask = (c_m[:, None] < M) & (c_n[None, :] < N)
    c_ptrs = c_ptr + stride_cm * c_m[:, None] + stride_cn * c_n[None, :]
    tl.store(c_ptrs, acc.to(tl.float16), mask=c_mask)


def custom_kernel(data: input_t) -> output_t:
    A, B, C = data
    M, K  = A.shape
    K2, N = B.shape
    assert K == K2
    grid = lambda meta: (triton.cdiv(M, meta['BM']) * triton.cdiv(N, meta['BN']),)
    _matmul_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),
    )
    return C