[3/5] autotune: offline config emit + runtime lookup (#770)#786
Open
jhinpan wants to merge 5 commits into
Open
[3/5] autotune: offline config emit + runtime lookup (#770)#786jhinpan wants to merge 5 commits into
jhinpan wants to merge 5 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Adds an “offline” autotuning consumption path to FlyDSL’s autotuner so kernels can be tuned once, emit a portable JSON artifact keyed by a filename convention, and later resolve the config at runtime with zero search. This extends the existing autotune flow with builder-mode support (rebuild-per-config), a two-track default-vs-search behavior, and adopts it for RMSNorm with both unit and GPU integration tests.
Changes:
- Extend
flydsl.autotune.Autotunerwith offline config emit/lookup, builder mode (build_fn), and a heuristicdefaultthat skips search unlessFLYDSL_AUTOTUNE=1. - Add RMSNorm tuning configuration + autotuned front-end that uses builder-mode autotune and offline artifacts.
- Add GPU-free unit tests and a GPU integration test covering search/no-search, caching, and offline artifact round-trip; add a user guide doc.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_autotune.py | Adds GPU-free tests for cache keys, restore/reset semantics, builder mode, and offline artifact filename + round-trip. |
| tests/kernels/test_rmsnorm_autotune.py | Adds GPU integration coverage for RMSNorm default path, forced search + cache reuse, and offline emit/lookup. |
| python/flydsl/autotune.py | Implements offline config artifact support, builder-mode execution, heuristic default gating, and expanded cache key axes. |
| kernels/rmsnorm_kernel.py | Adds a build-time BLOCK_THREADS knob and conditions known_block_size emission for >256-thread blocks. |
| kernels/rmsnorm_config.py | Introduces RMSNorm default heuristic and exhaustive config set for tuning (BLOCK_THREADS × waves_per_eu). |
| kernels/rmsnorm_autotune.py | Adds the autotuned RMSNorm wrapper using builder mode plus offline config key extraction. |
| docs/autotune_guide.md | Documents the three consumption paths (default / online search / offline artifacts), cache semantics, and correctness notes. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+413
to
+420
| if self.build_fn is None: | ||
| return self.fn | ||
| cache_key = (key, repr(config)) | ||
| built = self._build_cache.get(cache_key) | ||
| if built is None: | ||
| built = self.build_fn(config, *args, **kwargs) | ||
| self._build_cache[cache_key] = built | ||
| return built |
Comment on lines
308
to
316
| fn_name = op_name | ||
| if fn_name is None: | ||
| 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" | ||
| self._fn_name = fn_name | ||
| cache_dir = Path(os.environ.get("FLYDSL_AUTOTUNE_CACHE_DIR", os.path.expanduser("~/.flydsl/autotune"))) | ||
| self._cache_file = cache_dir / f"{fn_name}.json" |
Comment on lines
+388
to
+392
| for name in self.restore_value: | ||
| t = sig_args.get(name) | ||
| if t is not None and hasattr(t, "clone"): | ||
| snapshot[name] = (t, t.clone()) | ||
| return snapshot |
Comment on lines
+59
to
+70
| return Autotuner( | ||
| fn=None, | ||
| configs=[], # filled per-shape by _prune (needs N) | ||
| key=["input_t"], | ||
| warmup=10, | ||
| rep=50, | ||
| prune_configs_by=_prune, | ||
| build_fn=_build, | ||
| default=_default, | ||
| op_name="rmsnorm", | ||
| config_key_fn=_config_key, | ||
| ) |
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>
jhinpan
force-pushed
the
feat/autotune-offline-configs
branch
from
bd834da to
b7b3c25
Compare
jhinpan
force-pushed
the
feat/autotune-offline-configs
branch
2 times, most recently
from
554ce32 to
5f21d98
Compare
jhinpan
force-pushed
the
feat/autotune-offline-configs
branch
6 times, most recently
from
e369703 to
8e7d1e8
Compare
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>
jhinpan
force-pushed
the
feat/autotune-offline-configs
branch
from
8e7d1e8 to
689f570
Compare
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>
jhinpan
force-pushed
the
feat/autotune-offline-configs
branch
from
689f570 to
5e2cda0
Compare
Third in the series. Adds the offline path (aiter/SGLang model): tune once,
commit the config, look it up at serving with no search — complementing the
online JIT search from PR2.
A tuned Config is emitted to a committed tree (FLYDSL_AUTOTUNE_CONFIG_DIR) as a
self-describing JSON whose filename is the lookup key:
name,<spec axes>,device_name.json. The spec axes reuse the tuner's key_fn, so
the offline key is exactly the tuning key axes (shape + build-only scalars).
Lookup order in __call__: in-memory/disk cache -> offline artifact -> heuristic
default -> full search. FLYDSL_AUTOTUNE=1 bypasses all to force a fresh search
and re-emit.
Robustness (from two independent fresh reviews):
- spec is normalized through JSON on both emit and lookup, so a value JSON
turns into another type (e.g. a tuple -> list) still self-matches instead of
the artifact rejecting its own emitted config.
- A malformed artifact never raises: a non-dict spec (e.g. "spec": null),
missing fields, or unparseable JSON all warn and fall back to the default,
rather than crashing the serving call.
- name/spec/device_name are all validated against the call before use; a
stale/hand-edited mismatch is ignored with a warning.
- Filenames are sanitized to ASCII [A-Za-z0-9._-] and the resolved path is
confirmed under the config dir.
Because PR2's autotune_builder already carries name and key_fn, the offline path
is pure reuse — no new adopter-facing parameters. rmsnorm gets it for free.
Docs (docs/autotune_guide.md) now spell out the offline contract: specialize()
must include every axis the best config depends on (an omitted axis silently
collides — the offline tree is coarser than the scratch cache, which also keys
on toolchain/stride/env) and must use JSON-scalar values; an older-toolchain
artifact on the same device is served (retune is a review policy, not automatic).
Verified on MI350X: forced tune emits
rmsnorm,N=8192,dtype_str=bf16,device_name=AMD_Instinct_MI350X.json
and a normal run serves from it — the test asserts the served Config's
BLOCK_THREADS equals the emitted artifact's (proving offline was the source, not
a None->default fallback) with zero benchmarks.
Tests: offline unit tests assert served-config identity (offline value !=
default), plus tuple-spec round-trip, malformed-spec (no crash), corrupt-JSON,
and missing-field coverage; GPU test asserts the loaded Config came from the
artifact. ruff + black clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhinpan
force-pushed
the
feat/autotune-offline-configs
branch
from
5e2cda0 to
e7edf3b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Third in the #770 series: the offline path (aiter/SGLang model) — tune once, commit the config, look it up at serving with no search. Complements the online JIT search from #785.
Offline artifacts
A tuned
Configis emitted to a checked-in tree (FLYDSL_AUTOTUNE_CONFIG_DIR) as a self-describing JSON whose filename is the lookup key:{ "config": { "BLOCK_THREADS": 512 }, "device_name": "AMD_Instinct_MI350X", "dtype": "bf16", "op": "rmsnorm", "shape": { "N": 8192 } }This is SGLang's filename-as-key convention. Portable across identical GPUs, unlike the machine-local fingerprint cache from #783.
Lookup order in
__call__in-memory/disk cache → offline artifact → heuristic default → full search. With tuning off, a committed artifact is resolved first (no search);FLYDSL_AUTOTUNE=1forces the search and emits a fresh artifact.New on
@autotuneop_name— names the artifact / disk cache (required in builder mode wherefn=None; also fixes the builder-mode disk cache that wasunknown.json).config_key_fn(*args) -> (shape_key, dtype)— extracts the portable offline lookup axes.rmsnorm adopts it:
op_name="rmsnorm", keyed onN(the only tuned shape axis;Mjust scales the grid).Docs
docs/autotune_guide.md— the three consumption paths (heuristic default / online JIT / offline committed), when to tune, why not in CI, cache-key vs offline-artifact semantics, andrestore_valuecorrectness.Verification (MI350X / gfx950)
Forced tune emits the artifact above; a fresh tuner with tuning off and an empty scratch cache serves from it (no search), output matches the torch reference (max err 1.5e-2).
Tests
ruff + black clean.
Refs #770.
🤖 Generated with Claude Code