PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashInfer-Bench: Building the Virtuous Cycle for AI-driven LLM Systems

Shanli Xing, Yiyan Zhai, Alexander Jiang, Yixin Dong, Yong Wu, Zihao Ye, Charlie Ruan, Yingyi Huang, Yineng Zhang, Liangsheng Yin, Aksara Bayyapu, Luis Ceze, Tianqi Chen

arXiv:2601.00227v1cs.AI

TL;DR

AI-generated GPU kernels remain difficult to move from code generation into production LLM inference because workloads, validation, and engine integration are complex. FlashInfer-Bench closes this gap with a standardized trace, real-serving dataset, robust benchmark, leaderboard, and runtime substitution mechanism. Its evaluation identifies compilation as the dominant failure mode, hardware specialization as a persistent weakness, and a Triton–CUDA trade-off between usability and peak performance.

  • Problem

    AI-generated GPU kernels face workload-dependent performance, difficult real-world tracking, and extra integration effort before they can enter production LLM systems.

  • Method

    FlashInfer-Bench standardizes operator contracts, workloads, implementations, and evaluations, then combines real-serving traces, rigorous correctness and performance checks, and runtime substitution into engines.

  • Results

    Compilation is the dominant failure mode; models struggle with hardware-specific features; Triton generally offers higher correctness and usability, while CUDA can reach higher peak performance.

  • Takeaways & Limitations

    The framework provides a practical closed loop for evaluating frontier models and converting validated kernel-level gains into lower latency and higher throughput in LLM serving.

  • Takeaways & Limitations

    The current scope excludes multi-GPU and communication kernels, and support for models, hardware devices, and programming languages remains limited.

Abstract

from arXiv · show

Recent advances show that large language models (LLMs) can act as autonomous agents capable of generating GPU kernels, but integrating these AI-generated kernels into real-world inference systems remains challenging. FlashInfer-Bench addresses this gap by establishing a standardized, closed-loop framework that connects kernel generation, benchmarking, and deployment. At its core, FlashInfer Trace provides a unified schema describing kernel definitions, workloads, implementations, and evaluations, enabling consistent communication between agents and systems. Built on real serving traces, FlashInfer-Bench includes a curated dataset, a robust correctness- and performance-aware benchmarking framework, a public leaderboard to track LLM agents' GPU programming capabilities, and a dynamic substitution mechanism (apply()) that seamlessly injects the best-performing kernels into production LLM engines such as SGLang and vLLM. Using FlashInfer-Bench, we further evaluate the performance and limitations of LLM agents, compare the trade-offs among different GPU programming languages, and provide insights for future agent design. FlashInfer-Bench thus establishes a practical, reproducible pathway for continuously improving AI-generated kernels and deploying them into large-scale LLM inference.

1 INTRODUCTION

FlashInfer-Bench addresses the gap between AI-generated GPU kernels and production LLM inference by standardizing workloads, evaluation, and deployment. It combines a trace schema, real-world dataset, benchmarking workflow, leaderboard, and dynamic substitution mechanism while analyzing agent capabilities and language trade-offs.

  • Motivation: Production LLM inference depends on performance-sensitive GPU kernels, whose optimization requires expert-level engineering.These kernels execute operations such as attention, matrix multiplication, and sampling.
  • Challenges: AI-generated kernels face workload-dependent performance, real-world benchmarking, and integration challenges before deployment.Relevant factors include ragged distributions and data precision, while engine integration can require additional engineering effort.
  • Framework: FlashInfer Trace standardizes kernel tasks, workloads, solutions, and evaluations in a self-contained format for AI-driven LLM systems.The framework also curates real-world workloads and supports robust benchmarking and dynamic kernel substitution into SGLang and vLLM without code changes.
  • Evaluation: A public leaderboard evaluates frontier models on real-world workloads and tracks their GPU programming capabilities.It is built on the benchmarking stack and reports correctness, performance, workload latency, and end-to-end latency metrics.
  • Findings: The evaluation finds that compilation failures dominate correctness errors, hardware-specific optimization remains difficult, and Triton–CUDA language choice involves a performance trade-off.Triton generally performs better across tasks, while CUDA offers greater potential for specialized optimization.
  • Contributions: The paper contributes a pragmatic workflow for continuously generating and directly applying AI-generated kernels in production systems.It also provides a comprehensive analysis of how LLM-generated kernels perform on LLM systems.

