fla-optimization-loop

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

Updated Sep 11, 2026
One-click install
npx skills add https://github.com/swiss-ai/flash-linear-attention --skill fla-optimization-loop-swiss-ai
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: fla-optimization-loop
Source: https://github.com/swiss-ai/flash-linear-attention/tree/main/.agents/skills/fla-optimization-loop
Command: npx skills add https://github.com/swiss-ai/flash-linear-attention --skill fla-optimization-loop-swiss-ai

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill includes references (resource) components.

What problem does it solve? Optimizing GPU kernels in the Flash Linear Attention (FLA) library over multiple iterations risks silently breaking correctness, gaming benchmarks, or losing track of what was tried. This Skill imposes a disciplined loop — frozen test contract, three-phase plan, one-change-per-iteration protocol, and evidence logging — so speedups on fla/ops/** kernels are real, reproducible, and promotable. ## Core Features & Use Cases - Frozen correctness gate: Treats the op's pytest file and naive reference as immutable, banning tolerance loosening, shape dropping, and one-sided numeric-flag relaxation, with NaN-poisoned memory catching partial writes. - Structured iteration protocol: Enforces one change per iteration followed by gate run, benchmark, log entry, and commit, with stall handling after three non-improving rounds and explicit no-go criteria. - Evidence and trap catalogs: Provides references/TRAPS.md (TF32 reference inflation, autotune staleness, clock drift, int64 overflow) and log templates for iteration rows, kept-candidate headers, and shape-dispatch justification. - Use Case: While making chunk_gla faster in Triton, you write a task contract, profile with NCU, try three candidate directions, log each result in OPT_LOG.md, and promote only the candidate with a full green gate and a profiler-explained win. ## Quick Start Ask the agent to optimize the FLA operator chunk_gla for a 1.2x speedup on the target shapes using this optimization loop, starting with a task contract draft 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 pytest file and naive.py reference as frozen, and run the verify gate after every single kernel change. Never loosen tolerances, drop shapes, or relax numeric flags on one side; log each iteration and commit before trying the next direction.

How to benchmark Triton kernel speedups against a baseline?

Run the verify command with --base main so the candidate is compared against the unmodified baseline in one session. Use the solution's own runtime for fast iteration signal, but require the full comparison on all shapes before claiming a win.

Which kernel languages does this optimization loop support?

The loop applies to FLA kernels written in Triton, Gluon, TileLang, or CuTe DSL. Hand-written kernels in any of these languages are allowed, while delegating the whole compute to a vendor library call as the operator is banned.

Why does my kernel pass tests but produce wrong results on large shapes?

Program IDs and grid offsets multiplied by strides can silently overflow 32-bit integers at large sequence lengths. Cast to int64 before address arithmetic and always run the gate's large-T cases, not just small fast shapes.

When should I stop iterating on kernel optimization?

Stop when you hit a user-set iteration cap, when re-assessment shows hard evidence of a floor such as HBM bandwidth or launch overhead, or after documenting three distinct tried directions. A no-go verdict requires a baseline number, a reasoned attempt, gate status, bench evidence, and a named blocker.