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 indices: torch.Tensor,7 local_shard: torch.Tensor,8) -> torch.Tensor:9 rank = dist.get_rank()10 world_size = dist.get_world_size()11 shard_size = local_shard.shape[0]12 embed_dim = local_shard.shape[1]13 14 indices = indices.contiguous().to(torch.cuda.current_device())15 16 target_ranks = indices // shard_size17 18 send_indices_list = [indices[target_ranks == r] for r in range(world_size)]19 send_counts = torch.tensor([len(idx) for idx in send_indices_list], dtype=torch.long, device='cuda')20 21 recv_counts = torch.zeros(world_size, dtype=torch.long, device='cuda')22 dist.all_to_all_single(recv_counts, send_counts)23 24 non_empty_lists = [idx_list for idx_list in send_indices_list if len(idx_list) > 0]25 if non_empty_lists:26 flat_send_indices = torch.cat(non_empty_lists)27 else:28 flat_send_indices = torch.empty(0, dtype=torch.long, device='cuda')29 30 total_recv = recv_counts.sum().item()31 total_send = send_counts.sum().item()32 received_indices = torch.empty(total_recv, dtype=torch.long, device='cuda')33 34 if total_recv > 0 or total_send > 0:35 dist.all_to_all_single(36 received_indices, 37 flat_send_indices,38 output_split_sizes=recv_counts.tolist(),39 input_split_sizes=send_counts.tolist()40 )41 42 if total_recv > 0:43 local_lookup_indices = received_indices - (rank * shard_size)44 local_lookup_indices = torch.clamp(local_lookup_indices, 0, shard_size - 1)45 retrieved_vectors = local_shard[local_lookup_indices]46 else:47 retrieved_vectors = torch.empty((0, embed_dim), dtype=local_shard.dtype, device='cuda')48 49 output_vectors = torch.empty((len(indices), embed_dim), dtype=local_shard.dtype, device='cuda')50 51 input_split_sizes = recv_counts.cpu().tolist()52 output_split_sizes = send_counts.cpu().tolist()53 54 if len(indices) > 0 or retrieved_vectors.numel() > 0:55 dist.all_to_all_single(56 output_vectors,57 retrieved_vectors,58 output_split_sizes=output_split_sizes,59 input_split_sizes=input_split_sizes60 )61 62 return output_vectors63