Budget-Constrained Step-Level Diffusion Caching

[remyx-ai]

Drafted by an autonomous discovery loop — Remyx ranks recent arXiv papers against this team's research interest and shipping history; Claude Code selects the candidate most directly implementable against this repo from the lookback window and drafts it.

Recommended paper: Budget-Constrained Step-Level Diffusion Caching
Confidence: 🟢 high (Remyx relevance 0.92)
Research interest: diffusers
Implementation by: Claude Code as autonomous agent

---

Why this paper for this team

The diffusers team's work on 'Add LLaDA2 Discrete Diffusion Pipeline' and broad support for diffusion models makes inference efficiency a continuous priority. This paper directly bridges with the 'Model Efficiency' research interest by introducing BudCache, a method that optimizes caching policies to preserve output quality under a fixed compute budget. This is highly relevant for diffusers as it enables granular control over inference cost for iterative denoising processes, ensuring optimal performance for large models like Ideogram 4 or Cosmos 3 without sacrificing quality. It complements the existing performance and quantization efforts by offering a principled approach to resource allocation.

Why this candidate (selected from the lookback pool)

BudCache's contribution (a budget-aware step-level caching schedule that decides per-denoising-step whether to reuse cached activations) maps directly onto the existing caching-hook contract: every sibling hook (FasterCache, MagCache, TaylorSeerCache) already implements exactly this skip-or-compute decision over transformer blocks via ModelHook, so BudCache adds a new XxxConfig+apply_xxx pair wired into hooks/init.py with no contract changes. It is the only caching-family candidate with real, permissively-licensed (Apache-2.0) reference code, making it a clean drop-in addition where ReCache [10] (license missing, compat 0.00) and the no-code budget-cache abstracts cannot anchor an implementation. The repo's 'Deepening Focus on Performance and Quantization' direction makes inference-cost control directly relevant.

License & code availability

🟢 Permissive license — safe to adopt.

• Code: https://github.com/Westlake-AGI-Lab/BudCache
• License: Apache-2.0 (class: permissive, compat: 1.00, source: github)

Suggested experiment

Implement BudCache's Simulated Annealing with Hill Climbing for step selection within a diffusion pipeline in diffusers (e.g., FLUX.1-dev or a widely used DiT). Evaluate the generation quality (e.g., LPIPS) under varying fixed compute budgets, comparing its performance against existing heuristic caching baselines provided by the library.

---

What this PR delivers

Call site: Public package API: exported via src/diffusers/init.pyandsrc/diffusers/hooks/init.pyasapply_budget_cache/BudgetCacheConfig, the same surface through which sibling caching hooks (apply_faster_cache, apply_first_block_cache, etc.) are reached.

Delivers (from the paper):

• Adds BudCache's budget-constrained step-level caching to the diffusers caching-hook family: BudgetCacheConfig + apply_budget_cache wired into the public package namespace (diffusers.__init__, hooks/__init__.py, dummy objects), matching the existing FasterCache/TaylorSeerCache config+apply pattern.
• Implements the paper's core inversion of threshold caching: search_cache_schedule fixes a compute_budget in advance and searches for the compute/skip schedule that best preserves output, combining Simulated Annealing (budget-preserving swaps under a cooling temperature) with deterministic Hill Climbing local refinement over a surrogate trajectory-error cost — deterministic given a seed, run fully offline so inference has no thresholding/search overhead.
• Provides the runtime replay machinery via ModelHooks (head/tail/block hooks + BudgetCacheState): on scheduled-compute steps blocks run normally and the tail caches the input→output residual; on skipped steps the head reuses the cached residual (with shape-alignment fallbacks), wired across a module's transformer blocks by apply_budget_cache.
• Ships a _uniform_schedule budget-respecting heuristic baseline used when no step_errors profile is supplied, plus config validation (budget range, first-step-always-computes, schedule length).
• Adds unit tests covering schedule search (budget/first-step invariants, beating the uniform heuristic), config validation, and end-to-end hook replay (cached reuse vs. full compute) on a dummy transformer.

Intentionally out of scope (not needed for this contribution):

• Cache-aware schedule alignment (adapting the time discretization to the chosen cache policy under very tight budgets) — the paper's secondary contribution — is not implemented; it would require coupling the schedule into the scheduler's timestep spacing.
• Automatic derivation of the per-step error profile (step_errors) from a real model/calibration run is left to the caller; the search consumes a user-provided profile (e.g. MagCache-style residual ratios) or falls back to the uniform heuristic.
• No integration into any specific pipeline (FLUX.1-dev / Wan2.1) and no quality evaluation (LPIPS under fixed budgets vs. heuristic baselines) — the paper's experimental validation — since that needs model weights and a benchmark harness the repo doesn't run in-tree.
• The surrogate _schedule_cost is a staleness-weighted heuristic rather than the paper's exact trajectory-error objective (offline search infra to evaluate true final-output error is out of scope).

This delivers BudCache's central result — a budget-first (rather than threshold-first) step-level cache policy found by an offline SA+Hill-Climbing search and replayed at inference with zero online overhead — packaged as a first-class public diffusers caching hook (BudgetCacheConfig/apply_budget_cache) that mirrors the existing apply_faster_cache-style API exactly. Like every sibling caching hook, it is user-invoked through the public package namespace rather than auto-called by a pipeline, so no pre-existing pipeline auto-exercises it, but it is properly wired into the product's public optimization surface (not just reachable from tests). Intentionally scoped out as unnecessary for that core value: the paper's tight-budget schedule alignment, automatic step-error calibration, and the FLUX/Wan quality benchmarks — each of which needs model weights or scheduler-coupling infra not hosted in-tree.

Test results

ℹ️ Tests could not run in CI — the runner lacks this repo's dependencies (a collection/import error, not a code failure). Run the suite locally to validate.

ages have valid Python names.
Traceback:
/opt/hostedtoolcache/Python/3.11.15/x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
examples/controlnet/test_controlnet.py:23: in <module>
    from test_examples_utils import ExamplesTestsAccelerate, run_command  # noqa: E402
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
examples/test_examples_utils.py:23: in <module>
    from accelerate.utils import write_basic_config
E   ModuleNotFoundError: No module named 'accelerate'
=========================== short test summary info ============================
ERROR examples/advanced_diffusion_training/test_dreambooth_lora_flux_advanced.py
ERROR examples/consistency_distillation/test_lcm_lora.py
ERROR examples/controlnet/test_controlnet.py
!!!!!!!!!!!!!!!!!!!!!!!!!! stopping after 3 failures !!!!!!!!!!!!!!!!!!!!!!!!!!!
3 errors in 1.16s

---

Want eval-on-every-PR? Outrider Validate (coming soon, paid tier) runs your benchmark suite against this diff and posts the results as a PR comment. Design partner pilot is open — join the waitlist.

Opened by the Remyx Recommendation orchestrator.