Diffusion: auto TE/VAE precision, video DiT quant correctness, and HunyuanVideo-1.5 attention speedup#7021
Diffusion: auto TE/VAE precision, video DiT quant correctness, and HunyuanVideo-1.5 attention speedup#7021danielhanchen wants to merge 55 commits into
Conversation
The LoRA row div was keyed on sel.id, the editable repo-id field. Typing the first character changed the key from the fallback index to that character, which remounted the row and dropped focus on the input, so only one character was ever captured and manual entry of a repo id was impossible (paste still worked). The list is index-addressed (every mutation uses j === i) and rows are removed explicitly via the delete button, so the index is the stable key, matching the reference-image rows above.
The auto ladder still listed nvfp4 in the Blackwell tier even though the comment above it says nvfp4 must never be the auto pick for diffusion: at the DiT's real shapes it is slower (0.81x end-to-end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044) than fp8. It was unreachable in practice (fp8 precedes it and the only fp8-denied families also deny nvfp4), but the entry contradicted the stated intent and left a latent path to the worse scheme. Drop nvfp4 from the ladder so Blackwell auto is fp8 -> mxfp8 -> int8, keeping the previous line commented with a TODO to restore it once the FP4 GEMM wins at these shapes. nvfp4 stays fully available as an explicit transformer_quant="nvfp4" request.
The transformer already defaults to auto-quant (fp8/int8); the companion text encoder was opt-in and stayed dense bf16 unless a scheme was named, even though it is often the largest resident component. Add an auto policy mirroring the transformer's ladder: select_te_quant_scheme walks a per-capability ladder (data-center fp8-GEMM: fp8_dynamic -> int8 -> layerwise fp8; Ampere: int8 -> fp8), reorders int8 first on consumer GDDR parts, falls to layerwise fp8 under group offload (the only offload-safe cast), only picks int8 for a family with a measured keep-bf16 schedule, honors a per-family deny list, and smoke-probes the torchao kernel so a missing build degrades gracefully. The image + video loaders now map an unset text_encoder_quant to auto (explicit none/off stays dense; a named scheme is forced), so the shipped default quantizes the encoder to the fastest accurate scheme for the GPU. Records the engaged scheme in the image resolved-record too. Verified on a B200: auto -> fp8_dynamic, offload -> layerwise fp8.
for more information, see https://pre-commit.ci
The transformer and text encoder auto-quantize; the VAE stayed dense. VAEs are convolutional, so torchao int8 (Linear/2D-only) does not apply, but Float8DynamicActivationFloat8WeightConfig quantizes Conv2d/Conv3d weights with PerTensor granularity (auto-skipping convs whose channels are not a multiple of 16, so the 3-channel RGB head stays dense). New diffusion_vae_quant.py offers two schemes: fp8_dynamic (torchao conv compute fp8, cc>=8.9, resident) and fp8 (diffusers layerwise storage cast, any conv, survives offload); no int8 (no Conv3d int8 kernel). select_vae_quant_scheme walks (fp8_dynamic, fp8) with a live conv smoke probe, an offload gate, a per-family deny list, and a force_fp32 gate; the image + video loaders map unset vae_quant to auto, skip the vae_force_fp32 Wan families, and record the engaged scheme. Guards _align_vae_dtype to skip the img2img/inpaint re-cast when the VAE is quantized (its fp8 tensor subclasses reject .to(dtype=)). Verified on a B200: %16 Conv2d/Conv3d/Linear -> Float8Tensor, conv_out dense, forward runs.
for more information, see https://pre-commit.ci
…_dynamic opt-in)
A B200 decoded-image LPIPS/SSIM sweep vs the dense bf16 VAE (new
scripts/quant_accuracy_sweep.py) settles the two VAE schemes:
- Layerwise fp8 (storage-only) holds across families (SSIM >= 0.977 on all but
SDXL), so auto now engages layerwise fp8 ONLY. For a VAE decode (a few percent
of end to end) fp8_dynamic's fp8-matmul speedup over storage fp8 is negligible,
so auto never takes the accuracy risk.
- fp8_dynamic (torchao PerTensor conv compute) is in-bar on only FLUX.2 and
Hunyuan and out-of-bar or catastrophic elsewhere (Qwen-Image SSIM 0.46), so it
is now an explicit opt-in, re-gated by a per-family deny list derived from the
sweep. SDXL denies both schemes (its small VAE stays dense).
Also fixes a real decode-time crash: torchao 0.17's fp8 conv kernel rejects
pointwise (1x1 / 1x1x1) convs ("Activation and filter channels must match"), so
an explicit fp8_dynamic request cast fine then threw at the first decode on most
families. The conv filter now keeps 1x1 convs dense, the smoke probe uses a
spatial 3x3 conv (so it exercises the path that actually runs), and an explicit
fp8_dynamic request runs that probe before casting.
Tests updated for the fp8-only auto ladder, the 1x1 exclusion, the explicit
probe gate, and the shipped deny list.
Five diffusion source and test files landed without the standard two-line SPDX header the rest of studio/backend carries. Prepend it (matching the sibling convention) so the whole backend is uniformly licensed. Header-only, no code change.
… image VAEs A B200 speed/memory sweep (new scripts/quant_speedmem_bench.py) shows the VAE quant win is a video story. Image AutoencoderKLs are ~0.15-0.26 GB, so fp8 saves ~0.1 GB and only slows their tiny decode (+6-16%); the video Conv3d VAEs are ~2.5 GB and halve to ~1.2 GB at ~2% decode cost. So VAE auto now only engages above a ~1 GB size floor: small image VAEs stay dense (faster decode, no quant quality risk), video VAEs still quantize. An explicit fp8 / fp8_dynamic request skips the gate (opted in). The same sweep confirmed the text-encoder default is already right: fp8_dynamic is E2E-neutral (denoise per-step unchanged; +2% one-time encode) and, by hidden-state cosine vs bf16, marginally more accurate than layerwise fp8 -- so that default is left as is. Tests cover the gate (small skipped, large quantized, explicit bypasses).
The dense video default engages transformer auto-quant, and on Blackwell the auto ladder leads with fp8. On the Wan DiT the production per-row fp8 path (torch._scaled_mm) renders every frame black (mean luma 0.0 at 512x320 and 704x480, LPIPS ~0.80 vs bf16): Wan's activation outliers exceed per-row fp8's range, the same failure already denied for qwen-image. First-Block-Cache then over-caches the degenerate activations (per-step collapses to ~10ms), compounding it. Add the Wan families (wan2.2-ti2v-5b, wan2.2-t2v-a14b, same WanTransformer3DModel) to _FAMILY_SCHEME_DENY for fp8/mxfp8/nvfp4 so auto falls through to int8, which is clean on Wan (per-token, outlier-robust), saves the same weight memory on the DiT, and lets First-Block-Cache engage normally instead of over-caching. mxfp8/nvfp4 are denied alongside fp8 conservatively so auto lands on the battle-tested int8; they can be re-enabled per family once validated in-bar, like the nvfp4 auto-ladder TODO. Validated on B200: the shipped video default now selects int8 for the Wan DiT and renders clean frames (mean 172.6) at 15.6 GB resident (down from 24.2 GB dense), with First-Block-Cache engaged. Adds two deny tests; 48/48 transformer-quant tests pass.
video_speedmem_bench.py drives the real video-loader lever functions (quantize_transformer / quantize_text_encoders / quantize_vae / apply_speed_optims / apply_attention_backend / apply_step_cache) with the loader's own defaults, so each measured config reflects a real load. It decomposes the video speed/memory stack (compile, cuDNN attention, First-Block-Cache, DiT/TE/VAE quant) with per-step latency, peak resident GB, and per-frame LPIPS vs a bit-exact reference. This is the harness that surfaced and validated the Wan fp8 black-frame fix. quant_speedmem_bench.py gains the DiT-quant mode (dense vs fp8/int8/mxfp8 speed, peak memory, and LPIPS vs the dense render) plus the shared LPIPS(AlexNet) helper.
Measured the fp8 DiT auto-quant path across the remaining dense-pipeline video families on
B200 (production torch._scaled_mm per-row fp8, no MSLK):
- HunyuanVideo-1.5 (480p + 720p repacks): every frame black (mean luma 0.0, LPIPS 0.82);
int8 is clean (mean 102.7 vs dense 99.9). Same failure as Wan / qwen-image.
- LTX-2: fp8 renders clean (mean 153.7, matches int8's 157.7) -- NOT a black-frame family.
So deny fp8/mxfp8/nvfp4 for hunyuanvideo-1.5 and hunyuanvideo-1.5-720p (fall to int8), and
deliberately leave LTX-2 on fp8. The deny stays measured per family, not a blanket video rule:
a blanket deny would have wrongly forced LTX-2 off fp8. Adds a Hunyuan deny test that also
asserts LTX-2 keeps fp8; 49/49 transformer-quant tests pass.
video_speedmem_bench.py gains guidance_via_guider support (HunyuanVideo-1.5 sets CFG on a
guider component and its __call__ takes no guidance_scale / callback_on_step_end), so the
harness can drive Hunyuan the same way the loader does.
The Wan fp8 black frame was root-caused (scripts/fp8_layer_ablation.py, measured on B200 with the production torch._scaled_mm path): per-row fp8 scales each activation row by row_amax/448, and the text prompt is padded to 512 tokens (~all padding for a short prompt), so condition_embedder's text embedder divides a zero padding row by a zero scale, which infs and renders every frame black. That embedder's bias makes every downstream row non-zero, so the whole 30-block attn1/attn2/ffn stack is fp8-clean (fp8-except- condition_embedder measured cosine 0.9998 vs bf16, 0 non-finite; fp8- everywhere is 100% non-finite). So the blanket fp8 deny was heavier than needed for Wan. Remove fp8 from the Wan deny and keep only condition_embedder in bf16 via a new _FP8_FAMILY_EXCLUDE_NAME_TOKENS; auto now restores fp8 (the Blackwell ladder head) for Wan2.2-TI2V-5B and -T2V-A14B (shared DiT class and padded-text conditioning). Full-generation check (512x320, 25 frames, 30 steps, cache on and off): mixed-fp8 is non-black (mean luma 182.6 vs dense 181.2), more accurate than int8 (LPIPS 0.129 vs 0.180 no-cache, 0.224 vs 0.251 with FBCache), faster (49.9 vs 64.6 ms/step; int8 was a per-step regression vs the 59.8 ms/step dense), at the same memory (19.34 GB, both -20% vs dense). HunyuanVideo-1.5 keeps the fp8 deny: its MMDiT masks the padding text tokens to zero inside every block, so the per-block context stream (add_*_proj / to_add_out / ff_context) regenerates zero rows layer after layer (fp8 on only the main blocks is 100% non-finite) so no small exclude set exists and int8 stays. mxfp8 / nvfp4 remain denied for Wan (same per-row scaled_mm family, not separately validated). exclude_tokens_for_scheme now takes an optional family, threaded through the runtime quantiser and the offline prequant builder + validator so offline == runtime (a stale Wan fp8 checkpoint baked without the exclude is rejected and re-quantised rather than loaded). Adds scripts/fp8_layer_ablation.py (the per-layer ablation probe) and a mean-luma black-frame metric plus mixed-fp8 vs int8 configs to the video bench.
…ilies Adds wan2.2-t2v-a14b and hunyuanvideo-1.5-720p to the bench family table and makes _apply_levers quantize / compile BOTH experts of a dual-expert MoE (Wan2.2-A14B) via a _SecondExpertView proxy that mirrors the loader's _SecondDiTView, so A14B latency and accuracy are measured on the real two-DiT path instead of only the first expert. Used to validate that every video family is on the fastest DiT quant scheme that is not less accurate than its alternative (B200, 512x320, 25 frames, 30 steps, no cache, LPIPS vs dense bf16): - Wan2.2-TI2V-5B fp8 49.9 ms/step vs int8 64.6 vs dense 59.8; LPIPS 0.129 vs 0.180 - Wan2.2-T2V-A14B fp8 195.8 ms/step vs int8 193.5 vs dense 369.1; LPIPS 0.288 vs 0.349 (both experts) - LTX-2 / LTX-2.3 fp8 133.5 ms/step vs int8 138.4 vs dense 204.7; LPIPS 0.026 vs 0.027 - HunyuanVideo-1.5 fp8 is black (denied, not localizable); int8 21.2 s vs dense+compile 15.4 s -- a memory saving at a speed cost, so int8 stays only because fp8 is impossible there. fp8 is faster than dense AND more accurate than int8 on all three Wan/LTX families (the two Wan ones via the condition_embedder exclude; LTX-2 needs none -- its conditioning has no zero-amax padding rows). The auto-ladder + deny + exclude already select exactly these schemes, so no scheme-selection change was needed; this commit is the bench faithfulness improvement that let the campaign confirm it.
…lback Root-caused "HunyuanVideo-1.5 int8 is slower than dense" with a per-forward profiler (scripts/hunyuan_int8_profile.py, dynamo-reset, back-to-back on a clean B200): int8 compiles cleanly (0 recompiles, 0 graph breaks, steady 268.3 ms/forward) and is only ~7% slower than dense + regional compile (250.5 ms/forward), not the 38% a contended-GPU bench run suggested. int8 is also less accurate (LPIPS 0.085 vs dense+compile 0.037). So for a family where fp8 is denied (Hunyuan black-frames on per-row fp8), int8 is a MEMORY lever, not a speed win, yet the auto-quant default quantised it even when the dense DiT already fit resident. Fix: is_int8_memory_fallback(target, family) is True only when AUTO quant lands on int8 as a denied/black-frame fallback on a data-center, fp8-capable GPU (fp8 would be the arch pick but is denied for the family). The video loader now skips the auto-quant and runs dense+compile when that holds AND the bf16 memory plan already fits resident (offload_policy == none), so there is no new OOM risk. Scoped tightly: only an AUTO request (explicit int8/fp8 honored), only int8-fallback families (Wan / LTX resolve to fp8 -> keep quantising), only data-center fp8-capable parts (consumer GPUs and pre-Ada, where int8 is a genuine accelerator, keep int8), and only when dense provably fits; a memory-constrained plan still quantises. Result: Hunyuan on a resident-fit B200 now runs faster AND more accurate, quantising only when memory is the constraint. Also resets dynamo per config in the video bench (so compiled graphs cannot leak across configs in one process) and adds the per-forward profiler used for the diagnosis.
HunyuanVideo-1.5's DiT runs a joint [video; text] self-attention and, on every block and step, builds a dense [B,1,N,N] boolean mask so the video never attends to the padded text. A dense bool attn_mask disables every fused SDPA kernel (flash rejects it; cuDNN and memory-efficient fall back), so the attention runs the slow math-style path: at the production shape (121 frames, 480p, N about 50k) one attention call is ~421ms with the mask vs ~19ms with attn_mask=None. The text is ~99.5% padding (a t2v prompt fills ~9 of ~1985 slots), so nearly all of that cost is spent masking padding. install_hunyuan_attention_trim installs an eager forward pre-hook that drops the all-zero image stream (t2v) and trims the mllm/byt5 text streams to their globally-valid columns, plus a null-mask attention processor that runs attn_mask=None once no partially-padded column remains (the batch-1 / per-guidance-branch case) and otherwise delegates to the stock dense-mask processor. The model already zeroes and masks the padded text and discards its attention output (only the video split feeds proj_out), so removing it is exact for the video; the only numeric change is the SDPA kernel (masked fallback to fused). Measured on a B200: 23.3s to 1.3s per DiT forward at 121 frames (~18x with regional compile, 0 graph breaks); per-forward cosine 0.99998 vs stock; equal distance to an fp32 reference (LPIPS fp32-vs-stock 0.292, fp32-vs-trim 0.307), so it is not less accurate than the current bf16 default. Wired auto-on for HunyuanVideo-1.5 in the video loader, before the attention backend set so the requested kernel pins onto the new processors; a no-op for every other family and reversible (stock dense-mask path on any anomaly). Adds hermetic tests and the diagnostic/validation scripts.
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
for more information, see https://pre-commit.ci
![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: a5928064a0
ℹ️ 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".
|
|
||
| # ── HunyuanVideo-1.5 padded-text attention trim ───────────────────────────────────── | ||
| # _trim_stream / _hunyuan_trim_pre_hook use real torch tensor ops, so these run on CPU torch. | ||
| import torch # noqa: E402 |
There was a problem hiding this comment.
In environments that run these hermetic attention tests without torch installed (the module docstring explicitly says no torch/diffusers are needed), this top-level import aborts collection before pytest can run or skip anything, so all of the existing attention-backend tests fail to collect. Please use pytest.importorskip("torch") for just the new Hunyuan trim tests or move the import behind a skip fixture so the non-torch tests remain runnable.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. Moved the 15 torch-dependent trim tests into a new test_diffusion_attention_trim.py (which declares the torch dependency), so test_diffusion_attention.py stays collectable and runnable without torch as its docstring promises. Fixed in a4a03fe.
danielhanchen
force-pushed
the
image-generation
branch
from
42a4c07 to
e9db36a
Compare
danielhanchen
changed the title
…on tests torch-free) test_diffusion_attention.py documents itself as hermetic with no torch/diffusers needed, but the HunyuanVideo trim tests added a module-level 'import torch' that aborted collection of the whole file when torch is absent. Move those 15 tests to test_diffusion_attention_trim.py (which declares the torch dependency) so the attention-backend policy tests stay collectable and runnable without torch.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a4a03feca8
ℹ️ 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".
| memory_mode = request.memory_mode, | ||
| speed_mode = request.speed_mode, | ||
| text_encoder_quant = request.text_encoder_quant, | ||
| vae_quant = request.vae_quant, |
There was a problem hiding this comment.
When the image load is routed to the native sd.cpp engine (CPU/no usable GPU, MPS-enabled native, or a forced native engine), this new keyword is still passed to engine.begin_load, but SdCppDiffusionBackend.begin_load does not accept vae_quant (it only has the other diffusers-only knobs). That raises TypeError: unexpected keyword argument 'vae_quant' before any native GGUF load can start, so native image loading is broken even when the request leaves vae_quant unset. Please add the ignored parameter to the native backend signature, as was done for text_encoder_quant.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 4fd93b2: the native SdCppDiffusionBackend.begin_load now accepts (and ignores) vae_quant like the other diffusers-only knobs. The image route passes vae_quant=request.vae_quant to both engines uniformly, so without this a native (CPU-only / MPS / forced-native) image GGUF load raised TypeError on every request. Added test_begin_load_accepts_diffusers_only_vae_quant.
The image load route calls engine.begin_load(..., vae_quant=request.vae_quant, ...) uniformly for both engines, but the native SdCppDiffusionBackend.begin_load accepted every other diffusers-only knob except vae_quant and had no **kwargs, so a native (CPU-only / MPS / forced-native) image GGUF load raised TypeError on every request (vae_quant is always passed, defaulting to None). Accept and ignore it like the other diffusers-only knobs; sd.cpp has no torchao VAE quant.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4fd93b2c7a
ℹ️ 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".
| text_encoder_quant = TE_QUANT_AUTO | ||
| # vae_quant tri-state, same contract: UNSET -> auto (fp8_dynamic conv compute on resident | ||
| # fp8-GEMM silicon that passes the conv probe, else layerwise fp8, else dense); none/off -> | ||
| # dense; an explicit scheme forces it. | ||
| if vae_quant is None or str(vae_quant).strip() == "": |
There was a problem hiding this comment.
When an image load sets speed_mode="off" (the UI labels this bit-exact) and omits the companion quant fields, this still rewrites the unset text encoder request to auto and the following VAE default does the same. On CUDA/bf16 systems where the auto selector succeeds, quantize_text_encoders/quantize_vae later mutate those components even though the transformer path has a separate speed_off guard, so an explicit bit-exact/off load can still produce quantized, non-dense outputs. Please keep companion quant disabled under explicit speed off unless the user explicitly requested a scheme.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in be04ba0. An explicit Speed=off now pins the companions dense: the unset TE/VAE default resolves to off (not auto) when speed_mode is explicitly off, mirroring the transformer_quant suppression right below it. An unset speed still auto-quantises, and an explicit companion scheme still forces it. Test: test_speed_off_load_suppresses_auto_companion_quant.
| text_encoder_quant = TE_QUANT_AUTO | ||
| # vae_quant tri-state, same contract. The vae_force_fp32 families (Wan) keep the VAE dense | ||
| # regardless (quantize_vae's force_fp32 gate), so auto is safe as the shipped default. | ||
| if vae_quant is None or str(vae_quant).strip() == "": | ||
| vae_quant = VAE_QUANT_AUTO |
There was a problem hiding this comment.
For video loads with speed_mode="off" and no explicit companion quant request, these defaults still turn both text-encoder and VAE quantization on via auto. The transformer auto default is suppressed for speed-off later in this loader, but the companion defaults are not, so Hunyuan/LTX-style loads can still quantize resident encoders/VAEs while the user asked for the dense bit-exact/off path. Please gate these auto defaults the same way, while still honoring explicit text_encoder_quant/vae_quant requests.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in be04ba0, same as the image backend: speed_off (already computed here for the DiT) now also gates the unset TE/VAE default to off. Test: test_video_speed_off_suppresses_auto_companion_quant.
| # attention so it runs the fused (cuDNN/flash) SDPA kernel instead of the dense-mask | ||
| # fallback (~18x/DiT-forward at 121 frames, cosine ~1.0). Must precede the backend set | ||
| # so the requested kernel pins onto the new processors. No-op for every other family. | ||
| trim = install_hunyuan_attention_trim(view, fam, logger = logger) |
There was a problem hiding this comment.
For HunyuanVideo-1.5 loads with speed_mode="off", this still installs the padded-token trim unconditionally before attention selection. The comment here notes that the trim makes attention run the fused SDPA path instead of the dense-mask fallback, so an explicit Off/bit-exact load still gets a different attention kernel and rounding behavior (and reports hunyuan_attn_trim in speed_optims). Please skip this optimization when effective_speed == SPEED_OFF unless the user has explicitly opted into a speed tier.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in be04ba0. The trim is a speed lever (it swaps the joint attention onto the fused SDPA kernel), so it is now gated on effective_speed != off exactly like the attention-backend selection right below it -- the bit-exact off path keeps the stock dense-mask attention. Test: test_video_speed_off_skips_hunyuan_trim.
Two fidelity fixes to the video speed/mem bench so its numbers match production:
- _build_pipe loaded the pipeline with a scalar bf16 torch_dtype and then upcast the
VAE, which truncates the fp32-stored Wan VAE at load (a later .to(float32) only widens
the lossy values). Pin the VAE fp32 per-component like the production loader
({"vae": fp32, "default": bf16}) so the bench decodes the same weights production does.
- The persisted reference frames were written/read as a single unkeyed ref_frames.npz in
the fixed default --out dir, so a reference-less run of a different family/seed/steps/
frames/resolution scored LPIPS against a stale baseline. Key the cache by those
parameters so a run only reuses a reference computed for the same parameters.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
unsloth/studio/backend/core/inference/diffusion_precision.py
Lines 254 to 255 in 5eb64bd
Probe explicit torchao TE kernels before reporting quantized
When a user explicitly requests a torchao text-encoder mode such as fp8_dynamic, int8, or nvfp4, this path only checks the CUDA capability and skips the _te_scheme_probe smoke test that auto uses. On a torchao/PyTorch build where quantize_ can wrap the encoder but the actual GEMM kernel fails, the load will report text_encoder_quant as engaged and then crash on the first prompt-encoder forward instead of falling back to dense; please run the same probe for explicit torchao TE modes before selecting the caster.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # Attention (per expert). | ||
| backend = select_attention_backend(tgt, cfg["attn"], speed_active = speed_active) | ||
| for v in views: | ||
| engaged["attn"] = apply_attention_backend(v, backend, logger = logger) |
There was a problem hiding this comment.
For hunyuanvideo-1.5 benchmark runs, this helper applies the attention backend without first installing install_hunyuan_attention_trim, while the production video loader installs that trim before backend selection. As a result shipped/speedmax Hunyuan rows still benchmark the dense-mask attention path and miss the PR's main fused-attention speedup, producing misleading latency and LPIPS comparisons for the optimized family; add the trim install in this per-expert section before applying the backend.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in be04ba0. The bench now installs the Hunyuan trim per expert BEFORE the backend set (the loader's order), gated on an active speed tier, so the shipped/attention rows measure the fused-kernel path the loader runs.
| # DiT quant (pipeline kind, resident): mutates each expert's transformer in place. | ||
| if cfg["dit"] not in ("none", "off"): | ||
| schemes = [ | ||
| quantize_transformer(v, tgt, mode = cfg["dit"], family = fam_name, logger = logger) | ||
| for v in views |
There was a problem hiding this comment.
For HunyuanVideo-1.5 on a data-center fp8-capable GPU, the production loader skips an AUTO DiT quantization when the dense bf16 DiT already fits resident, because the int8 fallback is slower and less accurate there. This benchmark has already loaded the dense pipeline resident, but still unconditionally calls quantize_transformer(..., mode="auto"), so shipped/ditquant rows measure an int8 path that production would not run and the speed/LPIPS numbers are misleading; apply the same dense-fit is_int8_memory_fallback skip here.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in be04ba0. The bench now mirrors the loader's dense-fit skip: for an AUTO dit config on an int8-fallback family (is_int8_memory_fallback -> HunyuanVideo, where fp8 is denied so auto lands on int8, a memory-only lever), it runs the dense DiT instead of int8, since the bench always loads resident. Explicit int8/fp8 configs are still honored (the point of the sweep).
| except Exception: # noqa: BLE001 — optimisation only; never break the forward | ||
| for blk in getattr(module, "transformer_blocks", []): | ||
| attn = getattr(blk, "attn", None) | ||
| if attn is not None: | ||
| setattr(attn, _NULL_ATTN_FLAG, False) | ||
| return args, kwargs |
There was a problem hiding this comment.
If this best-effort hook hits an unexpected error after it has already mutated kwargs (for example after emptying all-zero image_embeds and then encountering a changed/malformed text mask), the except path only clears the null-mask flag and returns the partially trimmed inputs. That no longer falls back to the stock path promised by the comment and can leave the model running with image/text lengths that no longer match the dense attention mask; restore original in this exception path before returning.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in be04ba0. The fallback now restores the caller's original kwargs (kwargs.clear(); kwargs.update(original)) before taking the stock dense-mask path, so a failure mid-trim (after image_embeds was emptied / a stream trimmed) never hands the dense path a half-trimmed mix -- same restore the empty-prompt guard already does. Test: test_trim_pre_hook_restores_inputs_on_midtrim_failure.
…plicit TE kernels, bench fidelity Address the Codex review round on the video/quant work: - Companion auto-quant now honors an explicit Speed=off. Both loaders already pin the DiT dense under an explicit off (bit-exact reference), but the unset text-encoder / VAE quant still promoted to auto and silently fp8/int8'd the companions, breaking the bit-exact request. An UNSET speed still auto-quantises; an explicit companion scheme still forces it. - The HunyuanVideo joint-attention trim is a speed lever (it swaps to the fused SDPA kernel), so gate it on a non-off speed tier exactly like the adjacent attention-backend selection -- the off path keeps the stock dense-mask attention. - Explicit torchao text-encoder modes (int8 / fp8_dynamic / nvfp4) now run the same kernel smoke test the auto ladder uses. They could clear the capability gate yet fail the real GEMM on a build where quantize_ wraps the encoder but the kernel is broken; the caster's try/except only covers the cast, not the first forward, so the load would report engaged then crash at generation. Now it falls back to dense. Layerwise fp8 has no torchao GEMM, so the probe is a no-op for it. - The trim pre-hook's fallback restores the caller's original kwargs (it may have emptied the image stream / trimmed a text stream before failing), so the stock dense-mask path runs on exactly what it expects, matching the empty-prompt guard. - video_speedmem_bench mirrors the loader: installs the Hunyuan trim before the backend set (gated on an active tier) and skips the auto int8 quant when it is the fp8-denied memory fallback and dense fits resident, so the shipped/auto rows measure what the loader actually runs. Tests: TE explicit-mode kernel probe (+ layerwise-fp8 bypass), trim mid-trim restore, and loader-level speed=off companion suppression + trim skip for both backends. 262 backend tests pass; ruff clean.
|
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". |
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: 0d378bc496
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| normalize_cache_quality, | ||
| ) | ||
|
|
||
| if cache_mode == "auto": |
There was a problem hiding this comment.
When a benchmark config explicitly uses cache="magcache" (for example magcache_explicit) and --steps differs from the family default, _apply_levers installs MagCache using default_steps, but this generation-time path only rechecks cache_mode == "auto". With the script's default --steps 30 against 50-step calibrated defaults, the row keeps the #s50 curve/skip budget while timing a 30-step clip, unlike production which re-engages explicit MagCache for the actual step count. Please mirror that explicit-MagCache reconfiguration here before timing so the benchmark doesn't publish latency/LPIPS for a cache state users won't run.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The bench now mirrors the loader's explicit-MagCache re-engagement: when the actual step count differs from the configured marker, each expert's cache is disengaged and re-applied over the real schedule before timing, so magcache rows no longer publish latency or LPIPS for a curve users cannot run. Fixed in bd1c2e0.
| "te_scheme": engaged["te"] or "dense", | ||
| "vae_scheme": engaged["vae"] or "dense", | ||
| "attn": engaged["attn"] or "native", | ||
| "cache": engaged["cache"] or "off", |
There was a problem hiding this comment.
When cache="auto" is benchmarked with an actual --steps value that crosses the auto threshold (for example a Hunyuan/Wan default cache installed at 50 steps but a run with --steps 8), _timed_video toggles the transformer cache off/on for the real generation, but this row still records the load-time engaged["cache"]. That publishes timings/LPIPS under cache: magcache/fbcache even when the measured clip ran uncached (or vice versa); derive the field from the post-toggle marker/current transformer state, matching production status updates.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The row's cache field is now derived from the post-toggle transformer marker, so an auto cache toggled off below the step threshold (or a re-sized explicit magcache) is reported as it actually ran; the load-time engagement stays available as cache_at_load. Fixed in bd1c2e0.
| _empty() | ||
|
|
||
| # ── optimisation layers ── | ||
| snapshot_backend_flags() # process-wide flags; benchmark process is short-lived so no restore. |
There was a problem hiding this comment.
When --configs runs a speed-enabled row before a bit-exact row in the same process (for example speedmax,reference or epc_off,reference), this snapshot is discarded and the process-wide flags that apply_speed_optims mutates (cudnn.benchmark, TF32/fp16 accumulation, and emulate_precision_casts) are never restored. The later speed="off"/reference row can therefore inherit the previous row's backend settings, corrupting the baseline and LPIPS/speed comparisons; keep the snapshot and restore it after each config, as the production loader does on unload.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The snapshot is kept and restored after each config, like the production unload path, so a speedmax or epc_off row no longer leaks cudnn.benchmark, TF32/fp16 accumulation, or emulate_precision_casts into a later reference or off row in the same process. Fixed in bd1c2e0.
…he reporting, per-config flag restore Mirror production's explicit-MagCache re-engagement in the video bench: a magcache row installed at the family default step count now re-interpolates its curve, retention window, and skip budget when --steps differs, instead of timing a stale schedule users never run. Report the generation-time cache state: the row's cache field is derived from the post-toggle transformer marker rather than the load-time engagement, so an auto cache toggled off below the step threshold (or a re-sized explicit magcache) is published as it actually ran; the load-time value stays available as cache_at_load. Restore the process-wide backend flags (cudnn.benchmark, TF32 and fp16 accumulation, emulate_precision_casts) after each config, like the production unload path, so a speed-enabled row cannot poison a later reference or off row in the same --configs run.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 87054e7f94
ℹ️ 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".
| if mode == VAE_QUANT_FP8_DYNAMIC: | ||
| _cast_vae_fp8_dynamic(vae, target) |
There was a problem hiding this comment.
When an API caller explicitly requests vae_quant="fp8_dynamic" on a resident load, the image and video loaders call quantize_vae before apply_memory_plan, so the VAE is still the freshly loaded CPU module here. This path probes CUDA conv kernels, then casts the CPU VAE and relies on the later pipe.to(device) to move torchao FP8 tensor subclasses; the surrounding code already notes those subclasses reject Module.to(), so this can make the explicit fp8_dynamic load fail during placement instead of producing a usable resident VAE. Move the VAE to the target before this cast or run the dynamic cast after resident placement.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Could not reproduce the mechanism: on the installed torchao a single Module.to moves the fp8_dynamic Float8Tensor weights fine (verified CPU-quantize then to-cuda then forward), which is the resident path's one placement; the reject-Module.to note applies to offload hooks' repeated moves, and fp8_dynamic is already offload-denied. A build lacking the fp8 conv kernel is caught by the per-device, per-dimensionality probe before the cast.
| self._primary.enable_cache(config) | ||
| try: | ||
| self._replica.enable_cache(config) | ||
| _invalidate_registry(self._replica) |
There was a problem hiding this comment.
When CFG-parallel is installed and the step cache is re-engaged through the proxy, such as a MagCache generation with a different step count after a prior run, apply_step_cache calls this proxy method. It invalidates only the replica registry, while the later _invalidate_child_registry_cache(transformer) sees the proxy object rather than _primary, so the primary DiT can keep the stale child-registry list from the previous hooks and the new primary cache hooks may not receive a context. Invalidate _primary here as well after self._primary.enable_cache(config).
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The delegation goes the other way: the proxy's getattr forwards _diffusers_hook to the primary, so apply_step_cache's post-enable invalidation nulls the primary's child-registry cache through the proxy (verified by running that call against the real proxy), and the fan-out's explicit replica invalidation covers the one object delegation cannot reach. Both DiTs are invalidated on every enable.
There was a problem hiding this comment.
💡 Codex Review
unsloth/studio/backend/core/inference/diffusion.py
Lines 2905 to 2906 in 87054e7
Toggle step cache on the workflow pipe
When an auto-cached image load is reused for img2img/inpaint/upscale/ControlNet, the actual pipeline selected above is pipe, but this toggle still runs against state.pipe (the original txt2img pipeline). If the effective denoise count stays above FBCACHE_MIN_STEPS, the cache remains enabled because the txt2img pipeline opens cache_context, then the workflow pipe is invoked without that context (these workflows are exactly why _pipeline_opens_cache_context exists) and the cached transformer can fail with No context is set. Re-check/disable caching against the actual pipe that will be called.
unsloth/studio/backend/core/inference/diffusion_attention.py
Lines 328 to 334 in 87054e7
Roll back partial attention backend sets
When a multi-DiT pipe has one denoiser accept the requested backend and another denoiser's set_attention_backend raises, this leaves engaged true and returns the backend without restoring the successful module to native. That violates the helper's all-denoisers contract (_attention_dits includes transformer_2/unconditional_transformer) and makes status report a single attention backend while the clip runs with mixed kernels. Treat partial success as a fallback or roll back already-set denoisers when any setter fails.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self._primary.enable_cache(config) | ||
| try: | ||
| self._replica.enable_cache(config) | ||
| _invalidate_registry(self._replica) |
There was a problem hiding this comment.
When CFG parallelism is installed and an auto cache toggles off for a short generation, the uncached run still enters proxy.cache_context(...), which can populate the primary DiT's HookRegistry child cache without block hooks. On a later generation that re-enables caching through the proxy, apply_step_cache only sees the proxy object and this fan-out invalidates only the replica, so the primary's newly installed cache hooks can miss the context and fail with the same stale-registry No context is set path that _invalidate_child_registry_cache is meant to prevent. Invalidate _primary after enabling it as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Same as the sibling thread: getattr on the proxy forwards _diffusers_hook to the primary, so the post-enable invalidation reaches the primary through the proxy, and the fan-out explicitly invalidates the replica. Both registries are cleared on every enable, including the uncached-context-then-re-enable sequence.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 87054e7f94
ℹ️ 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".
| def enable_cache(self, config: Any) -> None: | ||
| self._primary.enable_cache(config) | ||
| try: | ||
| self._replica.enable_cache(config) | ||
| _invalidate_registry(self._replica) |
There was a problem hiding this comment.
When cfg_parallel is installed and the auto step cache is toggled off by a short generation, the primary transformer's HookRegistry can cache an empty child-registry list during that uncached cache_context. A later >=20-step generation re-enables cache through this proxy, but this path invalidates only the replica while apply_step_cache invalidates the proxy object, not _primary; the primary block hooks can then miss the cache context and fail on the first cached forward with the same “No context is set” issue the cache helper guards against. Please invalidate _primary as well after enable_cache.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Same premise as the two sibling threads and it does not hold: the proxy's getattr forwards _diffusers_hook to the primary, so apply_step_cache's post-enable invalidation nulls the PRIMARY's child-registry cache through the proxy (verified by executing that call against the real proxy object); the fan-out's explicit replica invalidation covers the one object delegation cannot reach. A comment now documents this invariant at the fan-out in ce32c1e.
|
@codex review |
1 similar comment
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ce32c1e1fd
ℹ️ 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".
| if text_encoder_quant is None or str(text_encoder_quant).strip().lower() in ("", "auto"): | ||
| text_encoder_quant = "off" if speed_off else TE_QUANT_AUTO |
There was a problem hiding this comment.
When text_encoder_quant is left unset on a resident CUDA image load, this now defaults to auto. For families where auto resolves to a torchao scheme such as fp8_dynamic or int8, quantize_text_encoders runs before apply_memory_plan, while the pipeline is still CPU-resident; the same helper skips torchao under offload because those tensor subclasses reject Module.to(). The resident path then calls pipe.to(device), so default loads can fail during placement unless the user disables TE quant. Please either quantize after the resident .to() or avoid torchao TE schemes before placement.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Verified on-GPU with the exact configs this path can resolve: int8 dynamic-activation, per-row fp8_dynamic and nvfp4 subclasses all accept the single resident Module.to() and forward correctly afterwards; the offload skip targets per-module hook moves, not the one-shot placement, and layerwise fp8 is not a torchao subclass. Repro at temp/repro_r31_cpu_quant_to_cuda.py.
| if text_encoder_quant is None or str(text_encoder_quant).strip().lower() in ("", "auto"): | ||
| text_encoder_quant = "off" if speed_off else TE_QUANT_AUTO |
There was a problem hiding this comment.
This makes an unset video text_encoder_quant default to auto, but the video load path still calls quantize_text_encoders before apply_memory_plan places the freshly loaded pipeline. On resident CUDA loads where auto picks a torchao TE scheme, the encoder is quantized while still CPU-resident and then the later resident placement calls pipe.to(device) on the torchao tensor subclasses that this code already treats as incompatible with Module.to() under offload. Default video loads on those GPUs can therefore fail during placement; quantize after resident placement or keep the default dense until then.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Same evidence as the image thread: every reachable TE scheme survives the video loader's single resident placement move on-GPU; the incompatibility note is about offload hooks' repeated per-module moves.
| quantize_( | ||
| vae, Float8DynamicActivationFloat8WeightConfig(granularity = PerTensor()), filter_fn = filter_fn | ||
| ) |
There was a problem hiding this comment.
For an explicit vae_quant="fp8_dynamic" request on a resident image or video load, both loaders call quantize_vae before the later apply_memory_plan(...).to(device) placement, but this caster ignores target and quantizes the VAE in its current CPU-resident state. That leaves torchao fp8 tensor subclasses in the VAE and the subsequent resident placement has to run Module.to() over them, the same operation this module already avoids under offload because it is unsupported. Move the VAE to the target device before this torchao cast or defer this mode until after resident placement.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The Float8Tensor conv weights survive the resident .to() on-GPU (Conv2d, Conv3d and Linear all verified); kernel availability is orthogonal to placement order and the per-device, per-dimensionality probe already keeps the VAE dense on builds without the fp8 conv path.
| mode = normalize_cfg_parallel(requested) | ||
| if mode == CFG_PARALLEL_OFF: | ||
| return None, "disabled by request" | ||
| explicit_on = mode == CFG_PARALLEL_ON |
There was a problem hiding this comment.
When speed_mode="off" and cfg_parallel is left at auto, speed_active is passed as False but never gates the auto path here, so a resident two-GPU HunyuanVideo-1.5 load can still build and install the CFG-parallel proxy. That makes the explicit Speed=off reference load reserve a second GPU and run a speed lever unless the user also knows to set cfg_parallel="off"; auto should return off when speed_active is false, while still allowing an explicit cfg_parallel="on" if desired.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch. The auto gate now returns off when speed_active is false, before any build or GPU reservation, so a Speed=off load stays a true single-device reference; an explicit cfg_parallel=on still overrides. Fixed in d7a5a01 with a test.
| pred_uncond = None, | ||
| **kw, | ||
| ): | ||
| return orig_forward(_resolve(pred_cond), _resolve(pred_uncond), **kw) |
There was a problem hiding this comment.
When CFG parallel engages for HunyuanVideo-1.5, the real denoise loop stores each branch output on guider_state_batch.noise_pred and then calls self.guider(guider_state), so pred_cond here is the guider-state container rather than the prediction tensor. _resolve therefore returns the container unchanged, leaving a _PendingPred in thread mode or a cuda:1 tensor in inline mode for the original guider to combine, which either breaks arithmetic on the placeholder or mixes devices. Resolve the nested noise_pred entries before delegating to the original guider.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
On the installed diffusers, BaseGuidance.call unpacks each batch's noise_pred before calling forward (guider_utils.py:176-184), so the patched forward receives the per-branch prediction values, never the container; call does no tensor arithmetic. The round-2 bit-identical parallel-vs-sequential validation exercised exactly this path.
Speed=off is the reference contract: the loaders pin every auto speed lever (transformer/TE/VAE quant tri-states) to off, but the cfg_parallel auto path never consulted speed_active, so a resident two GPU HunyuanVideo-1.5 load with Speed=off could still reserve a second GPU and install the CFG-parallel proxy. Auto now returns off when speed_active is false; an explicit cfg_parallel=on stays honored as a deliberate override (the install-failure test now exercises exactly that override path).
|
@codex review |
|
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". |
The pre-warmed torch.compile cache (diffusion_compile_cache.py: fingerprinted bundles over torch.compiler.save/load_cache_artifacts plus a persistent per-key TORCHINDUCTOR_CACHE_DIR) was wired into the image backend only, so every video load re-paid the full first-generation compile after every process restart (the stock inductor dir lives in /tmp). video.py now mirrors the image backend exactly: compile_cache.begin runs after the attention-backend set (the fingerprint keys on the engaged kernel) and before apply_speed_optims on a compile-eligible default/max tier, keyed on the same fullgraph decision as the compile itself (an engaged or still-toggleable step cache and a planned offload both drop it); the context is committed to _VideoLoadState and compile_cache.save persists the bundle after the first successful generation (env-gated distributor / first-run-warm mode); _teardown_state restores the inductor dir, and a token-scoped _rollback_precommit_compile_cache covers loads that die before the state commit, mirroring the globals and CFG-parallel rollbacks. Measured on HunyuanVideo-1.5-480p through the real VideoBackend (B200, 480x288/17f/30 steps): the first-generation extra drops from 107.5 s cold to 13.8 s when the 12.8 MB bundle loads into a fresh inductor dir (0.10 s load) and to 11.7 s from the persistent per-key dir alone; through the wired production path a restart lands at 10.5-10.8 s (bundle-only included) vs 86.5 s cold. Steady state is unchanged (2.4-2.6 s), and the loaded artifacts are the same bits a local compile would produce, so numerics are untouched. Tests: begin/save/restore lifecycle with the fullgraph keying, the Speed=off and compile-ineligible skips, and the token-scoped pre-commit rollback.
Every current release of the kernels package (0.13 through 0.16) builds its dependency tables with huggingface_hub >= 1.0's strict-dataclass API, and with an older hub the breakage is not contained to the requested backend: import kernels raises at module scope, and diffusers imports kernels whenever it is installed, so a single on-demand install (an explicit flash3/flash4 attention request on a stack pinned to hub < 1.0) permanently breaks every later diffusers pipeline import on the box until the package is removed. Reproduced against hub 0.36.2 with kernels 0.13.0 and 0.16.0: the HunyuanVideo-1.5 pipeline import dies in hub's strict-dataclass validator both times. _ensure_attention_backend_installed now checks the resident hub version before installing kernels (_kernels_hub_compatible) and refuses on < 1.0, logging why and falling back to the native default, which is the best-effort contract the installer already promises for an uninstallable wheel. The refusal is a policy decision, not a failed attempt, so it is not memoised and a later request on a fixed environment can still install. An undeterminable hub version keeps the previous permissive behaviour, and the gate applies only to the kernels package: sage/flash-attn/xformers wheels do not import hub at module scope. Tests: the refusal (nothing memoised), the hub >= 1.0 allow, the package-scoping, and the version-parse fallback.
|
@codex review |
1 similar comment
|
@codex review |
for more information, see https://pre-commit.ci
… prewarm
vLLM and SGLang finish every compilation at server startup (dummy batches
through each compiled shape) so no request ever pays a compile mid-serving.
The video backend's compiled tier instead paid a first-generation extra after
every restart: ~54 s cold and ~11.3 s even with a warm Mega-cache bundle (the
residual is dynamo tracing plus cudnn.benchmark autotune, which the bundle
cannot carry).
After a compiled DEFAULT-tier resident load commits, a daemon thread now runs
one tiny throwaway generation (192x128 snapped, 4k+1-lattice 9 frames, 2
steps) under the generate lock. The default tier compiles dynamic=True, so
the small trace serves every later resolution. Measured through the real
backend (HunyuanVideo-1.5-480p, B200, 480x288/17f/30 steps):
warm bundle: first-generation extra 11.3 s -> 2.1 s (9.6 s background warmup)
cold start: the full compile moves off the user's first request entirely
(14.5 s background; first generation extra 2.1 s), and the
warmup persists the Mega-cache bundle itself
steady state: unchanged (2.4-2.5 s per 30-step clip in every phase)
Exactly lossless by construction: the warmup only changes when compilation
work happens. It resets its step-cache residuals, the real generation seeds
its own generator, and no process-wide flag is touched.
The warmup registers itself as the active cancellable job, so unload, a new
load, or cancel_generate abort it at a step boundary (verified: unload 2 s
into a running prewarm returns in 6.7 s with the warmup cancelled). It yields
untouched when a real request arrived first and is token-scoped against
superseded loads. Gated per family (supports_compile_prewarm), skipped for
speed=max (static per-shape graphs a warmup shape cannot serve), offload
(every warmup forward would stream the DiT over PCIe), and CFG parallel (its
planner owns compile-sensitive runs); the UNSLOTH_DIFFUSION_COMPILE_PREWARM
kill switch disables it. The decision and reason land in the resolved record.
Tests: +4 hermetic (decision gates, engaged-load spawn with snapped tiny
shape, skip-without-compile, yield to generations / stale tokens); related
backend set 551 passed; ruff clean.
|
@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: 6adcc34e1b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| return # superseded/unloaded before the warmup could start | ||
| if self._generate_job_active or self._gen.get("active"): | ||
| return # a real request beat us; it absorbs the warmup itself | ||
| self._active_generate_cancel = cancel |
There was a problem hiding this comment.
On compiled video loads where the background prewarm has already entered _generate_lock, this only registers a cancel event but does not mark the warmup as the active generation. A user generation submitted during that window sees _generate_job_active == False, overwrites _active_generate_cancel with its own event, and then waits behind the warmup, so cancel/unload no longer signals the prewarm and the first real request can be delayed by the 9–54s warmup this path is meant to hide. Mark the prewarm as busy/preemptible, or have begin_generate cancel and wait for it before queuing the real job.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Genuine catch, fixed in 47c202e. The prewarm now registers its cancel event in a dedicated _prewarm_cancel slot (in addition to _active_generate_cancel, so unload/cancel keep working unchanged), and begin_generate signals that slot before installing the real job's event; direct generate() calls do the same. The warmup aborts at its next step boundary and the real job takes _generate_lock instead of waiting out the full warmup. Covered by test_begin_generate_preempts_running_prewarm.
The prewarm registered its cancel event in _active_generate_cancel, but a begin_generate arriving mid-warmup overwrote that slot with its own event and then queued its worker behind the full warmup on _generate_lock. From that point unload and cancel_generate signalled the wrong event, so the warmup could no longer be aborted and the first real request waited out the 9-54s the prewarm exists to hide. Track the prewarm's event in a dedicated _prewarm_cancel slot (cleared identity-checked alongside _active_generate_cancel) and signal it from begin_generate before registering the real job's event, and from direct generate() calls that skip begin_generate. The warmup then aborts at its next step boundary and the real job takes the lock, while unload/cancel keep working against whichever run is actually active.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Breezy! 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". |
1 participant
A batch of accuracy-gated precision and speed improvements for the Studio diffusion backend, split out of #6763 for focused review. Every change here builds on the diffusion loader / quant / attention code in #6763, so this PR is stacked on the
image-generationbranch. All defaults were chosen from B200 accuracy sweeps; anything that could regress an output is gated or opt-in, and every other model family is left on its existing path.Auto text-encoder and VAE precision
fp8_dynamicstays opt-in) and by size, so only the large video VAEs are quantized and the tiny image VAEs are left dense where the saving is not worth the risk.Video DiT quantization correctness
HunyuanVideo-1.5 attention speedup
Hunyuan's DiT runs a joint
[video; text]self-attention and, on every block and step, builds a dense[B,1,N,N]boolean mask so the video never attends to padded text. A dense boolattn_maskdisables every fused SDPA kernel (flash rejects it; cuDNN and memory-efficient fall back), so the attention runs the slow math path: at 121 frames / 480p (N about 50k) one attention call is ~421ms with the mask vs ~19ms withattn_mask=None. The text is ~99.5% padding (a t2v prompt fills ~9 of ~1985 slots).install_hunyuan_attention_trimadds an eager forward pre-hook that drops the all-zero image stream (t2v) and trims the mllm/byt5 text streams to their globally-valid columns, plus a null-mask attention processor that runsattn_mask=Noneonce no partially-padded column remains (the batch-1 / per-guidance-branch case) and otherwise delegates to the stock dense-mask processor. The model already zeroes and masks the padded text and discards its attention output (only the video split feedsproj_out), so removing it is exact for the video; the only numeric change is the SDPA kernel (masked fallback to fused).Measured on a B200: 23.3s to 1.3s per DiT forward at 121 frames (~18x with regional compile, 0 graph breaks); per-forward cosine 0.99998 vs stock; equal distance to an fp32 reference (LPIPS fp32-vs-stock 0.292, fp32-vs-trim 0.307), so it is not less accurate than the current bf16 default. Wired auto-on for HunyuanVideo-1.5 only (before the attention backend is set, so the requested kernel pins onto the new processors); a no-op for every other family and reversible (stock dense-mask path on any anomaly).
Benchmark tooling
Housekeeping
Tests and validation
Hermetic CPU tests in
studio/backend/tests/test_diffusion_attention.pycover the attention trim helper, the null-mask processor gating, idempotency, the t2v image-empty / i2v-keep paths, mixed-batch fallback, and the positional-call and empty-prompt guards. TE and VAE auto-precision and the black-frame denies have their own tests alongside the existing diffusion suites. Diagnostic and validation scripts underscripts/reproduce the findings, includingsdpa_mask_backend_probe.py(the dense-mask kernel tax),hunyuan_attn_diag.py(token breakdown),hunyuan_trim_validate.py(per-forward speed and accuracy),hunyuan_trim_e2e.py(end-to-end LPIPS and mean luma), andhunyuan_trim_fp32ref.py(accuracy vs fp32 ground truth).Stacked on #6763.