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 9@torch.no_grad()10def _block_int8_quant_dequant(x_flat: Tensor, block_size: int) -> Tensor:11 n = x_flat.numel()12 if n == 0:13 return x_flat.clone()14 flat = x_flat.contiguous().reshape(-1)15 pad = (-n) % block_size16 if pad:17 flat = F.pad(flat, (0, pad))18 nb = flat.numel() // block_size19 xv = flat.view(nb, block_size)20 scales = xv.abs().amax(dim=1).float().clamp(min=1e-8) / 127.021 q = (xv.float() / scales.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)22 out = (q.float() * scales.unsqueeze(1)).reshape(-1)23 return out[:n]24 25 26@torch.no_grad()27def solution(28 flat_grad: Tensor,29 block_size: int,30) -> Tensor:31 assert block_size >= 132 33 world_size = dist.get_world_size()34 orig_shape = flat_grad.shape35 x = flat_grad.reshape(-1)36 37 rec = _block_int8_quant_dequant(x, block_size)38 acc = rec.float()39 dist.all_reduce(acc, op=dist.ReduceOp.SUM)40 acc.div_(world_size)41 42 return acc.to(dtype=flat_grad.dtype).reshape(orig_shape)43