2 BACKGROUND

LLM inference relies on GPU kernels whose performance depends strongly on hardware, numerical format, and workload shape. FlashInfer-Bench combines generative program synthesis with rigorous evaluation to make structurally novel kernels usable while controlling correctness and reward-hacking risks.

  • LLM inference workloads: LLM serving engines combine batching, scheduling, parallelism, GPU kernels, and CPU logic, with GPU kernels dominating execution time.Common kernel families include GEMM, attention, fused MoE, and sampling or post-processing.
  • Optimization constraints: Kernel optimization is sensitive to SM count, memory hierarchy, tensor-core generation, numerical format, and workload shape.These dependencies make universal one-size-fits-all kernels elusive.
  • Existing approaches: Kernel libraries and search-based auto-scheduling provide strong optimization methods but remain limited by workload specificity, template expressiveness, or search cost.These constraints become more significant when hardware or shapes change.
  • Generative synthesis: Generative program synthesis lets LLMs write low-level GPU code and discover fusion or dataflow patterns beyond existing templates, but requires stringent validation.Without validation, generated kernels introduce correctness and security risks.
  • FlashInfer-Bench: FlashInfer-Bench combines generative synthesis with a production-grade evaluation harness that prevents regressions and reward hacking.The framework is designed to preserve generation’s ability to propose structurally new kernels while adding rigorous validation.
  • Production requirements: Production deployment additionally requires precise task specifications, reward-hacking defenses, robust validation, standardized exchange, and dynamic substitution without engine rewrites.These ingredients extend public generation evaluations toward system-level use.

3 FLASHINFER-BENCH DESIGN

FlashInfer-Bench uses a portable trace schema, real serving workloads, rigorous correctness and timing checks, a leaderboard, and runtime dispatch to connect kernel generation with deployment. Its design binds operator semantics, concrete workloads, implementations, and immutable evaluations into a reproducible workflow.

  • FlashInfer Trace: FlashInfer Trace is a self-contained schema whose Definition, Workload, Solution, and Evaluation components describe the task, input, implementation, and benchmark result.Auxiliary fields support grouping and filtering, while the combined Trace object provides portability and reproducibility.
  • FlashInfer Trace: Definitions specify operator I/O tensors, dtypes, dimension axes, reference semantics, and optional axis constraints.Axes can be static const values or workload-determined var values, enabling shape-specific optimization and support for ragged inputs.
  • FlashInfer Trace: Workloads bind concrete inputs to Definitions, while Solutions provide callable implementations with source files and compatibility metadata.Solutions can target specified GPU architectures and software versions, with extendable language or DSL support.
  • FlashInfer Trace: Evaluations immutably bind a Definition, Solution, and workload while recording execution state, correctness, performance, and environment information.This creates a concrete record for comparing implementations under reproducible conditions.
  • Dataset curation: The dataset pairs common serving-kernel Definitions with representative workloads collected from major models served against real-world traffic.It covers DeepSeek-V3, Llama-3.1-8B, and Qwen3-30B-A3B across GEMM, attention, normalization, sampling, and MoE operators.
  • Dataset curation: Workloads are deduplicated along performance-sensitive axes and tensor statistics while preserving diversity, leaving approximately 50 workloads per Definition.Full tensors are retained when values affect performance or correctness; otherwise seeded random tensors reduce storage.
  • Robust benchmarking: The benchmarking subsystem provides correctness validation, reproducible timing, and native multi-device execution.Timing uses per-GPU locking, warm-up runs, CUDA event-based device timing, and mean latency over measured runs.
  • Robust benchmarking: Deterministic kernels are checked elementwise against reference outputs, with non-finite values rejected and maximum observed error recorded.Low-precision kernels instead use a matched-ratio rule requiring a specified fraction of outputs to satisfy standard error bounds.

