From 15ff470472d32cbdbefd872f0e185f0c668c6b0e Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 8 Sep 2026 22:37:21 +0800 Subject: [PATCH] [diffusion] chore: key the VAE decode-dtype store by module layout (#38496) Co-authored-by: Mick Qian Co-authored-by: Claude Opus 5 --- .../loader/component_loaders/vae_loader.py | 13 ++++++-- .../test/unit/test_vae_decoder_store.py | 33 ++++++++++++++++--- 2 files changed, 39 insertions(+), 7 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py index 7f3460dc9..ddd427f42 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py @@ -159,10 +159,17 @@ def _should_use_channels_last_3d( def _decode_dtype_store_path( - component_model_path: str, component_name: str, dtype: torch.dtype + component_model_path: str, component_name: str, dtype: torch.dtype, vae=None ) -> str: + # The module layout is part of the key: two runtime versions that expose a + # different parameter set for the same checkpoint must not share a store. + layout = "" + if vae is not None: + layout = "|" + ",".join( + f"{name}:{tuple(tensor.shape)}" for name, tensor in vae.state_dict().items() + ) key = hashlib.sha1( - f"{os.path.realpath(component_model_path)}|{component_name}|{dtype}".encode() + f"{os.path.realpath(component_model_path)}|{component_name}|{dtype}{layout}".encode() ).hexdigest()[:16] return os.path.join( envs.SGLANG_DIFFUSION_CACHE_ROOT, "decode_dtype_store", f"{key}.safetensors" @@ -211,7 +218,7 @@ def _rehome_cast_weights_to_file( Returns (weights held, file-backed?). """ - path = _decode_dtype_store_path(component_model_path, component_name, dtype) + path = _decode_dtype_store_path(component_model_path, component_name, dtype, vae) try: if os.path.exists(path): mapped = _load_store(path) diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_decoder_store.py b/python/sglang/multimodal_gen/test/unit/test_vae_decoder_store.py index 13b3b431c..8a53b8dfa 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_decoder_store.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_decoder_store.py @@ -61,7 +61,7 @@ def test_the_cast_weights_end_up_file_backed(tmp_path): vae, _server_args(), "video_vae", str(model_path) ) - path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16) + path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16, vae) import os assert os.path.exists(path) @@ -101,7 +101,8 @@ def test_a_second_start_adopts_the_store_without_casting(tmp_path): def test_a_mismatched_store_is_discarded_and_the_cast_kept(tmp_path): model_path = tmp_path / "ckpt" model_path.mkdir() - path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16) + vae = _TinyVAE() + path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16, vae) import os os.makedirs(os.path.dirname(path), exist_ok=True) @@ -109,7 +110,6 @@ def test_a_mismatched_store_is_discarded_and_the_cast_kept(tmp_path): save_file({"blocks.0.weight": torch.zeros(4, 4, dtype=torch.float16)}, path) - vae = _TinyVAE() _hold_decoder_weights_in_decode_dtype( vae, _server_args(), "video_vae", str(model_path) ) @@ -127,8 +127,33 @@ def test_the_store_kill_switch_keeps_the_copies_in_memory(monkeypatch, tmp_path) vae, _server_args(), "video_vae", str(model_path) ) - path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16) + path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16, vae) import os assert not os.path.exists(path) assert all(b.weight.dtype == torch.float16 for b in vae.blocks) + + +def test_module_layouts_do_not_share_a_store(tmp_path): + """A runtime exposing a different parameter set must not evict another's store.""" + model_path = tmp_path / "ckpt" + model_path.mkdir() + first = _TinyVAE() + _hold_decoder_weights_in_decode_dtype( + first, _server_args(), "video_vae", str(model_path) + ) + first_path = _decode_dtype_store_path( + str(model_path), "video_vae", torch.float16, first + ) + other = _TinyVAE() + other.blocks.append(nn.Linear(8, 8)) + _hold_decoder_weights_in_decode_dtype( + other, _server_args(), "video_vae", str(model_path) + ) + other_path = _decode_dtype_store_path( + str(model_path), "video_vae", torch.float16, other + ) + import os + + assert first_path != other_path + assert os.path.exists(first_path) and os.path.exists(other_path)