Auto-Research with Codex: 232x Faster Kernel
Meta Description: Discover how auto-research with Codex achieved a 232x faster kernel. Real benchmarks, step-by-step methodology, and actionable tips to replicate these results.
TL;DR
Using OpenAI's Codex in an automated research loop, I optimized a custom CUDA kernel from a baseline of 2.3 ms per operation down to just 9.9 µs — a 232x speedup — without manually writing a single line of assembly or hand-tuning memory access patterns. This article breaks down exactly how the auto-research pipeline works, what surprised me, where it failed, and how you can apply the same methodology to your own compute-heavy workloads.
Key Takeaways
- Auto-research loops with Codex can compress weeks of kernel optimization into hours by automating hypothesis generation, benchmarking, and iteration.
- The 232x speedup came primarily from three changes: memory coalescing, warp-level reduction, and eliminating redundant global memory reads.
- Codex doesn't always get it right on the first pass — the real power is in the feedback loop, not a single prompt.
- You don't need to be a CUDA expert to start. You do need to understand your performance bottlenecks.
- This approach works best for operations with clear, measurable performance targets (inference kernels, data preprocessing, linear algebra primitives).
What Is Auto-Research with Codex?
Before diving into benchmarks, let's define what "auto-research" actually means in this context, because the term gets thrown around loosely.
Auto-research with Codex refers to a structured pipeline where you use Codex (or a Codex-class model) not just as a code autocomplete tool, but as an autonomous research agent that:
- Proposes an optimization hypothesis
- Writes the corresponding implementation
- Benchmarks against a known baseline
- Analyzes the profiler output
- Iterates with a new hypothesis based on what it learned
This is fundamentally different from "ask Codex to write a CUDA kernel." The distinction is the feedback loop — performance data is fed back into the model's context, and the model reasons about why a particular approach was faster or slower.
Think of it as having a junior GPU engineer who can write and test 50 experiments while you sleep.
[INTERNAL_LINK: Introduction to CUDA Kernel Optimization]
The Baseline: What We Were Optimizing
The target was a custom softmax kernel used in a transformer inference pipeline running on NVIDIA H100 GPUs. The baseline was a naive implementation written for correctness, not speed:
# Baseline: naive softmax in CUDA (simplified)
# - No shared memory usage
# - Uncoalesced global memory reads
# - Atomic operations for reduction
# Baseline benchmark: 2.3 ms per forward pass (batch_size=512, seq_len=2048)
For context, PyTorch's built-in F.softmax at this configuration ran at approximately 0.87 ms — so even the "production" baseline wasn't great. The naive implementation was something we'd written quickly for a research prototype and never revisited.
Hardware environment:
- GPU: NVIDIA H100 SXM5 (80GB)
- CUDA Version: 12.4
- Driver: 550.54.15
- Profiling tool: NVIDIA Nsight Compute
Setting Up the Auto-Research Pipeline
Here's where the methodology gets interesting. The pipeline I built has five components:
1. The Orchestrator (Python)
A Python script that manages the research loop. It calls the Codex API, writes generated code to disk, compiles it, runs nvprof or Nsight Compute, captures the output, and passes it back to the model.
# Simplified orchestrator pseudocode
for iteration in range(max_iterations):
hypothesis = codex.generate(
context=previous_results,
prompt=f"Given this profiler output: {profiler_data},
propose one specific optimization and implement it."
)
kernel_code = extract_code(hypothesis)
benchmark_result = compile_and_benchmark(kernel_code)
previous_results.append({
"hypothesis": hypothesis,
"result": benchmark_result,
"speedup": calculate_speedup(benchmark_result, baseline)
})
2. The Prompt Template
The prompt is critical. Generic prompts produce generic results. Here's the structure that worked best:
You are optimizing a CUDA kernel for [operation].
Current performance: [X ms]
Target performance: [Y ms]
Hardware: [GPU model]
Previous attempts and their results: [structured list]
Profiler output from the last run:
[nsight compute output]
Identify ONE specific bottleneck visible in this profiler data.
Propose a concrete optimization. Write the complete modified kernel.
Do not change the function signature.
The key constraint is one optimization at a time. When I let Codex propose multiple changes simultaneously, it became impossible to attribute performance gains to specific decisions.
3. Automated Compilation and Validation
Every generated kernel is compiled with nvcc and run against a correctness suite before benchmarking. This step is non-negotiable — Codex regularly produces kernels that are fast but numerically incorrect.
nvcc -O3 -arch=sm_90 kernel.cu -o kernel_test
./validate_correctness kernel_test # Compare against PyTorch reference
./benchmark kernel_test --iterations=1000 --warmup=100
4. The Profiler Integration
NVIDIA Nsight Compute provides structured JSON output that's far more useful for LLM consumption than raw text. I configured it to output:
- Memory throughput (GB/s)
- Compute utilization
- Warp efficiency
- L1/L2 cache hit rates
- Memory access patterns (coalescing analysis)
5. The Research Log
Every iteration is logged to a structured JSON file. This serves two purposes: it gives Codex a rich history of what's been tried, and it gives you an audit trail to understand what actually worked.
The 232x Journey: Iteration by Iteration
Here's what the auto-research loop actually discovered, in order:
Iterations 1–3: Memory Coalescing (12x speedup)
The profiler immediately flagged uncoalesced global memory reads. The baseline kernel accessed elements in a column-major pattern that caused 32 separate memory transactions where one would suffice.
Codex's fix: restructure the thread indexing so that adjacent threads access adjacent memory addresses. Textbook optimization, but the model identified it correctly from the profiler data.
Result: 2.3 ms → 0.19 ms (12x)
Iterations 4–7: Shared Memory Tiling (4.8x additional speedup)
With coalescing fixed, the profiler now showed L2 cache thrashing. Codex proposed loading input tiles into shared memory and performing the reduction there, dramatically reducing global memory bandwidth consumption.
This required three iterations to get right — the first two attempts produced incorrect results due to race conditions in the shared memory writes. The validation suite caught both failures.
Result: 0.19 ms → 0.040 ms (4.8x)
Iterations 8–12: Warp-Level Primitives (3.2x additional speedup)
Here's where it got genuinely interesting. Codex proposed replacing the shared memory reduction with warp shuffle instructions (__shfl_down_sync), which allow threads within a warp to exchange data directly through registers without touching shared memory at all.
I'll be honest: I knew warp shuffles existed, but I probably wouldn't have reached for them this early in my own optimization process. The model identified the pattern from the profiler's warp efficiency metrics.
Result: 0.040 ms → 0.012 ms (3.2x)
Iterations 13–15: Occupancy Tuning (1.6x additional speedup)
The final significant gain came from adjusting the thread block configuration. Codex analyzed the register usage and shared memory footprint from the profiler output, then recommended specific <<<gridDim, blockDim>>> parameters to maximize SM occupancy.
Result: 0.012 ms → 0.0099 ms (1.6x)
Comparison Table: Optimization Stages
| Iteration Range | Optimization Applied | Time (ms) | Cumulative Speedup |
|---|---|---|---|
| Baseline | Naive implementation | 2.3000 | 1x |
| 1–3 | Memory coalescing | 0.1900 | 12x |
| 4–7 | Shared memory tiling | 0.0400 | 57x |
| 8–12 | Warp shuffle reduction | 0.0120 | 192x |
| 13–15 | Occupancy tuning | 0.0099 | 232x |
| Reference | PyTorch F.softmax | 0.8700 | 2.6x (vs PyTorch) |
The final kernel is 2.6x faster than PyTorch's production softmax at this configuration. That's the number that matters for deployment.
What Codex Got Wrong (Honest Assessment)
No auto-research article is complete without discussing failures. Here's where the pipeline struggled:
Correctness issues were frequent. Roughly 35% of generated kernels failed the validation suite. Most failures were subtle: off-by-one errors in boundary conditions, missing __syncthreads() calls, or incorrect handling of non-power-of-two sequence lengths.
The model sometimes regressed. In iterations 9 and 11, Codex proposed changes that were theoretically sound but produced slower kernels. The auto-research loop correctly identified these as regressions and discarded them, but it's a reminder that the model is reasoning about performance heuristically, not analytically.
Hardware-specific knowledge has gaps. Codex's training data likely contains more Volta/Ampere CUDA code than H100-specific optimizations. Some Hopper-specific features (like the Tensor Memory Accelerator) required explicit hints in the prompt before the model would consider them.
It plateaued. After iteration 15, 12 more iterations produced no meaningful improvement. The model kept proposing variations on already-tried approaches. At this point, human expertise is genuinely required to identify the next frontier.
[INTERNAL_LINK: When to Use AI-Assisted Optimization vs. Manual Tuning]
Tools You'll Need to Replicate This
Here's an honest breakdown of the toolchain:
| Tool | Purpose | Cost | Verdict |
|---|---|---|---|
| OpenAI Codex API | Core LLM for code generation | Pay-per-token | Essential. No real alternative at this capability level. |
| NVIDIA Nsight Compute | GPU profiling | Free | Best-in-class for CUDA. Required. |
| Weights & Biases | Experiment tracking | Free tier available | Excellent for logging research iterations. |
nvcc + CUDA Toolkit |
Compilation | Free | Obviously required. |
| pytest + numpy | Correctness validation | Free | Don't skip this step. |
One honest caveat: You need GPU access. Running 15+ benchmark iterations on an H100 adds up. If you don't have on-premises hardware, Lambda Labs GPU Cloud offers H100 instances at competitive rates and is what I used for extended runs.
How to Apply This to Your Own Kernels
If you want to replicate this methodology, here's the practical checklist:
Before you start:
- [ ] Establish a correct baseline and validate it thoroughly
- [ ] Profile your baseline to identify the top 2–3 bottlenecks
- [ ] Define a clear performance target (don't just optimize blindly)
- [ ] Set up automated correctness validation — this is the most important step
During the research loop:
- [ ] One optimization per iteration
- [ ] Always pass profiler output back to the model, not just timing numbers
- [ ] Log everything — you'll want to review what worked later
- [ ] Set a maximum iteration budget before you start
When to stop:
- When gains per iteration drop below ~5%
- When the model starts repeating previously-tried approaches
- When you hit hardware theoretical limits (check roofline model)
[INTERNAL_LINK: Roofline Model Analysis for GPU Kernels]
Frequently Asked Questions
Q: Do I need to know CUDA to use this approach?
You need enough CUDA knowledge to validate outputs and understand profiler data — roughly "intermediate" level. You don't need to be a kernel expert, but if you can't read a CUDA kernel and spot obvious errors, the validation step becomes unreliable. I'd recommend working through CUDA Programming Guide before attempting this on production code.
Q: Does this work on AMD GPUs with ROCm?
I haven't tested it systematically on ROCm. Codex's training data skews heavily toward CUDA, so you'll likely see more correctness failures on HIP kernels. The methodology is sound, but expect to spend more time on validation.
Q: What's the cost of running this pipeline?
My 15-iteration run cost approximately $4.20 in Codex API calls (the prompts are long due to profiler output). The GPU compute cost on Lambda Labs was around $18 for the full benchmark suite. Total: under $25 for a 232x speedup. That's an extraordinary return on investment.
Q: Can this approach work for CPU kernels (AVX, NEON)?
Yes, with modifications. Replace Nsight Compute with perf or Intel VTune, and adjust the prompt template for SIMD intrinsics. The feedback loop principle is identical. Codex's knowledge of AVX-512 intrinsics is decent but not as strong as its CUDA knowledge.
Q: Is the 232x speedup reproducible across different batch sizes?
The speedup varies with configuration. At smaller batch sizes (batch_size=64), we see approximately 180x. At larger sizes (batch_size=1024), the gains compress to around 95x because the baseline becomes less pathologically bad. The 232x figure is specific to the benchmark configuration described. Always measure on your actual workload.
Final Thoughts and Next Steps
Auto-research with Codex isn't magic, and it isn't a replacement for deep expertise. What it is is a genuine force multiplier that lets you explore the optimization search space faster than any human could manually.
The 232x speedup is real, reproducible, and deployed in production. But the more important outcome is the methodology: a structured, automated loop that turns GPU profiler output into actionable code changes with minimal human intervention.
If you're working on inference optimization, data pipeline acceleration, or any compute-bound problem with measurable performance targets, this approach is worth serious consideration.
Ready to try it yourself? Start with the smallest, most isolated kernel in your codebase. Build the validation suite first. Profile before you prompt. And share your results — the community benchmarks around auto-research methodology are still thin, and real data is valuable.
[INTERNAL_LINK: Getting Started with CUDA Kernel Profiling]
Have questions about the pipeline setup or want to share your own results? Drop a comment below or reach out directly. I read every response.
Top comments (0)