--- name: fla-optimization-loop description: > Disciplined, reproducible loop for making an FLA kernel faster (Triton, Gluon, TileLang, CuTe) without ever breaking or gaming correctness. Synthesizes the task-contract / three-phase / iteration-protocol / silent-bug-catalog discipline of agent kernel-optimization frameworks (KDA, the MLSys FlashInfer contest workflow, AKO4ALL/AKO4X), and anchors all of it on FLA's frozen pytest (forward AND backward, under NaN poisoning) as the immutable correctness gate. Use when iterating on `fla/ops/**` performance over multiple rounds. --- # FLA Optimization Loop Skill Use this when you are making an existing `fla/ops/**` kernel faster — or bringing a new kernel from correct to fast — across more than one iteration, in any FLA backend language (Triton, Gluon, TileLang, CuTe DSL). This skill is the **search discipline** that ties the other skills together; it does not replace them: - **`fla-design-coverage`** — define reachable contracts, numerical budgets, and production workloads before the loop. - **`fla-nvidia-performance`** — how to profile (NCU), hardware baselines, PR-ready perf evidence. - **`fla-ascend-performance`** — how to profile (torch_npu), diagnose Cube/Vector/MTE/UB bottlenecks, and optimize Triton-Ascend kernels. - **`fla-correctness-coverage`** — how to design the test coverage matrix for an op. - **`fla-pr-readiness`** — how to package the promoted change into a PR. This skill covers the loop *around* those: what to lock, how to iterate, what to record, and when to stop. ## 0. The inviolable rule: the test file is a frozen contract The op's `tests/ops/test_.py` and its `fla/ops//naive.py` reference are **frozen** for the entire optimization loop. Add any missing regression or new-path coverage before freezing the gate; those additions belong in the final PR. Freezing the gate during optimization does not waive the contribution requirement to add or update matching tests. The whole point of "faster" only means something if correctness — forward **and** backward, under the conftest NaN-memory poisoning — is held fixed. During a perf loop you may **not**: - edit the test file, or its `naive.py` reference; - loosen an `assert_close` tolerance, or widen `rms_eps` / dtype to make a diff pass; - drop, narrow, or `skip` parametrized shapes; - special-case the kernel on values it only sees in the test; - cache module-level tensors so trials reuse warm/identical data. Run the gate with the unmodified test, every iteration: ```bash python -m benchmarks.ops.verify --op [--gate-k ] ``` `verify.py` runs the current pytest file as a black box and **stops before benchmarking on a red gate by default**. It does not freeze files or reject an all-skipped run. Keep the gate unchanged yourself, use strict local validation (`FLA_CI_ENV=0`), and check that the expected tests actually ran. `--gate-k` only *selects* a shape subset for a fast signal — it never edits the test; promote only on a full (no `-k`) green gate. **Banned vs. allowed implementation (anti-reward-hacking):** | Banned | Allowed | | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------ | | Making the op a thin wrapper that delegates the whole compute to a vendor lib (a plain `torch.matmul` / `F.scaled_dot_product_attention` standing in *as* the operator) just to win latency | Hand-written Triton / Gluon / TileLang / CuTe kernels; `torch` ops used as *glue* around a kernel you wrote | | Returning uninitialized / partially-written outputs that happen to pass | Fully initialized outputs (NaN poisoning will catch partial writes) | | Stream tricks / monkey-patching the bench to dodge timing | Genuine latency reduction measured by `verify.py` / `run.py` | | One-sided numeric relaxation the baseline doesn't get — flipping `allow_tf32` on, dropping the fp32 accumulator to bf16/tf32, a config that quietly changes the numeric path | Same accumulation precision and numeric flags on both sides; speed comes from the kernel, not from computing something less accurate | If you genuinely believe a test is **wrong**, that is a **separate PR** with its own justification — never bundled into a perf change. Stop and ask the user. ## 1. Write the task contract first (before any code) Put a short `docs/draft.md` in your scratch workspace (see §6) stating: - **Op + entry point** — e.g. `chunk_gla` in `fla.ops.gla`. - **Target** — production workloads derived from layer/model callers, plus applicable shapes from `benchmarks/ops/registry.py`, and a target speedup against a recorded baseline SHA. - **Allowed languages** — Triton / Gluon / TileLang / CuTe (state any constraint). - **Validation command** — `python -m benchmarks.ops.verify --op ` (the frozen gate). - **Benchmark command** — `python -m benchmarks.ops.verify --op --base "$FLA_BENCH_BASE"`, where `FLA_BENCH_BASE=$(git rev-parse origin/main)` is recorded before the loop. Pass a SHA to avoid checked-out-branch and slash-containing-ref worktree failures. - **Promotion criteria** — full green gate, a measured repeatable win, and a profiler reading that explains it (§7). - **Frozen scope** — the test file, `naive.py`, and the public op signature. Do not start editing kernels until the draft exists. (Borrowed from KDA: plan, then execute.) ## 2. Three phases Run these in order; repeat 2 and 3 with progressively higher targets. - **Phase 1 — correct baseline.** Confirm the current kernel passes the full gate, and record baseline numbers: `python -m benchmarks.ops.verify --op --base "$FLA_BENCH_BASE"`. For a brand-new kernel, get the gate green first; performance is secondary here. - **Phase 2 — profile-guided optimization.** Use `fla-nvidia-performance` for NCU evidence on NVIDIA backends, or `fla-ascend-performance` for Ascend NPU / Triton-Ascend backends. Enumerate candidate directions, rank them by expected benefit vs. implementation risk, and explore each for at most a few iterations. Keep, revise, or reject each with evidence — don't optimize blindly. - **Phase 3 — shape specialization.** Only when profiling shows *different* bottlenecks across shape regimes (short vs. long `T`, small vs. large `D`), add dispatch / specialized paths. Justify the added complexity with the measured win; validate on the full shape set, not the one you tuned on. Record each bucket (condition / entry point / per-bucket latency + speedup / reason) in `dispatch.md` (template in `references/opt-log-template.md`) — a specialized path without that evidence is unjustified complexity. ## 3. Iteration protocol Every iteration is exactly three steps, in order, with no telescoping into the next iteration between them: 1. Make **one** change to the kernel. 2. Run `verify.py` — gate must stay green; record the bench number. 3. Append one row to `OPT_LOG.md` (see `references/opt-log-template.md`) and checkpoint the code on the task's feature branch, following `AGENTS.md` and the user's git authorization. A failed or no-change iteration still needs a log entry before exploring the next direction. Do not create empty commits for unchanged code or commit ignored scratch records. (Borrowed from AKO4ALL: bench → log → commit, the most-skipped step in practice.) **Stall handling.** After 3 consecutive iterations with no improvement (≥ a few % over current best, above noise), stop and re-assess: re-profile, re-read `OPT_LOG.md` for which axes you've already tried, and search for known techniques for this op family before picking a new direction. **When to stop.** A user-set iteration cap is reached; or re-assessment produces hard evidence of a floor (bandwidth-bound at HBM limit, launch-overhead dominated, timer-resolution limited — cite it in `OPT_LOG.md`); or you've documented ≥3 distinct directions tried with evidence. Don't stop silently because a tool (e.g. NCU) was unavailable — that's a re-assessment input. **No-go bar.** If stopping means concluding there's no win to promote — a no-go — that verdict has its own bar: a first candidate losing doesn't clear it. A no-go needs a recorded baseline number, at least one reasoned candidate attempt (not a blind guess), the gate status, the bench evidence, and a *named* active bound or blocker (name which roofline/launch/timer limit, with the number). Without those five it's an unfinished loop. ## 4. Reproducibility - **Seed** — the tests already pin `torch.manual_seed(42)`; don't undermine it. - **Environment** — `verify.py` prints GPU / CUDA / PyTorch / Triton / commit SHA; paste that line into `OPT_LOG.md` so a number is interpretable later. - **Full-shape verdict before promotion** — a `--gate-k` subset is a signal only. - **Rank by the solution's own runtime** for fast iteration signal; pay for the full `--base "$FLA_BENCH_BASE"` comparison only at the verdict. Clock noise on unlocked GPUs can swing absolute speedup — see `references/TRAPS.md`. ## 5. Silent-bug & measurement traps Read `references/TRAPS.md` before trusting any number. It catalogs FLA-specific traps (NaN poisoning on partial writes, TF32 inflating fp32 reference diffs, `assert_close` relative tolerance, one-sided numeric-flag relaxation, autotune-cache staleness across edits, `int64` address arithmetic, implausible speedups that signal a silently-skipped path). When you hit a new one, add it there with **Fact / Why / How to apply** so the next session doesn't re-learn it. (Borrowed from AKO4X's `TRAPS.md`.) ## 6. Evidence records (scratch workspace) Keep your search artifacts in a git-ignored directory — `profile/-opt/` is already ignored: ```text profile/-opt/ docs/draft.md # the task contract (§1) OPT_LOG.md # one row per iteration (template in references/) dispatch.md # one row per shape bucket — only if Phase 3 specializes (template in references/) TRAPS.md # traps you hit this session (seed: references/TRAPS.md) trace/ # torch.profiler / NCU artifacts (kept out of git) ``` Record each kept candidate's Identity / Delta / Lessons / Dead-ends / Open-directions in the scratch log per `references/opt-log-template.md`. This history stays out of committed kernel comments and PR summaries. ## 7. Promotion → PR A candidate is promotable only on a **full green gate** plus a measured, repeatable win on the target shapes, **and** a profiler/roofline reading that *explains* the win (or, for a no-go, the blocker) — a speedup you can't account for is a silent-skip suspect, not a result (see §5). Then: - Keep the diff minimal — change only what the win needs, plus light cleanups (CONTRIBUTING "Protect battle-tested paths; keep diffs minimal"). - Collect perf evidence per backend skill — `fla-nvidia-performance` (before/after, NCU summary, dense + varlen coverage) or `fla-ascend-performance` (before/after, pipe/UB metrics, round summary template) — and a full final-claim stats block (median/mean/std/min/p10/p90 per shape, equal-weight geomean speedup, exact commands, baseline commit + candidate SHA, GPU id/model with idle-clock evidence — the last guards the clock-drift trap in §5). - Package the PR per `fla-pr-readiness`. The scratch workspace under `profile/-opt/` stays local; it is not part of the PR.