[2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770) by jhinpan · Pull Request #785 · ROCm/FlyDSL

jhinpan · 2026-07-01T07:05:40Z

Second in the #770 series. This PR adds a no-search default path to the existing @autotune, adopts it in an explicit RMSNorm tuning entry, and makes ROCm waves_per_eu candidates compile as distinct, effective binaries.

It keeps the existing direct-JIT and kernel-factory model: there is no builder-mode tuner, second kernel body, or separate tuner-owned build cache.

Autotuner behavior

On a normal call, a cached winner runs first. On a cache miss, an optional callable default= supplies a config without benchmarking.
FLYDSL_AUTOTUNE=1 (also true, yes, or on) bypasses both cache and default and performs a fresh search.
configs= may be callable and is evaluated only when a search is required.
Autotuners without default= retain the existing search-on-cache-miss behavior.
Benchmark setup and final execution use the launch stream for both PyTorch stream objects and raw stream pointers; raw 0 maps to the device default stream.

RMSNorm tuning entry

rmsnorm_autotuned is an explicit tuning entry; existing RMSNorm callers remain unchanged.
Its direct-JIT entry exposes N, dtype_str, and BLOCK_THREADS as Constexpr inputs and specializes the existing build_rmsnorm_module factory.
The factory remains the single owner of the kernel body, storage, tiling, and launch geometry. Existing call sites keep the same default block size.
known_block_size is derived from the specialized block size when it exceeds the AMDGPU default limit of 256.
Small-N RMSNorm keeps its existing geometry and does not sweep multiple configs. Large-N search covers seven BLOCK_THREADS and waves_per_eu combinations.
The adopter key includes row count, hidden size, and dtype. Compilation and launch are bound to the input device, and an omitted stream resolves to the current stream on that device.

Compile-hint behavior

waves_per_eu is a compile option, not a runtime kernel argument.

Persistent JIT hints form the base; nested thread-local hints shallow-merge over them, with later non-None values winning.
The fully merged hint snapshot participates in the JIT binary cache key and is passed consistently through compilation. Candidate compiler options remain part of the selected Config.
None inherits, 0 leaves a source kernel attribute unchanged, and a positive integer is an exact compile-time override.
Before ROCDL lowering, ROCm materializes a positive override on entry-point gpu.func operations carrying gpu.kernel as amdgpu-waves-per-eu=N,N. It replaces the native minimum-only rocdl.waves_per_eu form while preserving unrelated passthrough attributes.
The backend hook is generic and a no-op by default. Existing hand-authored kernel attributes and the existing ROCm pipeline option remain available.

Verification

Focused autotune, compiler, and RMSNorm suite: 70 passed, 2 skipped.
Existing RMSNorm coverage: 2 passed.
Black, Ruff, py_compile, git diff --check, and scripts/check_python_style.sh --include-local: passed.
PR CI passed Python and C++ style, documentation, MI325, MI355, and Navi: Fly DSL test run 29537829776.
The GPU adopter test covers forced search across all seven candidates on an explicit raw stream and verifies subsequent cached-winner reuse.

Refs #770.

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.

FlyDSL's autotuner exists but nothing uses it, and two gaps block real adoption. This is the first of a series making it a correct, adopted path. Cache key (_make_key) previously specialized on shape/dtype only. A config tuned under one compiler build, GPU arch, or memory layout would be silently reused under another. Fold in the axes Triton/quack rely on: - normalized stride pattern ({0,1,other}: broadcast vs contiguous vs strided) - device arch (get_rocm_arch) - toolchain fingerprint (reuses jit_function._flydsl_key) - cache-invalidating env vars (reuses _cache_invalidating_env_values) The dtype/stride axes are sorted by arg name so a call is keyed identically regardless of kwarg order (no duplicate tuning / cache files). restore_value (new) is the correctness soul of autotune: benchmarking runs the same kernel dozens of times, so an in-place / accumulating kernel (e.g. fused-add rmsnorm) corrupts its own inputs and picks a config on garbage. Snapshot the named tensors once and restore before every rep. reset_to_zero is now also re-applied on the real (non-benchmark) call — both the post-tune run and cache hits — via a shared _run_config, so an accumulate-into-zero kernel returns the single-clean-run result instead of carrying benchmark-rep state. (Was applied only inside the bench loop.) Also defer the CompilationContext import so the autotuner core stays importable and unit-testable without the compiled flydsl._mlir bindings. Adds tests/unit/test_autotune.py: 19 GPU-free tests covering Config serialization, every cache-key axis (incl. env-fingerprint change and kwarg-order insensitivity), restore_value/reset_to_zero semantics (incl. the final-run and cache-hit reset), pruning, and disk-cache round-trip. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Comment-only cleanup of the PR1 additions: keep the one key fact per helper, drop the Triton/quack background, redundant restatements, and by-example prose. No logic change; 19 unit tests still pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

