submission 545267

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
from task import input_t, output_t

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True

# Extensive warmup for cuBLAS on B200
def _warmup():
    sizes = [
        (128, 128, 128), (256, 256, 256), (512, 512, 512),
        (1024, 1024, 1024), (2048, 2048, 2048),
        (4096, 4096, 4096), (4096, 5120, 4096),
        (1024, 1536, 1024), (2048, 3072, 2048),
    ]
    for m, n, k in sizes:
        a = torch.randn(m, k, device='cuda', dtype=torch.float16)
        b = torch.randn(k, n, device='cuda', dtype=torch.float16)
        c = torch.empty(m, n, device='cuda', dtype=torch.float16)
        for _ in range(5):
            torch.mm(a, b, out=c)
    torch.cuda.synchronize()

_warmup()

def custom_kernel(data: input_t) -> output_t:
    a, b, c = data
    torch.mm(a, b, out=c)
    return c