willychan21/ParallelKernelBench_Problems
ParallelKernelBench (benchmark) Reference problems for ParallelKernelBench: a benchmark for LLM-generated multi-GPU CUDA kernels. This dataset contains 87 reference implementations in reference/ and the input tensor specification in utils/input_output_tensors.py. Files Path Description data/problems.parquet One row per problem (tabular access) reference/*.py Reference solution() implementations utils/input_output_tensors.py Input/output tensor… See the full description on the dataset page: https://huggingface.co/datasets/willychan21/ParallelKernelBench_Problems.
0266
1import torch2import torch.distributed as dist3 4 5@torch.no_grad()6def solution(X_hat: torch.Tensor, dY: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:7 d_beta = dY.sum(dim=0)8 d_gamma = (dY * X_hat).sum(dim=0)9 dist.all_reduce(d_beta, op=dist.ReduceOp.SUM)10 dist.all_reduce(d_gamma, op=dist.ReduceOp.SUM)11 return d_gamma, d_beta12 