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
1from typing import Optional2 3import torch4import torch.distributed as dist5 6 7def _broadcast_data(8 rank: int,9 world_size: int,10 data: torch.Tensor,11 group: dist.ProcessGroup,12) -> torch.Tensor:13 if world_size == 1:14 return data15 16 sizes = torch.zeros(world_size, dtype=torch.long, device=data.device)17 sizes[rank] = data.shape[0]18 dist.all_reduce(sizes, op=dist.ReduceOp.SUM, group=group)19 20 send_splits = [data.shape[0] for _ in range(world_size)]21 recv_splits = sizes.to("cpu").tolist()22 send = data.repeat(*([world_size] + [1] * (data.ndim - 1))).contiguous()23 recv = torch.empty(24 (int(sizes.sum().item()), *data.shape[1:]),25 dtype=data.dtype,26 device=data.device,27 )28 dist.all_to_all_single(29 recv,30 send,31 output_split_sizes=recv_splits,32 input_split_sizes=send_splits,33 group=group,34 )35 return recv36 37 38def _calc_ranking(pos_score: torch.Tensor, neg_score: torch.Tensor) -> torch.Tensor:39 scores = torch.cat([pos_score.view(-1, 1), neg_score], dim=1)40 _, indices = torch.sort(torch.sigmoid(scores), dim=1, descending=True)41 return torch.nonzero(indices == 0)[:, 1].view(-1).detach() + 142 43 44@torch.no_grad()45def solution(46 local_pos_scores: torch.Tensor,47 local_neg_scores: torch.Tensor,48 group: Optional[dist.ProcessGroup] = None,49) -> torch.Tensor:50 group = group or dist.group.WORLD51 rank = dist.get_rank(group)52 world_size = dist.get_world_size(group)53 54 pos_scores = _broadcast_data(rank, world_size, local_pos_scores, group)55 neg_scores = _broadcast_data(rank, world_size, local_neg_scores, group)56 return _calc_ranking(pos_scores, neg_scores)