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 List, Optional, Union2 3import torch4import torch.distributed as dist5 6 7def solution(8 local_tensor: torch.Tensor,9 input_split_sizes: Optional[Union[List[int], torch.Tensor]] = None,10 output_split_sizes: Optional[Union[List[int], torch.Tensor]] = None,11 group: Optional[dist.ProcessGroup] = None,12) -> torch.Tensor:13 group = group or dist.group.WORLD14 world_size = dist.get_world_size(group)15 if world_size == 1:16 return local_tensor.contiguous()17 18 local_tensor = local_tensor.contiguous()19 if output_split_sizes is None:20 output = torch.empty_like(local_tensor)21 else:22 out_size = sum(output_split_sizes) if isinstance(output_split_sizes, list) else int(output_split_sizes.sum().item())23 output = torch.empty(24 (out_size, local_tensor.size(1)),25 dtype=local_tensor.dtype,26 device=local_tensor.device,27 )28 dist.all_to_all_single(29 output,30 local_tensor,31 output_split_sizes=output_split_sizes,32 input_split_sizes=input_split_sizes,33 group=group,34 )35 return output36