[4/5] autotune: CI guard + committed-config regression check (#770)#788
Open
jhinpan wants to merge 6 commits into
Open
[4/5] autotune: CI guard + committed-config regression check (#770)#788jhinpan wants to merge 6 commits into
jhinpan wants to merge 6 commits into
Conversation
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>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds CI-safe guardrails around FlyDSL’s autotune “offline config” workflow: committed configs are validated in GPU-free unit tests, CI is explicitly documented to avoid running autotune search, and RMSNorm gains an end-to-end autotune integration test path.
Changes:
- Add a GPU-free regression test that validates every committed
configs/autotune/*.jsonis well-formed and filename/content-consistent. - Seed
configs/autotune/with an initial RMSNorm tuned config + add documentation for generating/committing configs. - Extend autotune + RMSNorm plumbing/tests to support offline emit/lookup and CI-safe behavior (no search unless explicitly enabled).
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_autotune.py | GPU-free unit tests covering autotune keying, restore/reset semantics, builder mode, and offline emit/lookup behavior. |
| tests/unit/test_autotune_configs.py | New regression guard ensuring committed offline configs under configs/autotune/ are parseable, loadable, and filename-consistent. |
| tests/kernels/test_rmsnorm_autotune.py | GPU integration tests validating RMSNorm default path, forced-search path, cache reuse, and offline emit/lookup behavior. |
| python/flydsl/autotune.py | Adds CI-safe tuning gating, hardened cache key axes, builder-mode support, restore/reset semantics, and offline config emit/lookup. |
| kernels/rmsnorm_kernel.py | Makes RMSNorm’s build-time BLOCK_THREADS knob explicit and wires known_block_size for >256 threads. |
| kernels/rmsnorm_config.py | Introduces RMSNorm default heuristic + exhaustive config space definition for tuning. |
| kernels/rmsnorm_autotune.py | Adds the autotuned RMSNorm entrypoint using builder-mode autotuning + offline key extraction. |
| docs/autotune_guide.md | New guide documenting default/search/offline modes and CI guidance (“do not tune in CI”). |
| configs/autotune/rmsnorm,N=8192,dtype=bf16,device_name=AMD_Instinct_MI350X.json | Seed committed offline tuned config for RMSNorm. |
| configs/autotune/README.md | Documents the committed-config directory purpose and generation workflow. |
| .github/workflows/flydsl.yaml | Documents that CI must not set FLYDSL_AUTOTUNE and relies on default/offline-lookup paths. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+651
to
+660
| shape_key, dtype_str = self.config_key_fn(*args, **kwargs) | ||
| expected_op = self.op_name or self._fn_name | ||
| mismatches = [] | ||
| if data.get("op") not in (None, expected_op): | ||
| mismatches.append(f"op {data.get('op')!r}!={expected_op!r}") | ||
| if data.get("dtype") not in (None, dtype_str): | ||
| mismatches.append(f"dtype {data.get('dtype')!r}!={dtype_str!r}") | ||
| if "shape" in data and dict(shape_key) != data["shape"]: | ||
| mismatches.append(f"shape {data['shape']}!={dict(shape_key)}") | ||
| if mismatches: |
Comment on lines
+320
to
+328
| # Disk cache. Prefer an explicit op_name (builder mode has fn=None); | ||
| # otherwise fall back to the wrapped fn's name. | ||
| 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 |
Comment on lines
+126
to
+133
| if device_name is None: | ||
| device_name = _device_name() | ||
| parts = [_safe_component(op_name)] | ||
| for k, v in shape_key: | ||
| parts.append(f"{_safe_component(k)}={_safe_component(v)}") | ||
| parts.append(f"dtype={_safe_component(dtype_str)}") | ||
| parts.append(f"device_name={_safe_component(device_name)}") | ||
| return ",".join(parts) + ".json" |
| deterministic config with **no search**. This is the aiter/SGLang "offline" | ||
| model. See `docs/autotune_guide.md`. | ||
|
|
||
| - **Filename is the key:** `op,shape...,dtype=...,device_name=...json`. |
jhinpan
force-pushed
the
feat/autotune-ci-guard
branch
from
9fcb683 to
d9f1e14
Compare
jhinpan
force-pushed
the
feat/autotune-ci-guard
branch
6 times, most recently
from
1654d98 to
c767d77
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-ci-guard
branch
from
c767d77 to
1c55a7f
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-ci-guard
branch
from
1c55a7f to
b5c2305
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>
Fourth in the series. Makes the autotune path CI-safe and guards committed offline configs, without ever running the (slow, noisy) search on CI. CI already runs tests/unit/ and tests/kernels/ via scripts/run_tests.sh, so the GPU-free autotune unit tests and the l2_device rmsnorm autotune test are picked up automatically. This PR: - Adds tests/unit/test_autotune_configs.py, a GPU-free regression guard (aiter check_tuned_op_regression.sh analog): every committed offline config under configs/autotune/ must be well-formed, filename-consistent with its content (name + sorted spec + device_name, matching the emit path), and have scalar spec/config values. Config.from_dict accepts any dict, so the guard validates value types rather than relying on a round-trip that always passes. No-op when nothing is committed. - Seeds configs/autotune/ with the MI350X rmsnorm bf16 N=8192 config emitted by PR3's offline path, plus a README documenting how to generate/commit configs. - Documents in .github/workflows/flydsl.yaml that CI must not set FLYDSL_AUTOTUNE (search stays off; only default + offline-lookup paths run). GPU-free autotune unit tests pass. ruff + black clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhinpan
force-pushed
the
feat/autotune-ci-guard
branch
from
b5c2305 to
92406ee
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.
Fourth in the #770 series: make the autotune path CI-safe and guard committed offline configs, without ever running the (slow, noisy) search on CI.
What this adds
CI already runs
tests/unit/andtests/kernels/viascripts/run_tests.sh, so the GPU-free autotune unit tests and thel2_devicermsnorm autotune test are picked up automatically. On top of that:tests/unit/test_autotune_configs.py— a GPU-free regression guard (the aitercheck_tuned_op_regression.shanalog): every committed offline config underconfigs/autotune/must be well-formed,Config-loadable, round-trip stable, and filename-consistent with its content (filename == key). No-op when nothing is committed.configs/autotune/— seeded with the MI350X rmsnorm bf16 N=8192 config emitted by [3/5] autotune: offline config emit + runtime lookup (#770) #786's offline path, plus a README on how to generate/commit configs..github/workflows/flydsl.yaml— documents that CI must not setFLYDSL_AUTOTUNE(search stays off; only the default + offline-lookup paths run). Seedocs/autotune_guide.md→ "Do not tune in CI".Tests
33 GPU-free autotune unit tests pass (incl. the new config guard). ruff + black clean.
Refs #770.
🤖 Generated with Claude Code