ako4all

name

ako4all

description

Drive an agentic loop that iteratively optimizes a GPU kernel for maximum speedup. Use this skill whenever the user wants to optimize / speed up / benchmark a GPU kernel (CUDA, Triton, TileLang, C++, Python), mentions AKO / AKO4ALL / AKO4X / agentic kernel optimization, asks to "make this kernel faster", or has a kernel they want measured against a PyTorch reference. The skill handles setup, profiling (ncu), correctness checking, iteration logging, and git commits. Bootstraps a workspace in any directory the user points at.

AKO4ALL — Agentic Kernel Optimization

Drive a profile → modify → benchmark → log → commit loop on a GPU kernel until it runs faster than the reference. The user provides at minimum a kernel; everything else (reference, inputs, bench script, hints) is optional.

When this skill applies

• "optimize this kernel" / "speed up this CUDA / Triton / TileLang kernel"
• "run AKO / AKO4ALL on ..."
• "benchmark this kernel against PyTorch"
• "iterate on this kernel until it's faster"
• mentions of ncu, kernel profiling, GPU speedup target

Does NOT apply when:

• User wants to write a new kernel from scratch with no optimization target — just write code, no loop.
• User wants Codex / GPT to review or implement — use codex:rescue instead.
• User wants generic performance advice for code that isn't a GPU kernel.

First action

Before doing anything else, establish the workspace — the directory the loop runs in. It is typically the user's CWD, or a subdirectory / path they name in the prompt.

Inventory the workspace + prompt

