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 Optional, Sequence2 3import torch4import torch.distributed as dist5 6 7_COEFFICIENTS: dict[str, Sequence[tuple[float, float, float]]] = {8 "simple": ((3.4445, -4.7750, 2.0315),),9 "quintic": (10 (4.0848, -6.8946, 2.9270),11 (3.9505, -6.3029, 2.6377),12 (3.7418, -5.5913, 2.3037),13 (2.8769, -3.1427, 1.2046),14 (2.8366, -3.0525, 1.2012),15 ),16 "polar_express": (17 (8.2051, -22.9019, 16.4607),18 (4.0664, -2.8612, 0.5184),19 (3.9096, -2.8234, 0.5250),20 (3.2856, -2.4647, 0.5074),21 (2.2779, -1.6447, 0.4162),22 (1.8726, -1.2307, 0.3585),23 (1.8564, -1.2132, 0.3568),24 (1.8750, -1.2500, 0.3750),25 ),26 "aol": (27 (4.0098, -7.0585, 2.4635),28 (3.4585, -5.5479, 2.5959),29 (2.7573, -3.2939, 1.4254),30 (2.7215, -3.0494, 1.3169),31 ),32}33 34 35def _coefficient_at(36 coefficients: Sequence[tuple[float, float, float]],37 step: int,38) -> tuple[float, float, float]:39 return coefficients[step % len(coefficients)]40 41 42def _distributed_normalize(43 x: torch.Tensor,44 group: dist.ProcessGroup,45 eps: float = 1e-7,46) -> torch.Tensor:47 norm_sq = (x * x).sum()48 dist.all_reduce(norm_sq, op=dist.ReduceOp.SUM, group=group)49 return x / torch.sqrt(norm_sq).clamp_min(eps)50 51 52def _newton_schulz_step(53 x: torch.Tensor,54 a: float,55 b: float,56 c: float,57 group: dist.ProcessGroup,58) -> torch.Tensor:59 gram = x @ x.mT60 dist.all_reduce(gram, op=dist.ReduceOp.SUM, group=group)61 update = torch.addmm(gram, gram, gram, alpha=c, beta=b)62 return torch.addmm(x, update, x, alpha=1.0, beta=a)63 64 65@torch.no_grad()66def solution(67 x: torch.Tensor,68 steps: int = 5,69 coefficient_type: str = "quintic",70 partition_dim: int = 1,71 group: Optional[dist.ProcessGroup] = None,72) -> torch.Tensor:73 group = group or dist.group.WORLD74 assert x.ndim == 275 assert coefficient_type in _COEFFICIENTS76 coefficients = _COEFFICIENTS[coefficient_type]77 assert steps % len(coefficients) == 078 79 in_dtype = x.dtype80 81 if partition_dim == 0:82 x_work = x.mT.contiguous()83 elif partition_dim == 1:84 x_work = x85 else:86 raise AssertionError("invalid partition_dim")87 88 x_work = x_work.to(torch.float32)89 x_work = _distributed_normalize(x_work, group)90 91 for step in range(steps):92 a, b, c = _coefficient_at(coefficients, step)93 x_work = _newton_schulz_step(x_work, a, b, c, group)94 95 x_work = x_work.to(in_dtype)96 if partition_dim == 0:97 return x_work.mT.contiguous()98 return x_work.contiguous()99