diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index f354c388c..a3d4a08f3 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: SGLANG_DIFFUSION_LOGGING_PREFIX: str = "" SGLANG_DIFFUSION_TRACE_FUNCTION: int = 0 SGLANG_DIFFUSION_DISABLE_EARLY_VAE_DECODER_CAST: bool = False + SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE: bool = False SGLANG_DIFFUSION_WORKER_MULTIPROC_METHOD: str = "fork" SGLANG_DIFFUSION_TARGET_DEVICE: str = "cuda" SGLANG_DIFFUSION_PLATFORM_OVERRIDE: str = "" @@ -272,6 +273,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "SGLANG_DIFFUSION_DISABLE_EARLY_VAE_DECODER_CAST": _lazy_bool( "SGLANG_DIFFUSION_DISABLE_EARLY_VAE_DECODER_CAST" ), + # Kill-switch: keep the decode-dtype VAE decoder weights in anonymous host + # memory instead of a file-backed cache mapping the page cache can drop. + "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE": _lazy_bool( + "SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE" + ), # ================== cache-dit Env Vars ================== # Enable cache-dit acceleration for DiT inference # CUDA-IPC transport for 2-rank Ulysses all-to-all (NVLink same-node) 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 8c752de9b..c1f92703e 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 @@ -1,9 +1,11 @@ +import hashlib import importlib.util import os import torch import torch.nn as nn from safetensors.torch import load_file as safetensors_load_file +from safetensors.torch import save_file as safetensors_save_file from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig @@ -131,8 +133,85 @@ def _should_use_channels_last_3d( return False +def _decode_dtype_store_path( + component_model_path: str, component_name: str, dtype: torch.dtype +) -> str: + key = hashlib.sha1( + f"{os.path.realpath(component_model_path)}|{component_name}|{dtype}".encode() + ).hexdigest()[:16] + return os.path.join( + envs.SGLANG_DIFFUSION_CACHE_ROOT, "decode_dtype_store", f"{key}.safetensors" + ) + + +def _assign_matching_store(vae, mapped: dict, dtype: torch.dtype) -> bool: + """Adopt a decode-dtype store if it matches the module, else refuse.""" + state = vae.state_dict() + for name, tensor in mapped.items(): + param = state.get(name) + if param is None or param.shape != tensor.shape or tensor.dtype != dtype: + return False + vae.load_state_dict(mapped, strict=False, assign=True) + return True + + +def _rehome_cast_weights_to_file( + vae, dtype: torch.dtype, component_model_path: str, component_name: str, prepare +) -> tuple[int, bool]: + """Hold the decode-dtype weights in a file-backed mapping. + + The cast copies are anonymous host memory the kernel cannot reclaim, and + on a budgeted host every one of those bytes comes out of the pin budget + the stepped components live on. Written once to a cache file and mapped + back, the same bytes become page cache: droppable under pressure, free to + re-fault, and absent from the anonymous accounting. safetensors round-trips + tensors byte-exactly, so the mapping holds the identical rounded values — + and a later start adopts the store without paying the cast at all. + + Returns (weights held, file-backed?). + """ + path = _decode_dtype_store_path(component_model_path, component_name, dtype) + try: + if os.path.exists(path): + mapped = safetensors_load_file(path) + if mapped and _assign_matching_store(vae, mapped, dtype): + return len(mapped), True + raise ValueError("existing decode-dtype store does not match the module") + converted = prepare(dtype) + if not converted: + return 0, False + cast_state = { + name: tensor + for name, tensor in vae.state_dict().items() + if tensor.dtype == dtype and tensor.device.type == "cpu" + } + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = f"{path}.tmp.{os.getpid()}" + safetensors_save_file({k: v.contiguous() for k, v in cast_state.items()}, tmp) + os.replace(tmp, path) + mapped = safetensors_load_file(path) + if set(mapped) != set(cast_state): + raise ValueError("decode-dtype store does not match the cast weights") + vae.load_state_dict(mapped, strict=False, assign=True) + return converted, True + except Exception as exc: + logger.warning( + "VAE: could not re-home %s decode-dtype weights to %s (%s); " + "keeping in-memory copies", + component_name, + path, + exc, + ) + try: + if os.path.exists(path): + os.remove(path) + except OSError: + pass + return prepare(dtype), False + + def _hold_decoder_weights_in_decode_dtype( - vae, server_args: ServerArgs, component_name: str + vae, server_args: ServerArgs, component_name: str, component_model_path: str = "" ) -> None: """Round decoder weights to their decode compute dtype at load. @@ -156,13 +235,19 @@ def _hold_decoder_weights_in_decode_dtype( return if not autocast_enabled(dtype, server_args.disable_autocast): return - converted = prepare(dtype) - if converted: + if component_model_path and not envs.SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE: + held, file_backed = _rehome_cast_weights_to_file( + vae, dtype, component_model_path, component_name, prepare + ) + else: + held, file_backed = prepare(dtype), False + if held: logger.info( - "VAE: %s holds %d decoder weights in %s from load", + "VAE: %s holds %d decoder weights in %s from load (%s)", component_name, - converted, + held, dtype, + "file-backed" if file_backed else "anonymous host memory", ) @@ -275,7 +360,9 @@ class VAELoader(ComponentLoader): logger.info( "VAE: converted %d Conv3d weights to channels_last_3d", n ) - _hold_decoder_weights_in_decode_dtype(vae, server_args, component_name) + _hold_decoder_weights_in_decode_dtype( + vae, server_args, component_name, component_model_path + ) vae = current_platform.optimize_vae(vae) return vae @@ -349,6 +436,8 @@ class VAELoader(ComponentLoader): if n > 0: logger.info("VAE: converted %d Conv3d weights to channels_last_3d", n) - _hold_decoder_weights_in_decode_dtype(vae, server_args, component_name) + _hold_decoder_weights_in_decode_dtype( + vae, server_args, component_name, component_model_path + ) vae = current_platform.optimize_vae(vae) return vae 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 new file mode 100644 index 000000000..543d5a1a9 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_vae_decoder_store.py @@ -0,0 +1,133 @@ +"""The decode-dtype weights live in a file-backed mapping, not anonymous memory. + +What matters: the store round-trips the exact rounded bytes, a later start +adopts it without paying the cast, a mismatched store is discarded rather than +adopted, and the kill switch keeps everything in memory. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.loader.component_loaders.vae_loader import ( + _decode_dtype_store_path, + _hold_decoder_weights_in_decode_dtype, +) + + +class _TinyVAE(nn.Module): + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([nn.Linear(8, 8) for _ in range(3)]) + self.head = nn.Linear(8, 8) # stays fp32, like the output projection + self.prepare_calls = 0 + + def prepare_decoder_autocast_weights(self, dtype) -> int: + self.prepare_calls += 1 + converted = 0 + for block in self.blocks: + if block.weight.dtype != dtype: + block.to(dtype=dtype) + converted += 1 + return converted + + +def _server_args(): + return SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_decode_precision="fp16", + vae_precision="fp32", + ), + disable_autocast=False, + ) + + +@pytest.fixture(autouse=True) +def _env(monkeypatch, tmp_path): + from sglang.multimodal_gen.runtime.utils import precision + + monkeypatch.setattr(precision.current_platform, "is_amp_supported", lambda: True) + monkeypatch.setenv("SGLANG_DIFFUSION_CACHE_ROOT", str(tmp_path / "cache")) + + +def test_the_cast_weights_end_up_file_backed(tmp_path): + vae = _TinyVAE() + model_path = tmp_path / "ckpt" + model_path.mkdir() + _hold_decoder_weights_in_decode_dtype( + vae, _server_args(), "video_vae", str(model_path) + ) + + path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16) + import os + + assert os.path.exists(path) + assert all(b.weight.dtype == torch.float16 for b in vae.blocks) + assert vae.head.weight.dtype == torch.float32 + + from safetensors.torch import load_file + + stored = load_file(path) + for name, tensor in stored.items(): + assert torch.equal(tensor, vae.state_dict()[name]) + + +def test_a_second_start_adopts_the_store_without_casting(tmp_path): + model_path = tmp_path / "ckpt" + model_path.mkdir() + first = _TinyVAE() + _hold_decoder_weights_in_decode_dtype( + first, _server_args(), "video_vae", str(model_path) + ) + + second = _TinyVAE() + second.load_state_dict( + { + k: v.to(torch.float32) if v.dtype == torch.float16 else v + for k, v in first.state_dict().items() + } + ) + _hold_decoder_weights_in_decode_dtype( + second, _server_args(), "video_vae", str(model_path) + ) + assert second.prepare_calls == 0 + for name in first.state_dict(): + assert torch.equal(second.state_dict()[name], first.state_dict()[name]) + + +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) + import os + + os.makedirs(os.path.dirname(path), exist_ok=True) + from safetensors.torch import save_file + + 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) + ) + assert vae.prepare_calls == 1 + assert all(b.weight.dtype == torch.float16 for b in vae.blocks) + assert not os.path.exists(path) + + +def test_the_store_kill_switch_keeps_the_copies_in_memory(monkeypatch, tmp_path): + monkeypatch.setenv("SGLANG_DIFFUSION_DISABLE_VAE_DECODER_STORE", "1") + model_path = tmp_path / "ckpt" + model_path.mkdir() + vae = _TinyVAE() + _hold_decoder_weights_in_decode_dtype( + vae, _server_args(), "video_vae", str(model_path) + ) + + path = _decode_dtype_store_path(str(model_path), "video_vae", torch.float16) + import os + + assert not os.path.exists(path) + assert all(b.weight.dtype == torch.float16 for b in vae.blocks)