feat(autotuner): reach both persistence backends through autotune()

[aleozlx]

📌 Description

Targets yanxu/autotune-cache-v2-mvp, not main — a proposal for flashinfer-ai#3861.

Reach v2 through the existing autotune() instead of a second public context manager. An explicit
v2_opt_in argument selects the implementation, and autotune() dispatches to autotune_v2()
rather than absorbing it:

autotune(True)                                          # legacy: memory only
autotune(True, cache="cfg.json")                        # legacy: JSON file cache
autotune(True, v2_opt_in=True)                          # v2 + managed store
autotune(False, v2_opt_in=True)                         # v2 replay, no profiling
autotune(True, v2_opt_in=True, persistent_cache=False)  # v2, no disk

if v2_opt_in:
    _reject_v1_only_args(cache)
    with autotune_v2(...):          # v2's implementation, unchanged, in autotune_cache.py
        yield
    return
_reject_v2_only_args(cache_root, persistent_cache, measure)
# ---- v1 implementation ----

Not the "subtle behaviour change" flashinfer-ai#3920 warned about. v2_opt_in defaults to False, so absent
the argument behaviour is byte-identical. The v2 branch hands off the whole context and returns, so
the v1 body never runs for it — the two implementations never share a function scope, and
interleaving is not merely prevented but unrepresentable. No recursion: autotune_v2() calls
autotune() without the opt-in for the shared machinery (bucketing, skip_ops, tuning-mode
refcount).

Branch selection is positive. if not v2_opt_in: asserts "everything that is not v2", an
unbounded set that would silently absorb a third implementation; each branch now names its own and
rejects only the arguments the other owns.

The RFC's two objections are answered rather than accepted:

Objection	Answer
cache=<file> would be silently redefined as a root dir	It isn't. cache keeps its exact v1 meaning; the store uses a separate cache_root, and combining them raises.
v1 is block-scoped, v2 attaches for the process	Attach happens only inside autotune_v2(), reached only via the opt-in. autotune(True) and autotune(True, cache=...) keep block scope.

Why one entry point. A shared front door is a structural brake on divergence, and it changes
what happens to future edits: with two entry points a change to shared machinery must be plumbed
twice and drifts by omission, and a v1-only fix (a ported TRT-LLM change, say) is invisible to v2
unless someone remembers to mirror it. With one signature both land in front of the same reviewer —
divergence is paid per-change as a review question rather than later, all at once, as a migration.

Why v2_opt_in and not a behaviour name. A version-named flag announces its own expiry and
retires along a path a function name cannot offer: flip the default → ignore the argument internally
once no caller passes it → delete it, with callers working at every step, each step independently
verifiable and revertible. autotune_v2() cannot be "ignored" — an import either resolves or raises
— so removing it is one breaking edit for every caller. The durable knobs (persistent_cache,
cache_root, measure) keep behaviour names because they outlive v1.

⚠️ Critical assumption, flagged not resolved. The graduation plan assumes v1 is eventually
removed. That is not decided — no owner, no date. It requires the frameworks to stay migrated,
new TRT-LLM work to land against the runner contract rather than v1's cache, and a major version
to spend. If any fails, the honest outcome is indefinite coexistence with a preferred default, and
§5 should say so instead of carrying a removal plan that never runs. This PR does not depend on
that assumption: one entry point is worth having under either outcome — permanent coexistence is
exactly when two front doors drift most.

User-facing docs, the gap flagged on flashinfer-ai#3861. docs/autotuning.rst described the legacy
implementation only, and two sections were wrong once v2 exists: the Config Caching note warned
that concurrent writes "may result in lost updates" — precisely what v2's atomic per-entry publish
fixes — and Config Lookup Priority omitted the managed store. Both corrected, plus an
Autotuner v2 (Experimental) section and API-reference entries for v2_opt_in, persistent_cache,
cache_root, measure, MeasurementPolicy, autotune_v2_reload and autotune_v2 — three of
which are exported from the top-level package and had no user-facing documentation anywhere.

Design doc, updated in the same change: §2.1 describes dispatch (it previously claimed v2 was
"deliberately disjoint from autotune()"); §4 explains why the two share an entry point; §5 opens
with the assumption callout; §5.1 carries the option-space and future-edits comparison; §5.4 records
the docs as written rather than owed; the Scope line narrows to the v2 branch of autotuner.py.

🔍 Related Issues

