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 7@torch.no_grad()8def solution(9 local_embedding_shard: torch.Tensor,10 input_node_ids: torch.Tensor,11 proj_matrix: torch.Tensor,12 num_total_nodes: int,13 group: Optional[dist.ProcessGroup] = None,14) -> torch.Tensor:15 group = group or dist.group.WORLD16 world_size = dist.get_world_size(group)17 rank = dist.get_rank(group)18 shard_size = (num_total_nodes + world_size - 1) // world_size19 embed_dim = local_embedding_shard.shape[1]20 num_queries = input_node_ids.shape[0]21 22 owner = (input_node_ids // shard_size).clamp(max=world_size - 1)23 sort_idx = torch.argsort(owner, stable=True)24 sorted_ids = input_node_ids[sort_idx]25 sorted_owner = owner[sort_idx]26 27 send_counts = torch.zeros(28 world_size, dtype=torch.long, device=input_node_ids.device29 )30 send_counts.scatter_add_(0, sorted_owner, torch.ones_like(sorted_owner))31 recv_counts = torch.empty_like(send_counts)32 dist.all_to_all_single(recv_counts, send_counts, group=group)33 34 send_splits = send_counts.to("cpu").tolist()35 recv_splits = recv_counts.to("cpu").tolist()36 recv_ids = torch.empty(37 int(recv_counts.sum().item()),38 dtype=input_node_ids.dtype,39 device=input_node_ids.device,40 )41 dist.all_to_all_single(42 recv_ids,43 sorted_ids,44 output_split_sizes=recv_splits,45 input_split_sizes=send_splits,46 group=group,47 )48 49 local_ids = (recv_ids - rank * shard_size).long()50 fetched = local_embedding_shard[local_ids]51 52 gathered = torch.empty(53 (num_queries, embed_dim), dtype=fetched.dtype, device=fetched.device54 )55 dist.all_to_all_single(56 gathered,57 fetched,58 output_split_sizes=send_splits,59 input_split_sizes=recv_splits,60 group=group,61 )62 emb = gathered[torch.argsort(sort_idx, stable=True)]63 64 return emb @ proj_matrix65