minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[diffusion] chore: key the VAE decode-dtype store by module layout (#38496)

Browse Source 
Co-authored-by: Mick Qian <mickqian@radixark.ai>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Mick
2026-09-08 22:37:21 +08:00
committed by GitHub
co-authored by Mick Qian Claude Opus 5
parent a6b542813f
commit 15ff470472
2 changed files with 39 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+10 -3
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_vae_decoder_store.py
+29 -4
View File
@@ -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)