RFC flashinfer-ai#3920 · MVP flashinfer-ai#3861 · design-doc policy flashinfer-ai#4334

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests added — tests/autotuner/test_autotune_unified_api.py, 12 GPU-free cases covering
  the dispatcher, which is the new code: argument ownership and its error messages, the nesting
  guard, rollback on a rejected argument, validation ordering, one smoke test that dispatch reaches
  v2, and the alias's argument translation. Deliberately thin on v2 behaviour — that is unchanged
  and already covered by test_autotune_cache_v2.py; six cases written against the earlier
  merged-body design were deleted once dispatch made the boundary structural.
• All tests passing — NOT VERIFIED. No torch/pytest in this environment, so the new tests
  and the existing tests/autotuner/ suite are unrun. Needs CI or a GPU box.

Reviewer Notes

1. measure= now requires v2_opt_in=True (behaviour change from an earlier revision of this
   PR). It is a v2 concept and the dispatcher has no path that could honour it on the v1 branch.
   Say the word if you would rather it stayed permissive.
2. persistent_cache without v2_opt_in=True raises — deliberate, but it is a new error for an
   argument name autotune_v2() already accepts.
3. Is v1 still worth keeping? It is largely a TRT-LLM integration. If the argument for keeping it
   is absorbing upstream changes, that weakens once those changes must satisfy the runner contract
   anyway. Worth checking against a real recent TRT-LLM port — it decides whether §5's timeline is
   conservative or simply never.

[aleozlx]

Investigation: how vLLM and SGLang call the autotuner today

New in this comment

• vLLM has a second autotune call site (flashinfer_sparse_mla_warmup.py) that framework_patches/vllm_autotune_v2.patch does not migrate. Because the draft's header says it deletes flashinfer_autotune_cache.py entirely, and that file's second caller imports two functions from it, applying the draft as written is an ImportError at import time — not a subtle regression. Analysis of why it was missed below.
• SGLang can only detect "am I inside autotune?" by grepping the Python stack (pynccl_allocator.py:393 matches _flashinfer_autotune in traceback.format_stack()). No public predicate exists; it breaks on rename.
• vLLM omits skip_ops from its cache key; SGLang includes it. Changing VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS replays tactics tuned under a different skip policy — the vllm#43119 class in a different field. Free fix under the managed store.
• Effort and sequencing estimate for us sending both migration PRs: diffs are ~half a day and a few hours; §5.2 gate 1 unpacks into four sequential steps, the first two ours and blocking (completeness check, export()/install()).
• Wording tension in [RFC] Autotuner v2: FlashInfer-managed persistent cache and deployment-matched measurement flashinfer-ai/flashinfer#3920 Evidence 1: the vLLM flow it describes (rank 0 tunes → broadcast → load_configs) executes only at world_size == 1 at HEAD, where the broadcast is a no-op.

Corrections to the first revision of this comment

• vLLM's multi-rank cache disable presented as a headline finding — already in [RFC] Autotuner v2: FlashInfer-managed persistent cache and deployment-matched measurement flashinfer-ai/flashinfer#3920 Exhibit A and removed line-by-line in the patch draft. Re-verification added currency only.
• "two hash functions, two path layouts … policy neither should have had to invent" — that is [RFC] Autotuner v2: FlashInfer-managed persistent cache and deployment-matched measurement flashinfer-ai/flashinfer#3920's thesis, restated as if new.

Re-verified, unchanged: every claim in flashinfer-ai#3920 Evidence 1 still holds at HEAD 2026-08-04, including the zero-external-use feature list.

---

Read against vllm-project/vllm@0b1c151 and sgl-project/sglang@5e6c37f (both HEAD, 2026-08-04), to check the migration surface this PR implies and re-verify the RFC survey at a newer HEAD.

Most of this was already established in flashinfer-ai#3920 and in framework_patches/vllm_autotune_v2.patch. Attribution is inline below so the genuinely new items are separable.

Correction to the first revision of this comment. It presented vLLM's multi-rank cache disable as a headline finding:
"This is the lasting damage from vllm#43119 and it is still in effect at HEAD — the RFC cites it, but it is worth stating that it is current, not historical."
That overclaimed. flashinfer-ai#3920's Exhibit A already says "still disabled on the general path today", and framework_patches/vllm_autotune_v2.patch removes the exact use_persistent_cache lines and calls out the win in its header. Re-verification added currency, nothing more. Also struck: "two projects, two hash functions ... all of it is FlashInfer policy that neither should have had to invent" — that is flashinfer-ai#3920's thesis, restated here as if new.

Confirmed, already documented

