togethercomputer/ParallelKernelBench_Problems

ParallelKernelBench（基准测试）是一个用于评估由大型语言模型（LLM）生成的多GPU CUDA内核的基准测试数据集。该数据集包含87个参考实现，位于reference/目录中，以及输入输出张量生成规范在utils/input_output_tensors.py文件中。输入张量是确定性的，可以通过提供的create_input_tensor函数重现，无需存储.pt文件。数据集以parquet格式存储问题信息，包括问题ID、描述、参考代码路径、默认评估设置（如8个H100 GPU、1024x1024基础形状、bfloat16数据类型和5次试验）等。用户可以使用HuggingFace的datasets库加载数据集，并通过相关工具进行本地输入重现和评估。ParallelKernelBench (benchmark) is a reference problem dataset 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. Inputs are deterministic and can be reproduced using the create_input_tensor function from that file without needing stored .pt files. The dataset is stored in parquet format with columns including problem_id, stem, reference_code, reference_path, input_tensor_spec_path, and default evaluation settings such as world_size, default_m, default_n, default_dtype, default_trials. It can be loaded via HuggingFace datasets and used for local input reproduction and evaluation.

搜集汇总

数据集介绍

构建方式

ParallelKernelBench_Problems数据集专为评估大语言模型生成多GPU CUDA内核的基准测试而设计，其构建过程融合了系统的工程化与科学严谨性。该数据集包含87个参考实现，存储于`reference/`目录下，每个问题对应一个独立的Python文件，其中定义了`solution()`函数作为标准答案。输入张量的规格说明被封装在`utils/input_output_tensors.py`模块中，通过`create_input_tensor`函数实现确定性生成，无需依赖存储的`.pt`文件。数据集的元信息以Parquet格式组织在`data/problems.parquet`中，每一行代表一个问题，包含问题标识符、参考代码路径、默认评估参数（如8×H100 GPU、1024×1024张量尺寸、bfloat16精度及5次试验）等关键字段，从而为自动化评测提供了一致且可复现的基础。

使用方法

使用ParallelKernelBench_Problems数据集时，用户首先通过HuggingFace Datasets库加载`togethercomputer/ParallelKernelBench_Problems`数据集，访问`train`分片中的问题信息。利用`hf_hub_download`函数可获取输入张量规格文件，并将其添加至Python路径或克隆仓库以导入`create_input_tensor`函数。该函数接受rank、world_size、problem_id、base_shape、dtype和trial参数，生成确定的输入张量。评估过程通过执行`run_local.py`脚本进行，指定问题编号、解决方案类型（cuda）、解决方案根目录、数据类型及试验次数。用户亦可直接读取Parquet文件中的`reference_code`字段，获取完整的参考实现代码。此流程为研究人员提供了一条从数据加载到性能测量的完整闭环途径。

背景与挑战

背景概述

ParallelKernelBench_Problems 数据集由 Together Computer 团队于 2024 年创建，旨在为大型语言模型生成的多 GPU CUDA 内核提供标准化的基准测试平台。该数据集包含 87 个参考实现及其输入张量规范，覆盖了分布式系统中常见的并行计算问题，如矩阵操作和通信模式。作为 ParallelKernelBench 项目的核心组件，它填补了 LLM 在高性能计算领域缺乏系统性评估的空白，推动了代码生成模型从单 GPU 向多 GPU 环境扩展的研究。通过提供确定性输入和可复现的评估流程，该数据集为度量 LLM 生成代码的正确性、效率及可扩展性奠定了重要基础，在系统与机器学习交叉领域产生了积极影响。

当前挑战

该数据集面临的核心挑战包括：首先，领域问题层面，多 GPU CUDA 内核生成需要协调通信、计算与内存访问，现有 LLM 难以处理同步、负载均衡及跨设备数据依赖等复杂语义，导致生成代码常出现死锁或性能退化。其次，构建过程中，参考实现需兼顾数值精度与硬件特性，而输入张量规范必须适应不同 GPU 拓扑和规模，确保 87 个问题的难度梯度与实际问题匹配。此外，评估指标的制定极具挑战性，单纯依赖正确性不足，还需纳入运行时间、内存占用及可扩展性等多维指标，同时保持跨运行的可重复性，这对测试框架的设计提出了严苛要求。

常用场景

经典使用场景

ParallelKernelBench_Problems数据集专为评估与优化大语言模型在多GPU环境下生成CUDA内核的能力而设计。其核心使用场景在于提供一个包含87个参考实现的标准化基准，涵盖从张量操作到分布式通信的典型并行计算问题。研究者可通过该数据集中的问题描述、输入张量规格及参考代码，系统性地测试LLM生成的CUDA内核在正确性、性能与可扩展性方面的表现。结合其提供的确定性输入生成函数和默认评估配置（如8×H100 GPU、bf16精度），该数据集成为衡量模型在多GPU编程领域代码生成能力的权威标杆。

解决学术问题

该数据集直击当前大语言模型在高性能计算领域面临的两大学术难题：一是如何验证LLM生成的多GPU CUDA内核在分布式环境下的功能正确性与数值稳定性，二是如何量化其生成代码相对于人工手写优化的性能差距。通过提供涵盖不同通信模式、计算负载与张量形状的参考实现，数据集使得研究者能够系统性地分析模型在并行编程任务中的错误模式与效率瓶颈。它为分布式系统与代码生成交叉领域的实证研究提供了可复现的基准，推动了对LLM在底层硬件编程中泛化能力的深入理解，并为未来更智能的自动化内核优化策略奠定了评估基础。

实际应用

在实际工程场景中，ParallelKernelBench_Problems数据集加速了基于LLM的GPU内核自动生成与调优工具的研发进程。例如，高性能计算团队可利用此基准筛选并微调模型，以生成更高效的通信算子或矩阵乘法核函数，从而降低多GPU应用的手动编码成本。深度学习框架开发者亦可借助该数据集验证模型能否生成适配特定硬件拓扑（如NVLink层级互联）的分布式执行代码。此外，该数据集在自动化代码审查与性能预测系统中也发挥着关键作用，其标准化的输入输出规范使得对LLM生成方案的系统性评估成为可能。

数据集最近研究

最新研究方向

ParallelKernelBench_Problems数据集聚焦于大语言模型在多GPU场景下自动生成CUDA内核的基准测试研究，近期研究热点包括利用该数据集评估前沿LLM（如GPT-4、Claude）编写分布式计算代码的能力，以及探索模型在跨GPU通信、并行策略优化等复杂系统任务中的表现。伴随AI辅助编程工具的爆发式增长，该数据集为衡量LLM对底层硬件抽象和并行计算范式理解程度提供了标准化测试场，其87个参考实现和确定性输入设计确保了评测的可复现性。该工作对推动高效GPU内核自动生成、降低高性能计算开发门槛具有标杆意义，尤其在多卡协同场景下，为评估LLM处理真实分布式系统难题的能力奠定了关键基础。