KernelForge

KernelForge is a research monorepo for improving LLM generation of Triton GPU kernels with constrained generation and validation. The repo combines benchmark tooling, exploratory notebooks, grammar experiments, a course Flex/Bison compiler component, inference backend experiments, and the current kernel-development agent loop.

The current project is research-first. Reusable Python code lives under src/kernelforge; notebooks and scripts are interfaces around that package.

Project goals

• Measure how different LLMs perform on niche Triton kernel generation tasks.
• Validate generated kernels against known-good Triton kernels and their expected values before treating results as usable benchmark data.
• Design constrained-generation grammars that make invalid or non-existent Triton APIs harder for the model to emit.
• Build and harden a differentiator: an agent loop for developing kernels. The current implementation uses Pi Agent tools under apps/agent/.
• Maintain a first MVP validation loop with a basic semantic checker before expanding the agent workflow.
• Maintain a Flex/Bison syntax validator for the course requirement.
• Keep benchmark, inference, evaluation, and notebook workflows reproducible.

Current direction

The near-term focus is not a broad product surface. It is hardening and measuring a validated kernel-development loop:

1. Generate candidate Triton kernels.
2. Run a basic semantic checker over obvious Triton mistakes.
3. Validate candidates against known-good kernels and values.
4. Record benchmark results in reproducible ledgers.

Pi Agent is the current agent-loop adapter. Its tools call KernelForge scripts for task loading, semantic checks, optional constrained decoding, Modal validation, and result ledger writes without forcing a large repo reshaping.

The constrained-decoding path is now centered on XGrammar through the Modal vLLM backend's structured-output support. grammar/triton.gbnf is the active grammar artifact for those experiments; plain llama.cpp GBNF remains useful for local debugging, but the production smoke tests should exercise the vLLM/XGrammar path.

Repository map

.
├── compiler/                # Course Flex/Bison Triton JIT syntax validator
│   └── report/              # Typst report for the compiler component
├── benchmarks/              # First-party benchmark datasets, including KAGBench
├── docs/                    # Onboarding, evaluation, architecture notes
├── grammar/                 # GBNF grammar assets and viewers
├── notebooks/               # Marimo exploration and visualization notebooks
├── scripts/                 # Operational one-file scripts, e.g. Modal inference
├── src/kernelforge/         # Installable reusable Python package
│   ├── agent/               # KAGBench agent workflow, prompts, ledgers, checks
│   ├── benchmark/           # TritonBench loaders, prompts, inference, results
│   └── grammar/             # Python grammar/constrained-generation utilities
├── tests/                   # CPU-first regression tests and opt-in corpus/GPU tests
├── apps/                    # Agent-loop adapters and multi-file surfaces
├── poster/                  # Presentation/poster assets
├── vendor/TritonBench/      # Vendored upstream benchmark data/scripts
└── runs/                    # Generated ledgers/results; git-ignored

For the rationale behind these boundaries, see docs/architecture.md.

Setup

uv sync

For local ROCm/PyTorch/Triton execution:

uv sync --extra rocm

If you use Nix/NixOS:

nix develop .#default
# or, for local ROCm execution
nix develop .#rocm

Common commands

# List configured inference models
uv run python -m kernelforge.benchmark.llm_inference --list-models

# One-row dry run against the configured benchmark prompt builder
uv run python -m kernelforge.benchmark.llm_inference \
  --model lightning-ai/gemma-4-31B-it \
  --limit 1 \
  --dry-run \
  --dry-run-preview 1

# Modal Gemma backend
uv run modal serve scripts/modal_vllm.py

# List KAGBench tasks for the reusable agent workflow
uv run python -m kernelforge.agent list-tasks

# Build one KAGBench workflow payload without calling models or GPUs
uv run python -m kernelforge.agent run --limit 1 --dry-run

# Course compiler front end
make -C compiler compiler scanner
make -C compiler test

# Interactive notebooks
uv run marimo edit notebooks/grammar.py
uv run marimo edit notebooks/tritonbench.py

# Static checks
uv run ruff check .
uv run basedpyright

Use kernelforge.benchmark for reusable benchmark imports and CLI entry points.

Credentials

Set only the provider credentials needed for the workflow you are running:

export LIGHTNING_API_KEY="..."
export MODAL_API_KEY="..."      # or MODAL_API_TOKEN
export MODAL_API_SECRET="..."
export ANTHROPIC_API_KEY="..."  # Pi Agent brain
export GOOGLE_API_KEY="..."     # older Google AI Studio notebook path only

Keep local .env* files uncommitted.

Current stable entry points

• benchmarks/KAGBench/: agentic Triton synthesis benchmark built from TritonBench cases, with public tests, hidden/unit tests, prompts, and PyTorch references. The current tree contains 48 case directories.
• kernelforge.agent: reusable teacher/implementer workflow over KAGBench, including prompt construction, grammar request shaping, static checks, local/Modal evaluation adapters, and JSONL ledgers.
• kernelforge.benchmark.tritonbench: TritonBench-T loading, prompt construction, generated-code cleanup, and lightweight local evaluation helpers.
• kernelforge.benchmark.llm_inference: resumable OpenAI-compatible batch inference CLI over TritonBench-T simple tasks.
• kernelforge.benchmark.llm_results: JSONL loading, syntax metrics, cost estimates, and notebook-friendly tables.
• scripts/modal_vllm.py: Modal deployment for the Gemma 4 E4B vLLM backend.
• apps/agent/: Pi Agent tools for generating, checking, validating, and logging kernels.
• compiler/: Flex/Bison course compiler front end for top-level @triton.jit blocks, with scanner/parser binaries, tests, and a Typst report.
• grammar/triton.gbnf: current XGrammar/vLLM constrained-decoding grammar experiment.

Generated inference and evaluation ledgers should go under runs/.

More documentation

• docs/onboarding.md: detailed setup, provider, inference, and troubleshooting workflow.
• docs/evaluate.md: Colab/local evaluation notes for generated kernels.
• docs/architecture.md: monorepo boundaries and migration rule.
• compiler/report/: course compiler report source.
• docs/compiler_oral_exam_answers.md: English oral-exam notes for the compiler component.