Finding	Already in
vLLM runs use_persistent_cache = False whenever world_size > 1 — every multi-GPU deployment re-tunes in memory each start	flashinfer-ai#3920 Exhibit A ("still disabled on the general path today"); framework_patches/vllm_autotune_v2.patch removes those exact lines
SGLang writes one file per rank (rank_tp{n}_pp{n}_dp{n}.json) and tunes on every rank	flashinfer-ai#3920 Evidence 1, SGLang bullet
Zero external use of tuning_buckets, round_up, tune_mode=False, direct save_configs, FLASHINFER_AUTOTUNER_LOAD_FROM_FILE	flashinfer-ai#3920 Evidence 1 — still true at this HEAD. (round_up greps 220× in vLLM; all unrelated arithmetic.)
vLLM reinvents hashing / atomic write / enable knob	flashinfer-ai#3920 Evidence 1
export()/install() must land before framework migration	flashinfer-ai#3920 roadmap ordering (item 4 before item 6)
vllm#43119 as the key-completeness exhibit	flashinfer-ai#3920 Exhibit A

The only thing re-verification adds is currency: all of it still holds three weeks later.

Where the code lives (for the migration PRs)

vLLM — warmup/kernel_warmup.py (general warmup), warmup/flashinfer_sparse_mla_warmup.py (a second pass), warmup/flashinfer_autotune_cache.py (hash + path + atomic write), config/kernel.py + config/vllm.py (enable_flashinfer_autotune), envs.py (VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR, ..._SKIP_OPS), utils/flashinfer.py (re-export).

SGLang — srt/model_executor/runner/flashinfer_autotune.py (all of it), runner/base_runner.py (call site), srt/server_args.py (--disable-flashinfer-autotune, --flashinfer-autotune-skip-ops), srt/environ.py (SGLANG_FLASHINFER_AUTOTUNE_CACHE).

Full used surface across both projects: autotune(True, cache=<path>, skip_ops=<set>) and AutoTuner.get().load_configs(<path>). Nothing else.

New — not found in flashinfer-ai#3920 or the flashinfer-ai#3861 thread

1. vLLM has a second autotune call site, and the migration draft misses it.

kernel_warmup.py:145 calls flashinfer_sparse_mla_decode_autotune_warmup(worker) before flashinfer_autotune(...) at line 168. Both resolve the same resolve_flashinfer_autotune_file(runner) path and both open autotune(True, cache=str(cache_path)). They coexist only because v1's save_configs does read/merge/write against the on-disk file.

framework_patches/vllm_autotune_v2.patch migrates kernel_warmup.py only — zero occurrences of sparse_mla. That turns into a concrete breakage rather than an omission, because the draft's header states it deletes:

flashinfer_autotune_cache.py entirely, ~56 lines

and flashinfer_sparse_mla_warmup.py:10-12 imports two functions from that module:

from vllm.model_executor.warmup.flashinfer_autotune_cache import (
    resolve_flashinfer_autotune_file,
    write_flashinfer_autotune_cache,
)

So applying the draft as written is an ImportError at module import, not a subtle regression — it fails loudly, which is the good version of this problem. The fix is to migrate both sites or keep the helper.

Why it was missed — worth recording, because the obvious explanations do not hold:

