[Feature] Make autotune a real, adopted tuning path (learn from aiter + existing flydsl autotune)

[jhinpan]

Context

Supersedes #612 (discussion-only) with an actionable plan, and feeds RFC #749 (downstream libraries — the PoC target is Tri Dao's quack — expect a usable tuning path to consume FlyDSL kernels). #612 established that FlyDSL has an autotune implementation but nothing actually uses it.

Current state of FlyDSL autotune (verified)

• python/flydsl/autotune.py (296 lines) already provides:
  • Config (constexpr params + waves_per_eu / maxnreg compiler hints),
  • Autotuner (benchmark configs -> pick fastest -> cache),
  • disk cache under FLYDSL_AUTOTUNE_CACHE_DIR / ~/.flydsl/autotune,
  • a _make_key (key args' shapes/dtypes), _prune, and the public @autotune decorator (exported from flydsl.__init__).
• But a repo-wide search shows no kernel actually uses @autotune — only the definition plus README / CLAUDE / SKILL docs. No dedicated tests, no real adoption. It is effectively dead code today.
• Prior history: Feat/autotune compile hints #241 (autotune compile-hint support), feat(jit): two-phase compilation for autotuning compile_hints #266 (two-phase compile for autotune hints — closed unmerged).

Reference comparison: how others do it

Verified by reading source. Two distinct questions are on this table: how to tune a kernel (Triton / quack / aiter / FlyDSL) vs how an application consumes tuned configs (SGLang). SGLang is listed because it's a real downstream FlyDSL consumer and shows what the "offline path" looks like at scale — it is not itself an autotuner, so it's flagged as a consumer view, not ranked against the autotuners. Triton is the template the others (quack, FlyDSL) imitate, so it's the most important one to understand deeply.

Aspect	Triton (the template)	quack (Tri Dao)	aiter	SGLang (consumer)	FlyDSL (today)
Role	JIT DSL autotuner	JIT DSL autotuner	offline kernel-lib tuner	inference engine consuming tuned configs	JIT DSL autotuner
Impl	python/triton/runtime/autotuner.py (511 ln)	quack/autotuner.py (849 ln)	csrc/.../*_tune.py + CSVs	benchmark/.../tuning_fused_moe_triton.py + get_moe_configs	python/flydsl/autotune.py (296 ln)
Adoption	ubiquitous	rmsnorm + gemm	GEMM family	fused MoE etc.	none
Mode	online JIT	online JIT + parallel precompile	offline CSV	offline tune -> JSON, runtime lookup	online JIT
Avoid-search path	@heuristics (compute params, skip search)	analytical get_default + exhaustive	per-shape CSV	get_default_config fallback when no JSON	none
In-place correctness	restore_value / reset_to_zero (snapshot/restore between bench reps)	inherited pattern	n/a (offline)	n/a	missing
Cache key	triton_key() (compiler ver) + backend.hash + fn.cache_key (src) + invalidating env + shape/dtype	source fingerprint + shape + stride	(per-shape CSV row)	filename: E,N,device_name,dtype + triton-version dir	shape/dtype only
Persisted as	*.autotune.json in cache dir	.o + result cache	CSV in repo	JSON keyed by device+shape, committed	best-config JSON
First-run cost	pays search (mitigated by disk cache)	parallel precompile	none at serving	none at serving (lookup)	pays full search
CI guard	unit tests	test_autotuner.py	check_tuned_op_regression.sh	auto-tune.yml	none

Deep-dive: Triton autotune (the design FlyDSL/quack imitate)

Three mechanisms in triton/runtime/autotuner.py that FlyDSL's port currently lacks and that matter most:

1. restore_value / reset_to_zero — the correctness soul of autotune. Benchmarking runs the same kernel dozens of times to pick the fastest. If the kernel writes its output in place or accumulates, repeated runs corrupt the data and the timing/selection is made on garbage. Triton snapshots and restores the affected tensors around each bench rep. FlyDSL has nothing equivalent — autotuning any in-place kernel (e.g. fused-add rmsnorm) would currently be incorrect.
2. @heuristics alongside @autotune. Params that can be computed are derived by a heuristic (zero search); only genuinely-uncertain params are searched. Same idea as quack's two-track get_default + exhaustive.
3. Thorough cache key. triton_key() (compiler version) + backend target hash + fn.cache_key (kernel source) + cache-invalidating env vars + shape/dtype — far more complete than FlyDSL's shape/dtype-only key.

SGLang shows the consumer end: tuned configs are committed as JSON whose filename encodes E, N, device_name, dtype under a triton-version directory; at runtime get_moe_configs looks up by device+shape and falls back to get_default_config on a miss (incl. AMD gfx configs). This is the shape a FlyDSL tuned-config artifact should take if an engine like SGLang is to consume it.

What to borrow (synthesis of best-of-breed)

1. Two-track configs (quack get_default + exhaustive, == Triton @heuristics + @autotune): heuristic default for zero-search normal runs, exhaustive space only when explicitly autotuning. Avoids FlyDSL's all-or-nothing.
2. In-place correctness (Triton restore_value/reset_to_zero): required before autotuning any kernel that mutates its inputs/outputs — otherwise tuning picks a config based on corrupted reruns. Prerequisite, not a nice-to-have.
3. Online + offline (FlyDSL online + aiter/SGLang offline): keep @autotune JIT for dev; add an offline "tune -> emit config -> load at runtime" path, opt-in via env/marker so normal tests don't pay search cost. (Discussion: make autotune a supported and broadly usable tuning path #612 Q3)
4. Thorough cache key (Triton + quack): compiler/version key + backend/arch + kernel source fingerprint + cache-invalidating env + shape/dtype + stride normalization. (Discussion: make autotune a supported and broadly usable tuning path #612 Q4)
5. Consumer-friendly config artifact (SGLang): if FlyDSL is to be tuned-and-consumed by an engine (ties to [RFC]: a contribution path for FlyDSL kernels into downstream OSS libraries (PoC: Dao-AILab/quack) #749), emit tuned configs keyed by device+shape+dtype that a downstream can look up, with a heuristic default fallback on miss.
6. Parallel precompile (quack): subprocess pool to warm the compile cache for all candidates before benchmarking — kills first-run latency; practical answer to feat(jit): two-phase compilation for autotuning compile_hints #266.
7. CI guard (aiter + quack + SGLang auto-tune.yml): GPU-free unit tests for Config/serialization/pruning; opt-in GPU integration test; a tuned-config regression guard.

Deliverables

• Decide the online + offline story and the cache-key spec (incl. source fingerprint + stride)
• Two-track config model (heuristic default + exhaustive search) for the first adopted kernel
• First kernel actually adopting @autotune (+ a committed tuned config for the offline path). Candidates: softmax / rmsnorm for API validation; a GEMM / MoE variant for perf. (Discussion: make autotune a supported and broadly usable tuning path #612 Q2)
• Parallel precompile path to bound first-run cost (revisit feat(jit): two-phase compilation for autotuning compile_hints #266)
• GPU-free unit tests + an opt-in GPU integration test + a tuned-config regression guard
• Docs: when to tune, cache behavior, when not to run in CI

Refs

• Closes Discussion: make autotune a supported and broadly usable tuning path #612 (carrying its questions forward into this actionable form).
• Related: [RFC]: a contribution path for FlyDSL kernels into downstream OSS libraries (PoC: Dao-AILab/quack) #749 (PoC target quack), Feat/autotune compile hints #241, feat(jit): two-phase compilation for autotuning compile_hints #266.
• Source refs: FlyDSL python/flydsl/autotune.py; aiter aiter/configs/*_tuned_gemm.csv, aiter/ops/flydsl/gemm_tune/, .github/scripts/check_tuned_op_regression.sh; quack quack/autotuner.py, quack/rmsnorm_config.py, AI/caching_and_parallel_compilation.md, tests/test_autotuner.py. Triton python/triton/runtime/autotuner.py (Autotuner, check_disk_cache, restore_value/reset_to_zero, @heuristics). SGLang benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py, srt/layers/moe/.../fused_moe_triton_config.py (get_moe_configs/get_default_config), committed JSON configs keyed by E,N,device_name,dtype, .github/workflows/auto-tune.yml.