4 DATASET OVERVIEW AND EVALUATION

The evaluation covers representative LLM-inference kernels and shows that agents remain limited by compilation errors, hardware-specific optimization, and low-level implementation complexity, while higher-level DSLs often perform better. Kernel substitution nevertheless transfers benchmark improvements to end-to-end serving with minimal overhead.

  • Dataset: The dataset covers eight representative LLM-inference kernel types, including GEMM, attention variants, MoE, normalization, and sampling.These include fused, non-deterministic, and low-bit kernels, producing 41 distinct kernel definitions under fixed parameter combinations.
  • Correctness: 30 of 32 correctness errors arise from compilation failures, while only 2 are runtime or numerical errors.The numerical errors mainly involve incorrect padding calculations for large input shapes.
  • Hardware optimization: LLMs struggle to exploit hardware-specific intrinsics, often defaulting to older primitives instead of newer instructions such as mma and tcgen05.The reported failures include incomplete CUDA attention optimization and insufficient use of tiling or tensor cores.
  • Language trade-offs: Triton achieves higher correctness and speed on most tasks, while CUDA exposes a higher ceiling for specialized optimization when agents can coordinate low-level details.Triton’s compiler can select advanced hardware implementations, whereas CUDA requires explicit management of hardware abstractions.
  • Case studies: CUDA agents can learn to call optimized libraries, with some cuBLAS-based GEMM solutions matching or exceeding the baseline.This demonstrates environment use but also indicates reliance on library calls for strong performance in that task.
  • Case studies: GPT-5’s Triton GEMM kernel runs in 0.11 ms versus 0.5 ms for its CUDA kernel, a 4.5× speedup.The difference is attributed to compiler-supported pipelining and automatic targeting of tcgen05 tensor-core instructions.
  • Case studies: Low-level CUDA optimization remains difficult because agents fail to coordinate memory management, tiling, asynchronous execution, and pipelining.Explicit prompting failed to produce correct kernels using the requested attention optimizations in 10 attempts.
  • End-to-end deployment: apply() adds 1–2 us per kernel call and keeps end-to-end overhead below 0.8% across batch sizes.Kernel-level gains also carry through to serving: reported end-to-end times were 934 ms, 939 ms, and 1055 ms for the compared substitutions.

5 RELATED WORK

Prior work evaluates kernel generation, trains kernel-generation agents, or optimizes within human-designed schedule spaces. FlashInfer-Bench instead connects direct kernel synthesis with strict validation, benchmarking, and deployment into production-oriented LLM systems.

  • Kernel-generation benchmarks: KernelBench and TritonBench primarily evaluate whether models can generate compiling kernels with reasonable speedups, while BackendBench studies integration into PyTorch.FlashInfer-Bench focuses on major workloads in LLM systems and connects candidate kernels to production systems.
  • Kernel-generation agents: Post-trained kernel models and specialized agents address generation, transpilation, or CUDA design, and are complementary to FlashInfer-Bench’s evaluation framework.The benchmark can further assess these models and agents on real-world inference workloads.
  • Compiler optimization: TVM, AutoTVM, Ansor, and Meta-Schedule optimize within fixed spaces defined by human-crafted templates.Their methods use tensor-program optimization, search-based tuning, design-space abstraction, or policy learning.
  • Framework positioning: FlashInfer-Bench directly generates candidate kernels that may lie outside predefined schedule spaces, then subjects them to functional validation and performance benchmarking.This shifts the effective frontier from template search toward synthesis of new implementation patterns.
  • Implementation ecosystems: Triton, CUTLASS, cuBLAS, and FlashInfer serve as implementation targets or strong baselines while FlashInfer-Bench emphasizes connecting candidates to production systems.The framework treats existing kernel ecosystems as targets for agents rather than as the sole optimization space.
  • LLM serving systems: FlashInfer-Bench targets modern LLM operations such as attention and MoE and supports rapid evaluation and deployment through serving frameworks including vLLM and SGLang.These systems provide scalable inference infrastructure and reference implementations.