Browse the workspace (don't run a fixed checklist — look around) and read the user's prompt to identify what the loop needs:

• Kernel (required) — the code to optimize
• Reference (optional) — correctness golden
• Input data (optional) — data files the kernel consumes (.npz, .bin, shape lists, custom formats, etc.)
• Knowledge (optional) — reference materials the user wants you to consult: algorithm notes, papers, design docs, prior PRs. Typically under knowledge/ but anywhere the user points at.
• Bench mode — user-provided bench script vs. default bench/kernelbench/ evaluator
• Scaffold presence — whether bench-wrapper.sh, HINTS.md, ITERATIONS.md, bench/kernelbench/ are already at workspace root

Whether the workspace follows AKO4ALL's source/ / knowledge/ / bench/ naming or some entirely different layout is not the signal. What matters is whether you can identify each item above with confidence.

Ask only when genuinely uncertain

If the user's prompt + filesystem give you confidence about every required item, don't ask — skip straight to presenting the plan. Ask only when a piece's role is genuinely ambiguous (a kernel-shaped file with no obvious reference, two files that could both be the kernel, an input data file in an unfamiliar format you need permission to wire up a custom way, etc.). When in doubt, asking is cheaper than guessing wrong.

Always present the resolved plan before running anything

Whether you asked the user or not, list back what you decided — so the user can correct you even when you didn't think you needed to ask.

Use the format below. Bold field labels + inline-code path values + the leading emoji marker make the plan visually scannable in any terminal theme (don't flatten to a wall of prose):

📋 Resolved Plan

• Workspace — <path>
• Kernel — <path>
• Reference — <path> (or none — will use original kernel)
• Input data — <path> (or inline in ref, or none)
• Knowledge — <path> (or none)
• Bench mode — default (KernelBench) (or custom: <path>)
• Scaffold to copy — <list of missing files> (or none — already present)

If anything still feels uncertain at this point, stop and ask. Otherwise proceed to Workflow.

Bringing in scaffold

When copying scaffold (bench-wrapper.sh, bench/kernelbench/, starter HINTS.md / ITERATIONS.md, workspace.gitignore → as .gitignore in the workspace) from this skill's own directory into the workspace, do not overwrite files that already exist — the user may have edited HINTS.md, or ITERATIONS.md may carry prior iteration history. Copy only what's missing.

Persisting user-supplied hints

The user may supply behavior directives in two ways:

• Inline in the prompt — e.g., "do not use shared memory" or "prefer Triton".
• External file reference — e.g., "follow rules in /tmp/x.md" or "see hints.md".

In both cases, merge those directives into HINTS.md. It's the persistence layer — directives that only live in the current session's plan are lost on resume.

Surfacing HINTS.md changes

Whenever you merge directives into HINTS.md, tell the user explicitly what happened. Example phrasings:

"I added your 'avoid shared memory' directive from the prompt to HINTS.md." "I added the 3 rules from /tmp/user-hints.md to HINTS.md."

Without this acknowledgment the user can't tell from your reply whether you added, replaced, or silently dropped their directives. Always name the source ("from your prompt" / "from /tmp/x.md").

Workflow

1. Analyze inputs. Building on the inventory above, confirm class Model and get_inputs() can be assembled for default bench mode; if not, stop and ask the user. See bench/kernelbench/GUIDE.md for the input assembly contract (KernelBench-format input / raw kernel / kernel + separate data file / external path patterns).

2. Create branch. git checkout -b opt/<kernel-name>. If the workspace isn't a git repo, init one first.

3. Initialize solution. Create solution/ and scripts/. Copy the kernel implementation files into solution/ (only the kernel itself — reference / inputs helper files stay at their resolved locations). Do not copy or mkdir canonical directories (source/, input/, etc.) when the user's files already exist elsewhere. Point bench.sh's --ref and --inputs flags at the resolved paths in place. solution/ is the only directory the loop owns.

4. Generate bench.sh. Build the bench command with adjusted paths, pipe through 2>&1 | tee _bench_output.txt. Replace {{BENCH_COMMAND}} in bench-wrapper.sh to produce scripts/bench.sh. For default bench mode the command is python bench/kernelbench/bench.py --ref <ref> --solution solution/<kernel> [--inputs <inputs-file>] --verbose — include --inputs only when inputs are defined outside the ref file. Do not hardcode --backend in the rendered command; bench.py auto-detects backend from solution source. Add --backend only to override the sniff (explicit HIP labelling or mixed-backend solutions). scripts/bench.sh is a starting template — when the bench env needs setup (conda activate, sub-env python paths, multi-CUDA toolkit selection), edit it freely; preserve only the trajectory section (LABEL/TIMESTAMP handling and cp -r solution/* "$TRAJ_DIR/").

   Common env friction: base shell often has no python on PATH when python lives in a sub-env (e.g. ~/anaconda3/envs/py312/bin/python). Tools that internally subprocess python (sol-execbench CLI, torch cpp_extension.load_inline) will then fail with command not found. Workaround: put PATH=<env-bin>:$PATH at the top of scripts/bench.sh (or source <conda>/etc/profile.d/conda.sh && conda activate <env>). Discover sub-envs via ls /home/*/anaconda3/envs/*/bin /root/*/envs/*/bin /opt/conda/envs/*/bin 2>/dev/null.

5. Verify baseline. Run bash scripts/bench.sh. Expect CORRECT=True. If not, diagnose and fix before iterating. Commit: git add -A && git commit -m "[baseline] Initialize solution and benchmark". Then run ncu once on the baseline to inform iter-1 direction.

Iteration protocol

Every modification to solution/ followed by a bench run = one iteration. Number sequentially (1, 2, 3, …). Each iter is exactly three steps:

1. bash scripts/bench.sh iter-N — label is required, must match iter-N format.
2. Append a structured entry to ITERATIONS.md (template inside that file).
3. git commit -m "[iter N] <short description of optimization direction>".

Steps 2 and 3 MUST be the next two tool calls after step 1 — no ncu, no probes, no reads, no planning the next iter between them. A failed or partial bench is still an iter; log + commit first, debug after. This is the most-missed step in practice: agents read the bench result and telescope into next-iter analysis (probes, ncu, hypothesis forming) without closing out the current one, leaving commit gaps with ITERATIONS.md entries written from memory later.

Backstop: if you catch yourself starting a new iter (Editing solution/, or running ncu/probes for the next direction) and git log -1 doesn't show [iter N] ..., stop and finish the prior iter's steps 2 and 3 first. Related experiments that belong together narratively get grouped in ITERATIONS.md analysis prose, not in batched git commits.

Profile to identify bottlenecks — see "ncu profiling" below for the ncu workflow and analytical fallback. Do not optimize blindly.

Keeping the iteration loop fast

A bench run must be cheap enough to iterate against (seconds to low minutes). When it isn't, the cause is almost always an expensive reference — it's re-run for every correctness trial and re-timed for the speedup denominator, yet it's invariant across solution edits, so most of that cost is wasted. This is an eval-time problem, not a metric problem: never change what you compare against to make the bench cheaper.

Separate the per-iteration signal from the final verdict:

• Signal (every iter): rank candidates by the solution's own RUNTIME (lower is better) — the reference contributes nothing to comparing two solutions.
• Verdict (before committing a winner / a final): a full run — full trial counts, reference measured, real SPEEDUP, reward-hack check.

Whose eval is it determines what you may touch:

• Default bench (bench/kernelbench/bench.py — the skill owns it): pull levers freely, cheapest-and-safest first — --no-ref (skip reference timing; REF_RUNTIME/SPEEDUP → -1, COMPILED/CORRECT/RUNTIME unaffected) → trim --num-perf-trials (e.g. 100→20; latency noise only) → trim --num-correct-trials (higher risk: weakens the fresh-input anti-cheat and still runs the ref once per trial, so keep ≥1 in the loop, full at the gate).
• User-provided eval (custom {{BENCH_COMMAND}}): the trial counts, correctness rounds, and reference handling are the user's contract — do not inject --no-ref or cut counts on a script you didn't author (it may have no such flag, break the interface, or silently invalidate the measurement / a leaderboard's required N). Use only the fast-iteration switches the user exposed (flags / env vars documented in the prompt or HINTS.md). If iteration is too slow and none exist, raise it with the user — don't fabricate one.

Caching the reference's runtime across iterations is sound only on a clock-locked GPU; on unlocked clocks (ref and solution timed in different clock states) prefer ranking by the solution's own latency.

Stall handling

When 3 consecutive iterations show no improvement (≥3% over current best), pause the loop and re-assess before iter N+1. Re-assessment combines:

• Re-profile with ncu if available, or re-read runtime stats from ITERATIONS.md (min vs mean, distribution shape) if not.
• WebSearch for op-specific best-known techniques / numbers on the same hardware class.
• Review ITERATIONS.md for patterns (which axes have been tried, which haven't, where prior wins came from).

Default outcome: pick a new direction and continue. Only escalate to stop (see next section) if re-assessment produces concrete evidence the current state is at a physical floor.

When to stop

Legitimate triggers:

1. User-specified iteration cap reached (in prompt or HINTS.md).
2. Stall re-assessment produced hard evidence of a floor — e.g., min runtime at cuda-event timer resolution, kernel arithmetic at HBM bandwidth limit, launch overhead dominating compute. Cite the evidence in ITERATIONS.md.
3. All viable directions exhausted: document at least 3 distinct directions tried (with their iteration numbers) in ITERATIONS.md before invoking this trigger, to prevent premature stops.

Do not stop silently because tooling is unavailable — that's a re-assessment input, not a stop reason.

HEAD handling on stop

After deciding to stop, leave HEAD at the best-performing iter — not necessarily the latest. Procedure:

1. Identify the best iter by reading ITERATIONS.md Summary, the bench output for each iter under trajectory/, and your own reasoning notes. Useful signals from KernelBench output: mean speedup, runtime std (consistency), min runtime (tail), CORRECT flag. Other bench harnesses expose different shapes — use what's available. Justify your pick in the commit message (e.g., "iter 4: best mean AND lowest min, while iter 6 ties on mean but has higher std").

2. If best iter ≠ latest iter:

   • git checkout <best-iter-sha> -- solution/ — verbatim copy, do NOT hand-reconstruct from memory or earlier notes.
   • bash scripts/bench.sh final to sanity-verify on a fresh run.
   • git commit -m "[final] Restore iter-K (X.XXx) — <one-sentence why>".

The git checkout step is mandatory. Manual reconstruction risks introducing silent drift from the actually-benched code.

ncu profiling — best effort, not a gate

Probe ncu once after baseline. If it fails (driver mismatch, missing toolkit, user opt-out via free-text HINTS.md directive), proceed analytically for the rest of the loop without re-probing within this session. Don't gate iteration progress on ncu availability; analytical reasoning + runtime stats from the bench harness are a valid substitute for the optimizer's direction picking.

Gotchas

• Pursue genuine latency reduction, not reward hacking. No CUDA stream injection to evade timing, no monkey-patching the benchmark, no returning uninitialized results. The built-in evaluator flags >10× speedups for a reason — investigate before celebrating.
• The solution file must not contain get_inputs / get_init_inputs. The bench script strips the solution's module-level tail before eval as an anti-cheat boundary. Inputs come from the reference or --inputs file, never the solution.
• get_inputs() must produce fresh data each call. Bench calls it 5+ times across trials. Module-level cached tensors make correctness checks trivially pass and let timing measure cache-warm performance. Use torch.randn or reload from disk on every call.
• Don't be lazy. Stay-in-PyTorch, only-tune-configurations, skip-profiling — these are the default low-effort failure modes for agents. The point of the loop is to rewrite the implementation — switch languages (Triton → CUDA, etc.) when it helps.

Reference files

• bench/kernelbench/GUIDE.md — full input assembly patterns, CLI flags, timing methods, tolerances. Read this before writing the bench command if anything about input shape, precision, or backend is non-obvious.
• HINTS.md (workspace) — user-editable behavior directives. Read at session start; respect any constraint named there.
• ITERATIONS.md (workspace) — your own iteration log. Write to it every iteration.