[3/5] autotune: add offline config artifacts (#770) - #786
Conversation
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.
bd834da to
b7b3c25
Compare
5e2cda0 to
e7edf3b
Compare
Resolve merge conflicts with upstream/main (through ROCm#818) for PR ROCm#786. Conflicts, both in the shared autotune feature series (ROCm#786 is [3/5], the conflicting upstream commit is [1/5] ROCm#783 of the same series): - python/flydsl/autotune.py - tests/unit/test_autotune.py Both were resolved by taking this branch's version: PR ROCm#786 is a strict superset of upstream's autotune changes (restore_value, _run_config, _snapshot/_restore_tensors, env/toolchain/device cache-key axes) and evolves them further (builder mode, offline config emit/lookup). Taking ours also drops the auto-merge duplications (doubled snapshot line, duplicate _run_config). Verified: tests/unit/test_autotune.py passes 44/1-skipped. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
/rerun-ci |
…fline path
Blockers:
- rmsnorm adopter imported kernels.rmsnorm_kernel, which no longer exists after
the upstream merge relocated the kernel to kernels/norm/. Fix the imports in
rmsnorm_autotune.py / rmsnorm_config.py to kernels.norm.rmsnorm_kernel; the
adopter was dead-on-import (hidden by the GPU test's pre-import skip). Add a
GPU-free regression test asserting the adopter's import path resolves.
- Offline load accepted any JSON scalar for a config knob, so a wrong-typed
value ("256" vs 256) passed validation and crashed build_fn at serve time.
Type-check knob values against the heuristic default's config, and wrap the
offline-sourced run so a bad artifact falls back to default/search instead of
crashing the serving call.
Robustness / serving hot path:
- Negative-cache the resolved decision (default/offline) so a missing or
rejected artifact isn't re-stat/parsed and re-warned on every serving call.
- Route serving-path warnings through logging (warn-once per artifact) instead
of unconditional print to stdout; keep tuning progress on stdout.
- Atomic writes (tempfile + os.replace) for offline artifacts and the disk
cache, so a concurrent reader never sees a torn file.
- Persist only searched bests to the disk cache (not memoized default/offline
configs), so they can't shadow a committed artifact next process.
- Serialize the compiler-hint mutate/restore on the shared launch fn with a lock
(thread-safe serving of hinted configs).
- Reject non-finite float knob values; sanitize the scratch-cache filename;
guard offline path ops against OSError (e.g. ENAMETOOLONG); warn on filename
sanitization collisions at emit; warn when builder mode drops non-structural
config kwargs; invalidate memoized offline decisions if CONFIG_DIR changes;
gate restore_value snapshot on copy_ support.
rmsnorm:
- get_all_configs no longer emits a single-config "search" for f32; it returns
the heuristic default explicitly (the BLOCK_THREADS sweep is 16-bit only) and
reuses the kernel's VEC_WIDTH as the single source of truth for tile width.
GPU-free unit tests extended (48 total): mistyped-scalar rejection, no-reread/
no-respam on rejected artifacts, and the adopter import-path guard. ruff + black
clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
|
/rerun-ci |
|
/rerun-ci |
1da82fc to
1633bf7
Compare
1633bf7 to
ff067a9
Compare
ff067a9 to
09dd110
Compare
jhinpan
left a comment
There was a problem hiding this comment.
Two non-blocking environment-configuration nits; neither is a correctness issue.
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
|
/rerun-ci |
Keep the RMSNorm autotune tests accepted by the Python style gate. Co-authored-by: Cursor <cursoragent@cursor.com>
|
@jhinpan test failed. |
|
@coderfeli Thanks — fixed in c4da52c, and CI is green now (mi325, mi355, navi all pass). Root cause: the The assertion is intentionally an exact literal: the artifact filename digest is derived from this key, so any change to the declared key axes invalidates artifacts already on disk and should fail loudly here. Updated it to the full four-axis key rather than deriving it from Verified locally on MI355X before pushing: 4544 passed, 2098 skipped, 0 failed across |
Adds one optional offline-config layer to the direct
@autotuneflow from #785.artifact_name=enables lookup underFLYDSL_AUTOTUNE_CONFIG_DIR.FLYDSL_AUTOTUNE=1bypasses those serving decisions, searches the existing configs, updates the scratch winner, and atomically emits an artifact.key, plus the call device descriptor. Adopters must put every dtype/layout/mode that changes the winner inkey.RMSNorm opts in with
artifact_name="rmsnorm"; its #785 direct-JIT kernel,m_in/N/dtype_strkey, default, and seven BLOCK_THREADS/waves-per-EU candidates are unchanged.There is no builder path, artifact-specific key callback, compiler/backend change, new export, or second RMSNorm factory.
Local verification on current
main: 99 passed, 3 skipped across autotune, cache-key, compile-hint, external-LLVM, autotuned RMSNorm, and legacy RMSNorm tests; Python style,py_compile, diff check, and the Sphinx build pass.Refs #770.