Ako4all

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.