AIAny
Icon for item

CUDA-Agent-Ops-6K

Provides 6,000 runnable, operator-level PyTorch tasks for training and evaluating CUDA kernel generation models; each sample includes executable code, operator descriptors, and provenance tags, with execution-driven filtering to ensure reproducibility and contamination control.

Introduction

Why this matters

High-quality, reproducible training data for CUDA kernel generation is scarce, and manual expert annotation is costly. This dataset provides a middle ground: a synthesized, execution-verified corpus of operator-level PyTorch tasks designed to train RL and code-generation models to produce high-performance CUDA kernels.

What Sets It Apart
  • Synthesized + verified samples: tasks are created by combining seed operators (from torch and transformers) via an LLM-based synthesis step and then filtered by actual execution. This yields diverse fused-operator problems (up to 5 ops) while ensuring each sample runs correctly.
  • Execution-driven filtering: every sample must run in both eager and torch.compile modes, excludes stochastic or degenerate outputs, enforces eager runtime between ~1ms–100ms, and removes near-duplicates to KernelBench to reduce evaluation contamination. This makes rewards and profiling signals more reliable for RL training.
  • Practical data format: the release contains 6,000 rows in a parquet file with three string columns — ops (operator descriptor), data_source (provenance tag like torch#N or transformers), and code (runnable Python/PyTorch implementation) — so samples can be executed or integrated into training pipelines with minimal preprocessing.
Who it’s for and tradeoffs

Great fit if you are training or researching models that generate or optimize CUDA kernels (especially RL-based agentic workflows), benchmarking kernel synthesis approaches, or building tooling around torch.compile. Look elsewhere if you need real-world end-to-end model training workloads (this is operator-level synthesis, not full production workloads) or if you require coverage of extremely long-running kernels—samples are curated to a controlled runtime range. Also note that synthesized problems may not capture every hardware-specific idiom present in hand-optimized kernels.

Where it fits

Use this dataset to bootstrap or scale RL/LLM training for CUDA codegen, to run contamination-controlled ablations against KernelBench, or to prototype skill-augmented environments that rely on runnable operator tasks. It complements, rather than replaces, hand-crafted benchmark suites and real-world production kernels.

Information

  • Websitehuggingface.co
  • OrganizationsBytedTsinghua-SIA, Tsinghua University Intelligent Industry Research Institute (AIR), ByteDance Seed SIA-Lab
  • Published date2026/02/26

Categories

More Items

Hugging Face

Provides manually curated Japanese instruction pairs (questions and safe reference answers) for improving LLM output safety, covering broad harm categories and regionally sensitive cases. Includes English meta-tags and standard splits for benchmarking and fine-tuning.

Hugging Face

A 16 GB, 507-file PhD‑level cybersecurity knowledge base for training and evaluating security-focused LLMs and automation. Covers offensive/defensive/forensics/cloud/iot and AI-security across 30+ domains with real-world labs and framework mappings.

Hugging Face

Structured dataset for training and evaluating LLM agentic behavior: function-calling conversations, JSON-mode structured outputs, and extraction samples for teaching models to generate tool calls and strict structured responses. Includes single-turn and multi-turn scenarios across several configs.