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 __future__ import annotations2 3import torch4import torch.distributed as dist5import torch.nn.functional as F6from torch import Tensor7 8 9def _make_tp_fsdp_groups(n_tp: int, n_fsdp: int, rank: int):10 tp_group = None11 fsdp_group = None12 for j in range(n_fsdp):13 ranks = [j * n_tp + ii for ii in range(n_tp)]14 g = dist.new_group(ranks)15 if rank in ranks:16 tp_group = g17 for i in range(n_tp):18 ranks = [jj * n_tp + i for jj in range(n_fsdp)]19 g = dist.new_group(ranks)20 if rank in ranks:21 fsdp_group = g22 assert tp_group is not None and fsdp_group is not None23 return tp_group, fsdp_group24 25 26def _gather_fsdp_concat_dim0(shard: Tensor, fsdp_group, parts: int) -> Tensor:27 lst = [torch.empty_like(shard) for _ in range(parts)]28 dist.all_gather(lst, shard.contiguous(), group=fsdp_group)29 return torch.cat(lst, dim=0)30 31 32def _gather_fsdp_concat_dim1(shard: Tensor, fsdp_group, parts: int) -> Tensor:33 lst = [torch.empty_like(shard) for _ in range(parts)]34 dist.all_gather(lst, shard.contiguous(), group=fsdp_group)35 return torch.cat(lst, dim=1)36 37 38@torch.no_grad()39def solution(40 x_local: Tensor,41 W1_shard: Tensor,42 W2_shard: Tensor,43 W3_shard: Tensor,44 n_tp: int,45 n_fsdp: int,46) -> Tensor:47 world_size = dist.get_world_size()48 rank = dist.get_rank()49 assert world_size == n_tp * n_fsdp, (50 f"world_size ({world_size}) must equal n_tp * n_fsdp ({n_tp} * {n_fsdp})"51 )52 53 tp_group, fsdp_group = _make_tp_fsdp_groups(n_tp, n_fsdp, rank)54 55 W1 = _gather_fsdp_concat_dim0(W1_shard, fsdp_group, n_fsdp)56 W2 = _gather_fsdp_concat_dim0(W2_shard, fsdp_group, n_fsdp)57 W3 = _gather_fsdp_concat_dim1(W3_shard, fsdp_group, n_fsdp)58 59 x1 = x_local @ W160 x2 = x_local @ W261 z = F.silu(x1) * x262 y_partial = z @ W363 64 y = y_partial.clone()65 dist.all_reduce(y, op=dist.ReduceOp.SUM, group=tp_group)66 return y67