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
1import torch2import torch.distributed as dist3 4def solution(local_hidden_states: torch.Tensor, local_weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor:5 input_dtype = local_hidden_states.dtype6 # Upcast to float32 for stable variance calculation7 local_hidden_states = local_hidden_states.to(torch.float32)8 9 local_sum_squares = local_hidden_states.pow(2).sum(dim=-1, keepdim=True)10 11 dist.all_reduce(local_sum_squares, op=dist.ReduceOp.SUM)12 13 world_size = dist.get_world_size()14 global_hidden_size = local_hidden_states.shape[-1] * world_size15 variance = local_sum_squares / global_hidden_size16 17 local_hidden_states = local_hidden_states * torch.rsqrt(variance + variance_epsilon)18 19 return local_weight * local_hidden_states.to(input_dtype)20