# Optimizer Agent

Runs an autonomous optimization loop over an existing Liger kernel, generating multiple variants, benchmarking each against the baseline, and selecting the best performer while maintaining correctness.

## Inputs

- `target_kernel`: kernel name (e.g., "rms_norm", "cross_entropy", "swiglu")
- `optimization_goal`: speed / memory / balanced
- `scope`: general / forward / backward / specific input regime
- `max_variants`: maximum number of variants to try (default 8)
- `target_metric`: optional concrete target (e.g., "forward under 0.3ms at hidden_size=4096")
- `optimization_profile`: the profile produced by the Profiler Agent at `optimization/{kernel}/profile.md`
- `original_kernel`: the original kernel snapshot at `optimization/{kernel}/original_{kernel}.py`

## Pre-Requisites

Before starting the optimization loop:

1. Verify the workspace exists: `optimization/{kernel}/` with `original_{kernel}.py` and `profile.md`
2. Verify the benchmark script exists: `benchmark/scripts/benchmark_{kernel}.py`
3. Verify the test file exists: `test/transformers/test_{kernel}.py`
4. Read the optimization profile thoroughly — extract bottleneck classification, tier, current config, and recommended strategy order
5. Read the original kernel code thoroughly — understand every `@triton.jit` function, the `torch.autograd.Function`, and the launcher logic
6. Read `optimization-strategies.md` for the full catalog of available techniques

## Variant Numbering

- v0 = baseline (original kernel, already benchmarked by the Profiler)
- v1, v2, ... = optimization variants generated by this agent
- Each variant gets two files:
  - `optimization/{kernel}/{kernel}_vN.py` — the modified kernel code
  - `optimization/{kernel}/{kernel}_vN_notes.md` — the lab notebook entry

## The Optimization Loop

```
variant_number = 1
best_variant = 0  (baseline)
consecutive_no_improvement = 0

while variant_number <= max_variants:
    1. Read ALL prior variant notes (v1 through v{variant_number-1})
    2. Decide what to try next (see Strategy Selection below)
    3. Generate variant code
    4. Write variant notes (hypothesis section)
    5. Run smoke test → if FAIL, mark REJECTED, increment variant_number, continue
    6. Run full benchmark → record results
    7. Check guardrails → if FAIL, mark REJECTED, increment variant_number, continue
    8. Update variant notes with actual results and verdict
    9. If improvement over best_variant: update best_variant, reset consecutive_no_improvement
   10. Else: increment consecutive_no_improvement
   11. If consecutive_no_improvement >= 2: STOP (diminishing returns)
   12. If target_metric is set and met: STOP (goal achieved)
   13. Increment variant_number
```

## Strategy Selection

### Always Try Parameter Tuning First (v1)

The first variant MUST be a parameter sweep over `BLOCK_SIZE`, `num_warps`, and `num_stages`. This is the cheapest optimization with the highest expected payoff.

**Why NOT @triton.autotune**: Liger kernels use `torch.autograd.Function` where the forward pass stores `BLOCK_SIZE` and `num_warps` in `ctx` and passes them to the backward kernel. With `@triton.autotune`, these become compile-time constants that cannot be stored in `ctx` — the backward kernel would not know which config the forward used. Additionally, NPU/Ascend backends do not support `num_warps` or `num_stages` at all. For these reasons, `@triton.autotune` is intentionally not used in Liger production kernels.

**Instead, do a manual parameter sweep**: Replace the `calculate_settings(n_cols)` call with hardcoded values and benchmark each configuration:

```python
# In the variant, replace:
#   BLOCK_SIZE, num_warps = calculate_settings(n_cols)
# With a specific config to test:
BLOCK_SIZE = 2048   # Try: 256, 512, 1024, 2048, 4096, 8192
num_warps = 8       # Try: 2, 4, 8, 16, 32
num_stages = 2      # Try: 1, 2, 3, 4
```

**Sweep procedure for v1**:
1. Read the current `calculate_settings(n_cols)` output for the benchmark's representative input sizes
2. Create a variant that replaces `calculate_settings` with a custom function that returns a better config
3. Test 3-5 promising configs (not the full cartesian product — use the architecture guidance below to prune)
4. Pick the best config across ALL benchmark x_values (not just one size)
5. The variant should implement a smarter `calculate_settings` replacement, not just a single hardcoded value

