unsloth/scripts
Daniel Han efc7a44747 Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT
Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).

Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.

OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.

- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
  (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
  quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
  when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
  transformer_cache_threshold through begin_load -> load_pipeline and report the
  engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
  transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
  selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
  LPIPS vs the compiled no-cache baseline).
2026-06-26 12:35:50 +00:00
..
data CI: scope GITHUB_TOKEN permissions, add MLX CI, unblock ~60 skipped tests (#5312) 2026-05-11 03:19:13 -07:00
build_prequant_checkpoint.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 11:24:35 +00:00
check_frontend_dep_removal.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
check_new_install_scripts.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
compare_engines.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 11:24:35 +00:00
compile_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 03:20:02 +00:00
diffusion_bench.py Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override 2026-06-26 08:01:15 +00:00
diffusion_quality.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-25 13:56:37 +00:00
enforce_kwargs_spacing.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
fbcache_flux_probe.py Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT 2026-06-26 12:35:50 +00:00
fp8_overflow_check.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 07:38:27 +00:00
install_gemma4_mlx.sh Update Install Scripts (#5968) 2026-06-03 05:39:42 -07:00
install_qwen3_6_mlx.sh Update Install Scripts (#5968) 2026-06-03 05:39:42 -07:00
install_rocm_wsl_strixhalo.sh Windows/WSL installer: fix winget msstore cert failure, amd-smi DiskPart prompt, and enable AMD GPU (Strix Halo gfx1151) (#5940) 2026-06-10 04:24:49 -07:00
leverage_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 11:24:35 +00:00
lint_workflow_triggers.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
lockfile_supply_chain_audit.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
notebook_to_python.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
notebook_validator.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
nvfp4_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 08:24:58 +00:00
nvfp4_t211_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 12:03:36 +00:00
perf_levers_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 12:03:36 +00:00
perf_verify.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 03:20:02 +00:00
prequant_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 11:24:35 +00:00
quant_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 06:17:18 +00:00
run_ruff_format.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
scan_npm_packages.py Package scanners: close fail-open gaps in the sdist fallback and hidden-payload paths (#6359) 2026-06-18 06:50:16 -07:00
scan_npm_packages_baseline.json Package scanners: close fail-open gaps in the sdist fallback and hidden-payload paths (#6359) 2026-06-18 06:50:16 -07:00
scan_packages.py Package scanners: close fail-open gaps in the sdist fallback and hidden-payload paths (#6359) 2026-06-18 06:50:16 -07:00
scan_packages_baseline.json Studio: remove OpenEnv and other unused packages (#6585) 2026-06-23 07:20:47 -07:00
sd_cpp_smoke.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-25 16:08:55 +00:00
sparse_accum_probe.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 06:48:40 +00:00
stamp_studio_release.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
sync_allow_scripts_pins.py Studio: auto-sync allowScripts pins after dependency bumps (#6136) 2026-06-10 02:35:37 -07:00
uninstall.ps1 Windows installer: fix DiskPart UAC mid-install, drive-root cache, and spurious unsloth.exe rename warning (#6296) 2026-06-22 03:09:08 -07:00
uninstall.sh Windows installer: fix DiskPart UAC mid-install, drive-root cache, and spurious unsloth.exe rename warning (#6296) 2026-06-22 03:09:08 -07:00
verify_comment_only_diff.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
verify_import_hoist.py CI: fix import-hoist false positive, vision-cache test cwd, llama.cpp CLI smoke (#6598) 2026-06-23 01:16:47 -07:00
verify_prequant_backend.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-26 11:24:35 +00:00