EvanOLeary/pallasbench-robust

PallasBench是一个包含45个JAX Pallas GPU内核的基准数据集，具有稳健的评估结果、编译工件和GPU性能指标。它采用了来自SakanaAI的Towards Robust Agentic CUDA Kernel Benchmarking中的稳健评估方法，并针对Pallas/JAX编译管道进行了适配，提供了首个专注于Pallas内核的GPU基准。每个内核都经过GPU兼容性修复（通过块大小钳位解决Triton的100万元素限制），并通过五个稳健性过滤器进行评估，以区分真正正确的实现与退化实现。PallasBench is a benchmark dataset containing 45 JAX Pallas GPU kernels, with robust evaluation results, compiled artifacts, and GPU performance metrics. It adopts the robust evaluation methodology from *Towards Robust Agentic CUDA Kernel Benchmarking* by SakanaAI, and adapts it to the Pallas/JAX compilation pipeline, providing the first GPU benchmark exclusively focused on Pallas kernels. Each kernel has undergone GPU compatibility fixes (resolving Triton's 1 million element limit via block size clamping) and is evaluated through five robustness filters to distinguish truly correct implementations from degenerate ones.

搜集汇总

数据集介绍

构建方式

PallasBench-Robust数据集基于开源项目PallasBench构建，其核心工作是将原本面向TPU的45个JAX Pallas GPU内核系统地移植至NVIDIA A100 GPU环境。构建过程中，研究团队针对Triton编译器在GPU上的1M元素限制，对所有内核实施了块尺寸箝制修复，涉及35个文件、106行代码的修改。此外，数据集引入了源自SakanaAI鲁棒性评估方法学的五重过滤器，包括输出范围、输出标准差、坐标轴变化、输入影响及源码分析，用以甄别内核实现的正确性与鲁棒性。每一条数据均记录了完整的编译中间产物（Jaxpr DAG与StableHLO MLIR）以及详尽的GPU性能指标，最终以KernelBook兼容的JSONL格式存储。

使用方法

该数据集主要服务于大语言模型驱动的内核优化、编译器研发以及鲁棒性评估方法学研究。使用者可通过HuggingFace Datasets库加载JSONL格式数据，每条记录以kernel_name为唯一标识符。研究人员可借助source_original与source_fixed字段对比TPU与GPU内核实现的差异，利用jaxpr与stablehlo字段深入分析编译流水线中的IR转换过程。result字段中的鲁棒性过滤器结果可用于训练自动化的内核验证模型，而wall_time_seconds与jit_compile_time_seconds则为性能建模提供了基准参考。鉴于首次运行包含显著的JIT编译开销（约60-100分钟），建议在预热后收集稳态性能数据。

背景与挑战

背景概述

PallasBench-Robust数据集由Evan O'Leary于2025年创建，聚焦于JAX Pallas GPU内核的鲁棒性评估与基准测试。该数据集源于SakanaAI的鲁棒性评估方法论，首次将Pallas/JAX编译管道中的GPU内核性能评估系统化。数据集包含45个涵盖激活函数、矩阵运算、注意力机制等类别的内核，并提供了完整编译中间表示与硬件性能指标。作为面向Pallas语言的GPU内核基准，它填补了可移植领域特定语言与CUDA内核基准之间的空白，为编译器研究与内核优化提供了标准化评估框架。

当前挑战

当前挑战主要集中于三方面：其一，领域问题层面，Pallas内核从TPU迁移至GPU时面临Triton编译器1M元素上限的约束，需通过块大小钳制解决兼容性；其二，构建过程中，JIT编译开销巨大，完整基准套件的首次运行需60-100分钟，且性能结果受限于NVIDIA A100单一硬件平台；其三，数值精度与鲁棒性验证的平衡问题，某些内核需放宽容忍度，而轴变异性过滤器对稀疏输出存在误判风险，限制了评估的普适性。

常用场景

经典使用场景

在GPU计算与编译器研究领域，PallasBench-Robust数据集为评估和优化基于JAX Pallas框架的GPU内核提供了标准化基准平台。该数据集涵盖45个精心设计的Pallas内核，从单算子操作到融合模式再到完整架构组件，横跨激活函数、归一化、矩阵乘法、注意力机制、基因组分析等多个类别。每个内核都经过GPU兼容性修复，并附带完整的编译产物——包括Jaxpr中间表示、StableHLO MLIR以及PTX汇编指令，使得研究者能够深入剖析从高级语言到底层机器码的完整编译流水线。经典的用法聚焦于驱动大语言模型进行内核自动优化：研究者可以利用该数据集提供的原始与修复版本内核、正确性验证结果以及鲁棒性过滤器的输出，训练模型理解合法内核的构造模式，进而实现自动化的内核生成、缺陷修复与性能调优。

解决学术问题

该数据集直面GPU内核性能评测中长期存在的核心痛点：如何区分一个看似正确但实则退化的内核实现与真正正确的实现。传统评测往往只关注最终输出结果是否与参考一致，却容易忽略那些运气好输出正确但算法设计不当的退化内核。PallasBench-Robust通过引入五层鲁棒性过滤器——输出范围检查、输出标准差分析、轴变化敏感度测试、输入扰动影响评估以及源码级结构审查，系统性地解决了这一验证难题。这一方法论使得研究者能够可靠地度量内核实现的真实质量，为编译器后端优化策略的评估提供了可信依据。数据集的发布显著推进了可移植领域特定语言（DSL）内核基准测试研究，将原本集中在CUDA生态的性能评测框架拓展至JAX/Pallas这一日益重要的深度学习编译器栈。

实际应用

在实际工程应用中，该数据集为GPU计算流水线的自动化工具链建设提供了扎实的测试床。基于这个基准，AI驱动的代码辅助工具（如GitHub Copilot、代码生成模型）能够获取人类专家的推理过程和调试策略，提升生成代码的正确性和性能。在健康医疗领域，数据集中的基因组分析内核（如k-mer计数、反向互补序列、汉明距离计算）可直接用于加速生物信息学管道的运行速度，帮助研究人员更快地分析遗传测序数据。对于云计算服务提供商而言，该数据集中的注意力机制、多层感知器块以及完整Transformer编码器内核的鲁棒性评测结果，为在云端部署高性能推理服务提供了关键的代码质量控制手段。此外，内存使用和编译时间的详细记录也为资源受限场景下的模型部署提供了宝贵的性能预测依据。

数据集最近研究

最新研究方向

PallasBench-Robust数据集聚焦于面向GPU的Pallas内核稳健性评估与优化，其研究前沿在于引入五重鲁棒性过滤机制以甄别退化内核实现，并结合JAX/Triton编译流水线的中间表示（Jaxpr、StableHLO）与A100实测性能指标，为LLM驱动的内核自动生成与编译器优化提供首个可移植的基准测试。该工作呼应了SakanaAI关于CUDA内核稳健基准化的热点方法论，将稳健性验证从CUDA扩展到Pallas这一可移植领域专用语言，显著推动了GPU编程抽象层的高效验证与自动化调优研究，对降低GPU内核开发中因编译差异与数值误差导致的错误风险具有重要实践意义。