GPU Kernel Code Generation Benchmark — Multi-Agent vs Zero-Shot LLMs

A domain-specific coding benchmark for CUDA/C++ GPU kernel programming, plus a 3-agent ReAct framework that measurably beats zero-shot generation — evaluated on 25 real-world tasks with an extended CodeBLEU metric and unbiased Pass@1.

PythonMulti-Agent SystemsReActCUDALLM EvaluationCodeBLEUNVIDIA NIMBenchmarking

---

Why this project

Writing correct GPU kernels is one of the hardest code-generation tasks for LLMs: it demands thread-hierarchy reasoning, memory-layout awareness, synchronization correctness, and CUDA API fluency. It's a perfect stress test for whether agentic frameworks actually improve code quality over plain zero-shot prompting — not just for toy problems, but on 25 tasks collected from 13 real production repositories.

This project answers three questions with data, not vibes:

1. How well do small/medium open-weight models generate CUDA kernels zero-shot?
2. Does a multi-agent ReAct framework (Coder → Reviewer → Refiner) improve their scores?
3. Can a domain-weighted evaluation metric (extended CodeBLEU) capture what makes GPU code correct?

Key Results

Configuration	Pass@1	CodeBLEU	Δ Pass@1 (agent vs zero-shot)
Bielik-11B (zero-shot)	0.888	0.3998	—
Bielik-11B (agent)	0.960	0.3397	+0.072
Nemotron-Nano-9B (zero-shot)	0.156	0.2548	—
Nemotron-Nano-9B (agent)	0.400	0.2269	+0.244 (2.6×)
GPT-OSS-20B (zero-shot)	0.652	0.2951	—
GPT-OSS-20B (agent)	0.640	0.2461	−0.012 (stable)

The agent harness improves Pass@1 for 2 of 3 models — the Reviewer catches missing __syncthreads, unbalanced memory usage, and wrong indexing; the Refiner fixes them. CodeBLEU drops slightly because agent output is longer and structurally complete, reducing n-gram overlap with compact reference snippets — while actually improving correctness. This gap between surface similarity and functional correctness is exactly why we report both.

---

Architecture

Three specialized agents collaborate in a Thought → Action → Observation loop, following the multi-agent cross-team collaboration pattern from arXiv:2408.08927:

┌─────────────────────────────────────────────────────────────────────┐
│ Iteration 1 │
│ CoderAgent Thought : analyse task (kernel, memory, sync) │
│ Action : generate initial CUDA code │
│ Observation: code → ReviewerAgent │
├─────────────────────────────────────────────────────────────────────┤
│ Iterations 2–3 │
│ ReviewerAgent Thought : check against 7-point CUDA correctness list │
│ Action : structured JSON review │
│ Observation: review → RefinerAgent │
│ │
│ RefinerAgent Thought : plan fixes for every reported issue │
│ Action : rewrite code addressing all issues │
│ Observation: improved code → next iteration / output │
├─────────────────────────────────────────────────────────────────────┤
│ Early stop: quality = "excellent" AND zero critical issues │
│ AND zero missing components │
└─────────────────────────────────────────────────────────────────────┘

Every step records {iteration, agent, thought, action, observation} — full ReAct traces are saved to results/ for inspection and audit.

Design Decisions

• Multi-agent, not single-agent. A single "fix your code" loop collapses into echo-chambering. Separating generation, critique, and refinement into distinct agents with different system prompts creates real adversarial pressure — the Reviewer has no incentive to agree with the Coder.
• Structured JSON reviews. The Reviewer emits a 7-point checklist verdict (critical / major / minor, missing components) instead of free text. This makes refinement deterministic, debuggable, and measurable.
• Heuristic validity checks instead of nvcc. Compiling 750 generated kernels (25 tasks × 10 samples × 3 models × 2 modes) would need GPU toolchains on every runner. The 6-criterion validator (kernel qualifier, blockIdx+threadIdx, balanced braces, minimum size, no stubs, control flow) is compiler-free and portable — with documented tradeoffs.
• Domain-weighted CodeBLEU. Standard CodeBLEU treats all tokens equally. GPU code has load-bearing keywords (__syncthreads, __shared__, atomicAdd). We reweighted 42 CUDA keywords at 2.0–5.0× and added a missing-keyword penalty so the metric actually rewards kernel semantics, not just n-gram overlap. Sanity-checked: compute_codebleu(ref, ref) = 1.0000 on all 25 samples.
• Unbiased Pass@1. Uses the Chen et al. (2021) estimator Pass@1 = c/n with n=10 generations per task — no optimistic bias from sampling without replacement.
• One API key per model, per mode. Six isolated NIM keys with an RPM limiter so zero-shot and agent phases never contend for rate limits — results stay comparable.
• Parallel execution. All six phases run in independent screen sessions (launch_all.sh) — a 750-generation benchmark completes in hours, not days.

Benchmark Dataset

25 tasks from 13 real GitHub repositories across five GPU programming categories:

Category	Repos	Tasks
Core CUDA	NVIDIA/cuda-samples, moderngpu/moderngpu	9
Image Processing	opencv/opencv, NaderAlAwar/image-processing-cuda, NajatN/Parallel-Image-Processing-CUDA-, rpgolshan/CUDA-image-processing	6
Deep Learning	a-hamdi/GPU, SartajBhuvaji/Cuda	5
Physics / Simulation	Baey/N-Body-CUDA, niteya-shah/Fluid-Simulation-CUDA, vlvovch/lennard-jones-cuda	3
Sparse Matrix	shreyansh26/SparseMatrix-Computation-CUDA, bsampson1/SpMV-CUDA	2