@autotune

Second in the series. Gives the autotuner an "avoid-search" path and puts it to work on a real kernel, so it stops being dead code. Builder mode. Every current FlyDSL kernel bakes its structural knobs at module-build time rather than exposing them as jit Constexpr params. The existing @autotune, which injects config kwargs into one jit call, can't tune those. Add builder mode: build_fn(config, *args) -> launch_callable rebuilds the module per config. autotune_builder(): one-call adoption. Instead of a hand-rolled wrapper per kernel, a kernel opts in with just its kernel-specific pieces: rmsnorm_autotuned = autotune_builder( name="rmsnorm", build=build_rmsnorm_module, specialize=lambda inp, g, out, m, dtype_str="bf16", **kw: { "N": inp.shape[-1], "dtype_str": dtype_str}, configs=get_all_configs, default=get_default, structural=("BLOCK_THREADS",)) The helper owns cache naming, callable config generation, the structural-vs- compiler knob split, build caching, and launch-kwarg filtering. rmsnorm's adopter dropped from ~79 lines of boilerplate to a single declaration. Correctness (surfaced by two independent fresh reviews): - Build cache keys only on the STRUCTURAL knobs + spec, not repr(config), so configs differing only in a compiler hint (waves_per_eu) reuse one built module. Verified on MI350X: a 12-config sweep now builds 4 modules, not 12. - A build-only scalar (dtype_str) passed positionally is rejected with a clear error instead of silently binding to the wrong launch slot (e.g. stream). - autotune_builder requires a non-empty name (empty/None would fall back to unknown.json and defeat the per-kernel cache identity). - Build-only scalars enter the cache key via the specialize() axes. - Compiler hints (waves_per_eu / maxnreg) are folded into the JitFunction's compile_hints (restored after) so each distinct hint compiles a distinct binary instead of reusing a cached one. - FLYDSL_AUTOTUNE=1 bypasses cache + default to force a fresh search. - Disk cache is re-loaded when FLYDSL_AUTOTUNE_CACHE_DIR changes (a module- level tuner isn't pinned to the import-time dir for loads or saves). - Config.pre_hook is documented as non-persistable (not serialized). Two-track config == Triton @Heuristics + @autotune: default= gives zero-search normal runs; FLYDSL_AUTOTUNE=1 forces the sweep. First adopter rmsnorm: build_rmsnorm_module gains a BLOCK_THREADS build knob (+known_block_size when >256; default arg keeps the old signature). config space in rmsnorm_config.py; small-N (imported SMALL_N_THRESHOLD) emits a single config since that kernel ignores BLOCK_THREADS. VEC_WIDTH stays pinned (128b copy). Verified on MI350X (gfx950), M=4096 N=8192 bf16: default and forced-search match the torch reference (max err 1.5e-2); sweep picks BLOCK_THREADS 256-512 and a subsequent normal call reuses the cache (zero benchmark invocations asserted). Tests: GPU-free unit tests for builder mode + autotune_builder (rebuild/cache, build-cache-ignores-hints, default skip/force, scalar-in-key, name required, positional-scalar rejected) = 31; tests/kernels/test_rmsnorm_autotune.py (l2_device) covers default, forced-search, and no-re-tune cache reuse. ruff + black clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

The autotuner routes Config(waves_per_eu=...) through compile_hints -> --amdgpu-waves-per-eu on gpu-module-to-binary opts=, which the AMDGPU backend silently ignores (a known limitation). So the rmsnorm _WAVES_PER_EU_CHOICES sweep tuned nothing. Lower the waves_per_eu compile-hint onto the kernel's rocdl.waves_per_eu gpu.func attribute in MlirCompiler.compile (before convert-gpu-to-rocdl) -- the same path the MoE kernels use. convert-gpu-to-rocdl turns it into the LLVM amdgpu-waves-per-eu function attribute, which the backend honors. Verified on MI355X (gfx950): the compiled LLVM IR now carries "amdgpu-waves-per-eu"="N" (previously absent) and the backend acts on it (ISA changes; occupancy warning at aggressive values). The JIT compile cache key already includes compile_hints, so each waves_per_eu recompiles a distinct binary. Co-authored-by: Cursor <cursoragent@cursor.com>

