Tlx Kernel Optimization Agent

Run The TLX Kernel Optimization Agent

The executable is third_party/tlx/tools/agents/kernel_optimization/cli.py. Follow the layers below in order. Target-specific profiling rules supplement, but never replace, the generic workflow.

Layer 0: Invocation And Safety

1. Invoke the public CLI. Do not manually optimize the kernel or substitute a generic subagent before the first CLI attempt.
2. Resolve the repository root, absolute kernel path, and target bundle. Record initial source-control status and treat existing bytes as the user baseline. Never clean or revert a dirty worktree.
3. Keep candidate generation isolated from the live checkout. Candidate source may exist only in the provider's temporary workspace and output artifacts until final promotion.
4. Run stdout and stderr separately. Tee stderr to a stable absolute live log; write stdout JSON to a separate artifact. When starting a run, respond only with the absolute log path unless the user asks for more.
5. The CLI commits a revalidated winner by default. Use --no-commit-winner only for an explicitly requested artifact-only run. Submission remains a separate explicit action.
6. If the live kernel changes concurrently, do not overwrite it. Preserve the winner artifact and report the conflict.
7. When continuing a completed optimization, pass --prior-run <output-dir>. This imports prior evidence and source hashes for cross-run deduplication but never adopts the old winner or replaces validation of the current kernel.
8. Keep production kernel authoring and heuristic tuning independent:
   • Use --task authoring --op <op> --arch <arch> --suite <suite> to author against the production kernel, harness, and cases. This does not tune the heuristic policy.
   • Add --tune-after-authoring only when source changes may alter the winning configurations across the production suite and the user wants both steps. The flag is opt-in and tuning runs only after authoring succeeds.
   • Use --task tuning --op <op> --arch <arch> --suite <suite> when the kernel implementation is already ready and only the heuristic policy needs work.

Reading agent implementation is allowed only after the CLI reports an internal failure that requires diagnosis. The first attempt must use the public contract.

Layer 1: Inputs And Standard Loop

The CLI needs:

kernel.py                  complete source file
bundle/harness.py          build, verify, benchmark, optional profile
bundle/cases.json          workloads, weights, protected cases
bundle/target.json         backend, architecture, device, environment
output/                    fresh artifact directory

Optional inputs are reference_kernel.py and budget.json. The higher-level coding agent owns target-bundle preparation; the TLX Agent consumes the bundle as a frozen trust boundary and must never generate or modify it during the optimization loop.

Before invoking the CLI, the higher-level agent must:

1. Search for an existing bundle that exactly matches the kernel entry point, workload, mode, backend, and architecture. Do not silently reuse a nearby shape or provider.
2. If no exact bundle exists, create a run-specific bundle from the nearest authoritative correctness test and production benchmark. Keep it outside the candidate workspace and do not modify the user's kernel permanently.
3. Encode all requested workloads in cases.json, hardware and environment in target.json, and kernel-specific invariants, known failed experiments, and evidence-to-knob guidance in target.json as optimization_guidance.
4. Run the bundle against the untouched kernel before launching the Agent. Confirm build succeeds, every protected case passes, repeated benchmark samples are stable, and summary profiling selects the intended kernel.
5. Negative-test the trust boundary with an intentionally invalid or incorrect temporary candidate and confirm build or verification rejects it.
6. Freeze the validated bundle for the duration of the run. Record its absolute paths and content hashes in the run log, and pass all paths explicitly to the public CLI.

Read references/input-contract.md for the complete construction and validation checklist. Do not start the optimization loop until the bundle is validated.

Candidate generation automatically receives trusted source-optimization skills owned by third_party/tlx/tools/agents/kernel_optimization/skills/: every target receives layout-conversion efficiency guidance, while CUDA/NVIDIA targets also receive async TMA output publication and warp-barrier efficiency guidance. Known Hopper and Blackwell targets additionally receive NVIDIA persistent pipeline efficiency guidance, and Blackwell targets also receive persistent CLC scheduling guidance. Unknown NVIDIA architectures receive the common NVIDIA guidance but require an explicit allowlist update before receiving persistent pipeline guidance. Canonical profiling workflow documentation lives under third_party/tlx/tools/agents/kernel_optimization/docs/profiling/ and is not injected as source guidance. Keep workload-specific invariants and exclusions in target.json.optimization_guidance; they are applied after the built-in target skills.

For every candidate, the standard loop is:

build -> verify -> benchmark -> profile -> decide -> repeat

The live log must include:

• hypothesis and evidence;
• one coherent source/configuration change;
• expected effect and risk;
• correctness, median, p95, CV, and speedup;
• Proton attribution plus target summary profile deltas and a concise decision.

Do not print kernel source or private chain-of-thought. Continue through the configured round budget after an unpromoted round. Reject incorrect, duplicate, unstable, and materially regressing candidates.

Layer 2: Profiling Request And Artifacts

The optimizer sends a profile request for the baseline (deep), every correctness-passing candidate (summary, escalated to deep near the promotion threshold), and the finalist (deep). Profiling is always on; the --profile flag is vestigial. Every request carries tools=["proton_launch", "native_profiler"], where native_profiler resolves to ncu on CUDA/NVIDIA targets automatically. Pass --diagnostic-proton-intra-kernel to additionally collect warp-granularity Proton instrumentation traces for the baseline and final winner only (diagnostic-only: never benchmark, promote, or commit instrumented source or timing).

These requests produce data only if the bundle's harness.py::profile() implements them. A stub that repackages endpoint timings yields ncu=unavailable and zeroed proton.* fields on every line, and all hypotheses degrade to endpoint latency plus source inspection. Before launching, verify the harness profile path end to end:

• Proton launch attribution: follow third_party/tlx/tools/agents/kernel_optimization/docs/profiling/proton.md and return profile()["proton"] with nonzero main_kernel_us for the expected kernel launch.
• Target counters: follow third_party/tlx/tools/agents/kernel_optimization/docs/profiling/nvidia-ncu.md (CUDA/NVIDIA) and return profile()["ncu"] with non-null summary duration. Report unsupported counters as JSON null with a diagnostic, never as zero.

Harnesses that implement profile should accept a structured request with:

{
  "level": "summary",
  "tools": ["proton_launch", "native_profiler"],
  "experiment_id": "stable candidate or baseline id",
  "artifacts_dir": "/absolute/path/to/profile-artifacts",
  "reason": "why this profile was requested",
  "diagnostic_only": false
}

The legacy two-argument profile(build_artifact, case) contract remains supported and means summary profiling for the default tools. Return compact, normalized JSON inline. Store raw .hatchet, .chrome_trace, .ncu-rep, CSV, mapping, and command files as artifacts, and reference them with absolute paths.

Profiling has three distinct layers:

• Proton wrapper/launch attribution: read third_party/tlx/tools/agents/kernel_optimization/docs/profiling/proton.md and collect for every correctness-passing candidate.
• Target summary/deep profiling: request native_profiler; each target harness maps it to its platform tool. For CUDA/NVIDIA, this is NCU and the harness should follow third_party/tlx/tools/agents/kernel_optimization/docs/profiling/nvidia-ncu.md.
• Diagnostic-only Proton intra-kernel instrumentation: use only for attribution questions that cannot be answered from wrapper timelines or target counters.