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 math4 5import torch6from torch import Tensor7 8 9@torch.no_grad()10def solution(11 flat_param_shard: Tensor,12 flat_grad_shard: Tensor,13 exp_avg_shard: Tensor,14 exp_avg_sq_shard: Tensor,15 lr: float,16 beta1: float,17 beta2: float,18 eps: float,19 weight_decay: float,20 step: int,21) -> tuple[Tensor, Tensor, Tensor]:22 assert step >= 123 24 m = exp_avg_shard.clone()25 v = exp_avg_sq_shard.clone()26 g = flat_grad_shard27 theta = flat_param_shard.clone()28 29 m.mul_(beta1).add_(g, alpha=1.0 - beta1)30 v.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)31 bc1 = 1.0 - math.pow(beta1, step)32 bc2 = 1.0 - math.pow(beta2, step)33 m_hat = m / bc134 v_hat = v / bc235 denom = v_hat.sqrt().add(eps)36 37 theta.add_(m_hat.div(denom), alpha=-lr)38 theta.add_(flat_param_shard, alpha=-lr * weight_decay)39 40 return theta, m, v41