• Not because it postdates the survey. The file was created 2026-06-22 (vllm#43477, "Enable DeepSeek V4 and GLM-5.1 on SM120"); the draft is dated 2026-07-24 and the RFC survey says "HEAD, 2026-07".
• Not a narrow grep. flashinfer_autotune(, with .*autotune(, resolve_flashinfer_autotune_file and load_configs all match both files. Only the very narrow fi_utils.autotune matches just kernel_warmup.py — the two sites import the same symbol differently (import vllm.utils.flashinfer as fi_utils vs from vllm.utils.flashinfer import autotune as flashinfer_autotune).
• Most likely: the survey characterised a usage pattern, and the draft migrated the canonical instance of it. The sparse-MLA site follows the identical pattern — one-shot warmup, cache=<path>, leader broadcast, load_configs — so it changes no conclusion in [RFC] Autotuner v2: FlashInfer-managed persistent cache and deployment-matched measurement flashinfer-ai/flashinfer#3920. It is a missed instance, not a missed pattern, and it only becomes a defect because the draft claims to delete a module a second caller depends on.

The finding also illustrates §2.2's filename-as-identity argument concretely: two independent tuning passes were never really distinct cache identities, and only v1's read/merge/write hid it.

2. SGLang can only detect "am I inside autotune?" by sniffing the Python stack.

srt/distributed/device_communicators/pynccl_allocator.py:393:

traces = traceback.format_stack()
# Skip autotune stack traces
if any("_flashinfer_autotune" in trace for trace in traces):
    return

A symmetric-memory debug check suppresses warnings during tuning by matching a function name in the stack — it breaks the day that function is renamed. There is no supported predicate; is_tuning_mode is not public. Cheap addition to v2's observability surface.

3. vLLM omits skip_ops from its cache key; SGLang includes it.

SGLang appends "skip_ops=" + ",".join(sorted(skip_ops)) to its key material. vLLM's key is aot_compile_hash_factors(vllm_config) only, so changing VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS reuses tactics tuned under a different skip policy — the same key-incompleteness class as vllm#43119, in a different field. The patch draft passes skip_ops through as a kwarg but does not touch keying. Under the managed store this disappears because we own the key, which is a concrete "you get a bug fix for free" line for the vLLM PR.

4. A wording tension in flashinfer-ai#3920's Evidence 1.

The vLLM bullet reads "rank 0 runs autotune(tune_mode=True, cache=<path>), broadcasts the raw JSON file bytes, every rank atomically rewrites the file and calls load_configs(path)." At HEAD that path executes only when world_size == 1, where the broadcast is a no-op. Exhibit A corrects it two bullets later, but Evidence 1 is the bullet most likely to be quoted in isolation, and read alone it implies multi-rank works that way. Suggest a half-sentence fix in the RFC.

Divergence between the two consumers

	vLLM	SGLang
Who tunes	rank 0 (single-GPU only); all ranks uncached (multi-GPU)	every rank, always
Cache file	one shared file per config hash	one file per rank
Hash inputs	aot_compile_hash_factors(vllm_config)	model path, dtype, quant, moe backend, tp/pp/dp/ep, hf config class, skip_ops
Distribution	broadcast_object of raw bytes + load_configs	none
Disable knob	enable_flashinfer_autotune	--disable-flashinfer-autotune
Reuse knob	—	SGLANG_FLASHINFER_AUTOTUNE_CACHE=0 → timestamped runs/ dir

---

Effort estimate: us sending the migration PRs

Framing: neither project will prioritise migrating someone else's caching layer, and the patch draft in flashinfer-ai#3861 already implies we send both. So this is our cost plus two external review queues.

Diffs are small, mostly deletion

vLLM ≈ 150 lines removed / 30 added — flashinfer_autotune_cache.py deletes entirely (except mapping VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR → cache_root); flashinfer_autotune() goes ~70 → ~15 lines; flashinfer_sparse_mla_warmup.py sheds ~40 lines of broadcast/barrier/load_configs (once finding 1 is folded in).

SGLang ≈ 60 removed / 10 added — flashinfer_autotune_cache_path() (~45 lines) deletes, per-rank filenames go, SGLANG_FLASHINFER_AUTOTUNE_CACHE=0 maps to persistent_cache=False. Their all-ranks pattern is already what the store is designed for.

Coding: ~half a day (vLLM), a few hours (SGLang). That is not the cost.

What actually costs

export()/install() is unimplemented, and vLLM's broadcast does not need a shared filesystem. flashinfer-ai#3920 already orders item 4 before item 6; the point here is that it is a hard prerequisite, not a nice-to-have — without it we migrate vLLM's single-node case and they keep broadcast for multi-node, retaining most of what the migration was meant to delete.

vLLM's multi-GPU path is a behaviour change, not a refactor. Persistence is currently off there, so migrating turns caching on in the configuration everyone runs. Needs a rank-consistency and perf gate, not a green test suite.

The hard sell is key completeness, and it is the one thing still unguarded. vLLM disabled this because a key omitted a field. Re-enabling requires showing the class is structurally fixed — contract rule 5 plus the debug-mode completeness check, which flashinfer-ai#3920 lists as "proposed", not in flashinfer-ai#3861. The thing that got the feature disabled is the thing not yet mechanically guarded. Build the check before opening the vLLM PR, or the review stalls on exactly the point that killed it last time.

Pre-empt one conceptual objection in the PR body. vLLM's cache identity is model-level; ours contains no model identity at all, separation coming from the op key. Correct, but it looks wrong at a glance and is better answered in the description than in a review round.

Estimate

	Coding	Validation	Wall-clock to land
SGLang	~4h	light — already tunes every rank	1–3 weeks
vLLM	~1d	heavy — multi-GPU perf + rank consistency	4–8 weeks, after export/install

Implication for §5.2 gate 1

"Frameworks migrated and released" unpacks into four sequential steps, the first two ours and blocking:

1. debug-mode key-completeness check (proposed only)
2. export() / install() (roadmap item 4, unimplemented)
3. SGLang PR → review → release
4. vLLM PR → review → release

Realistically multi-quarter. That does not change what this PR does — one entry point is right either way — but it should calibrate how the graduation timeline is written, and it reinforces the §5 assumption callout: worth making coexistence good rather than treating removal as imminent.

Line numbers read at HEAD 2026-08-04; both projects move fast.