togethercomputer/ParallelKernelBench_Problems

1---2license: apache-2.03task_categories:4- text-generation5tags:6- code7- cuda8- distributed-systems9- gpu-kernels10- benchmark11size_categories:12- n<1K13---14 15# ParallelKernelBench (benchmark)16 17Reference problems for [ParallelKernelBench](https://github.com/togethercomputer/ParallelKernelBench): a benchmark for LLM-generated multi-GPU CUDA kernels.18 19This dataset contains **87** reference implementations in `reference/` and the input tensor specification in `utils/input_output_tensors.py`. Inputs are **deterministic** — reproduce them with `create_input_tensor(rank, world_size, problem_id, base_shape, dtype, trial)` from that file; you do not need stored `.pt` files.20 21## Files22 23| Path | Description |24|------|-------------|25| `data/problems.parquet` | One row per problem (tabular access) |26| `reference/*.py` | Reference `solution()` implementations |27| `utils/input_output_tensors.py` | Input/output tensor generation for every problem |28 29## Columns (`data/problems.parquet`)30 31- `problem_id`, `stem` — problem identity32- `reference_code` — full Python source33- `reference_path` — path to the same file in this repo34- `input_tensor_spec_path` — path to `utils/input_output_tensors.py` (same on every row)35- `world_size`, `default_m`, `default_n`, `default_dtype`, `default_trials` — default eval settings (8× H100, 1024×1024, bfloat16, 5 trials)36 37## Usage38 39```python40from datasets import load_dataset41from huggingface_hub import hf_hub_download42 43ds = load_dataset("togethercomputer/ParallelKernelBench_Problems", split="train")44print(ds[0]["stem"], ds[0]["reference_code"][:200])45 46# Fetch the input tensor spec (same file on disk in this dataset repo)47spec_path = hf_hub_download("togethercomputer/ParallelKernelBench_Problems", "utils/input_output_tensors.py", repo_type="dataset")48```49 50Reproduce inputs locally (add the downloaded `utils/` folder to PYTHONPATH, or clone this repo):51 52```python53from utils.input_output_tensors import create_input_tensor54import torch55 56x = create_input_tensor(57rank=0, world_size=8, problem_id=17,58base_shape=(1024, 1024), dtype=torch.bfloat16,59)60```61 62## Related63 64Net-new LLM-generated kernels live in a separate dataset repo containing only `solutions/<run_id>/*.py`.65 66## Eval67 68```bash69python run_local.py --mode eval --problem 17 --solution cuda \70--solutions-root path/to/solutions_dir --dtype bfloat16 --trials 571```72