6 CONCLUSION

FlashInfer-Bench closes the loop from AI-generated kernel creation to production LLM serving through standardized traces, benchmarking, and dynamic substitution. Its evaluation identifies agent and language trade-offs while showing that validated kernel gains can improve serving performance, within a currently limited scope.

  • FlashInfer-Bench standardizes operator contracts, real serving workloads, candidate implementations, and immutable evaluations in FlashInfer Trace.
  • The benchmark evaluates deterministic, low-precision, and sampling kernels, then substitutes validated kernels into SGLang and vLLM without code changes.
  • Compilation is the dominant failure mode, models struggle to exploit hardware features, and language choice trades Triton usability and correctness against CUDA peak performance.
  • Dynamic substitution adds negligible overhead and converts kernel-level gains into lower latency and higher throughput in LLM serving.
  • The current scope excludes multi-GPU and communication kernels and remains limited in supported models, hardware devices, and programming languages.

A FLASHINFER TRACE EXAMPLES

This appendix introduces concrete examples of the FlashInfer Trace format.

  • The appendix provides concrete examples illustrating the FlashInfer Trace format.

A.1 GEMM

The GEMM example defines a matrix-multiplication kernel, instantiates a workload, records a generated Triton solution, and stores correctness and performance evaluation results.

  • The GEMM definition describes C = A @ B.T with variable M, constant N = 128 and K = 2048, float16 inputs, and float16 output.
  • The generated solution object records a named Claude Opus Triton implementation targeting NVIDIA B200 with a main.py::run entry point.
  • A sample workload instantiates gemm_n128_k2048 with M = 6 and random A and B inputs.
  • The trace records 0.023046740692633086 latency_ms against 0.025240250456929125 reference_latency_ms.

A.2 Attention

The attention example specifies a paged grouped-query attention decode operator with constrained interfaces, a PyTorch reference, a generated Triton solution, and a verified evaluation record.

  • The operator is batched grouped-query attention decode with a paged KV cache, using 32 query heads, 4 KV heads, head dimension 128, and page size 1.
  • The reference implementation gathers paged keys and values, maps query heads to KV heads, applies scaled softmax attention, and produces the output.

B.1 Triton GEMM Kernel

The Triton GEMM implementation computes C = A @ B.T for fixed N=4096 and K=4096 inputs, with validation, device management, and configurable tiled kernel launches. It is correct across evaluated workloads but slower than the torch.nn.functional baseline.

  • Results: Correctness was achieved across all evaluated workloads, while speedup was 0.2× overall and 0.6× at best versus torch.nn.functional.The generated implementation targets the GEMM definition with Triton.
  • Kernel: Triton launch configurations vary BLOCK_M, BLOCK_N, and BLOCK_K, including 128×128×64 and asymmetric 64×256×64 tiles.The configurations use eight warps and four pipeline stages.
  • Kernel: The kernel uses program IDs and tiled offsets over M and N, iterating across K in BLOCK_K chunks with masked loads and stores.The launch grid covers ceil(M/BLOCK_M) by ceil(N/BLOCK_N) tiles.
  • Interface: The wrapper computes C = A @ B.T for A shaped [M, 4096] and B shaped [4096, 4096], returning [M, 4096] float16 output.It validates tensor types and shapes before execution.
  • Runtime: The runtime moves inputs to a CUDA compute device, makes them contiguous, allocates device output, launches the kernel, and preserves the original output-device convention.No CPU fallback is provided when CUDA is unavailable.

