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.
0263
1from typing import List, Optional2 3import torch4import torch.distributed as dist5 6 7def _shift(chunks: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]:8 cutoff = len(chunks) - dist.get_rank(group)9 return chunks[cutoff:] + chunks[:cutoff]10 11 12def _all_to_all(13 outputs: List[torch.Tensor],14 inputs: List[torch.Tensor],15 group: dist.ProcessGroup,16) -> None:17 outputs = _shift(list(outputs), group)18 inputs = _shift(list(inputs), group)19 if outputs and outputs[0].is_cuda:20 dist.all_to_all(outputs, inputs, group=group)21 return22 23 output_splits = [out.size(0) for out in outputs]24 input_splits = [inp.size(0) for inp in inputs]25 flat_out = torch.cat(outputs) if outputs else torch.empty(0)26 flat_in = torch.cat(inputs) if inputs else torch.empty(0)27 dist.all_to_all_single(28 flat_out,29 flat_in,30 output_split_sizes=output_splits,31 input_split_sizes=input_splits,32 group=group,33 )34 for out, temp in zip(outputs, flat_out.split(output_splits)):35 out.copy_(temp)36 37 38@torch.no_grad()39def solution(40 grad_output: torch.Tensor,41 seed_inverse_ids: torch.Tensor,42 seed_size: int,43 counts_sent: List[int],44 counts_received: List[int],45 group: Optional[dist.ProcessGroup] = None,46) -> torch.Tensor:47 group = group or dist.group.WORLD48 out = grad_output.new_empty((sum(counts_received),) + grad_output.shape[1:])49 _all_to_all(50 list(torch.split(out, counts_received)),51 list(torch.split(grad_output, counts_sent)),52 group,53 )54 55 idx = torch.empty((2, out.shape[0]), dtype=torch.int64, device=grad_output.device)56 idx[0] = seed_inverse_ids57 idx[1] = torch.arange(out.shape[0], device=grad_output.device)58 coo = torch.sparse_coo_tensor(59 idx,60 torch.ones(idx.shape[1], dtype=grad_output.dtype, device=idx.device),61 size=(seed_size, idx.shape[1]),62 )63 return torch.sparse.mm(coo, out)64