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 __future__ import annotations2 3import torch4import torch.distributed as dist5from torch import Tensor6 7_FP8_E4M3_MAX = 448.08 9 10@torch.no_grad()11def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor:12 out = torch.roll(amax_history, shifts=-1, dims=0)13 out[-1] = cur_abs_max.to(dtype=out.dtype)14 return out15 16 17@torch.no_grad()18def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor:19 xf = x.float()20 qs = xf / scale21 q = qs.to(torch.float8_e4m3fn)22 return (q.float() * scale).to(dtype=x.dtype)23 24 25@torch.no_grad()26def solution(flat_grads: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]:27 world_size = dist.get_world_size()28 n = flat_grads.numel()29 shard_elems = n // world_size30 31 cur_abs_max = flat_grads.abs().max().to(torch.float32)32 updated_hist = _update_amax_history(amax_history, cur_abs_max)33 34 scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX35 recon = _fp8_round_trip_bf16(flat_grads, scale)36 37 out_shard = torch.empty(shard_elems, dtype=flat_grads.dtype, device=flat_grads.device)38 dist.reduce_scatter_tensor(out_shard, recon.contiguous(), op=dist.ReduceOp.SUM)39 out_shard.div_(world_size)40 41 return out_shard, updated_hist42