Studio: NVFP4 flashinfer backend kernel items (device guard, persistent barrier, bias path, cached dispatch) - #10731
danielhanchen wants to merge 56 commits into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ 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. |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7e87ca6da5
ℹ️ 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".
| from .diffusion_nvfp4_linear import reset_nvfp4_state | ||
| reset_nvfp4_state() |
There was a problem hiding this comment.
This unconditional process-global reset is unsafe when the sibling image backend still has an NVFP4 CUDA graph resident—for example, a GPU image model can coexist with a CPU video pipeline because the GPU arbiter only excludes competing GPU owners. Unloading that CPU video state clears the shared barrier tensor and then calls clear_gpu_cache(), while the image graph still contains a recorded zero_ kernel targeting the barrier's address; a later replay can therefore access freed or reused memory. Scope these caches to the owning model/backend, or reset them only after every graph that references them has been uninstalled.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Not reachable. Ownership in the GPU arbiter is exclusive and a video load evicts the image backend; both media routes gate on the host-wide device target, so on any host where a flashinfer NVFP4 graph can exist the video target is never cpu, and a cpu target means no image CUDA graph exists. Within each backend uninstall_all runs before reset_barriers, and the video reset sits inside the committed-state branch.
|
@codex security review |
🛡️ Codex Security ReviewSecurity review completed. No security issues were found in this pull request. Reviewed commit: Only the user who started this review can view the report in Codex. ℹ️ About Codex security reviews in GitHubThis is an experimental Codex feature. Security reviews are triggered when:
Once complete, Codex will leave suggestions, or a comment if no findings are found. |
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
Checked the device guard coverage in studio/backend/core/inference/diffusion_nvfp4_ops.py and the persistent barrier and cached dispatch paths, and the bit-identity numbers line up with the switches described. Will review once #10730 lands and this retargets to main. |
|
Codex Review: Didn't find any major issues. Keep it up! 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". |
|
Codex Review: Didn't find any major issues. Nice work! 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". |
for more information, see https://pre-commit.ci
|
Codex Review: Didn't find any major issues. Chef's kiss. 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". |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep them coming! 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". |
for more information, see https://pre-commit.ci
…loth into r10731 # Conflicts: # studio/backend/core/inference/diffusion_nvfp4_ops.py # studio/backend/tests/test_nvfp4_diffusion_flag.py
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
|
@codex review |
1 similar comment
|
@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
2 participants
Summary
Kernel-side items on the flashinfer NVFP4 backend introduced by the per-layer image policy PR #10730. Stacked on
studio-nvfp4-image; retarget tomainonce that one merges. Four items, each behind its own switch, none changing what a render looks like:Device guard audit (
diffusion_nvfp4_ops.py,diffusion_nvfp4_linear.py). Every flashinfer entry point sits insidewith torch.cuda.device(t.device): the quantiser body, the GEMM body spanning barrier and GEMM, the plan builder includingAutoTuner.choose_oneand the cutlass module and cache-buffer getters, the preflight, and the warm-up entry. Skipped inside the traced layer forward underis_compiling()because inductor's wrapper already opens one and a live context manager there is a graph break. A source test banstorch.cuda.set_streamin the subtree (it silently sets the current device; that is how three cards were lost on this host). A CUDA test loads the layer on card 1 with card 0 current.Persistent PDL barrier. flashinfer launches the CUTLASS FP4 GEMM with programmatic dependent launch while the kernel's
griddepcontrolwait is compiled out, so the GEMM can start before the quantiser has finished writing (torch.emptyoutput: 50 of 50 non-finite in the harness; a one-element fill: 0 of 50). The image PR paid atorch.zeros(1)allocation per GEMM. This PR keeps one process-global one-element buffer per device and zeroes it inside themmop body before the GEMM; never allocated during capture (falls back to uncached zeros), reset with the graph pool on unload,UNSLOTH_NVFP4_ZERO_BUFFER=1keeps the full memset for comparison. Not an op argument: a mutable input would go throughauto_functionalizedand clone. A comment above the buffer names what to delete when flashinfer ships the griddepcontrol fix.Eager bias fast path (
diffusion_nvfp4_bias.py).mm_fp4has no bias argument. The fused CUTLASS epilogue from the investigation is dropped: not bit-identical to the unfused path (max abs 2 to 4), sm_100a only, a 104 s nvcc build at runtime, and inductor already fuses the bias into the compiled path where a custom op would block that fusion. What ships is a Triton in-place add on the eager path only, bit-identical toTensor.add_on seven verified shapes, declined underis_compiling(), on non-contiguous input, on a dtype mismatch, without Triton, below 12 M output elements and above 2**31. The floor is measured: the kernel's launch is 20 to 28 us against 5 to 10 us foradd_, so it runs 0.26x to 0.37x below 4 M elements and 0.84x at 1024x10240, then 1.84x at 4096x3840 and 3.4x to 3.8x from 4096x10240 up; flux and qwen-image route only M = 1 GEMMs to nvfp4, where it was pure launch cost.UNSLOTH_NVFP4_FAST_BIAS=0disables it.Cached flashinfer dispatch state (
diffusion_nvfp4_dispatch.py). flashinfer's publicmm_fp4rebuilds its plan and re-transposes the weight on every call: 58 us of host per GEMM. All private-internal touching lives in one file behind an exact version allowlist ("0.6.6") plus a one-shot runtime bit-identity self-check (verify(), run inside the preflight, off the request path).gemm_planreturns None when capturing with a cold key so no profiling launch is baked into a graph; the transpose cache is keyed on(data_ptr, shape), weights only, bounded at 4096, reset on unload.UNSLOTH_NVFP4_FAST_DISPATCH=auto|0|1(1 skips only the version check, neververify). Any other flashinfer release silently uses the public path.Measurements
B200, torch 2.12.1+cu130, flashinfer 0.6.6. The PR 2 arm is the PR 2 worktree, not an env setting: the barrier has no switch back to the per-call allocation, so arms are whole trees and a "knobs off" arm on this tree isolates the barrier.
Per-kernel (commits 2 to 4): persistent barrier saves 1.7 us host per GEMM; cached dispatch 58 to 19 us host per GEMM,
torch.equalon every checked shape; Triton bias 1.6x to 3.7x overadd_, bit-identical.End to end, z-image, 9 steps, 3 warm-ups + 20 timed renders, median of 3 processes per arm:
This PR over the PR 2 tree: graphed 1.016x and 1.011x, un-graphed 1.128x at 512 and 0.992x at 1024. The un-graphed 512 gain is entirely the cached dispatch (the barrier arm shows no end-to-end effect and the compiled path declines the Triton bias). The 1.45x to 1.55x the investigation prototype measured does not reproduce and should not: that prototype had every admitted linear (239) on the flashinfer path, the shipped z-image policy puts 34 there. Wan2.2-TI2V-5B at 1280x704x121, 50 steps, 304 flashinfer linears: 48.15 s to 47.72 s (1.009x), the same with graphs requested or forced off since Studio does not capture a graph for that family.
Bit identity against the PR 2 tree: with deterministic algorithms pinned, every z-image pair is exactly 0 in all four cells. As shipped, the un-graphed compiled path is not reproducible against itself (two runs of the same tree and seed differ by up to 136 levels on 72 percent of pixels when one compiled cold and the other loaded a compile-cache bundle), every cross-tree pair lands inside that spread and several are exactly 0; Wan pairs differ by at most 2 levels on under 0.5 percent of pixels, within the same-arm range. No knob in this PR breaks identity. The compile-cache drift is pre-existing and independent of NVFP4, tracked separately.
Capture inertness: barrier and dispatch workspace pointers unchanged across a
GraphedForwardcapture and 20 replays, one barrier per device, no plan built during a cold capture, replays bit-identical to eager.verify()costs 0.9 ms cold inside a 110 to 130 ms preflight inside a 20 to 35 s load.Tests
test_diffusion_nvfp4_speed.py,test_diffusion_nvfp4_dispatch.py: barrier allocated once per device, not cached during capture,zero_precedesmm_fp4in the stub call log,ZERO_BUFFERenv; every flashinfer call sees the tensor's device with tensors on device 1 and device 0 current; noset_streamin the subtree; dispatch refuses an unknown version, a missing symbol, env off; no plan during capture; transpose cache bounded; fused bias falls back on each condition;reset()clears state. CUDA-gated: barrier bit identity over 50 iterations, fast vs public dispatchtorch.equalat five shapes, fused bias identity on seven shapes, multi-GPU guard, capture and replay identity with a stable barrier pointer.Full
studio/backend/testson this branch (test_blender_managed.pyexcluded), 828 files in 10 batches: 41593 passed, 124 failed, 262 skipped, 15 errors. Every failing id reproduces on the merge base (95feb6979) in the same 19 files none of which this PR touches (see the image PR for the list).hub/tests: 683 passed.Not measured
FLUX.1 and Qwen-Image end to end (the eager path where the Triton bias is live is not a shipped configuration for either); other flashinfer releases (the allowlist makes them the public path by construction); the 1.6 percent graphed z-image gap, three runs per arm do not resolve it and no claim is made for it.
UI evidence
No before/after pair is attached, because nothing a user sees changes, and a pair of identical screens is not evidence of that. What was checked against this PR's own base (
studio-nvfp4-image):studio/frontendis empty,studio/backend/main.pyandstudio/backend/models/are untouched, and no response payload gains or loses a field.reset_nvfp4_state()calls are teardown on paths that already tore down. The rest is kernel internals behindtorch.ops.unsloth_nvfp4, plus four environment switches listed indiffusion_speed.py's module docstring, all unset by default.transformer_quant_backendon/api/inference/images/statusfrom Studio: per-layer NVFP4 image policies, flashinfer FP4 backend and a gated auto row #10730, cannot move here.nvfp4_preflightsetsokfrom the GEMM probe BEFORE the new fast-dispatch check runs, and that check sits inside the sametrywhoseexceptrewrites onlyreason. Adispatch.verifythat fails or raises changes which dispatch path runs, never which backend is reported, so the loaded-models row reads the same on both sides by construction.The claim this PR does make is speed at bit identity, and that is the table above: a picture cannot show that two renders are equal to the last level.
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().This PR adds no new entry point. The fast dispatch is only reached through
nvfp4_preflightand the NVFP4 layers, both behind the switch. A test here checks that neither the dispatch probe nor its verify pass runs with the switch off. The gates themselves are on #10729 and #10730.