Honor memory_mode over legacy cpu_offload and prefetch dense-quant transformer shards

plan_diffusion_memory only applies the legacy cpu_offload override when no
memory_mode was supplied, matching the documented API contract that
memory_mode overrides cpu_offload when set; an explicit fast request now
stays resident even if the old flag is also enabled.

The transformer-quant dense path fetches the base repo's transformer/
shards inside the locked finalize phase, where unload and cancellation
cannot preempt the multi-GB download. The load worker now widens the
preemptible prefetch to include those shards when that path can actually
run: quant requested and supported for the device, scheme resolvable, and
no pre-quantized checkpoint shortcutting the dense build.
This commit is contained in:
Daniel Han 2026-07-02 03:51:53 +00:00
commit a4277a01e4
4 changed files with 122 additions and 7 deletions

View file

@ -270,6 +270,21 @@ def test_explicit_cpu_offload_overrides_resident_auto_choice():
assert any("explicit cpu_offload" in r for r in plan.reasons)
def test_explicit_memory_mode_wins_over_legacy_cpu_offload():
# The API documents memory_mode as overriding cpu_offload when set: fast +
# the legacy flag must stay resident, not silently downgrade to offload.
plan = plan_diffusion_memory(
target = _target(),
device_memory = _discrete(80000),
model_dense_mib = 4000,
runtime_headroom_mib = 2000,
requested_mode = MEMORY_MODE_FAST,
explicit_offload = True,
)
assert plan.offload_policy == OFFLOAD_NONE
assert not any("explicit cpu_offload" in r for r in plan.reasons)
def test_explicit_cpu_offload_ignored_on_cpu_target():
plan = plan_diffusion_memory(
target = _target(device = "cpu", backend = "cpu", supports_offload = False),