Datasetpublic

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.

sourceHugging Faceapache-2.0updated 4mo agoView on Hugging Face

0likes263downloads

rootcodereference33_ulysses_all_to_all_tensor_primitive.py

33_ulysses_all_to_all_tensor_primitive.py23 linesDownload Raw Back to reference

1from typing import Optional2 3import torch4import torch.distributed as dist5 6 7def solution(8    x: torch.Tensor,9    scatter_dim: int,10    gather_dim: int,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 x.contiguous()17 18    x = x.contiguous()19    input_list = [t.contiguous() for t in torch.tensor_split(x, world_size, scatter_dim)]20    output_list = [torch.empty_like(input_list[0]) for _ in range(world_size)]21    dist.all_to_all(output_list, input_list, group=group)22    return torch.cat(output_list, dim=gather_dim).contiguous()23