**Architecture-specific guidance for config selection**:
- **Ampere**: `num_stages=2`, `num_warps` up to 16, BLOCK_SIZE 1024-4096
- **Hopper**: `num_stages=3-5`, `num_warps` up to 32, BLOCK_SIZE 2048-8192
- **Blackwell**: `num_stages=4-5`, similar to Hopper

**Important**: Keep the same forward/backward ctx pattern. Both kernels must use the same `BLOCK_SIZE` and `num_warps` for a given input. The variant must still pass these through `ctx` just like the original.

### Subsequent Variants: Diagnosis-Driven

After autotuning, choose the next technique based on:

1. **The bottleneck classification from the profile** (memory-bound vs compute-bound)
2. **Learnings from all prior variants** (what worked, what did not, what was surprising)
3. **The recommended strategy order from the profile**

Consult [optimization-strategies.md](optimization-strategies.md) for the full technique catalog. Common patterns by bottleneck type:

**Memory-bound kernels** (Tier 1, simple element-wise):
- Fuse multiple memory-bound operations into a single kernel
- Add cache eviction hints (`tl.store(..., eviction_policy="evict_first")`)
- Increase BLOCK_SIZE to amortize memory latency
- Use `tl.load(..., eviction_policy="evict_last")` for reused data
- Reduce unnecessary memory traffic (eliminate redundant loads/stores)
- In-place operations where safe

**Compute-bound kernels** (Tier 3, heavy math):
- Reduce precision of intermediate computations where safe
- Use fast-math intrinsics (`tl.math.fast_expf`, etc.)
- Restructure computation to reduce total FLOPs
- Improve instruction-level parallelism via loop unrolling
- Increase occupancy by reducing register pressure

**Balanced / Tier 2 kernels** (reductions):
- Optimize the reduction pattern (two-pass vs single-pass)
- Tune BLOCK_ROW for multi-row processing
- Balance between register usage and parallelism
- Consider split-K or hierarchical reductions for large dimensions
- Block-row mode threshold tuning (the `n_rows < 4096 * 8` style checks)

### Combining Techniques

A single variant may combine multiple techniques if they are complementary (e.g., autotuning + cache hints). However, prefer isolated changes so that the lab notebook clearly attributes improvements to specific techniques. When combining, explain which part of the improvement is expected from each technique.

### When All Strategies Exhausted

If all recommended strategies have been tried and the consecutive_no_improvement counter has not triggered, the agent may:
- Try creative combinations of previous partial improvements
- Attempt architecture-specific tricks (Hopper warp specialization, Blackwell TMA)
- Revisit a rejected variant with a modified approach based on learnings

## Generating a Variant

For each variant, produce a complete, self-contained kernel file at `optimization/{kernel}/{kernel}_vN.py`.

Rules:
1. Start from the original kernel (not a previous variant) unless explicitly building on a prior variant's approach. If building on a prior variant, note which one and why.
2. The variant file must be a drop-in replacement for `src/liger_kernel/ops/{kernel}.py` — same public API, same class names, same function signatures.
3. Do not change the `torch.autograd.Function` class name or the public API.
4. Do not remove safety checks, dtype handling, or multi-device support (XPU, NPU paths) unless the change is specifically about those areas.
5. **Preserve ALL existing comments** — function docstrings, mathematical formula explanations, inline rationale comments, and section headers. Do NOT delete or shorten comments that explain what the code does or why.
6. **Add explanatory comments for every optimization change** — each modified section should have a comment explaining what was changed and why it improves performance. For example:
   ```python
   # Optimization: compute dW before dX to allow register reuse.
   # The original order kept 8 vectors live simultaneously (115 regs/thread).
   # This reordering reduces peak liveness to 5-6 vectors (95 regs/thread).
   ```
7. **Preserve function-level docstrings** — if the original kernel functions have docstrings or parameter descriptions, keep them. Update them if the function signature changes.
8. Add a comment block at the top of the file documenting what was changed:

```python
"""
Optimization Variant {N}: {Short Title}
Based on: original (or v{M} if building on a prior variant)
Changes: {brief description of modifications}
"""
```

## Variant Lab Notebook

For each variant, write `optimization/{kernel}/{kernel}_vN_notes.md` using the format in [templates/variant-notes.md](templates/variant-notes.md).

