From 2f00c77e75f679e27ef625d1796ad09aa8fd7159 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 25 Jun 2026 12:51:21 +0000 Subject: [PATCH] Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. --- studio/backend/core/inference/diffusion.py | 14 +- .../core/inference/diffusion_memory.py | 154 +++++++++++++----- studio/backend/models/inference.py | 3 +- .../backend/tests/test_diffusion_backend.py | 20 ++- studio/backend/tests/test_diffusion_memory.py | 100 ++++++++++-- 5 files changed, 226 insertions(+), 65 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 75b47657f5..45104bd1c3 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -458,9 +458,12 @@ class DiffusionBackend: plan = self._plan_memory( target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload ) - # apply_memory_plan returns the policy ACTUALLY engaged (it may fall - # back from sequential to whole-module offload), so status stays honest. - effective_policy = apply_memory_plan(pipe, plan, device = device, logger = logger) + # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it + # may fall back to whole-module offload, and tiling is a no-op on a + # pipeline with no tiling control), so status stays honest. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) self._state = _LoadState( pipe = pipe, @@ -471,13 +474,13 @@ class DiffusionBackend: dtype = str(dtype).replace("torch.", ""), cpu_offload = effective_policy != OFFLOAD_NONE, offload_policy = effective_policy, - vae_tiling = plan.vae_tiling, + vae_tiling = effective_tiling, memory_mode = plan.requested_mode, ) logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", - repo_id, base, device, effective_policy, plan.vae_tiling, "; ".join(plan.reasons), + repo_id, base, device, effective_policy, effective_tiling, "; ".join(plan.reasons), ) return self.status() @@ -511,6 +514,7 @@ class DiffusionBackend: target = target, device_memory = device_memory, model_dense_mib = model_dense_mib, + companion_dense_mib = companion_mib, runtime_headroom_mib = runtime_headroom, requested_mode = memory_mode, explicit_offload = cpu_offload, diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index cbea98a7a8..665c21fb9e 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -36,15 +36,23 @@ MEMORY_MODES = ( ) # ── offload policies (what the loader actually does) ────────────────────────── -# none -> all weights resident on the device (fastest; fits only if there -# is room). model -> diffusers enable_model_cpu_offload(): one -# top-level module on the GPU at a time (modest VRAM cut, small -# speed cost). sequential -> enable_sequential_cpu_offload(): -# submodule-level streaming (lowest VRAM, large speed cost). +# none -> all weights resident on the device (fastest; fits only if there is +# room). model -> diffusers enable_model_cpu_offload(): one top-level +# module on the GPU at a time (modest VRAM cut, small speed cost). +# group -> apply_group_offloading() on the transformer: stream it a few blocks at +# a time with a prefetch stream (lowest practical VRAM for the dominant +# module, less penalty than submodule sequential). sequential -> +# enable_sequential_cpu_offload(): submodule-level (broken for GGUF on +# diffusers 0.38, kept only as an explicit escape hatch). OFFLOAD_NONE = "none" OFFLOAD_MODEL = "model" +OFFLOAD_GROUP = "group" OFFLOAD_SEQUENTIAL = "sequential" +# Blocks of the transformer kept resident per group under group offloading: fewer +# = lower peak VRAM, more host<->device traffic. One is the lowest-VRAM setting. +DEFAULT_GROUP_BLOCKS = 1 + DEFAULT_IMAGE_WIDTH = 1024 DEFAULT_IMAGE_HEIGHT = 1024 # A flat allowance for the pipeline's fixed costs (scheduler, embeddings, the @@ -312,29 +320,47 @@ def plan_diffusion_memory( device_memory: DeviceMemory, model_dense_mib: Optional[int], runtime_headroom_mib: int, + companion_dense_mib: Optional[int] = None, base_overhead_mib: int = DEFAULT_BASE_OVERHEAD_MIB, requested_mode: Optional[str] = None, explicit_offload: bool = False, ) -> MemoryPlan: """Pick an offload policy + VAE memory savers for the current load. - ``model_dense_mib`` is the estimated resident device size of the weights - (transformer + companion text-encoder / VAE). ``explicit_offload`` is the - back-compat ``cpu_offload=True`` request: it forces at least whole-module - offload wherever offload is meaningful.""" + ``model_dense_mib`` is the estimated resident device size of all weights + (transformer + companion text-encoder / VAE); ``companion_dense_mib`` is just + the companions, which stay resident under streamed (group) offload while the + transformer is streamed block by block. ``explicit_offload`` is the back-compat + ``cpu_offload=True`` request: it forces whole-module offload. + + Policy meanings, ordered by measured speed/VRAM tradeoff: + none - everything resident: fastest, highest VRAM. + group - stream the transformer, companions resident: near-resident speed, + moderate VRAM cut (the balanced tradeoff). + model - offload every component incl. the text encoder: lowest VRAM, slow. + """ mode = normalize_memory_mode(requested_mode) or MEMORY_MODE_AUTO can_offload = bool(getattr(target, "supports_model_cpu_offload", False)) budget = _safe_device_budget_mib(device_memory) required = _sum_required(model_dense_mib, runtime_headroom_mib, base_overhead_mib) + # The resident floor under group offload: companions stay, the transformer streams. + group_floor = _sum_required(companion_dense_mib, runtime_headroom_mib, base_overhead_mib) reasons: list[str] = [] estimates: dict[str, Optional[int]] = { "safe_device_budget_mib": budget, "model_dense_mib": model_dense_mib, + "companion_dense_mib": companion_dense_mib, "runtime_headroom_mib": runtime_headroom_mib, "base_overhead_mib": base_overhead_mib, "resident_required_mib": required, + "group_floor_mib": group_floor, } + def _group_fits() -> bool: + # Group offload only helps if the resident remainder (companions) fits; when + # the text encoder itself is too big, only whole-module offload will do. + return group_floor is not None and budget is not None and group_floor <= budget + if not can_offload or device_memory.is_unified: # MPS / CPU can't stream to a separate device, and on unified / system # memory CPU offload just shuffles bytes within the same pool. @@ -346,28 +372,29 @@ def plan_diffusion_memory( elif mode == MEMORY_MODE_FAST: policy = OFFLOAD_NONE if budget is not None and required is not None and required > budget: - policy = OFFLOAD_MODEL - reasons.append("fast requested but weights do not fit resident; whole-module offload") + # Doesn't fit resident: the fastest offload is the streamed transformer. + policy = OFFLOAD_GROUP if _group_fits() else OFFLOAD_MODEL + reasons.append("fast requested but weights do not fit resident; offloading") else: reasons.append("fast requested; weights resident on device") elif mode == MEMORY_MODE_BALANCED: - policy = OFFLOAD_MODEL - reasons.append("balanced requested; whole-module CPU offload") + policy = OFFLOAD_GROUP + reasons.append("balanced requested; streamed block-level transformer offload") elif mode == MEMORY_MODE_LOW_VRAM: - policy = OFFLOAD_SEQUENTIAL - reasons.append("low_vram requested; submodule sequential CPU offload") + policy = OFFLOAD_MODEL + reasons.append("low_vram requested; whole-module offload of every component") elif budget is None or required is None: policy = OFFLOAD_NONE reasons.append("device budget or model size unknown; staying resident") elif required <= int(budget * 0.85): policy = OFFLOAD_NONE reasons.append("weights fit resident with headroom") - elif required <= budget: - policy = OFFLOAD_MODEL - reasons.append("tight fit; whole-module CPU offload") + elif _group_fits(): + policy = OFFLOAD_GROUP + reasons.append("tight fit; stream the transformer, companions resident") else: - policy = OFFLOAD_SEQUENTIAL - reasons.append("weights exceed device budget; submodule sequential CPU offload") + policy = OFFLOAD_MODEL + reasons.append("companions exceed budget; whole-module offload of every component") if explicit_offload and policy == OFFLOAD_NONE and can_offload and not device_memory.is_unified: policy = OFFLOAD_MODEL @@ -393,24 +420,31 @@ def plan_diffusion_memory( # ── apply to a built pipeline ───────────────────────────────────────────────── -def apply_memory_plan(pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = None) -> str: +def apply_memory_plan( + pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = None, +) -> tuple[str, bool]: """Apply ``plan`` to a freshly built diffusers pipeline: enable the VAE memory - savers (best-effort; not every pipeline exposes them) then place / offload the - weights. Exactly one placement call runs, so the pipeline ends up either fully - resident or wired for CPU offload, never both. + savers then place / offload the weights. Exactly one placement call runs, so the + pipeline ends up either fully resident or wired for offload, never both. - Returns the offload policy actually engaged, which can differ from the plan: - submodule sequential offload is unreliable for GGUF-dequantised transformers on - some diffusers versions, so it falls back to the robust whole-module offload - (the proper low-VRAM GGUF path is the CPU-resident cache, a later phase).""" + Returns ``(offload_policy, vae_tiling)`` ACTUALLY engaged, which can differ from + the plan: VAE tiling is a no-op on a pipeline that exposes no tiling control, and + block-level / sequential offload fall back to the robust whole-module offload if + the transformer doesn't support them (e.g. submodule sequential is broken for + GGUF on diffusers 0.38). Status then reflects what really happened.""" + tiling_engaged = False if plan.vae_tiling: - _try_call(pipe, "enable_vae_tiling", logger) + tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger) if plan.vae_slicing: - _try_call(pipe, "enable_vae_slicing", logger) + _enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger) policy = plan.offload_policy if policy == OFFLOAD_MODEL: pipe.enable_model_cpu_offload() + elif policy == OFFLOAD_GROUP: + if not _apply_group_offload(pipe, device, logger): + pipe.enable_model_cpu_offload() + policy = OFFLOAD_MODEL elif policy == OFFLOAD_SEQUENTIAL: try: pipe.enable_sequential_cpu_offload() @@ -424,15 +458,59 @@ def apply_memory_plan(pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = policy = OFFLOAD_MODEL else: pipe.to(device) - return policy + return policy, tiling_engaged -def _try_call(pipe: Any, method: str, logger: Any) -> None: - fn = getattr(pipe, method, None) - if not callable(fn): - return +def _enable_vae_saver(pipe: Any, pipe_method: str, vae_method: str, logger: Any) -> bool: + """Turn on a VAE memory saver, trying the pipeline-level shortcut first and the + VAE submodule directly otherwise (some pipelines, e.g. Z-Image, only expose it + on ``pipe.vae``). Returns whether it actually engaged.""" + for owner, method in ((pipe, pipe_method), (getattr(pipe, "vae", None), vae_method)): + fn = getattr(owner, method, None) + if not callable(fn): + continue + try: + fn() + return True + except Exception as exc: # noqa: BLE001 — a VAE saver is an optimisation, never fatal + if logger is not None: + logger.warning("diffusion.memory: %s() failed: %s", method, exc) + return False + + +def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: + """Stream the transformer through the device a few blocks at a time via + diffusers group offloading, keeping the smaller components resident. Returns + False (so the caller falls back to whole-module offload) on any failure.""" + transformer = getattr(pipe, "transformer", None) + if transformer is None: + return False try: - fn() - except Exception as exc: # noqa: BLE001 — a VAE saver is an optimisation, never fatal + import torch + from diffusers.hooks import apply_group_offloading + + onload = torch.device(device) + use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA + apply_group_offloading( + transformer, + onload_device = onload, + offload_device = torch.device("cpu"), + offload_type = "block_level", + num_blocks_per_group = DEFAULT_GROUP_BLOCKS, + use_stream = use_stream, + ) + # Place the remaining (smaller) components resident; the streamed + # transformer manages its own placement via the offloading hooks. + for name, comp in getattr(pipe, "components", {}).items(): + if name == "transformer": + continue + if isinstance(comp, torch.nn.Module): + comp.to(onload) + return True + except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if logger is not None: - logger.warning("diffusion.memory: %s() failed: %s", method, exc) + logger.warning( + "diffusion.memory: group offload failed (%s); falling back to " + "whole-module offload", exc, + ) + return False diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index d0b5daec2a..f30f8a31f6 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1704,7 +1704,8 @@ class DiffusionLoadRequest(BaseModel): memory_mode: Optional[Literal["auto", "fast", "balanced", "low_vram"]] = Field( None, description = "Memory policy: auto (measured), fast (resident), balanced " - "(model CPU offload), low_vram (sequential CPU offload). " + "(stream the transformer, near-resident speed, moderate VRAM " + "cut), low_vram (offload every component, lowest VRAM, slower). " "Overrides cpu_offload when set.", ) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index f06778466b..57be08fb4d 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -798,28 +798,34 @@ def _force_cuda_target(backend, monkeypatch): monkeypatch.setattr(backend, "_pick_device_and_dtype", lambda: ("cuda", torch.bfloat16)) -def test_load_memory_mode_balanced_engages_model_offload(fake_runtime, tmp_path, monkeypatch): +def test_load_memory_mode_balanced_streams_or_falls_back(fake_runtime, tmp_path, monkeypatch): + # balanced requests streamed block-level (group) offload. Under the stub there is + # no real diffusers.hooks, so group can't engage and the applier falls back to + # whole-module offload, reporting the policy actually engaged (the real "group" + # path is GPU-verified in the bench). (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) status = backend.load_pipeline( str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", memory_mode = "balanced" ) - assert status["offload_policy"] == "model" and status["cpu_offload"] is True + assert status["offload_policy"] in ("group", "model") and status["cpu_offload"] is True assert status["memory_mode"] == "balanced" - pipe = backend._state.pipe - assert pipe.offloaded is True and pipe.moved_to is None # offload owns placement + assert backend._state.pipe.offloaded is True # model-offload fallback engaged -def test_load_memory_mode_low_vram_uses_sequential_offload(fake_runtime, tmp_path, monkeypatch): +def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path, monkeypatch): + # low_vram offloads every component (lowest VRAM); whole-module offload is the + # robust path and engages directly (no streaming, so no diffusers.hooks needed). (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() _force_cuda_target(backend, monkeypatch) status = backend.load_pipeline( str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", memory_mode = "low_vram" ) - assert status["offload_policy"] == "sequential" and status["cpu_offload"] is True - assert backend._state.pipe.sequential_offloaded is True + assert status["offload_policy"] == "model" and status["cpu_offload"] is True + pipe = backend._state.pipe + assert pipe.offloaded is True and pipe.moved_to is None # offload owns placement def test_load_explicit_cpu_offload_engages_model_offload_on_cuda(fake_runtime, tmp_path, monkeypatch): diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 505b306623..357425c6c5 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -20,10 +20,12 @@ from core.inference.diffusion_memory import ( MEMORY_MODE_BALANCED, MEMORY_MODE_FAST, MEMORY_MODE_LOW_VRAM, + OFFLOAD_GROUP, OFFLOAD_MODEL, OFFLOAD_NONE, OFFLOAD_SEQUENTIAL, DeviceMemory, + MemoryPlan, apply_memory_plan, estimate_gguf_dense_mib, estimate_image_runtime_mib, @@ -166,14 +168,41 @@ def test_auto_model_offload_on_tight_fit(): assert plan.vae_tiling is True # offloading -> device is tight -> tile -def test_auto_sequential_when_model_exceeds_budget(): +def test_auto_group_offload_when_transformer_overflows_but_companions_fit(): + # Big transformer pushes the resident total over budget, but the companions + # (text encoder + VAE) still fit -> stream the transformer (fast, moderate cut). + plan = plan_diffusion_memory( + target = _target(), + device_memory = _discrete(8000, 8000), + model_dense_mib = 40000, + companion_dense_mib = 1500, + runtime_headroom_mib = 1000, + base_overhead_mib = 1000, + ) + assert plan.offload_policy == OFFLOAD_GROUP + + +def test_auto_model_offload_when_companions_exceed_budget(): + # The text encoder itself is too big to stay resident -> offload everything. + plan = plan_diffusion_memory( + target = _target(), + device_memory = _discrete(8000, 8000), + model_dense_mib = 40000, + companion_dense_mib = 30000, + runtime_headroom_mib = 4000, + ) + assert plan.offload_policy == OFFLOAD_MODEL + + +def test_auto_model_offload_when_companion_size_unknown(): + # Without a companion estimate the planner can't prove group fits -> safest cut. plan = plan_diffusion_memory( target = _target(), device_memory = _discrete(8000, 8000), model_dense_mib = 40000, runtime_headroom_mib = 4000, ) - assert plan.offload_policy == OFFLOAD_SEQUENTIAL + assert plan.offload_policy == OFFLOAD_MODEL def test_auto_stays_resident_when_budget_unknown(): @@ -199,11 +228,11 @@ def test_explicit_modes_force_policy_regardless_of_budget(): assert plan_diffusion_memory( target = _target(), device_memory = roomy, model_dense_mib = 1000, runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_BALANCED, - ).offload_policy == OFFLOAD_MODEL + ).offload_policy == OFFLOAD_GROUP assert plan_diffusion_memory( target = _target(), device_memory = roomy, model_dense_mib = 1000, runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_LOW_VRAM, - ).offload_policy == OFFLOAD_SEQUENTIAL + ).offload_policy == OFFLOAD_MODEL def test_fast_falls_back_to_model_offload_when_it_does_not_fit(): @@ -313,31 +342,74 @@ def _plan(policy, *, tiling): runtime_headroom_mib = 1000, requested_mode = { OFFLOAD_NONE: MEMORY_MODE_FAST, - OFFLOAD_MODEL: MEMORY_MODE_BALANCED, - OFFLOAD_SEQUENTIAL: MEMORY_MODE_LOW_VRAM, + OFFLOAD_GROUP: MEMORY_MODE_BALANCED, + OFFLOAD_MODEL: MEMORY_MODE_LOW_VRAM, }[policy], ) +def _manual_plan(policy, *, tiling): + """Build a plan for a policy the auto/explicit modes no longer emit (sequential).""" + return MemoryPlan( + requested_mode = "manual", + offload_policy = policy, + vae_tiling = tiling, + vae_slicing = tiling, + device_memory = _discrete(4000, 8000), + estimates = {}, + ) + + def test_apply_none_places_resident(): pipe = _RecordingPipe() - effective = apply_memory_plan(pipe, _plan(OFFLOAD_NONE, tiling = False), device = "cuda") + effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_NONE, tiling = False), device = "cuda") assert pipe.calls == ["to:cuda"] # no tiling on a roomy resident run - assert effective == OFFLOAD_NONE + assert effective == OFFLOAD_NONE and tiled is False def test_apply_model_offload_engages_offload_and_tiling(): pipe = _RecordingPipe() - effective = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda") + effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda") assert "model_offload" in pipe.calls assert "to:cuda" not in pipe.calls # offload owns placement; never both assert "vae_tiling" in pipe.calls and "vae_slicing" in pipe.calls - assert effective == OFFLOAD_MODEL + assert effective == OFFLOAD_MODEL and tiled is True + + +def test_apply_vae_tiling_falls_back_to_vae_submodule(): + # Z-Image-style pipeline: no pipeline-level enable_vae_tiling, only pipe.vae. + class _VaeOnly: + def __init__(self): + self.vae = types.SimpleNamespace( + tiled = False, sliced = False, + enable_tiling = self._tile, enable_slicing = self._slice, + ) + + def _tile(self): + self.vae.tiled = True + + def _slice(self): + self.vae.sliced = True + + def enable_model_cpu_offload(self): + self.offloaded = True + + pipe = _VaeOnly() + effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_MODEL, tiling = True), device = "cuda") + assert tiled is True and pipe.vae.tiled and pipe.vae.sliced + + +def test_apply_group_falls_back_to_model_without_transformer(): + # The recording pipe has no .transformer, so group offload can't engage and the + # applier falls back to whole-module offload, reporting the real policy. + pipe = _RecordingPipe() + effective, _ = apply_memory_plan(pipe, _plan(OFFLOAD_GROUP, tiling = True), device = "cuda") + assert effective == OFFLOAD_MODEL and "model_offload" in pipe.calls def test_apply_sequential_offload(): pipe = _RecordingPipe() - effective = apply_memory_plan(pipe, _plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda") + effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda") assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls assert effective == OFFLOAD_SEQUENTIAL @@ -350,7 +422,7 @@ def test_apply_sequential_falls_back_to_model_offload_when_unsupported(): raise RuntimeError("sequential offload not supported for this transformer") pipe = _NoSeqPipe() - effective = apply_memory_plan(pipe, _plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda") + effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda") assert effective == OFFLOAD_MODEL assert "model_offload" in pipe.calls @@ -365,5 +437,5 @@ def test_apply_tolerates_pipe_without_vae_savers(): self.moved = device bare = _Bare() - apply_memory_plan(bare, _plan(OFFLOAD_NONE, tiling = False), device = "cpu") - assert bare.moved == "cpu" + _, tiled = apply_memory_plan(bare, _plan(OFFLOAD_NONE, tiling = False), device = "cpu") + assert bare.moved == "cpu" and tiled is False