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.
0146
1from typing import List, Optional, Tuple2 3import torch4import torch.distributed as dist5 6 7def _preprocess_impl(8 expert_mask: torch.Tensor,9 num_experts: int,10 ep_group: dist.ProcessGroup,11) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]:12 ep_size = ep_group.size()13 num_local_experts = num_experts // ep_size14 rank = dist.get_rank(ep_group)15 num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2))16 17 input_splits = num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1).tolist()18 19 num_global_tokens_per_expert = torch.empty(20 ep_size,21 num_local_tokens_per_expert.size(0),22 dtype=num_local_tokens_per_expert.dtype,23 device=num_local_tokens_per_expert.device,24 )25 dist.all_gather_into_tensor(num_global_tokens_per_expert, num_local_tokens_per_expert, group=ep_group)26 27 start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts28 num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous()29 30 output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist()31 32 num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(33 torch.device("cpu"), non_blocking=True34 )35 num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to(36 torch.device("cpu"), non_blocking=True37 )38 39 return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert40 41 42def solution(43 expert_mask: torch.Tensor,44 num_experts: int,45 group: Optional[dist.ProcessGroup] = None,46) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]:47 group = group or dist.group.WORLD48 return _preprocess_impl(expert_mask=expert_mask, num_experts=num_experts, ep_group=group)49