Key requirements:
- The Hypothesis must reference the profile's bottleneck classification and learnings from prior variants
- Changes must include rationale for each modification with code snippets
- Expected Outcome must be specific and quantitative ("10-20% forward speed improvement" not "should be faster")
- Actual Results must include ALL x_values if improvement varies across sizes
- Verdict: ACCEPTED (improves target without violating guardrails), REJECTED (fails smoke test/guardrails/no improvement), or MIXED (improves some metrics, regresses others within limits)
- Learnings must include insights to carry forward to the next variant

### Writing Notes Across the Loop

- **Before smoke test**: Write the Hypothesis, Changes, and Expected Outcome sections
- **After benchmark**: Fill in Actual Results, Verdict, and Learnings
- **On smoke test failure**: Fill Actual Results with "SMOKE TEST FAILED: {error}", set Verdict to REJECTED, and write Learnings about what went wrong

## Smoke Test Procedure

The smoke test validates basic correctness before investing time in a full benchmark. Execute it by temporarily swapping the variant into the source tree.

### Step 1: Swap In the Variant

```bash
cp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py
```

### Step 2: Run a Minimal Correctness Check

```bash
python -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k "float32" 2>&1 | tail -80
```

This runs only float32 tests (the most numerically stable) with a timeout to catch hangs. The `-x` flag stops at the first failure.

### Step 3: Interpret Results

- **All tests pass**: Smoke test PASSES. Proceed to full benchmark.
- **Any test fails**: Smoke test FAILS. Read the error carefully.
  - If the failure is a **tolerance issue** (allclose fails but values are close): Note this in the variant notes. The optimization may have changed numerical behavior slightly. Try adjusting the variant to preserve precision, or note the tolerance difference.
  - If the failure is a **crash or shape error**: The variant has a bug. Mark REJECTED immediately. Record the error in the notes for learning.
  - If the failure is a **hang/timeout**: The variant has a deadlock or infinite loop. Mark REJECTED immediately.

### Step 4: Restore the Original

**CRITICAL**: Always restore the original kernel after the smoke test, regardless of outcome:

```bash
cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py
```

This must happen even if the smoke test crashes. Use a pattern that ensures restoration:

```bash
cp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py && \
python -m pytest test/transformers/test_{kernel}.py -x --timeout=120 -k "float32" 2>&1 | tail -80; \
SMOKE_EXIT=$?; \
cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py; \
exit $SMOKE_EXIT
```

## Full Benchmark Procedure

After a variant passes the smoke test, run the full existing benchmark suite to get precise measurements.

### Step 1: Swap In the Variant

```bash
cp optimization/{kernel}/{kernel}_vN.py src/liger_kernel/ops/{kernel}.py
```

### Step 2: Run the Full Benchmark

```bash
cd benchmark/scripts && python benchmark_{kernel}.py --overwrite
```

This runs speed benchmarks (forward, backward, full) and memory benchmarks across all configured x_values, comparing the "liger" provider against the baseline ("huggingface" or "torch").

### Step 3: Capture Results

The benchmark writes results to `benchmark/data/all_benchmark_data.csv`. Copy the relevant rows:

```bash
cp benchmark/data/all_benchmark_data.csv optimization/{kernel}/benchmarks/vN_results.csv
```

### Step 4: Parse Results

Read the CSV and extract:
- Speed (ms) at the median (p50) for forward, backward, and full modes
- Memory (MB) at the median (p50) for full mode
- Values at each x_value to detect size-dependent behavior

Compare against:
- v0 baseline: the absolute improvement
- Current best variant: whether this is the new best

### Step 5: Restore the Original

```bash
cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py
```

## Guardrail Checks

After benchmarking each variant, apply these guardrails. All comparisons are against the v0 baseline.

### Guardrail 1: Non-Target Metric Regression (>5%)

If the optimization goal is "speed", check that memory does not regress by more than 5%. If the goal is "memory", check that speed does not regress by more than 5%. If the goal is "balanced", check both.

```
if goal == "speed":
    reject if memory_vN > memory_v0 * 1.05
elif goal == "memory":
    reject if speed_vN > speed_v0 * 1.05  (higher ms = slower = worse)
elif goal == "balanced":
    reject if memory_vN > memory_v0 * 1.05 OR speed_vN > speed_v0 * 1.05
```

Check this at EVERY x_value. If the regression occurs only at certain sizes, note this in the variant notes but still reject — partial regressions are not acceptable.

### Guardrail 2: Cross-Pass Regression (>10%)

Do not sacrifice one pass (forward or backward) to improve the other beyond a 10% threshold:

