unsloth/studio/backend/core/inference
Daniel Han 395816cf7e Studio diffusion (Phase 7): accuracy-preserving speed pass
Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy
bug and a dead-on-arrival speed path; this fixes both and adds the lossless /
near-lossless wins, all measured on a B200.

Correctness:
- TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32
  process-wide and never restored them, so a later `off` load silently inherited
  TF32 and was no longer bit-identical. Added snapshot_backend_flags /
  restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer
  runs and restored on unload. Verified: load max -> unload -> load off is now
  byte-identical (PSNR inf) to a fresh off.
- sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and
  only checked the timeout after EOF, so a child stuck in model load / GPU init
  with no output ignored the timeout. Drained stdout on a reader thread with a
  wall-clock deadline. Added a silent-hang regression test.

Speed (diffusers path), near-lossless, opt-in tiers:
- Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and
  Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks
  compiles and runs ~2.2x faster on the GGUF Z-Image transformer on
  torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the
  block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB
  vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move
  output quality. Gate relaxed; default tier delivers it.
- cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs).
- torch.inference_mode() around the pipeline call (lossless, strictly faster than
  the no_grad diffusers uses internally).

Memory path:
- VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers;
  the balanced (group) tier keeps exact slicing only, so it is now bit-identical to
  the resident image (verified PSNR inf) and slightly faster.
- Group offload adds non_blocking + record_stream on the CUDA stream path to
  overlap each block's H2D copy with compute (lossless; gated on the installed
  diffusers signature so older versions still work).

Native (sd.cpp) path:
- native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a
  near-lossless CUDA win that was previously only added on offload tiers; max also
  -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so
  it is never auto-on. Engine generate() merges it, de-duped against offload flags.

Default profile: a GGUF model with no explicit speed_mode now resolves to the
`default` profile (resolve_speed_mode), since compile's perturbation sits below the
quantisation noise floor and so does not reduce quality versus the dense reference;
out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models
stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is
always honored, so the byte-identical path remains one flag away and is the
regression reference.

Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/
perf_verify.py (the B200 verification above), and diffusion_bench.py gains
--speed-mode so the speed tiers are benchmarkable.

Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore,
GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags +
the engine de-dup, and the sd-cli silent-hang timeout.
2026-06-26 03:18:44 +00:00
..
__init__.py Studio: make code comments and docstrings more succinct (#6029) 2026-06-08 23:07:28 -07:00
_html_to_md.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
anthropic_compat.py Studio: improve OpenAI- and Anthropic-compatible API spec compliance (#6010) 2026-06-09 17:13:25 +02:00
api_monitor.py Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
audio_codecs.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
chat_template_helpers.py Studio: make code comments and docstrings more succinct (#6029) 2026-06-08 23:07:28 -07:00
chat_templates.py Studio: bundle Gemma 4 chat templates (E2B/E4B + larger) and auto-apply to unsloth/gemma-4-*-GGUF (#6245) 2026-06-12 05:49:39 -07:00
defaults.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
diffusion.py Studio diffusion (Phase 7): accuracy-preserving speed pass 2026-06-26 03:18:44 +00:00
diffusion_device.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-25 11:30:41 +00:00
diffusion_families.py Studio diffusion (Phase 7): accuracy-preserving speed pass 2026-06-26 03:18:44 +00:00
diffusion_memory.py Studio diffusion (Phase 7): accuracy-preserving speed pass 2026-06-26 03:18:44 +00:00
diffusion_precision.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-25 14:43:45 +00:00
diffusion_speed.py Studio diffusion (Phase 7): accuracy-preserving speed pass 2026-06-26 03:18:44 +00:00
external_provider.py Studio: ignore unsupported env proxy during Studio startup (#6102) 2026-06-11 05:13:27 -07:00
gpu_arbiter.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-25 03:21:55 +00:00
image_gallery.py [pre-commit.ci] auto fixes from pre-commit.com hooks 2026-06-25 03:21:55 +00:00
inference.py fix(studio): handle multimodal list content in inference text paths (#4383) (#6480) 2026-06-23 01:26:11 -07:00
key_exchange.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
llama_cpp.py Studio: honor stream=false on the GGUF agentic tool path (#6570) (#6618) 2026-06-24 15:37:08 +01:00
llama_http.py Studio: serialize non-streaming responses once and pool the proxy client (#6393) 2026-06-17 22:38:02 -07:00
llama_server_args.py studio: deterministic VRAM auto-fit for GGUF (MTP reserve, compute buffer, total-based budget) (#6312) 2026-06-17 03:10:22 -07:00
llama_stats.py Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
mcp_client.py Studio: enable stdio MCP servers on a loopback bind (#6295) 2026-06-15 03:02:32 +01:00
mcp_config_import.py studio: show MCP "Import config" on the add-server form (#6030) 2026-06-11 16:17:22 +01:00
message_content.py fix(studio): handle multimodal list content in inference text paths (#4383) (#6480) 2026-06-23 01:26:11 -07:00
mlx_inference.py Expose runtime context length for hub models (#6154) 2026-06-11 22:13:53 +03:00
orchestrator.py Studio: persistent per-user trust_remote_code approval cache (#6551) 2026-06-22 05:12:49 -07:00
pricing.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
providers.py Studio: Add custom provider option to Connections (#6112) 2026-06-12 13:09:35 +02:00
runtime_context.py Expose runtime context length for hub models (#6154) 2026-06-11 22:13:53 +03:00
safetensors_agentic.py Studio: show tool-call progress for large GGUF tool arguments (#6484) 2026-06-22 05:50:10 -07:00
sd_cpp_args.py Studio diffusion (Phase 7): accuracy-preserving speed pass 2026-06-26 03:18:44 +00:00
sd_cpp_engine.py Studio diffusion (Phase 7): accuracy-preserving speed pass 2026-06-26 03:18:44 +00:00
tensor_fallback.py studio: deterministic VRAM auto-fit for GGUF (MTP reserve, compute buffer, total-based budget) (#6312) 2026-06-17 03:10:22 -07:00
tool_call_parser.py Fix Gemma 4 GGUF OpenAI API streams (#6476) 2026-06-23 06:13:56 -07:00
tool_loop_controller.py Studio: clean-room compact RAG (knowledge bases, hybrid search, fast indexing) (#5910) 2026-06-09 21:17:04 -07:00
tools.py Use UTF-8 for Python code-execution subprocess I/O (#6489 class) (#6548) 2026-06-22 09:06:03 -07:00
worker.py Generalize transformers tier selection by probing AutoConfig (#6550) 2026-06-22 08:20:06 -07:00