KernelBench

repository·main·Indexed 22 days ago

https://github.com/scalingintelligence/kernelbench

A benchmark for evaluating the ability of Large Language Models (LLMs) to generate efficient GPU kernels. It provides baseline timing results for various GPU architectures and PyTorch configurations to measure speedup between reference PyTorch modules and LLM-generated CUDA kernels. The framework supports both offline baseline timing using pre-computed JSON files and real-time live timing, with specific support for PyTorch 2.9.0 and Blackwell architecture.

Tokens
13.3K
Snippets
37
Records
55
Agent score
78%

What's inside kernelbench

  1. Understand the difference between baseline timing and live timing

    main

    KernelBench measures speedup by comparing the runtime of a reference architecture (PyTorch) against an LLM-generated architecture (the candidate kernel). There are two ways to handle the reference runtime:

    1. Baseline Timing (Offline Analysis): Uses pre-computed JSON files from the results/timing/ directory. This is used for offline scoring to ensure a consistent reference runtime without re-running the PyTorch reference on every evaluation.
    2. Live Timing (Real-time): Computes the reference runtime on the fly. This is used in workflows like scripts/run_and_check.py or via the kernelbench.eval.eval_kernel_against_ref(...) API. Live timing is necessary when running on hardware different from the provided baselines.

    Note: The provided baseline results use PyTorch 2.5.0+cu124 and CUDA 12.4.

  2. Generate baseline timing results on your own cluster

    main

    If your hardware (GPU, power settings, etc.) differs from the provided baselines, you can generate your own reference timing results using the provided script. This ensures your speedup measurements are accurate for your specific environment.

    Run the following command to generate baseline timings:

    uv run python scripts/generate_baseline_time.py
  3. KernelBench Task Format

    main
    KernelBench tasks follow a specific structure designed for in-context learning and benchmarking. While the exact schema is part of the task definition, the tutorial indicates that tasks are managed via a temporary directory for storing kernels and follow a standardized format used for benchmarking PyTorch modules paired with CUDA kernels.
  4. How the static checker identifies reward hacking patterns

    main

    The kernel_static_checker uses regex-based pattern matching to identify common ways models might 'cheat' during kernel benchmarking:

    1. Bypass Hacks: Detecting try-except blocks (to fall back to PyTorch if the kernel fails) or pass statements (to inherit and do nothing).
    2. Timing Manipulation: Detecting the use of CUDA streams, threading, or monkey-patching timing functions like time.perf_counter to fake benchmark results.
    3. Lazy Evaluation: Detecting the creation of fake tensors via torch.Tensor.__new__ or custom subclasses that pass correctness tests without performing actual math.
    4. Precision Downgrading: Detecting explicit casts (like __float2half in CUDA or tl.astype(..., tl.float16) in Triton) that make a kernel faster by reducing precision while appearing to meet the target precision requirements.
    5. High-level Op Usage: Detecting the use of torch.nn.functional or standard torch.* operations inside what should be a custom kernel implementation.
  5. How KernelBench datasets are abstracted

    main

    KernelBench uses a unified BaseDataset abstraction to allow users to interact with problems identically, regardless of whether they are stored on the local filesystem or hosted on HuggingFace.

    Core Components

    • Problem: A dataclass representing a single task, containing the source code, ID, and metadata.
    • BaseDataset: An abstract base class defining the interface for all datasets. It ensures that any dataset implementation provides methods for length, iteration, and ID-based access.
    • LocalKernelBenchDataset: An implementation that loads .py files from a local directory structure (e.g., KernelBench/level1/*.py).
    • HuggingFaceKernelBenchDataset: An implementation that pulls data from the ScalingIntelligence/KernelBench repository using the datasets library.

    Identity and Hashing

    To track problem identity across different versions or formatting changes, the Problem class provides a .hash property. This hash is computed by stripping all multi-line comments, inline comments, and whitespace from the code before hashing, ensuring that functionally identical code produces the same hash.

  6. Structure of a KernelBench dataset entry

    main

    KernelBench data is structured as a collection of Python files, where each file represents a specific kernel problem. When converted to a HuggingFace dataset, each entry contains the following fields:

    • code: The full source code of the Python file (containing a Model class and helper functions like get_inputs and get_init_inputs).
    • level: An integer representing the difficulty level (1 through 4).
    • name: The filename without the extension (e.g., 1_Square_matrix_multiplication_).
    • problem_id: An integer extracted from the prefix of the filename (e.g., 1 from 1_MLP.py).
  7. Understand the KernelBench Task Format

    main

    KernelBench is a benchmark designed to evaluate the ability of Large Language Models (LLMs) to generate efficient GPU kernels.

    In a typical task, the model is provided with a reference PyTorch model and must generate an optimized version. The goal is to replace specific PyTorch operations within the reference model with custom, inline CUDA kernels to improve performance.

    Task Input/Output Structure:

    • Input: A Model implemented as a standard PyTorch reference.
    • Output: A ModelNew which consists of the original PyTorch model augmented with inline CUDA kernels for optimized operations.