longcat: a cache hit follows the pipeline's device when none is given

encode_prompt's signature defaults device to None. Upstream's own callers always
pass it, so this never fires today, but a cached tuple returned on the CPU would
meet a DiT on the card — and the fallback is one expression, where discovering it
later is a device-mismatch traceback from inside a generation.

Falls back the way _get_t5_prompt_embeds does.
Co-Authored-By: 's avatarClaude Opus 5 <noreply@anthropic.com>
parent bfcd3f22
...@@ -369,6 +369,15 @@ def test_a_lazily_realised_quantised_encoder_is_never_swapped(SVC): ...@@ -369,6 +369,15 @@ def test_a_lazily_realised_quantised_encoder_is_never_swapped(SVC):
assert "_encoder_is_quantised(pipe)" in wrap assert "_encoder_is_quantised(pipe)" in wrap
def test_a_hit_without_an_explicit_device_follows_the_pipeline(SVC):
"""encode_prompt's signature defaults device to None. A cached tuple left on the
CPU would then meet a DiT on the card."""
wrap = SERVICE.read_text(encoding="utf-8")
wrap = wrap[wrap.index("def _wrap_encode_prompt("):]
wrap = wrap[:wrap.index("\n@contextlib")]
assert 'kwargs.get("device") or getattr(pipe, "device", None)' in wrap
def test_the_disk_cache_is_consulted_before_the_encoder_is_touched(SVC): def test_the_disk_cache_is_consulted_before_the_encoder_is_touched(SVC):
"""Order matters: in-memory, then disk, then encode. Realising the encoder before """Order matters: in-memory, then disk, then encode. Realising the encoder before
the disk lookup would spend the 11 GB this feature exists to save.""" the disk lookup would spend the 11 GB this feature exists to save."""
......
...@@ -683,7 +683,11 @@ def _wrap_encode_prompt(pipe, swap: bool, cache=None, lazy=None): ...@@ -683,7 +683,11 @@ def _wrap_encode_prompt(pipe, swap: bool, cache=None, lazy=None):
log(f"prompt cache: cannot key this call ({type(exc).__name__}: {exc})") log(f"prompt cache: cannot key this call ({type(exc).__name__}: {exc})")
disk_key = None disk_key = None
if disk_key is not None: if disk_key is not None:
hit = cache.get(disk_key, device=kwargs.get("device"), # Upstream's callers always pass device, but its signature defaults it to
# None, and a cached tuple left on the CPU would then meet a DiT on the
# card. Fall back the way _get_t5_prompt_embeds would.
hit = cache.get(disk_key,
device=kwargs.get("device") or getattr(pipe, "device", None),
dtype=kwargs.get("dtype")) dtype=kwargs.get("dtype"))
if hit is not None: if hit is not None:
log("prompt cache hit — the text encoder is not needed for this request") log("prompt cache hit — the text encoder is not needed for this request")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment