Liger Kernel Perf

Liger Kernel Perf

Optimizes existing Liger Kernel Triton kernels through a 3-stage pipeline: Profile, Optimize, Finalize. Supports interactive mode (human checkpoints between stages) and autonomous mode (runs end-to-end). NVIDIA GPUs only.

Mode Detection

• Interactive mode (default): Human checkpoints between each stage
• Autonomous mode: User says "just optimize it", "run without asking me", "optimize autonomously" → all stages run end-to-end, user sees only the final report

Input Parsing

Extract from the user's request:

Pre-Flight Validation

Before starting the pipeline, validate:

1. Kernel file exists: src/liger_kernel/ops/{kernel}.py
2. Benchmark script exists: benchmark/scripts/benchmark_{kernel}.py
3. Test file exists: test/transformers/test_{kernel}.py
4. GPU is available and CUDA works
5. Project is installed in dev mode (pip install -e ".[dev]")

If any validation fails, report clearly and stop.

Pipeline

Stage 1: Profile

Follow the Profiler workflow in profiler.md. If the host runtime supports parallel subagents, this stage may be delegated to one; otherwise execute the workflow directly.

This stage:

1. Creates the workspace directory optimization/{kernel}/
2. Copies the original kernel as a snapshot
3. Runs baseline benchmarks using the existing benchmark script
4. Detects GPU architecture (or uses user-specified target)
5. Optionally runs NCU profiling (if ncu is available)
6. Analyzes the kernel code (tier classification, patterns, optimization opportunities)
7. Classifies the bottleneck: memory-bound vs compute-bound
8. Produces an optimization profile with a recommended strategy order
9. Saves profile to optimization/{kernel}/profile.md

Human checkpoint (interactive mode): Present the optimization profile with bottleneck diagnosis and proposed strategy order. Confirm before proceeding.

Stage 2: Optimize

Follow the Optimizer workflow in optimizer.md.

This stage runs an autonomous optimization loop:

1. Read the optimization profile and original kernel
2. Always try parameter tuning first (BLOCK_SIZE, num_warps, num_stages manual sweep -- NOT @triton.autotune)
3. Then apply diagnosis-driven techniques from optimization-strategies.md
4. For each variant: a. Generate the variant code → optimization/{kernel}/{kernel}_vN.py b. Write the variant lab notebook → optimization/{kernel}/{kernel}_vN_notes.md c. Run quick smoke test (single shape, float32, forward+backward) → discard on failure d. Run the full existing benchmark script → optimization/{kernel}/benchmarks/vN_results.csv e. Check guardrails (no catastrophic regressions) f. Update the variant notes with actual results
5. Read all prior variant notes before generating the next variant
6. Stop when: budget exhausted, 2 consecutive variants with <1% improvement, or target metric met
7. Produce a comparison table of ALL variants

Human checkpoint (interactive mode): Present the comparison table across all variants. User approves the winner (or skill picks best if autonomous).

Stage 3: Finalize

Follow the Finalizer workflow in finalizer.md.

This stage:

1. Applies the winning variant in-place to src/liger_kernel/ops/{kernel}.py
2. Runs the full test suite: python -m pytest test/transformers/test_{kernel}.py -xvs (hard gate)
3. Runs checkstyle: make checkstyle (auto-fix with ruff check . --fix && ruff format .)
4. Generates 3-way comparison plots (original liger vs optimized liger vs huggingface baseline) using benchmarks_visualizer.py
5. Generates the final optimization report → optimization/{kernel}/report.md
6. Creates a PR with only the kernel code changes (no plots or optimization workspace files)
7. Presents the before/after summary with plots

Human checkpoint (interactive mode): Present the final report with before/after numbers, comparison plots, and test results.

Guardrails

These apply to EVERY variant, regardless of mode:

Reference Files

• profiler.md -- Profiler Agent specification
• optimizer.md -- Optimizer Agent specification
• finalizer.md -- Finalizer Agent specification
• optimization-strategies.md -- Catalog of optimization techniques
• templates/optimization-profile.md -- Profiling output format (cross-stage contract)
• templates/variant-notes.md -- Per-variant lab notebook format