B.2 CUDA GEMM Kernel

The CUDA GEMM implementation computes C = A * B.T through cuBLAS, adapting row-major PyTorch tensors to cuBLAS’s column-major view. It is correct across evaluated workloads but nearly matches, rather than clearly exceeds, the torch.nn.functional baseline.

  • Results: Correctness was achieved across all evaluated workloads, while speedup was 0.97× overall and 1.03× at best versus torch.nn.functional.The implementation invokes cuBLAS matmul.
  • Interface: The launcher computes C = A * B.T for A [M, 4096], B [4096, 4096], and half-precision output C [M, 4096].The PyTorch extension validates CUDA, contiguity, dtype, dimensionality, and compatible shapes.
  • Execution: A singleton cuBLAS handle is reused and configured for Tensor Core math, while the current PyTorch CUDA stream is assigned before launch.This avoids repeated handle creation and ties execution to the active stream.
  • Layout transformation: The cuBLAS call rephrases the row-major operation as column-major multiplication with op1 = B_cm.T and op2 = A_cm.The resulting [N, M] column-major layout matches the desired [M, N] row-major output memory.
  • Runtime: The extension allocates output on A’s device, obtains the current CUDA stream, passes raw half pointers to the launcher, and checks asynchronous CUDA errors.Exceptions from CUDA or cuBLAS are propagated through the wrapper.

B.3 Triton GQA Paged Decode Attention

The Triton GQA paged-decode kernel processes query heads against paged K/V caches using streaming softmax and returns both attention output and log-sum-exp values. It is correct but substantially slower than FlashInfer on the evaluated workloads.

  • Results: Correctness was achieved on all evaluated workloads, while speedup was 0.19× overall and 0.98× at best versus FlashInfer.The implementation targets h32, kv8, d128, page-size-1 grouped-query attention decode.
  • Attention computation: Streaming softmax maintains a running maximum, exponential sum, and output vector, merging each block with numerically stabilized scaling factors.The final output is normalized by the accumulated denominator and paired with a base-2 log-sum-exp value.
  • Paged access: Paged cache indices select K and V pages, which are gathered and processed in token blocks while masking positions beyond each sequence range.The wrapper flattens the page dimension before launching the kernel.
  • Runtime: The entry point moves inputs to CUDA when needed, makes them contiguous, launches the Triton kernel, and returns output and LSE tensors on the original device convention.CUDA is required for Triton execution.
  • Mapping: The kernel maps one program to each batch-and-query-head pair and derives the corresponding KV head using the 32-to-8 grouped-query ratio.The launch grid is batch_size × num_qo_heads.

B.4 CUDA GQA Paged Decode Attention

The CUDA GQA paged-decode implementation assigns four warps to query heads associated with each KV head and performs streaming softmax over paged BF16 caches. It is correct across evaluated workloads but has only 0.02× speedup overall against FlashInfer.

  • Results: Correctness was achieved across all evaluated workloads, while speedup was 0.02× overall and 1.02× at best versus FlashInfer.The implementation uses four warps per block and scalar FP32 operations for streaming softmax.
  • Data movement: The kernel loads BF16 query, key, and value data, converts values to FP32, and stages one token’s K/V vectors in shared memory.Paged indices are flattened into cache locations for each KV head.
  • Mapping: Each block handles one batch element and one KV head, while its four warps process the four query heads assigned to that KV head.The mapping covers 32 query heads and 8 KV heads.
  • Attention computation: For each token, the kernel computes scaled QK logits, updates streaming softmax accumulators, normalizes the output, and writes base-2 log-sum-exp.The running maximum stabilizes the exponential updates.
  • Runtime: The host wrapper validates CUDA placement, dimensions, dtypes, and index types before launching the kernel with output and LSE tensors.The interface expects BF16 q, K/V caches and FP32 LSE output for the specified shapes.
Loading 2601.00227v1…