Apex: Instructions file for Codex

AGENTS.md

This file provides guidance to the Codex agent when working with code in this repository.

Project Overview

Apex is an RL environment for GPU kernel optimization on AMD ROCm hardware. It trains LLM agents to optimize GPU kernels through a pipeline:

prompt constructor → LLM agent → output/ → grader (Magpie) → score

The agent receives a baseline kernel, writes an optimized version to output/<task_id>/solution.{py,hip}, and is scored on compilation (+20 pts), correctness (+100 pts), and speedup (×100 pts).

Default target: MI355X / gfx950 (CDNA4). Also supports gfx942 (MI300X), gfx940 (MI300A), gfx90a (MI250X).

Environment (always set first)

Ensure you are in the Apex repo root (the directory containing this AGENTS.md), then activate the venv and set MAGPIE_ROOT:

cd "$(dirname "$(readlink -f AGENTS.md)" 2>/dev/null || pwd)"
source .venv/bin/activate
export MAGPIE_ROOT=$(cd ../Magpie && pwd)

All commands below assume this working directory (APEX_ROOT).

CRITICAL: Read-Only Codebase

Do NOT modify any files in this repository or in the Magpie directory ($MAGPIE_ROOT). This includes:

• This repo (Apex) — all .py, .sh, .md, .yaml files
• The Magpie repo ($MAGPIE_ROOT) — all files including benchmark scripts
• The Python virtual environment

You may only write to:

• Your designated results directory (passed via -r)
• Temporary directories under /tmp/
• The output/<task_id>/ directories inside your results directory

If a benchmark script or pipeline step fails, do NOT patch the source code. Instead, work around it via environment variables, command-line flags, or report the issue.

GPU cleanup (only if needed)

Before launching a GPU workload, check if GPUs are free. Skip cleanup entirely if rocm-smi already shows ~0% VRAM on the target GPU(s).

Only if GPUs are busy, run:

pkill -9 -f 'vllm serve' 2>/dev/null
pkill -9 -f 'VLLM::EngineCore' 2>/dev/null
pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null
sleep 3