Difficulty distribution: 4 easy · 9 medium · 12 hard

Each sample pairs a natural-language prompt (used verbatim for zero-shot generation) with the ground-truth reference_code from its source repository. Dataset sources are documented per-repo with URLs in data/.

Metrics

Extended CodeBLEU

CodeBLEU = 0.30 × N-gram + 0.40 × Keyword + 0.15 × Dataflow + 0.15 × Syntax

The keyword component (40% weight) uses 42 CUDA-specific keywords:

Category	Keywords	Weight
Kernel qualifiers	__global__, __device__, __shared__	4.0–5.0
Thread indexing	threadIdx, blockIdx, blockDim, gridDim	3.5–4.5
Synchronisation	__syncthreads, __syncwarp	3.5–4.5
Memory management	cudaMalloc, cudaFree, cudaMemcpy	3.5
Atomic operations	atomicAdd, atomicCAS, etc.	3.0–3.5
Warp primitives	__shfl_down_sync, __shfl_xor_sync	4.0
Kernel launch	<<<, >>>	4.0

A missing-keyword penalty of −0.10 per absent critical keyword (max −0.50) is applied. Sanity check:compute_codebleu(reference, reference) = 1.0000 for all 25 samples.

Pass@1

Unbiased estimator from Chen et al. (2021): Pass@1 = c / n with n = 10 generations per task, where a generation is valid if it passes all 6 heuristic validity criteria.

Models Evaluated

All served via NVIDIA NIM, one dedicated API key per model:

Model	NIM Identifier	Roles
Bielik 11B	speakleash/bielik-11b-v2.6-instruct	Zero-shot + Agent
Nemotron Nano 9B	nvidia/nvidia-nemotron-nano-9b-v2	Zero-shot + Agent
GPT-OSS 20B	openai/gpt-oss-20b	Zero-shot + Agent

Findings

• Bielik-11B is the strongest model: 0.888 Pass@1 zero-shot → 0.960 with the agent.
• The agent delivers its largest gains on the weakest model — Nemotron-Nano-9B nearly triples (0.156 → 0.400). Critique-and-refine is a force-multiplier where raw capability is thin.
• GPT-OSS-20B is stable — the agent neither helps nor hurts significantly (0.652 → 0.640), suggesting its zero-shot output already saturates the validator.
• CodeBLEU and Pass@1 disagree (agent CodeBLEU ↓ while Pass@1 ↑) — a concrete argument for evaluating generated code by behavior, not just token overlap.

Project Structure

gpu-kernel-agent-benchmark/
├── agent/
│ └── multi_agent_system.py # ReAct: CoderAgent, ReviewerAgent, RefinerAgent
├── benchmark/
│ ├── phase1_bielik_zeroshot.py # Phase 1 — Bielik zero-shot
│ ├── phase2_nemotron_zeroshot.py# Phase 2 — Nemotron zero-shot
│ ├── phase3_gptoss_zeroshot.py # Phase 3 — GPT-OSS zero-shot
│ ├── phase4_bielik_agent.py # Phase 4 — Bielik agent
│ ├── phase5_nemotron_agent.py # Phase 5 — Nemotron agent
│ ├── phase6_gptoss_agent.py # Phase 6 — GPT-OSS agent
│ └── run_benchmark.py # Parallel runner (all 6 phases)
├── data/
│ ├── benchmark_data.py # 25 tasks + metadata
│ ├── samples_11_15.py # Deep learning + image tasks
│ ├── samples_16_20.py # Physics + sparse matrix tasks
│ └── samples_21_25.py # Image processing + simulation tasks
├── evaluation/
│ ├── extended_codebleu.py # Extended CodeBLEU metric
│ └── model_evaluator.py # NIM API + Pass@k + RPM limiter
├── validation/
│ └── basic_validator.py # Heuristic CUDA validator
├── results/ # Full ReAct traces + scores (48 JSON files)
├── demo.py # Self-contained demo (no API keys)
├── launch_all.sh # Launch all 6 screen sessions
├── requirements.txt
├── .env.example
└── README.md

Getting Started

# 1. Install
pip install -r requirements.txt
python -c "import nltk; nltk.download('punkt'); nltk.download('punkt_tab')"# 2. Configure NIM keys
cp .env.example .env # add your 4 NIM API keys# 3. Run the no-API-key demo (self-contained, deterministic)
python demo.py
# 4. Full benchmark — 6 parallel screen sessions
bash launch_all.sh
# 5. Monitor
tail -f results/phase4_bielik_agent_*.log
screen -r p1_bielik_zs

.env format

NIM_BASE_URL=https://integrate.api.nvidia.com/v1
RPM=40
NIM_KEY_NEMOTRON=nvapi-... # Key for Bielik
NIM_KEY_MINISTRAL=nvapi-... # Key for Nemotron
NIM_KEY_MISTRAL=nvapi-... # Key for GPT-OSS
NIM_KEY_AGENT=nvapi-... # Key for agent phases

References

• Multi-Agent Software Development through Cross-Team Collaboration (arXiv:2408.08927)
• Evaluating Large Language Models Trained on Code (Chen et al., 2021)
• CodeBLEU: a Method for Automatic Evaluation of Code Synthesis (Ren et al., 2020)
• NVIDIA NIM

License

MIT — see LICENSE.