Point prequant_repos for flux.1, flux.2-klein, flux.2-dev, qwen-image
(int8 only there; fp8 is family-denied), z-image and krea-2 at the
unsloth/<Model>-FP8 Hub repos carrying gate-validated int8 and fp8
transformer checkpoints, so the fast quant path loads the small
pre-quantized file instead of materialising the dense bf16 transformer
and quantising on device. Measured on FLUX.2-dev int8: build peak drops
from 60.7 GB (dense + quantize) to 30.7 GB (hosted prequant), identical
30.7 GB resident after either path since loading a checkpoint is
bit-identical to on-the-fly quantisation.
The hosted repos name files <Model>-<SCHEME>.pt, so resolve_prequant_source
now derives that model-name filename from the repo id (scheme suffix
stripped case-insensitively) and carries the legacy transformer_<scheme>.pt
as a fallback the resolver tries when the primary 404s, keeping older
repos loadable.
Wiring a repo also exposed a fallback hazard: with a prequant source
present, the dense-fit preflight used to be skipped entirely, so a failed
prequant download would fall through to the dense bf16 load the memory
plan never budgeted, OOMing after eviction. The preflight now always runs
and gates an allow_dense_fallback flag through _load_dense_quant_pipeline:
a dense misfit still skips the fast path when no prequant exists, but with
one it proceeds and a prequant failure raises to the GGUF build instead of
loading dense. The same flag is set when the auto-policy replans an
offloaded GGUF against a prequant-sized transient.
Tests updated to the new filename convention plus new coverage for the
derivation and the legacy-name fallback; the prequant-skips-refit test now
asserts the re-check runs and forbids the dense fallback. Verified end to
end on GPU: z-image int8 resolves the hosted repo, downloads the
model-name file and renders (6.8s load, 5.9 GB peak).
An all-zero activation token row makes the dynamic per-row fp8 scale 0,
which turns the quantized data to NaN and the render to black frames on
torchao's plain-torch kernel path. The fused fbgemm/mslk quantize kernels
clamp zero rows internally, so the bug only reproduces on machines without
them, which is most user environments. Zero rows are real inputs, not a
corner case: Wan 2.2 zero-pads its text conditioning, and Hunyuan-1.5 and
Qwen-Image regenerate zero rows inside their transformer blocks every step.
Pass activation_value_lb=1e-12 to Float8DynamicActivationFloat8WeightConfig
whenever the installed torchao supports the kwarg (Float8Tensor rework,
0.13+), checked via inspect.signature so older torchao keeps exactly the
current behaviour; the existing Float8MMConfig fallback chain is unchanged.
Verified on GPU: with the forced plain-torch kernel path a zero-row input
NaNs without the floor and stays finite with it, and end to end on
HunyuanVideo-1.5 fp8 goes from a solid black frame (LPIPS 1.00) to a normal
render (LPIPS 0.225); on Wan the floor matches the condition_embedder
exclusion (LPIPS 0.211 vs 0.206). This is defense in depth on top of the
family excludes and deny list, which stay as-is: it changes the failure
mode of any future zero-row family from black frames to graceful
degradation. Same-seed renders with fused kernels present are unaffected,
and pre-quantized fp8 checkpoints stay valid since weight scales are
untouched.
- diffusion_attention: arch-gate FlashAttention 2 to Ampere (SM80)+ in both the
primary selector and the heterogeneous-replica guard (it crashed on pre-Ampere).
- diffusion_cfg_parallel: convert boolean attn masks to additive bias before the direct
cuDNN op so partial masks match F.scaled_dot_product_attention; make proxy disable_cache
transactional (clean both branches, mark broken, surface a reload-required error).
- diffusion_cache: fail closed when a magcache step-count resize or below-threshold
disable cannot remove the old cache; surface a failed enable+cleanup instead of a false
uncached None.
- video: roll back earlier experts when a later expert raises in the all-or-none step-cache
loop; fail the load when the primary-only cache cannot be re-engaged through the
CFG-parallel proxy; validate transformer_cache_quality and cfg_parallel before the worker.
- scripts: place the fp8 ablation pipeline on CUDA; fail closed on a failed magcache resize
in the speedmem bench; label OOM distinctly in the SDPA mask probe.
- tests: regressions for the FA2 arch gate, transactional proxy disable, all-or-none
exception rollback, magcache fail-closed transitions, and enable+cleanup failure.
- _scan_models_dir: admit a scan folder that is itself a diffusers pipeline
(root model_index.json, weights in transformer/ vae/ subdirs). _is_model_directory
rejects such a root, so the child scan would list the component subdirs as bogus
models and hide the real pipeline; treat the root as one model via _local_pipeline_index.
- _local_is_diffusers / _local_model_task: include the sole checkpoint filename in the
family-detection needles (_local_family_needles, resolved via resolve_local_single_file).
A generically named folder holding one loadable qwen-image-*.safetensors / ltx-*.safetensors
identifies its family only from the filename; the load route already resolves that file, so
tag it or the task-scoped picker (which rejects task=null) hides the on-device model.
- list_cached_models: mark a companion-only base snapshot partial. A GGUF image load prefetches
the base repo's VAE / text-encoder / model_index.json but skips the transformer (the GGUF
supplies it); the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline, and
_cached_repo_partial misses it. _repo_pipeline_missing_denoiser flags a pipeline snapshot whose
transformer/ or unet/ component carries no weight, so the picker drops it instead of advertising
it as fully on-device.
Publish native sd.cpp generate progress (_gen) before LoRA resolution so a reload probe reads active during setup, matching the diffusers path.
Register the diffusion/video GPU load under the arbiter lock (acquire_for now takes a register callback) so a competing acquire cannot evict an owner before its load is marked in-flight and let two loaders allocate VRAM at once.
Admit local diffusers pipeline folders (root model_index.json, weights in component subdirs) in the local model scan so they reach task tagging and the On Device picker.
When the install target already exists, is non-empty and lacks the
.unsloth-studio-owned marker (a user's own stable-diffusion.cpp checkout,
or unrelated files beside a custom Studio root), install() previously still
extracted the release into it. Skipping the ownership marker only stopped the
uninstaller from deleting the directory; extraction still merged binaries into
the user's working tree and could overwrite same-named files.
Fail up front with a clear message pointing the user at a fresh/empty location
before any download or extraction, leaving their directory untouched. Update
the ownership test suite to assert the refusal.
Under speed=max the video DiT compiles dynamic=False, so inductor
artifacts are per (width, height, frames). The save path never called
compile_cache.register_shape, so after a bundle hit ctx.saved stayed
true and later resolutions/frame counts never re-dirtied the bundle,
leaving those shapes to recompile on every restart. Register the actual
generation shape before saving, gated on the static tier, mirroring the
image backend.
- diffusion_attention: clear the HunyuanVideo-1.5 null-mask flag with an always_call
post-hook so it is scoped to one hooked forward and never latches across an
exception; add attention_backend_supported_on_device to arch-gate an
already-resolved backend on a specific (heterogeneous) CUDA device.
- video: make the explicit MagCache resize transactional via _step_cache_all_or_none
(refuse to stack a fresh cache over one that could not be disabled; roll a mixed
resize back and report the true state); raise on a failed all-or-none rollback
instead of falsely reporting an uncached pipeline.
- diffusion_cfg_parallel: re-validate the attention backend on the replica device
and pin native there when unsupported; mirror the primary's max tier on the
replica (max-autotune compile + direct QKV fusion) via a new speed_mode arg;
prefer a viable heterogeneous secondary GPU over an unusable identical one; clear
the const cache at each plan_generation.
- diffusion_vae_quant / diffusion_precision: detect a partial diffusers
layerwise-fp8 mutation (leftover casting hooks the torchao detector cannot see)
and fail the load closed, while a clean failure still falls back to dense.
- video_speedmem_bench: engage the dual-expert cache all-or-none like the loader.
- frontend video api: add text_encoder_quant / vae_quant and the auto/off literals
to VideoLoadRequest so typed callers match the backend contract.
install_sd_cpp_prebuilt: only write the .unsloth-studio-owned marker when the
install created the target directory or it was empty. Adopting a pre-existing,
unowned, non-empty directory (a user's own stable-diffusion.cpp checkout) made
it eligible for the uninstaller's recursive delete.
routes/training upload: make the multi-file promotion transactional. Back up
each displaced original and roll every destination back on any failure, so a
mid-loop rename error can no longer partially overwrite the live dataset.
routes/training _resolve_dataset_folder: reject a symlinked dataset directory
and prove the resolved folder stays under the datasets root, so image
read/caption/delete cannot escape the root through a link.
routes/training delete: escape glob metacharacters in the thumbnail filename so
deleting an image named like [ab].png removes only its own thumbnails.
image_gallery / video_gallery listing: filter records against the response
schema inside the pager via a valid callback, so offset/limit/has_more all count
over accepted records. A leading schema-invalid record no longer returns an
empty page with has_more=true and stalls infinite scroll at offset 0.
image_gallery / video_gallery save: publish via a temp file plus atomic rename
(the sidecar is the video pair's commit marker) and clean up on failure, so a
partial write never surfaces a truncated PNG or strands an orphan MP4.
diffusion_train_common discovery: treat an empty caption sidecar as a metadata
tombstone that still falls through to the dreambooth instance prompt, so
clearing every metadata caption no longer fails with no captioned images found.
diffusion backend unload: wait for an in-flight denoise to exit before tearing
down process-wide patches and state, mirroring the load path.
diffusion_engine_router: serialize the whole check/unload/publish transition so
a concurrent selection cannot return the engine being unloaded.
uninstall.ps1: gate the default sd.cpp process stop on the owner marker so a
user's own sd-server is not terminated for a directory we then keep.
generate() assigned self._gen only at the pipe() call, after deferred
compile, LoRA resolution/application, and ControlNet download/build had
run. Across that setup window generate_progress() reported inactive even
though _generate_lock was held, so a reloaded page's mount probe showed
idle and let a second generate queue behind the first.
Publish an active step-0 _GenState the moment the generation lock is
acquired, before the setup work, and clear it in the outer finally so a
setup-time error cannot leave the UI stuck active. Mirrors the video
backend's queued phase and the training start guard.
Route on-device single-checkpoint video folders through the single_file loader:
a bare local .safetensors directory (no model_index.json) is advertised as a
pipeline with no filename, so validation rejected it before it could load.
Reinterpret the pick as a single_file load of the sole checkpoint, mirroring the
image load route.
Treat a reserved-but-not-yet-spawned LLM training start as active in
is_training_active() so /images/load, /video/load, and /diffusion/start cannot
race the reserved run for VRAM during the pre-spawn free window. Mirrors the
diffusion training service reservation.
Resume an in-flight image generation on the Images page mount: probe
generate-progress, re-enter the poll loop, and refresh the gallery on completion
so a run started elsewhere is reflected and its saved image appears without a
manual refresh. Seed resident image defaults from the resolved base_repo rather
than a possibly path-shaped repo_id so the first resident generation uses the
right recipe.
The previous commit introduced the single-device budget but wired it
into only the downloaded-group and Hub search sites. The LM Studio,
custom-folder, local-dir, live-search and exported-GGUF expanders
reachable from the Images/Video pickers still measured against the
summed multi-GPU total, as did the size-based GGUF row badge, so those
paths could still recommend a quant that OOMs on a single device. All
GgufVariantExpander call sites in HubModelPicker now share
expanderGpuGb, and the row badge derives the same task-scoped budget.
The generate callback omitted loraCapable from its dependencies, so when
an auto-compile flips supports_lora off mid-session the memoized handler
still sent the previously selected adapters and the next generation
failed with the backend's LoRA-not-supported error instead of omitting
adapters the UI had already hidden.
HubModelPicker's GGUF variant expanders, format lists and Hub row fit
hints measured against the summed multi-GPU total. When the picker is
task-scoped (Images/Video) the loaders place the whole pipeline on one
device, so a variant could be recommended as fitting and then OOM at
load; those sites now share the single-device budget the group fit gate
already uses, while chat pickers keep the summed total since llama.cpp
splits layers across devices.
The dataset labeling grid's Remove button relied on group-hover with no
group parent, leaving it permanently invisible to mouse users; the image
wrapper now carries the group class.
Register the dims the forward actually compiled with: image-conditioned
workflows (img2img, inpaint, upscale, edit) run at the input image's size,
not the slider's, so recording the slider values marked never-compiled
shapes as covered and warm restarts kept paying compile for the real one.
Validate a request-supplied transformer_prequant_path (existence plus the
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH allowlist) before treating prequant as
available at the resident-fit re-check: an unusable path skipped the dense
fit check up front and then fell back to materializing dense bf16 after
the previous pipeline was evicted, recreating the post-eviction OOM path.
Shared as usable_prequant_source, also used by the auto-policy planner.
A generation runs on a backend daemon thread and survives a page reload,
but the mount effect only probed the load progress, so reloading during a
generate showed an idle page that never picked up the finished clip until
a manual refresh. Hoist the generate poll loop out of handleGenerate and
re-enter it on mount when generate-progress reports an active job; merge a
terminal completed record into the gallery to cover the race with the
mount gallery fetch.
* Auto-detect completion masking markers with template table fallback
Studio's train_on_completions previously relied only on the hardcoded
MODEL_TO_TEMPLATE_MAPPER / TEMPLATE_TO_RESPONSES_MAPPER tables and
silently disabled masking when a model was not in the table, so unmapped
models (LFM2-8B-A1B, DeepSeek, and others) trained on full sequences
without telling the user. Several mapped templates (glm, mistral, llama,
starling, zephyr, qwen3-thinking) also carried markers that mask every
assistant token, which made every row drop in the post-masking filter.
Both training callsites (CUDA trainer.py and MLX worker.py) now share
utils.datasets.completion_masking.apply_completion_masking:
- Try unsloth_zoo chat template auto-detection first; it raises loudly
when the template cannot be parsed and never masks the EOS token.
- gpt-oss models keep their manual markers so non-final assistant
<|end|> tokens stay trained, matching current behavior.
- If auto-detection raises, fall back to the template table exactly as
before.
- If the table also misses, emit an explicit user-visible warning that
completion masking could not be applied and full-sequence training
will occur, instead of a quiet log line.
The >30 percent dropped-rows safety net in trainer.py now guards the
auto path as well. Table consumers for inference and chat templates are
unchanged. Validated against one representative tokenizer for every
template in TEMPLATE_TO_RESPONSES_MAPPER plus the unmapped models:
no template regresses; unit tests cover the four decision paths.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Restrict masking fallback to marker detection failures
The auto branch wrapped the whole train_on_responses_only call, so a real
failure while applying the masking (dataset map, tokenization) was treated
as a detection miss and training silently proceeded on full sequences.
Detect markers separately via get_chat_template_parts (test seam via
detect_fn), then apply them with errors propagating, matching the manual
path. Tokenizers with preset unsloth marker attrs skip detection and call
bare so zoo reuses the stored parts.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fail the run when applying completion masking raises
The helper already falls back internally on detection failures and returns
applied=False on a double miss, so an exception reaching the callsites is a
real failure applying the masking. Remove the callsite catches that
downgraded it to full-sequence training; the run now fails visibly instead.
Also use the explicit re-export alias form in utils/datasets/__init__.py for
the two new names, satisfying the import-hoist source lint.
* Import completion masking from its submodule
The import-hoist source lint counts only real name loads, so package-level
re-exports of the two new names cannot satisfy it. Import
apply_completion_masking from utils.datasets.completion_masking directly at
both callsites and leave utils/datasets/__init__.py untouched.
* Completion masking: gpt-oss renames and MLX raw/alpaca parity
Renamed or private gpt-oss checkpoints are name-detected as gpt-oss but miss
the exact-name table; default them to the gpt-oss template markers instead of
falling through to full-sequence training.
Gate the MLX masking call on not raw_text_mode and format_type != alpaca,
mirroring the CUDA path: raw/CPT text has no chat turns to mask and
Alpaca-rendered text lacks the tokenizer's chat markers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Define raw_text_mode outside the MLX feature-detect block
With an older zoo lacking the append_eos config field, the masking
gate referenced raw_text_mode before assignment. Hoist the assignment
above the feature detection so both consumers see it.
* Gate MLX masking on the formatter's resolved format
format_type auto can resolve to alpaca or raw text; the masking skip
checked only the requested value, so auto-detected Alpaca data got
chat-template markers applied to rendered prompt text. Track the
final_format returned by format_and_template_dataset and gate on it,
matching the CUDA path.
* Unwrap the mlx-lm TokenizerWrapper before marker checks
The wrapper delegates plain reads to the wrapped HF tokenizer but hides
underscore attrs, so preset unsloth markers were invisible and detection
relied on the loader's call patch. Unwrap to the real tokenizer first,
as the zoo MLX resolver does.
* Tighten masking comments
* gpt-oss: auto-detect markers first like every other template
The quantized and BF16 gpt-oss checkpoints ship a chat template without
the channel final header, so the pinned manual markers match nothing
there and masking trained zero tokens. Auto-detection derives markers
from whichever template the checkpoint ships and keeps the final
terminator trained; the manual gpt-oss markers remain the detection
failure fallback, including for renamed checkpoints.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Studio: route models by CONFIG_MAPPING_NAMES instead of hardcoded tables
A model whose model_type is absent from an overlay's transformers cannot load
there, so a new MoE arch not yet in the tier tables gets routed to default and
fails (e.g. lfm2_moe, deepseek_v4). Add a static resolver that parses each
overlay's CONFIG_MAPPING_NAMES straight from source (AST only, no import, no
network, no trust_remote_code) and picks the lowest tier that ships the
model_type. Runs after the existing checks and only ever upgrades default, so
no existing routing changes and new archs no longer need a table edit.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio router: harden the CONFIG_MAPPING_NAMES resolver
- Resolve the default tier map from the base install, skipping any .venv_t5_*
sidecar on sys.path, so an in-process 5.x activation cannot make a 5.x-only
model look loadable by 4.x.
- Do not cache an overlay whose sidecar dir is absent, so a later call re-reads
it once provisioned instead of serving a stale empty map.
- Also collect model types added via CONFIG_MAPPING_NAMES.update({...}) and
**{...} unpacking, not just the literal assignment (5.10 uses both).
- Wrap the AST walk in the try/except so a malformed source can never crash tier
resolution.
- Feed the mapping fallback from _load_config_json so a config served from the
hub cache during a transient outage still routes new architectures.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Tighten comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
LFM2-8B-A1B and any other lfm2_moe checkpoint were missing from the
transformers tier tables, so they fell through to the default 4.57.x
sidecar, which does not register lfm2_moe and errors with
"not supported yet in transformers==4.57.6". Only lfm2_vl was listed.
Add Lfm2MoeForCausalLM / lfm2_moe to the 5.3.0 tier (lfm2_moe is
registered in transformers 5.3.0). get_transformers_tier now returns
530 for LFM2-8B-A1B and the model loads and trains as expected.
- Wan2.2-A14B step cache: pin the balanced FBCache threshold to 0.08 even when
quant is active (per-family override in diffusion_cache.py). Auto-fp8 made the
generic quant promotion (0.12) the family's effective default at pairwise LPIPS
0.128, over the 0.08 quality gate the balanced preset is held to. Measured
operating point with fp8 actually engaged (1280x720/81f/50 steps, B200):
fb@0.08 = 1.08x at 0.129 vs the old fb@0.12 = 2.58x at 0.181; documented in
the preset table. Explicit thresholds and the fast preset are unaffected.
- MagCache curves: validated the shipped 33-frame calibrations at the production
121-frame default for hunyuanvideo-1.5-720p, hunyuanvideo-1.5 (480p) and
wan2.2-ti2v-5b. Fresh 121-frame calibrations differ by <= 0.024 max abs entry
and produce byte-identical frames at the auto presets (hv720 quality 1.69x at
LPIPS 0.042, hv480 quality 1.66x at 0.018, wan5b balanced 1.74x at 0.026, all
pairwise vs the same-load uncached stack), so the curves ship unchanged with
the frame-count transfer documented next to them.
- Dual-GPU CFG parallelism: the secondary-device pick now prefers a device whose
name and compute capability match the primary, and the gate declines a
mismatched pair in auto mode (eager kernel selection is arch-dependent, so the
advertised bit-identity cannot hold across different GPU models); an explicit
cfg_parallel=on proceeds but is downgraded to lossless=False with a warning.
- A14B expert step cache is now all-or-none, mirroring the transactional quant
loop: a mixed outcome (cache engaged on one expert but not the other) is
rolled back and reported uncached with the failure reason, on both the load
path and the generation-time auto toggle.
- Partial torchao quantization is no longer reported as dense: after an
in-place quantize_/caster failure, the DiT / text encoder / VAE is scanned
for leftover torchao tensor-subclass parameters and the load fails with a
clear error when any are found (a half-quantized module cannot run as dense,
and offload's Module.to() crashes on torchao tensors). Failures that swapped
nothing keep the best-effort dense fallback.
- Cleanup: apply_attention_backend / apply_speed_optims / the attention trim
are called once on the pipe (they already fan out over every DiT internally),
so the second A14B expert no longer passes through them twice; the stale
dual-DiT helper comment is rewritten to match the two helper shapes.
Tests: device-identity picker/gate/lossy-plan coverage, per-family threshold
pin scoping, all-or-none rollback in both failure directions, and partial-quant
detection for all three quant modules.
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.
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.
diffusion_compile_cache: auto mode now saves the Mega-cache bundle after the
first compiled generation (UNSLOTH_DIFFUSION_COMPILE_CACHE_SAVE=0 opts out), so
users get warm restarts without the distributor env; a bundle hit starts clean
(no pointless rewrite of the just-loaded artifacts) and explicit mode 1/on keeps
the distributor-style re-save. New register_shape + manifest shape coverage: a
STATIC compile produces new artifacts per (width, height, batch), so the
generate path registers each generation's shape and an uncovered shape
re-dirties the context, growing the bundle to cover every shape the session
used. Measured (B200, real backend): Qwen-Image deferred gen-3 hitch 29.1 ->
22.2 s warm with bit-identical output (7.9 MB bundle, ~0.5 s save); SDXL gen-3
115.7 -> 24.7 s and a mid-session 768px recompile 65.8 -> 12.6 s (bundle 63.6 ->
98.7 MB after the 768 re-save).
diffusion_speed: U-Net denoisers (UNet2DConditionModel; no _repeated_blocks, so
the regional compile never reached them) now get a whole-module STATIC
torch.compile on the default tier, plus fused QKV projections and a compiled VAE
decode. Measured on SDXL (30 steps / 7.0 / 1024px, 4 prompts, LPIPS vs the
bit-exact reference): 6.16 -> 3.14 s end to end (1.96x) at LPIPS 0.035, steady
state 0.70-0.88 s/image through the real backend. Rejected on measurement:
dynamic=True whole-module (366 s compile for 39.3 ms/step vs static's 73 s for
26.9), regional BasicTransformerBlock only (45.0 ms/step; ResNet convs stay
eager), max-autotune + inductor flags (25.9 ms/step for a 445 s warmup),
channels-last UNet alone (neutral). DiT tiers unchanged: fused QKV measured
exactly neutral under the regional compile (Qwen-Image 6.53 vs 6.52 s), so it
stays max-only there, and the DiT VAE decode stays eager (a few % of a DiT
generation). compiled_shapes_are_static tells the cache layer which loads are
per-shape (max tier, U-Net whole-module).
diffusion: register each generation's shape with the compile cache before the
save, pass pipe.unet to the cache fingerprint when the pipe has no transformer,
and correct the transformer_quant resolved reason on dense loads (it claimed a
GGUF transformer was loaded on every non-quantized pipeline load).
Tests: 333 passing across the related suites (speed 42, compile_cache 27, cache
40, precision 20, backend, base_precision, transformer_quant, memory); ruff
clean. Full measurement record: outputs/image_optim_round2_audit.md.
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.
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.
