Datasetpublic

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.

sourceHugging Faceapache-2.0updated 4mo agoView on Hugging Face

0likes266downloads

rootcodereference16_gemm_reducescatter.py

16_gemm_reducescatter.py25 linesDownload Raw Back to reference

1import torch2import torch.distributed as dist3 4@torch.no_grad()5def solution(6    A_local: torch.Tensor,7    B_local: torch.Tensor,8) -> torch.Tensor:9    rank = dist.get_rank()10    world_size = dist.get_world_size()11    12    M, K_local = A_local.shape13    K_B, N = B_local.shape14    15    M_local = M // world_size16    17    A_local = A_local.contiguous()18    B_local = B_local.contiguous()19    C_partial = torch.matmul(A_local, B_local)20    21    C_local = torch.empty((M_local, N), dtype=C_partial.dtype, device=C_partial.device)22    dist.reduce_scatter_tensor(C_local, C_partial, op=dist.ReduceOp.SUM)23    24    return C_local25