Follow-up to the waves_per_eu fix: the same "silently dropped compile hint" bug hit maxnreg, and num_warps had a related builder-mode trap. - maxnreg: Config(maxnreg=...) reached codegen only via --amdgpu-num-vgpr on gpu-module-to-binary opts=, which the AMDGPU backend ignores. Unlike waves_per_eu, amdgpu-num-vgpr has no native ROCDL translation, so lower it onto the gpu.func LLVM `passthrough`; convert-gpu-to-rocdl copies that to the llvm.func where the emitter turns it into the amdgpu-num-vgpr function attribute (same approach as FlyROCDLClusterAttrPass for amdgpu-cluster-dims). Generalize _apply_occupancy_compile_hints to a hint -> func-attr lowering. Verified on MI355X: the IR now carries "amdgpu-num-vgpr"="N" and the backend acts on it (ISA changes when the cap binds, e.g. 40 < the kernel's 60 VGPRs; 64/128 don't bind); outputs stay correct. - num_warps: in builder mode the block size is baked into build_fn, so a jit-kwarg num_warps can't be routed to the launch call and was silently dropped. Reject it with a clear error instead of tuning a no-op. - Tests: assert the func-attr lowering and the builder-mode num_warps rejection. Co-authored-by: Cursor <cursoragent@cursor.com>

…cache, f32) - Thread-safety: _run_with_hints no longer mutates the shared, cached fn.compile_hints (a race across concurrent tuned/served calls); it sets only the thread-local CompilationContext, and JitFunction folds the thread-local hints into its cache key + compile push, so each distinct waves_per_eu/maxnreg still compiles a distinct binary. - Search loop no longer silently swallows failures: unroutable configs (num_warps in builder mode) are rejected up front so they fail loudly even in a forced sweep, and "all configs failed" chains the last underlying error. - Disk cache: atomic write (tmp+rename); a single malformed entry is skipped rather than discarding the whole cache; a stale cached entry degrades to the default instead of crashing a normal call. - Cache key folds a source fingerprint of the adopter build/config functions, so editing the kernel/config invalidates a stale tuned best. - maxnreg lowering (_set_passthrough) replaces an existing amdgpu-num-vgpr entry instead of appending a duplicate. - restore_value/reset_to_zero run as an untimed per-rep setup (do_bench gains a setup param), so the restore copy no longer inflates the measurement. - get_all_configs sweeps f32 too (scalar path is BLOCK_THREADS-strided; it was silently collapsing to the single default). Verified on MI355X: f32 configs build and match reference (max_err 2.9e-6). - Builder mode rejects config kwargs not in structural (would route nowhere). Adds unit tests for each behavior. Co-authored-by: Cursor <cursoragent@cursor.com>

jhinpan · 2026-07-07T06:26:20Z

/rerun-ci

…#783/ROCm#807) Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # python/flydsl/autotune.py # tests/unit/test_autotune.py

jhinpan · 2026-07-07T06:50:54Z

/rerun-ci

…obustness) Follow-ups from the second multi-model review (round-1 fixes verified): - Cache-hit fallback catches only config-incompatibility errors (ValueError/ TypeError/KeyError); genuine compile/launch/runtime errors now propagate instead of being masked. Logs and drops the stale entry from disk too. - _call_do_bench passes `setup` only when the benchmarker explicitly declares it -- a **kwargs catch-all no longer gets setup passed-and-silently-dropped (which would skip restore/reset). - _save_disk_cache is best-effort: a write failure (read-only FS, full disk) is logged, never crashes an otherwise-successful tune. - _iter_gpu_kernel_funcs applies occupancy hints only to entry-point kernels (gpu.kernel attr), skipping non-kernel device helpers. - _source_fingerprint hashes the source file (transitive over module-level helpers/constants), not just the function body. - Tests: **kwargs do_bench folds setup; cache-hit degrades to default on a stale config; f32 get_all_configs sweeps BLOCK_THREADS. Co-authored-by: Cursor <cursoragent@cursor.com>

