Studio: install flashinfer on demand for NVFP4 without moving torch - #11730
danielhanchen wants to merge 56 commits into
Conversation
An explicit NVFP4 image or video load (or an nvfp4 prequant seed) now installs FlashInfer when it is missing, instead of silently falling back to torchao. - Pins flashinfer-python 0.6.6, the version the dispatch allowlist and the op fakes were checked against, plus the flashinfer-jit-cache built for the running CUDA (cu128 / cu129 / cu130, newest same-major build not newer than the runtime). Without a jit-cache the host must have nvcc, else it refuses. - Every installed distribution is pinned to its exact version through a constraints file, so the resolver can only add packages. Afterwards the metadata set, torch.__version__ / torch.version.cuda (fresh interpreter) and `import flashinfer` are verified; any drift or failure uninstalls what was added and restores anything moved. - Linux + CUDA torch + a device in the flashinfer NVFP4 set only. Refuses with a reason on ROCm, macOS, Windows, CPU, other archs, offline mode, an unreachable index, UNSLOTH_NVFP4_FLASHINFER_INSTALL=0, a broken existing flashinfer, or a stray jit-cache/cubin of another version. - Once per process under a lock; policy refusals are not memoised. - Runs outside the load locks, next to the attention-kernel pre-install. - Image and video status gain transformer_quant_backend_reason, set when an NVFP4 load runs torchao.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
Security findingsAdvisory findings (1)ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…lback reason per backend A local_files_only load downloads nothing, so the installer now refuses there instead of probing indexes and fetching wheels. The torchao fallback reason is kept per loading backend object, so the image and video status routes report the reason for their own loaded model rather than the last call anywhere.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
for more information, see https://pre-commit.ci
…rd it only after the swap
|
@codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: edc1b41fa4
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _nvfp4_install_outcome = ensure_flashinfer_for_nvfp4( | ||
| device, logger = logger, local_files_only = local_files_only | ||
| ) |
There was a problem hiding this comment.
When an image NVFP4 load is superseded or unloaded while this call is downloading the roughly 1.5 GB JIT cache, cancellation is not checked until ensure_flashinfer_for_nvfp4 returns, and the installer receives no cancellation event. The obsolete load therefore continues for up to 30 minutes and mutates the shared environment even while a newer load proceeds; the video path at video.py:4417 has the same behavior. Make the install subprocess cancellable or poll the load token during installation.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Leaving this as is. The install runs once per process under the install lock, rolls back on failure, and only adds the package the next NVFP4 load wants anyway, so a superseded load finishing it does no harm; the newer load's own ensure call sees the finished outcome. Killing pip mid-install would be the riskier path for the environment.
💡 Codex Review Honor UV_OFFLINE before probing package indexes
When an operator sets only ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Fixed the UV_OFFLINE item in b301c4a: the installer now treats UV_OFFLINE like HF offline mode and refuses before probing PyPI or the flashinfer index, with a test that no probe runs. |
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…e check (unslothai#11763) * Pin the legacy store outcome in the failed-delete data settings check The failed-delete check asserted that the Clear chats confirmation closes when the backend refuses the clear. The app closes it only when the legacy IndexedDB clear succeeds: with both stores failed, clearStoredChats throws and the confirmation stays open, armed for a retry, by design since unslothai#7029. Which of the two happens is decided by the legacy store gate (unslothai#9446). Any legacy read slower than 1 s shuts it for the life of the page, after which a clear reports the legacy store as failed. Firefox on the Windows runner occasionally crosses that line, and the check then fails at line 216 with the confirmation still on screen. Seen on main at 288b1a3 and on unslothai#11730 and a87d180, with Chromium passing in the same job. The fixture now refuses legacy writes on request, so the check drives the outcome it asserts instead of inheriting whatever the gate did earlier: the clear settles, the confirmation is still open with its action and switch re-enabled, no chat is lost, and Cancel closes it with no further delete. It is pinned to the confirmation by its switch, not `.last`, which slides onto Settings once the confirmation goes. * Record the dialogs still mounted when the data settings check fails A Firefox run on the Windows runner failed in reset() with one hidden dialog left 5 s after closing Settings, and the only evidence was a blank screenshot. The report now lists every [role=dialog] still in the DOM with its data-state, display, opacity and running animations, so the next occurrence says whether an exit animation held the dialog or something else kept it mounted.
…light fallback The jit-cache step now passes its index through _installer_prefix, which drops a pip-only mirror instead of adding a second --index-url that uv rejects. A transient (allocation) preflight failure is still not memoised for backend selection, but is kept per device so the status route can say why the resident model fell back to torchao. A later memoised preflight on that device replaces it.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 915d0f7095
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| present, detail = _import_flashinfer() | ||
| if present: | ||
| return _finish(True, f"flashinfer {detail} already installed", None, None) |
There was a problem hiding this comment.
When flashinfer-python is already installed but flashinfer-jit-cache is absent, import flashinfer can succeed and this return prevents the installer from adding the required cache. On Studio hosts without nvcc, the first NVFP4 preflight then attempts an uncached JIT build and falls back to torchao; the repository's Docker installation separately installs the matching cache for this reason (docker/Dockerfile:169-180). Check for a matching cache before treating an importable package as fully ready.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Leaving this as is, on purpose. The installer only covers a missing flashinfer. A flashinfer that is already present (usually pulled in by vLLM, at whatever version vLLM resolved) belongs to the user and is not changed, the same rule that keeps a present but broken install from being repaired. If that copy cannot JIT, the preflight fails safely to torchao and the status route reports the preflight reason. That matches the behaviour before this PR.
…stall Without uv, pip ignores UV_DEFAULT_INDEX / UV_INDEX_URL, yet the preflight skipped the pypi.org probe because a mirror was configured. Pass the uv index as --index-url to pip when PIP_INDEX_URL is unset.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…cache index flashinfer fixes its jit-cache directory once, at import. A second Studio process that imported it between another process's flashinfer-python and jit-cache steps kept JIT-compiling for its whole life. When the pinned flashinfer-python is on disk without its jit-cache and flashinfer is not yet imported, wait on the shared install lock before importing; nothing is installed or modified, and every other state returns without touching the lock. The jit-cache step now runs without UV_INDEX, UV_EXTRA_INDEX_URL, UV_FIND_LINKS and the pip equivalents in its environment. Those outrank uv's --index-url, so under uv's first-index strategy a mirror listing any flashinfer-jit-cache hid the pinned flashinfer index and the install rolled back.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
@codex security review |
…o-nvfp4-flashinfer-autoinstall
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7b56e24720
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| found: set[str] = set() | ||
| pending = [name for name in roots if name in candidates] |
There was a problem hiding this comment.
When pip times out or exits after installing one or more new dependencies but before installing flashinfer-python itself, the failed step has no success summary and none of the roots in this seed list may be present in candidates. _requirement_closure then returns nothing for those dependencies, so _rollback leaves packages added by this transaction in the environment despite claiming to roll back failed installs. Capture the resolved dependency closure before launching the installer, or otherwise identify partially installed dependencies without requiring the root distribution's metadata to exist.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
What can be left is additions only: the constraints keep every existing package where it was, and the next load's install reuses them. Not worth a resolve pass before every install.
| if Version is not None: | ||
| try: | ||
| Version(version) | ||
| except InvalidVersion: | ||
| continue |
There was a problem hiding this comment.
If an installed dependency has a non-PEP-440 version, such as a vendor-built torch distribution, this branch silently omits it from the constraints file. The resolver may then replace that dependency while adding FlashInfer; drift detection notices the replacement, but rollback constructs the same invalid name==old_version requirement and cannot restore it. Refuse the transaction before running either install step when an existing dependency cannot be represented as a constraint, rather than leaving it unprotected.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
pip 24.1 and uv both reject non-PEP 440 versions outright, and the dependencies here (torch, numpy, requests and the like) do not ship such versions. The torch re-probe and drift check still fail the install if torch moves.
| try: | ||
| with urllib.request.urlopen(urllib.request.Request(url, method = "HEAD"), timeout = 10): | ||
| return True | ||
| except urllib.error.HTTPError: | ||
| return False |
There was a problem hiding this comment.
When PyPI, the FlashInfer index, or an intervening proxy permits package downloads but rejects HEAD with HTTP 405, this branch treats the endpoint as unreachable and skips an installation that pip or uv could complete with its normal GET request. This leaves eligible NVFP4 loads on the slower torchao backend solely because of the probe method. Treat a method-not-allowed response as evidence of reachability or retry the probe with GET before refusing the install.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Both probed endpoints answer HEAD with 200 today, and a configured proxy or mirror skips the probe. Not changing it.
…is rolled back A step that dies without its summary gave the rollback no report to filter by, so it restored every drifted package, including one the built-in terminal upgraded meanwhile. The constraints pin every installed package, so this transaction cannot have moved those: restore only what the constraints could not pin.
…o-nvfp4-flashinfer-autoinstall
for more information, see https://pre-commit.ci
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…o-nvfp4-flashinfer-autoinstall
…o-nvfp4-flashinfer-autoinstall
…o-nvfp4-flashinfer-autoinstall
…r the rollback test by count tomllib is 3.11 but already falls back to tomli or None; mark the import novermin so the floor lint sees the guard. The rollback wait test compared two clock stamps under -n 4; it now records how many imports ran before the rollback, which is the ordering it asserts.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e6acf2fd9a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| and ( | ||
| _pipeline_prequant_planned == TQ_NVFP4 | ||
| or self._nvfp4_checkpoint_will_load( |
There was a problem hiding this comment.
When an explicit image transformer_quant="nvfp4" has _pipeline_prequant_planned == PIPELINE_SEED_DECLINED (for example, the artifact-sized prefetch plan still requires offload), the explicit request satisfies the earlier tuple check and this Hub probe can return true, so FlashInfer and its roughly 1.5 GB cache are installed. The loader later explicitly excludes PIPELINE_SEED_DECLINED from pipeline_seed_scheme at lines 4575-4578 and either performs the on-the-fly torchao conversion or declines quantization under offload, meaning FlashInfer cannot serve this load. Short-circuit this sentinel as the video gate already does before invoking the installer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct: a declined seed loads the released denoiser, so the install bought nothing. Fixed in eac82e1, which skips the install when the plan declined the seed, with a regression test.
…nsloth into r11730 # Conflicts: # studio/backend/core/inference/diffusion_nvfp4_ops.py # studio/backend/tests/test_nvfp4_diffusion_flag.py
…he NVFP4 seed A declined seed loads the released denoiser, so no pre-quantised checkpoint opens and FlashInfer serves nothing.
…all' into studio-nvfp4-flashinfer-autoinstall
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
This branch has not been deployed
1 participant
Stacked on #10731 (base
studio-nvfp4-kernels).Why
The NVFP4 FlashInfer backend from #10730 / #10731 only engages when
flashinferis already installed. Studio does not ship it, so a user who picks NVFP4 lands on the slower torchao path without knowing why. This installs it on demand, the first time an NVFP4 load needs it. It never moves torch, triton, CUDA or any other installed package.What it does
core/inference/diffusion_nvfp4_install.pyaddsensure_flashinfer_for_nvfp4(device, ...) -> (ok, reason), which never raises. The image loader (diffusion.py) and the video loader (video.py) call it for an explicitnvfp4request or an NVFP4 prequant seed, before the load locks, and only when a hosted NVFP4 checkpoint will actually load. FlashInfer only serves a pre-quantised checkpoint; an on-the-fly NVFP4 build always runs on torchao, so installing for it would buy 1.2-1.8 GB the load cannot use. The install goes ahead when the download plan already settled NVFP4, when a local override or cached checkpoint exists, or when the Hub listing shows the checkpoint. A private, gated or unpublished repo, a base with no hosted checkpoint, a LoRA bake or a declined video seed installs nothing. Offline loads check only the cache. The video gate also covers the MiniMax-H3 modular path.flashinfer-python==0.6.6, the version the fast-dispatch allowlist and op fakes were checked against.flashinfer-jit-cache==0.6.6+cuXYZ(the full local version, as flashinfer's own installer pins it) fromflashinfer.ai/whl/cuXYZwith--no-deps. This step dropsUV_INDEX,UV_EXTRA_INDEX_URL,UV_FIND_LINKSand the pip index variables from its environment, and under uv also passes the pinned index as--index, so neither an environment nor a uv.toml mirror can outrank it; the flashinfer-python step still uses the user's mirrors. It picks the newest cu128 / cu129 / cu130 build with the running torch's CUDA major and a minor no newer than it. With no matching build it installs flashinfer-python alone only ifnvccexists, and otherwise refuses, since every kernel would otherwise compile on first use.--only-binary :all:(uv, with pip as the fallback). The resolver can only add packages.torch.__version__/torch.version.cudaand import flashinfer 0.6.6. Anything else rolls back everything the install added.UNSLOTH_NVFP4_FLASHINFER_INSTALL=0UNSLOTH_NVFP4_BACKEND=torchaosys.prefix. A process that finds the pinned flashinfer-python on disk without its jit-cache, before importing flashinfer, waits on that file lock so it does not import halfway through another process's install; it never modifies that install. If that install outlasts the wait, the load reports flashinfer not ready rather than importing it without the jit-cache. Refusals are not remembered, so a later load can still install once the condition clears.transformer_quant_backend_reason. A failed flashinfer preflight on the loaded device is reported there too, including a transient out-of-memory one, which is still not cached for backend selection.diffusion_nvfp4_dispatch.forget_availability()drops the cached "flashinfer missing" answer once an install succeeds.Installer environment handling
The installer that will run (uv, else pip) is asked what it is configured to do before anything is probed or installed.
UV_INDEX_URL,UV_DEFAULT_INDEX,UV_INDEX,UV_EXTRA_INDEX_URL,PIP_INDEX_URL) and from configuration files: pip's effective settings viapip config list(environment over[install]over[global], every pip.conf andPIP_CONFIG_FILE), uv's viaUV_CONFIG_FILE,UV_NO_CONFIGand the project, user and systemuv.toml. A configured mirror skips the pypi.org probe; a mirror only one installer reads is handed to the other unless that one has its own.UV_OFFLINEor uvoffline, pip or uvno-index(so find-links-only setups), pip or uvno-deps, piptarget/prefix/rootand uv[pip] target/prefix(they install outside this interpreter's environment), and piprequire-virtualenvoutside a venv.HTTPS_PROXY/NO_PROXYare used by the probe itself. With onlyALL_PROXY, or pip's ownproxysetting, the probes are skipped and the installer decides. A failed TLS handshake counts as reachable, so pip'scertor uv's native TLS decide trust, not the probe.UV_PYTHON,UV_PROJECT_ENVIRONMENTandVIRTUAL_ENVlose to the explicit--python.PIP_CONSTRAINT/UV_CONSTRAINTthat conflicts with the pins fails the resolve before anything moves, and the reason is reported.UV_OVERRIDEthat moves a pinned package is caught by the drift check and rolled back, at the cost of the download.PIP_USERis left to pip: it refuses inside a venv, and outside one the user site is importable.UV_EXCLUDE_NEWER,UV_INDEX_STRATEGY, prerelease settings) either resolve 0.6.6 or fail cleanly before installing.Evidence
A real install ran on a scratch venv with torch 2.12.1+cu130 and triton 3.7.1, stubbing only the device-capability check:
uv pip freezediff.torchread2.12.1+cu130 13.0before and after, andimport flashinferreported 0.6.6. A second call found flashinfer already installed and ran nothing.packaging==23.2installed (flashinfer needs 24.2 or newer), the installer refused instead of upgrading it, and nothing changed.Tests
test_diffusion_nvfp4_install.pyhas 95 hermetic tests, which cover:test_diffusion_backend.py/test_video_backend.pydrive image, GGUF, Wan-A14B, Wan-5B and MiniMax-H3 NVFP4 loads up to the install hop. When the checkpoint repo returns 401 /RepositoryNotFoundError/GatedRepoError, lacks the file or is not hosted, nothing is installed; those tests fail on the previous head. When the checkpoint is reachable, the install still runs. Offline loads check only the cache.tests/test_diffusion*.py tests/test_video*.py tests/test_nvfp4*.py(CPU): 4437 passed. The failure list matches the base branch exactly: 15 environment-only failures intest_diffusion_dataset_api.py.Not covered
mm_fp4/nvfp4_quantizeon sm_100 / sm_120 withoutnvcc, or that flashinfer-cubin is not needed for the cutlass path.NVFP4 kill switch
NVFP4 for image and video generation is off by default until the hosted
unsloth/*-NVFP4repos are public.UNSLOTH_NVFP4_DIFFUSION=1(alsotrue,yes,on) restores the full NVFP4 behaviour of this stack, including private-repo access through the user's HF token; anything else, or unset, disables it. The default lives in one constant,NVFP4_DIFFUSION_DEFAULTinstudio/backend/core/inference/diffusion_nvfp4_flag.py, so turning it on for everyone later is a one-line change. The switch is read at call time throughnvfp4_diffusion_enabled().Added on this branch, with the switch off:
ensure_flashinfer_for_nvfp4returns before importing flashinfer, taking the install lock or running any subprocess, and records no install reason.With the switch on, the install runs exactly as before.