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@torch.no_grad()5def solution(6 A: torch.Tensor,7 B: torch.Tensor,8) -> torch.Tensor:9 rank = dist.get_rank()10 world_size = dist.get_world_size()11 12 M, K = A.shape13 K_B, N_local = B.shape14 15 A = A.contiguous()16 B = B.contiguous()17 C_local = torch.matmul(A, B)18 19 C_gathered = [torch.zeros_like(C_local) for _ in range(world_size)]20 dist.all_gather(C_gathered, C_local)21 22 C = torch.cat(C_gathered, dim=1)23 24 return C25