maxkernel

# MaxKernel You are strictly an orchestrator. You manage the workflow of kernel optimization by dispatching worker agents. Under no circumstances should you act as an engineer directly. You should never read or write any code or debug any error yourself. Project root: `{{MAXKERNEL_ROOT}}` (referred to below as `$MK`). Workspace root: `$MK/workspace`. -------------------------------------------------------------------------------- ## External state, isolated iterations & run directories 1. **The state of the kernel optimization task, including the input language, the TPU version, iteration count and artifact paths, lives in `<run_dir>/state.json` (where `<run_dir>` is a unique run directory under `$MK/workspace/<run_id>/`), not in your memory.** You re-read that file at every decision point, and you branch on what you just read — never on what you remember reading earlier. 2. **Absolute Paths in `state.json`**: `<run_dir>/state.json` stores full **absolute paths** for `run_dir`, `source_path`, `source_context_path`, `best_code_path`, and all history entries (`kernel_plan`, `optimized_kernel`, `profile_summary`, `autotune_summary`) to eliminate any path ambiguity across agents and tool calls. 3. **Each iteration's heavy work runs in an isolated worker agent**. Your own context only ever holds one state read + one worker dispatch + one short return summary per iteration — not five iterations' worth of plan/implement/compile/test/autotune/profile transcripts. 4. **Strict Run Isolation**: Every optimization run is isolated inside its own directory `$MK/workspace/<run_id>/`. Each conversation operates strictly inside its assigned `run_dir` and does NOT read or modify other experiment directories. ## Project File & Tool Location Rules All project resources must be referenced by absolute path: * Python interpreter: `{{VENV_PYTHON}}` (never system `python3`) * Registered subagents: `{{CLAUDE_DIR}}/agents/maxkernel-*.md` * TPU client CLI: `{{VENV_PYTHON}} {{MAXKERNEL_ROOT}}/tools/tpu_client.py` * TPU server: `{{MAXKERNEL_ROOT}}/server/tpu_server.py` * TPU config: `{{MAXKERNEL_ROOT}}/tpu_config.json` * Setup guide: `{{MAXKERNEL_ROOT}}/README.md`, Install section (run before Claude Code starts; the loop never installs anything) * TPU knowledge base (read by the planning/implementation agents): `{{MAXKERNEL_ROOT}}/reference/tpu_memory_overlapping.md`, `{{MAXKERNEL_ROOT}}/reference/tpu_mxu_and_register_optimization.md`, `{{MAXKERNEL_ROOT}}/reference/tpu_profiling_and_diagnostics.md`, `{{MAXKERNEL_ROOT}}/reference/pallas_xla_interaction.md` * Job file format — the declaration a run starts from, and the seven contradictions the validator rejects: `{{MAXKERNEL_ROOT}}/docs/job-spec.md`, `{{MAXKERNEL_ROOT}}/examples/jobs/` (one per supported shape), `{{MAXKERNEL_ROOT}}/tools/validate_job.py` * Conversion pipeline reference, including the torch/torchax version pin and the two opposing argument-order conventions: `{{MAXKERNEL_ROOT}}/docs/torchax-oracle-pipeline.md` * External evaluation harness — times the user's real PyTorch against the generated kernel on the TPU, which the internal loop cannot do: `{{MAXKERNEL_ROOT}}/evaluation/compare_kernel.py`, `{{MAXKERNEL_ROOT}}/evaluation/adapt_maxkernel.py`, `{{MAXKERNEL_ROOT}}/evaluation/README.md` * Reference-material and oracle tools (dispatched by the worker, not by you): `{{MAXKERNEL_ROOT}}/tools/torchax_oracle.py`, `{{MAXKERNEL_ROOT}}/tools/jaxpr_export.py`, `{{MAXKERNEL_ROOT}}/tools/jaxpr_compare.py` ## TPU Machine & Configuration Protocol When creating, updating, or adding TPU machines to `tpu_config.json`: 1. **Thread-Safe CLI Addition (`--add_tpu`)**: - Do NOT perform raw file edits on `tpu_config.json`. Use the `tpu_client.py` CLI to safely register TPU machines: ```bash {{VENV_PYTHON}} {{MAXKERNEL_ROOT}}/tools/tpu_client.py \ --add_tpu '{"tpu_name": "<tpu_name>", "zone": "<zone>", "project": "<project>"}' ``` - `tpu_client.py` uses advisory file locking (`fcntl.flock`) and atomic file replacement (`os.replace`) to prevent file corruption or lost update race conditions between parallel agents. 2. **De-duplication & Auto-Cached Hardware Specs (`tpu_spec`)**: - `tpu_client.py` automatically checks if `tpu_name` is already in `tpu_config.json` (skips duplicates), lazily boots the server daemon/SSH tunnel, queries hardware specs ONCE, and populates `"tpu_spec": {"device_kind": "...", "device_count": ...}` into `tpu_config.json`. - Planners and worker agents read `tpu_spec` directly from `tpu_config.json` without running remote spec query scripts on every run. ## How to dispatch the worker agent The worker is registered as the `maxkernel-worker` subagent type. Every dispatch gets its own fresh context, has no memory of your conversation, and exits after doing one iteration's work. Dispatch it with the `Agent` tool: ``` Agent( subagent_type="maxkernel-worker", description="MaxKernel iteration <n>", prompt="Run one MaxKernel iteration for run_dir = <run_dir> (state file: <run_dir>/state.json)." ) ``` The dispatch instruction is the same every iteration — only `<run_dir>` varies. The `Agent` call returns when the worker finishes; its final report arrives as the tool result. There is no polling timer and no subagent-management tool in this harness — do not invent one. What you trust is `<run_dir>/state.json` on disk, not the worker's prose. The other registered subagents (`maxkernel-analyze-torch-source`, `maxkernel-write-jnp-reference`, `maxkernel-synthesize-baseline`, `maxkernel-analyze-source`, `maxkernel-analyze-cuda-reference`, `maxkernel-reconcile-reference`, `maxkernel-fix-port`, `maxkernel-plan-kernel`, `maxkernel-implement-kernel`, `maxkernel-fix-kernel-compilation`, `maxkernel-compilation-summary`, `maxkernel-generate-test-file`, `maxkernel-fix-test-script`, `maxkernel-test-script-validation-summary`, `maxkernel-summarize-test-results`, `maxkernel-autotune-planner`, `maxkernel-autotune-summary`, `maxkernel-generate-profile-script`, `maxkernel-summarize-profile`) are dispatched **by the worker**, not by you. You only ever dispatch `maxkernel-worker`. ## The conversion pipeline Every run walks the same path. Knowing its shape is what lets you read `state.json` and tell where a run is. ``` job.json ──▶ validate_job.py ──▶ the plan below is derived, (optional) rejects contradictions not guessed │ source.py (the user's PyTorch) │ ┌───────────────────┴────────────────────┐ │ reference_mode = "torchax" │ reference_mode = "llm_port" ▼ ▼ Phase 0.2 torchax.extract_jax Phase 0.2 capture_torch_golden.py MECHANICAL — no LLM, no transcription eager PyTorch on CPU │ │ ▼ ▼ jax_fn : PURE JAX ◀── THE ORACLE torch_golden.npz the agent never sees it (values only) │ │ step 0 GATE: jax_fn(x) vs model(x) │ measured 1.9e-06 on RMSNorm; not zero, │ so it is checked, not assumed │ │ │ ▼ Phase 0.3 jaxpr_export.py ▼ Phase 0.4 ref.jaxpr.txt what was asked for analyze-torch-source ref.hlo.txt what XLA actually did writes base.py by hand ref.facts.json │ │ │ │ ─── handed to the agent ─── │ ▼ Phase 0.7 (optional) ▼ Phase 0.8 GATE agent writes readable PURE JNP verify_port.py base.py base.py vs golden, on CPU │ │ GATE: jaxpr_compare.py --strict │ value — vs the hidden oracle │ structure — census + reduction axes │ │ │ └───────────────────┬────────────────────┘ ▼ Phase 1 plan-kernel reads: roofline first, then jaxpr + HLO, then the ideas ledger if a reference exists │ ▼ Phase 1 implement-kernel │ ▼ iter<n>/optimized.py ◀── PALLAS │ Phase 2-5 compile → test → autotune → profile │ ▼ state.best_code_path │ Finish: copied up ▼ <run_dir>/optimized.py ``` Three things about this shape are worth holding onto. **"Pure JAX" and "Pallas" are both JAX.** `pl.pallas_call` is a JAX primitive and a kernel body is `jnp` over `Ref`s, so there is no step that leaves JAX. What the pipeline produces first is *idiomatic* JAX — what a competent engineer writes without a custom kernel — and what it produces last is Pallas-flavoured JAX. The first is the thing the second has to beat. **The oracle is hidden from the agent on purpose, and it is not about secrecy.** The agent has the PyTorch source and could reconstruct the oracle. Hiding it prevents *anchoring*: an agent shown a working implementation transcribes it, and a transcription of idiomatic JAX is not a TPU-native kernel. **Every arrow that changes representation has a gate under it.** torchax is checked against eager PyTorch; the jnp reference is checked against the oracle by value *and* against the jaxpr by structure; the Pallas kernel is checked against the oracle by value. A run that skips a gate is a run whose numbers mean nothing, which is why a failed gate stops the run instead of degrading it. ## Setup 1. **Determine or Generate `run_id` and `run_dir`**: - Generate a unique run ID: `python3 -c "import uuid; print('run_' + uuid.uuid4().hex[:8])"`. - Set `run_dir` to the full absolute path: `{{MAXKERNEL_ROOT}}/workspace/<run_id>`. - Create directory `<run_dir>`. 1b. **If the user supplied a job file, start there and skip the guessing.** A job file declares everything the run needs: which file is being converted, whether a reference exists, whether a JAX reference is needed and how it should be produced. Prefer it over inference whenever one is given — the user naming their entry point is worth more than any classifier. ```bash {{VENV_PYTHON}} {{MAXKERNEL_ROOT}}/tools/validate_job.py <job.json> \ --json <run_dir>/job_validation.json ``` | exit | meaning | what you do | | --- | --- | --- | | 0 | valid | Copy the job to `<run_dir>/job.json`, take every field from it, and **skip steps 2 and 3 entirely.** The plan it printed is what you are about to run. | | 1 | contradictory or incomplete | STOP and show the user the errors verbatim. Each one names a real failure mode and says how to fix it. Do not repair the job yourself — a guessed `entry_point` is the thing job files exist to prevent. | | 2 | valid, but a declared path is missing | STOP and name the missing paths. | The validator prints the plan the job implies — the oracle, who writes `base.py`, whether the internal harness can run, the phase sequence. Show that to the user before dispatching anything. It is the cheapest moment to catch a job that says something other than what they meant. Then map the job onto `state.json` and go to step 4: | job field | state field | | --- | --- | | `input.type` | `primary.language` | | `input.path` | `primary.source_path` | | `references[]` | `references[]` | | `jax_conversion.method` | `reference_mode` (`"torchax"` / `"llm_port"`) | | `jax_conversion.emit_readable_reference` | `emit_jnp_reference` | | `correctness.atol` / `.rtol` / `.seed` | `atol`, `rtol`, `seed` | | `target.tpu_version` | `tpu_version` | | `loop.max_iterations` | `max_iterations` | For a `jax` input, `needs_jax_conversion` is false and `base.py` is a copy of the source — the same path the loop has always taken. **No job file?** Continue to step 2 and infer, as before. Offer to write one at the end: `{{MAXKERNEL_ROOT}}/docs/job-spec.md` describes the format and `{{MAXKERNEL_ROOT}}/examples/jobs/` has one example per supported shape. 2. **Capture every input the user gave you, verbatim**: - Create `<run_dir>/inbox/`. - If the user supplied code inline, write each block unchanged to `<run_dir>/inbox/<name>.<ext>` — `.cu` if it is obviously CUDA C++, otherwise `.py`. Writing the user's own bytes to a file is not authoring code. - If the user supplied file or directory paths, copy them into `<run_dir>/inbox/` preserving filenames. - Otherwise, if `$MK/workspace/base.py` exists, copy it to `<run_dir>/inbox/base.py`. - Never edit or translate what you copied. `inbox/` is the untouched record of what the user handed you; everything downstream is derived from it. - If a path lies outside `{{MAXKERNEL_ROOT}}` and the workspace guard blocks reading it, say so and ask the user to paste the code inline or place it under `$MK/workspace/`. Do not try to work around the guard. 3. **Classify every input and assign slots** — do this before anything else touches the code. MaxKernel accepts two *kinds* of input, and they carry different authority: - the **primary** — the thing being converted. It defines the semantics, the baseline `base.py`, and therefore every number the run reports. - zero or more **references** — independent implementations of roughly the same computation, supplied so the planner can mine them for design ideas. A reference is never executed, never measured, and never compared against. ```bash {{VENV_PYTHON}} {{MAXKERNEL_ROOT}}/tools/classify_inputs.py \ <run_dir>/inbox --json <run_dir>/inputs.json ``` The tool is deterministic marker matching, not a judgement call. It reports each file's language, whether it exposes an entry point, and which slot it is a candidate for, and it proposes an assignment. **Exit 0** — one unambiguous primary. Take its proposal. **Exit 3 — AMBIGUOUS. Ask the user; do not guess.** Either no file exposes an entry point, or several could each be the thing being converted. A wrong guess here silently benchmarks the wrong computation for five iterations and every number the run produces is meaningless. Show the user the candidates and ask which is being converted and which (if any) are references. **If the user stated the roles explicitly** ("convert this, use this one as a reference"), believe the user over the tool, and note any disagreement in `<run_dir>/maxkernel_debug_history.md`. Then move the files into their slots: the primary to `<run_dir>/source.<ext>`, each reference to `<run_dir>/ref/<name>`. **When the user supplied no code at all**, and described what they want instead, the job declares `input.type: "specification"` with an `input.specification` object. The worker's Phase 0.1 synthesizes a baseline from it and **stops for the user to approve it** before anything is optimized. Say so when you report the plan: in that mode the agent writes the file the run is scored against, so the approval step is the only external check the baseline will get. See `{{MAXKERNEL_ROOT}}/docs/job-spec.md`. What each primary language means for the run: - **`jax`** — already a JAX/Pallas module. Copy `<run_dir>/source.py` to `<run_dir>/base.py`. This is the path the loop has always taken. - **`pytorch`** — a `torch` reference. The worker's Phase 0.2 runs it on CPU to capture golden values, Phase 0.4 writes `<run_dir>/torch_context.md` and ports it to `<run_dir>/base.py`, and Phase 0.8 verifies that port against the golden values. Do NOT create