[2/5] autotune: two-track config + first real adopter (rmsnorm) (#770)

[jhinpan]

Second in the #770 series: give the autotuner an avoid-search path and put it to work on a real kernel, so it stops being dead code.

Stacked on #783 (1/5). Targets main (cross-fork PRs can't base on a fork branch), so it currently shows PR1's commit too — review only eb94a0ac, or merge #783 first and this will shrink to its own diff.

Builder mode

Every current FlyDSL kernel bakes its structural knobs at module-build time — they size shared storage, the vectorized tile stride, and the launch block, and known_block_size is fixed at @kernel decoration — rather than exposing them as jit Constexpr params. The existing @autotune, which injects config kwargs into a single jit call, can't tune those.

So this adds builder mode: build_fn(config, *args, **kwargs) -> launch_callable rebuilds the module per config; built modules are cached per (key, config). Structural config kwargs route to build_fn (not the launch call); build-only kwargs like dtype_str are filtered out of the launch call by signature.

Two-track config

== Triton @heuristics + @autotune, quack get_default + exhaustive. New default= heuristic on @autotune: normal runs take the analytic default and skip the search entirely (zero-search); set FLYDSL_AUTOTUNE=1 to force the full sweep. Autotuners without a default always search.

First adopter: rmsnorm

• build_rmsnorm_module gains a BLOCK_THREADS build knob (adds known_block_size when >256); the default arg keeps the old signature working, so the existing rmsnorm test path is unchanged.
• kernels/rmsnorm_config.py — get_default (analytic BLOCK_THREADS) + get_all_configs (BLOCK_THREADS × waves_per_eu, fast-path only). VEC_WIDTH is intentionally not tuned (pinned to 128/elem_bits by the 128-bit buffer copy).
• kernels/rmsnorm_autotune.py — rmsnorm_autotuned, the tuned front end.

Verification (MI350X / gfx950, M=4096 N=8192 bf16)

Default and forced-search runs both match the torch reference (max err 1.5e-2). The 12-config sweep runs cleanly and picks BLOCK_THREADS 256–512; the result is cached (second call doesn't re-tune):

[autotune] tuning 12 configs...
  [4/12] Config(BLOCK_THREADS=256) -> 0.081 ms
  [7/12] Config(BLOCK_THREADS=512) -> 0.081 ms
  ...
[autotune] best: Config(BLOCK_THREADS=256) (0.081 ms)

Tests

• +6 GPU-free unit tests (builder rebuild/cache, default skip/force, kwarg filtering, env gate) → 22 total in tests/unit/test_autotune.py.
• tests/kernels/test_rmsnorm_autotune.py (l2_device): default run, forced-search + tuned-config persisted, cache reuse.

ruff + black clean.

Refs #770.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

This PR extends FlyDSL’s autotuning system to support a two-track “default heuristic vs forced search” flow and introduces a first real adopter (RMSNorm) that exercises builder-mode autotuning (rebuilding modules per structural config).

Changes:

• Add builder-mode support to Autotuner (build_fn-based rebuild + in-process build cache) and introduce a heuristic default= path that skips search unless FLYDSL_AUTOTUNE=1.
• Extend autotune machinery with stride normalization + extra cache-key axes, and add unit tests covering builder/default behavior.
• Add RMSNorm autotune integration (rmsnorm_config.py + rmsnorm_autotune.py) and a GPU integration test.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
python/flydsl/autotune.py	Adds builder-mode execution, two-track default/forced-search behavior, and additional cache key axes; updates benchmark/run path accordingly.
kernels/rmsnorm_kernel.py	Adds a build-time BLOCK_THREADS knob and applies known_block_size when needed.
kernels/rmsnorm_config.py	Introduces heuristic default + exhaustive config generation for RMSNorm tuning.
kernels/rmsnorm_autotune.py	Provides the tuned RMSNorm front-end using builder-mode autotuning.
tests/unit/test_autotune.py	Adds GPU-free unit tests for builder mode, default skip/force behavior, kwarg filtering, and env gating.
tests/kernels/test_rmsnorm_autotune.py	Adds an end-to-end GPU test validating default correctness, forced search, persistence, and cache reuse.

Comments suppressed due to low confidence (1)

python/flydsl/autotune.py:250

• In builder mode (fn=None), the disk-cache filename falls back to "unknown.json", so different tuners in the same cache dir can clobber each other’s tuned results. Consider deriving the cache filename from fn when present, otherwise from build_fn (and include the module name) so each tuner has an isolated cache file.

        # Disk cache
        fn_name = getattr(fn, "__name__", None) or getattr(fn, "func", None)
        if fn_name is not None and not isinstance(fn_name, str):
            fn_name = getattr(fn_name, "__name__", "unknown")
        fn_name = fn_name or "unknown"

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zhiding512]

waves_per_eu does any effect in FLYDSL?