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.

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.

Layer 3: Proton Attribution

Use Proton to attribute wrapper, benchmark phase, launch, and profiler overhead. An ordinary Proton timeline with hook='triton' that shows one kernel launch does not prove tlx.async_task overlap; it only shows launch/wrapper attribution around the compiled kernel. Do not infer per-task overlap from that timeline.

Diagnostic intra-kernel attribution must use Proton instrumentation mode with backend='instrumentation', data='trace', granularity='warp', and explicit Triton semantic enabled. Group warp lanes by known async-task warp ranges from the source, generated metadata, or a saved mapping artifact. Do not request warp_group granularity because the runtime rejects it today.

Instrumentation changes are diagnostic-only. Compiler transforms may move or merge scopes, so instrumented source and timing must never be benchmarked, promoted, committed, or used as speedup evidence.

Layer 4: Target Profiling

Select target profiling guidance from target.json:

• CUDA/NVIDIA: read third_party/tlx/tools/agents/kernel_optimization/docs/profiling/nvidia-ncu.md.
• Other backends: use a sibling target guide when present; otherwise use the vendor-neutral profile() contract without inventing NVIDIA requirements.

Every correctness-passing candidate receives Proton attribution and the target guide's summary profile. Escalate to the target guide's deep profile for:

• the baseline before candidate generation;
• a candidate within one percentage point of the promotion threshold;
• disagreement between endpoint benchmark, Proton attribution, and target summary profile;
• repeated rounds with no promoted candidate;
• schedule or WS hypotheses that need stall, occupancy, spill, or memory-hierarchy evidence;
• the final promoted candidate before commit.

Do not spend deep-profile time on incorrect or clearly slow candidates. Unsupported counters must be reported as unavailable with JSON null and a diagnostic, never as zero.

Layer 5: Evidence-To-Action Policy

Each candidate hypothesis must cite measured evidence and change one subsystem or tightly coupled invariant-preserving pair. Feed failed hypotheses and their metric regressions into subsequent prompts as exclusions.

Generic interpretation rules:

• Higher kernel duration with lower compute and memory utilization indicates scheduling, serialization, synchronization, or insufficient parallelism.
• A benchmark-only change inside the noise floor with unchanged kernel metrics is not optimization evidence.
• Lower utilization without lower work, traffic, and duration is not a win.
• Register-budget or buffering changes require spill/occupancy evidence and a producer-consumer/barrier proof.
• Main-kernel profile time and public-wrapper benchmark time cover different scopes; do not subtract them to manufacture a bottleneck.

Use the target guide for vendor-specific counter interpretation.

Layer 6: Promotion, Revalidation, And Commit

Promote only when:

• every protected case passes;
• weighted speedup meets the configured threshold;
• measurement variance is within budget;
• target profiling shows no material main-kernel regression.

Revalidate the finalist with benchmark, Proton attribution, summary target profile, and required deep target profile. Only then may the default commit occur. Auto-commit must detect Git or Mercurial from the kernel path, preserve unrelated dirty/staged work, include TLX agent authored in the commit body, and log VCS, revision, repository, target, subject, and failure diagnostics. If commit fails, keep all artifacts and return the distinct commit-failure status.

Layer 7: Command And Completion

From the repository root:

PYTHONPATH=<repo-root> python -m third_party.tlx.tools.agents.kernel_optimization.cli \
  --kernel <absolute-kernel.py> \
  --harness <absolute-harness.py> \
  --cases <absolute-cases.json> \
  --target <absolute-target.json> \
  --output-dir <absolute-output-dir> \
  --prior-run <optional-previous-output-dir> \
  --provider codex \
  --max-rounds 5 \
  --candidates-per-round 2 \
  --max-candidate-seconds 600 \
  --max-total-seconds 3600 \
  --min-speedup 1.01 \
  --max-cv 0.10 \
  --benchmark-repetitions 10 \
  --profile

Add --reference-kernel, --budget, --arch, --model, --vcs, or --commit-message only when required. Do not invent a model name. Never run optimizer.py directly.

The task is complete only after the actual CLI run reaches final revalidation or a diagnosed blocking failure. Report the exact workload, GPU, baseline and final latency, speedup, CV, correctness, stopping reason, commit result, and absolute output directory. best_kernel.py by itself is not completion.