fla-optimization-loop

Guides reproducible multi-round optimization of FLA kernels under a frozen pytest correctness gate.

Updated Aug 18, 2026
One-click install
npx skills add https://github.com/weichengz0616/fla --skill fla-optimization-loop-weichengz0616
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: fla-optimization-loop
Source: https://github.com/weichengz0616/fla/tree/main/.agents/skills/fla-optimization-loop
Command: npx skills add https://github.com/weichengz0616/fla --skill fla-optimization-loop-weichengz0616

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill includes references (resource) components.

What problem does it solve? Optimizing GPU kernels over many iterations risks silently breaking correctness, gaming benchmarks, or losing track of what was tried. This Skill imposes a disciplined loop for making FLA (Flash Linear Attention) kernels faster in Triton, Gluon, TileLang, or CuTe while keeping the frozen pytest gate — forward and backward, under NaN poisoning — green on every iteration. ## Core Features & Use Cases - Frozen correctness contract: Treats the op's test file and naive reference as immutable, banning tolerance loosening, shape dropping, one-sided numeric-flag relaxation, and vendor-library wrappers that fake speedups. - Three-phase iteration protocol: Baseline validation, profile-guided optimization, and shape specialization, with a strict one-change-per-iteration rule and mandatory bench-log-commit after every attempt. - Evidence and audit trail: Templates for OPT_LOG.md, dispatch.md, and a TRAPS.md catalog of silent-bug and measurement pitfalls (TF32 reference diffs, int64 overflow, stale autotune caches, clock drift). - Use Case: You are speeding up chunk_gla in fla/ops/gla across several sessions. The Skill has you write a task contract, run python -m benchmarks.ops.verify --op gla each iteration, log every result, and only promote on a full green gate with profiler evidence explaining the win. ## Quick Start Ask the AI to optimize the FLA operator you name using the fla-optimization-loop discipline, starting with a task contract and the frozen verify gate.

Frequently Asked Questions about fla-optimization-loop

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I optimize an FLA kernel without breaking correctness?

Treat the op's test file and naive.py reference as frozen, and run python -m benchmarks.ops.verify --op <op> after every single change. Each iteration makes one change, checks the gate stays green, logs the result, and commits before the next attempt.

How to benchmark FLA kernel speedups against the main branch?

Run python -m benchmarks.ops.verify --op <op> --base main to compare against main in a single session. Use --gate-k only for fast iteration signal on a shape subset, and promote only on a full green gate with a repeatable measured win.

Which kernel languages does the FLA optimization loop support?

The loop supports hand-written kernels in Triton, Gluon, TileLang, and CuTe DSL. Torch operations are allowed only as glue around a kernel you wrote, not as a wholesale replacement that delegates the compute to a vendor library.

Why does my FLA kernel pass ad-hoc tests but fail the pytest gate?

The conftest replaces torch.empty allocations with NaN-filled memory, so partially written outputs surface as NaN in assert_close. An ad-hoc script often sees coincidentally zeroed memory and hides the partial-write bug; trust the gate and fully initialize every output element.

When should I stop an FLA kernel optimization loop?

Stop when the iteration cap is reached, when re-assessment shows hard evidence of a floor such as HBM bandwidth or launch-overhead limits, or after documenting three distinct tried directions. A no-go verdict requires a recorded baseline, a reasoned candidate attempt, gate status, bench evidence, and a named blocker.

Can I loosen assert_close tolerance to make a faster kernel pass?

No. Tolerance is part of the frozen contract, and widening it converts a correctness regression into a silent one. If you believe a test or tolerance is genuinely wrong, that must be a separate justified PR, never bundled into a performance change.