```
reject if speed_fwd_vN > speed_fwd_v0 * 1.10 AND speed_bwd_vN < speed_bwd_v0
reject if speed_bwd_vN > speed_bwd_v0 * 1.10 AND speed_fwd_vN < speed_fwd_v0
```

An optimization that makes forward 15% slower to make backward 5% faster is rejected.

### Guardrail 3: Smoke Test Failure

Already handled in the smoke test procedure. Any correctness failure means immediate rejection.

### Handling Guardrail Failures

When a variant is rejected by guardrails:
1. Set Verdict to REJECTED in the variant notes
2. Record exactly which guardrail was violated and by how much
3. Record the Learnings — this information is valuable for future variants
4. Do NOT increment consecutive_no_improvement (guardrail failures are not "stagnation")
5. Proceed to the next variant

## Stopping Conditions

The loop terminates when ANY of these conditions is met:

### Condition 1: Budget Exhausted

`variant_number > max_variants`

Report: "Optimization budget of {max_variants} variants exhausted."

### Condition 2: Diminishing Returns

`consecutive_no_improvement >= 2`

Two consecutive variants that pass guardrails but improve the target metric by less than 1% over the current best. This indicates the kernel is near its optimization ceiling for the current approach.

Report: "Stopping due to diminishing returns — last 2 accepted variants showed <1% improvement."

### Condition 3: Target Metric Met

The user-specified `target_metric` has been achieved.

Report: "Target metric achieved by variant v{N}: {metric_description}."

## Reading Prior Variant Notes

Before generating each new variant, the agent MUST read ALL prior variant notes. This is the accumulated learning mechanism that makes later variants smarter.

When reading prior notes, extract:
1. **What techniques were tried**: Avoid repeating the exact same approach
2. **What worked**: Build on successful directions
3. **What failed and why**: Avoid repeating mistakes
4. **Surprising results**: These often reveal the true bottleneck
5. **Size-dependent behavior**: Some optimizations help at large sizes but hurt at small sizes

Use this information to form a more informed hypothesis for the next variant.

## Final Comparison Table

After the loop terminates (for any reason), produce a comparison table of ALL variants and save it to `optimization/{kernel}/comparison.md`.

The comparison must include:
- A master table with ALL variants (accepted and rejected): Variant | Strategy | Speed Fwd/Bwd/Full (ms) | Memory (MB) | Verdict
- Values at the representative x_value (typically the largest or most common) with percentage delta from baseline
- Winner declaration with justification
- Cumulative learnings summary (what worked, what didn't, why)

## Error Recovery

### Benchmark Script Crashes

If the benchmark script crashes (not a test failure — an actual crash):
1. Restore the original kernel immediately
2. Read the crash output
3. If it is a Triton compilation error: the variant has a syntax or semantic error in the kernel. Mark REJECTED with the error.
4. If it is a CUDA OOM: the variant uses too much memory for the GPU. Mark REJECTED. Note the memory constraint.
5. If it is an unrelated infrastructure error: retry once. If it fails again, report to the user.

### Workspace Corruption

If at any point `src/liger_kernel/ops/{kernel}.py` is found to differ from `optimization/{kernel}/original_{kernel}.py` when it should be the original:

```bash
diff src/liger_kernel/ops/{kernel}.py optimization/{kernel}/original_{kernel}.py
```

If they differ unexpectedly, restore immediately:

```bash
cp optimization/{kernel}/original_{kernel}.py src/liger_kernel/ops/{kernel}.py
```

### Agent Failure Mid-Loop

If the agent encounters an unrecoverable error:
1. Restore the original kernel
2. Save all variant notes written so far
3. Write a partial comparison table with whatever data is available
4. Report to the user with the current state and what went wrong

## Important Constraints

1. **Never leave a variant installed in the source tree.** The original kernel must always be restored after each smoke test and benchmark. This is the most critical invariant.
2. **Never modify test files.** The existing tests are the correctness oracle. If a variant fails tests, the variant is wrong, not the tests.
3. **Never modify benchmark files.** The existing benchmarks are the measurement standard. Variants must be measured by the same benchmark to enable fair comparison.
4. **Always write notes before running experiments.** The hypothesis must be written before seeing results to avoid post-hoc rationalization.
5. **Always read all prior notes before generating a new variant.** This is what makes the optimization loop learn from itself.
6. **Prefer the simplest change that achieves the goal.** A 5-line autotuning change that yields 15% speedup is better than a 200-line rewrite that yields 20%.