jhinpan · 2026-07-07T07:25:14Z

/rerun-ci

Sync to main c1b5218 (ROCm#800: fused-add/residual rmsnorm backward). # Conflicts: # kernels/norm/rmsnorm_kernel.py main (ROCm#800) moved rmsnorm's shared constants (EPS, BLOCK_THREADS, WARP_SIZE, VEC_WIDTH) and the vector/reduction helpers out into kernels/norm/rmsnorm_common.py and imports them; this branch still defined them inline next to its new SMALL_N_THRESHOLD knob. Adopt main's refactor -- the inline copies were byte-identical to rmsnorm_common's -- dropping the duplicates and keeping only the branch-only SMALL_N_THRESHOLD (used by build_rmsnorm_module, imported by rmsnorm_config). build_rmsnorm_module merged cleanly, retaining both main's store_rstd/eps params and this branch's BLOCK_THREADS build knob. ruff + black clean; 44 autotune unit tests pass (8 GPU skips). Co-authored-by: Cursor <cursoragent@cursor.com>

jhinpan · 2026-07-13T04:59:22Z

/rerun-ci

Make ROCm occupancy hint lowering use an explicit compile_hints snapshot so the backend hook documents where hints come from and why they must be materialized before ROCDL lowering. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep the occupancy hint documentation concise while preserving the key ROCm lowering invariants for reviewers. Co-authored-by: Cursor <cursoragent@cursor.com>

jhinpan · 2026-07-14T20:39:36Z

/rerun-ci

Resolve persistent, thread-local, and autotune hints through one detached snapshot shared by cache identity and codegen. Keep source attributes as the occupancy baseline while ROCm owns target-specific lowering and scalar compatibility options.\n\nExpose a generic Config compile_hints envelope with typed occupancy aliases, make builder-only inputs explicit, and harden autotune cache/error boundaries and environment handling.

Resolve RMSNorm autotuning on the input device/current stream, benchmark explicit torch and raw streams on the launch stream, and include row count in winner identity.\n\nMake the GPU adopter test deterministic and strict about candidate success and cache reuse, tighten ROCm hint boundary/lowering coverage, remove the external-backend reference, and sync main's CI workflow fixes.

jhinpan · 2026-07-16T21:54:18Z

/rerun-ci

Copilot AI review requested due to automatic review settings July 1, 2026 07:05

Copilot started reviewing on behalf of jhinpan July 1, 2026 07:06 View session

Copilot AI reviewed Jul 1, 2026

View reviewed changes

Comment thread python/flydsl/autotune.py Outdated

Comment thread python/flydsl/autotune.py Outdated

Comment thread python/flydsl/autotune.py Outdated

Comment thread python/flydsl/autotune.py Outdated

jhinpan mentioned this pull request Jul 1, 2026

[3/5] autotune: add offline config artifacts (#770) #786

Open

jhinpan force-pushed the feat/autotune-rmsnorm-adopt branch from eb94a0a to 0aa86e0 Compare July 1, 2026 07:56

jhinpan mentioned this pull request Jul 1, 2026

[4/5] autotune: CI guard + committed-config regression check (#770) #788

Open

Merge branch 'main' into feat/autotune-cache-key-restore

50a95c2

jhinpan force-pushed the feat/autotune-rmsnorm-adopt branch 5 times, most recently from e96dad5 to db07800 Compare July 2, 2026 06:54

jhinpan force-pushed the feat/autotune-rmsnorm-adopt branch from db07800 to 95ea147 Compare July 2, 2026 07:08

jhinpan force-pushed the feat/autotune-rmsnorm-adopt branch from 95ea147 to 1bc2ee7 Compare July 2, 2026 19:15

zhiding512 reviewed Jul 4, 2026

View reviewed changes

Comment thread kernels/rmsnorm_config.py Outdated

jhinpan and others added 2 commits July 6, 2026 23:30

Merge branch 'main' into feat/autotune-rmsnorm-adopt (sync to 0.2.3)

c6610e0

jhinpan force-pushed the feat/autotune-rmsnorm-adopt branch from 1e07aff to a19f7e6 Compare July 6, 2026 23:31

jhinpan and others added 3 commits July 7, 2026 00:07

Merge branch 'main' into feat/autotune-rmsnorm-adopt

2cb96df

Merge branch 'main' into feat/autotune-rmsnorm-adopt (sync after ROCm…

a8b87e4

…#783/ROCm#807) Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # python/flydsl/autotune.py # tests/unit/test_autotune.py

Merge branch 'main' into feat/autotune-rmsnorm-adopt

dd2180f

This was referenced Jul 10, 2026

[Feature] Extract reusable event-based benchmarking helper #829

Draft

[Feature]: Move benchmark/perf test from test folder to FlyDSL python source #304

Open

[Feature] Extract reusable event-based benchmarking helper jhinpan/FlyDSL-lab#2

Draft

jhinpan and others added 2 commits July 13, 2026 04:48

Merge branch 'main' into feat/autotune-rmsnorm-adopt

173ab72

Merge branch 'main' into feat/autotune-rmsnorm-adopt

b377161

jhinpan mentioned this pull request Jul 13, 2026

[CI] run_benchmark: treat runtime pytest.skip as skipped, not failed #842

Merged

3 tasks

Merge branch 'main' into feat/autotune-rmsnorm-adopt

8fbae45

zhiding512 reviewed Jul 14, 2026

View reviewed changes

Comment thread python/flydsl/compiler/backends/rocm.py Outdated

jhinpan and others added 3 commits July 14, 2026 20:03

refactor: clarify occupancy hint lowering

18ddf6a

Make ROCm occupancy hint lowering use an explicit compile_hints snapshot so the backend hook documents where hints come from and why they must be materialized before ROCDL lowering. Co-authored-by: Cursor <cursoragent@cursor.com>

docs: trim occupancy lowering comments

4d4a5a1

Keep the occupancy hint documentation concise while preserving the key ROCm lowering invariants for reviewers. Co-authored-by: Cursor <cursoragent@cursor.com>

Merge branch 'main' into feat/autotune-rmsnorm-adopt

b29c2e7

jhinpan mentioned this pull request Jul 14, 2026

[Test] Expand RMSNorm backward configs, cache reuse, and Torch benchmarks #854

Merged

jhinpan added 9 commits July 15, 2026 17:04

autotune: model waves per eu as a compile option

79d3ece

Merge branch 'main' into feat/autotune-rmsnorm-adopt

1bd314d

compiler: define occupancy hint precedence

011dacf

autotune: make rmsnorm config tests arch-stable

c4237a4

autotune: specialize rmsnorm through direct JIT

a6de6a2

autotune: trim rmsnorm adopter to essential path

f011227

autotune: fold rmsnorm config into adopter

81baca5

jhinpan changed the title ~~[2/5] autotune: two-track config + first real adopter (rmsnorm) (#770)~~ [2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770) Jul 17, 2026

zhiding512 previously approved these changes Jul 17, 2026

View reviewed changes

coderfeli reviewed Jul 17, 2026

View reviewed changes

Comment thread python/flydsl/compiler/jit_function.py Outdated

compiler: align compile-hint helper naming

154688c

jhinpan dismissed zhiding512’s stale review via 154688c July 17, 2026 23:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770)#785

[2/5] autotune: add opt-in search and RMSNorm waves-per-EU tuning (#770)#785
jhinpan wants to merge 38 commits into
ROCm:mainfrom
jhinpan:feat/autotune-rmsnorm-adopt

jhinpan commented Jul 1, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jhinpan commented Jul 7, 2026

Uh oh!

jhinpan commented Jul 7, 2026

Uh oh!

jhinpan commented Jul 7, 2026

Uh oh!

jhinpan commented Jul 13, 2026

Uh oh!

Uh oh!

jhinpan commented Jul 14, 2026

Uh oh!

jhinpan commented Jul 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

jhinpan commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Autotuner behavior

RMSNorm tuning entry

Compile-hint behavior

Verification

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jhinpan commented Jul 7, 2026

Uh oh!

jhinpan commented Jul 7, 2026

Uh oh!

jhinpan commented Jul 7, 2026

Uh oh!

jhinpan commented Jul 13, 2026

Uh oh!

Uh oh!

jhinpan commented Jul 14, 2026

Uh oh!

jhinpan commented Jul 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jhinpan commented Jul 1, 2026 •

edited

Loading