diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 5ceac76e47a3..3d5322d8e693 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -297,6 +297,34 @@ values. SGLang rejects a request outside that coverage instead of silently changing conditioning. Cache mode supports the matching unquantized checkpoint only. +### Advanced: online AdaLN rebuild with a host cache + +`--minimax-h3-adaln-online true` needs no prebuilt artifact: the server drops +the 24.2 GiB of `adaln_proj` weights from the GPU and computes each request's +AdaLN outputs from the checkpoint on demand, bit-exact with the resident-weight +path. Because a rebuild pass streams the whole 24.2 GiB, the first request of +every new `(task shape, num_inference_steps, flow_shift, audio_flow_shift)` +combination pays several seconds; after that the plans are served from a +64-slot GPU slab (LRU per plan) backed by a pinned host cache +(`--minimax-h3-adaln-host-cache-gb`, LRU per schedule, default 8 GB per +rank), so mixed-schedule serving does not re-read the checkpoint. Plan sets +that exceed the host budget are simply recomputed on their next occurrence. +Expert escape hatches live in environment variables: +`SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS` resizes the GPU slab (needed +only beyond 65 inference steps) and `SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32` +computes the one-time projections in fp32 (experimental; not bit-comparable +to resident weights, validate end-to-end before production). LoRA adapters +that modify `adaln_proj` are rejected in both cache modes rather than +silently ignored. + +Both cache modes hold values derived from `adaln_proj`, so a runtime weight +update is only accepted when the cache can follow it: online mode takes a disk +update whose target directory carries native `adaln_proj` safetensors, and +rejects everything else (tensor updates, and directories without those +tensors) before a single weight is written. A sidecar is built offline and +cannot be regenerated in the server, so weight updates are rejected outright; +rebuild the sidecar against the new weights and restart. + ### Serve MiniMax-H3 on Ascend NPUs For Ascend NPU, follow the diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 9820e42c5cb0..924a08fbb01b 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -78,6 +78,9 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--served-model-name {NAME}`: stable model name exposed by serving APIs. Defaults to `--model-id` when set, otherwise `--model-path`. - `--model-variant {NAME}`: semantic checkpoint variant to load when one model repository contains multiple weight partitions. The pipeline maps this stable name to the repository layout before loading; for example, MiniMax-H3 accepts `fl2va` and `ref2va`. This is a server/load-time choice, unlike a request's `task`. - `--minimax-h3-adaln-cache-path {FILE}`: advanced MiniMax-H3-only inference cache. It replaces the checkpoint's AdaLN projection weights with precomputed outputs and only accepts requests whose exact FP32 timestep plan is included in the cache. It requires unquantized weights and the matching model variant. +- `--minimax-h3-adaln-online {true,false}`: rebuild MiniMax-H3 AdaLN outputs from the checkpoint on demand instead of keeping the 24.2 GiB of `adaln_proj` weights resident. Works with any step count or schedule; requires the unquantized native-layout checkpoint. Built plans live in a GPU slab with per-plan LRU eviction and, by default, a pinned-host cache so previously seen schedules swap back in over PCIe instead of re-reading the checkpoint. +- `--minimax-h3-adaln-plan-width {N}`: widest timestep plan the online slab is sized for (default 4 covers every task; t2va needs 2, fl2va 3). +- `--minimax-h3-adaln-host-cache-gb {GB}`: pinned host memory per rank caching built AdaLN plans (default 8; 0 disables). A 50-step schedule needs about 0.9 (t2va) / 1.33 (fl2va) / 1.77 (ref2va) GB; over-capacity plan sets simply recompute. Expert escape hatches (GPU slot count, experimental fp32 rebuild) are the `SGLANG_DIFFUSION_MINIMAX_H3_ADALN_*` environment variables. - `--model-subfolder {PATH}`: advanced direct override for a component subfolder inside the model repository. Prefer `--model-variant` when the pipeline exposes semantic routing. If both are supplied, they must resolve to the same weight partition. - `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter from a local path, Hugging Face repo/subfolder, or exact Hub file URL - `--lora-weight-name {FILE}`: select one adapter file from a repository that contains multiple LoRA revisions. The Hub download is filtered to that file plus JSON metadata, so unused weights are not downloaded. diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index d67500b72877..3b0616787534 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/python/sglang/multimodal_gen/envs.py @@ -37,6 +37,8 @@ SGLANG_DIFFUSION_TEST_FORCE_HOST_AVAILABLE_GIB: float | None = None SGLANG_DIFFUSION_TEST_CAP_DEVICE_MEMORY_GIB: float | None = None SGLANG_DIFFUSION_STAGE_LOGGING: bool = False + SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS: int = 64 + SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32: bool = False SGLANG_DIFFUSION_CFG_GATE_STEP: float = 1.0 # cache-dit env vars (primary transformer) # on by default; engages only on 2 ranks with peer-to-peer access and falls @@ -260,6 +262,19 @@ def _getter(): # If set, sgl_diffusion will enable stage logging, which will print the time # taken for each stage "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"), + # Plan slots in the MiniMax-H3 --minimax-h3-adaln-online GPU slab + # (9.25 MiB per slot-timestep; 64 x width 4 = 2.31 GiB). A request needs + # up to num_inference_steps - 1 slots; the default covers the 50-step + # serving schedule, so this is an escape hatch, not a deployment knob. + "SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS": _lazy_int( + "SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS", 64 + ), + # Experimental: compute the online AdaLN rebuild projections once in fp32 + # (TF32 off) before the bf16 store. Not bit-comparable to resident + # adaln_proj weights; keep off until an e2e trajectory gate clears it. + "SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32": _lazy_bool( + "SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32" + ), # Fraction of denoising steps that run both CFG branches before reusing the # last conditional-minus-unconditional residual. Keep 1.0 to disable. "SGLANG_DIFFUSION_CFG_GATE_STEP": _lazy_float( diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index 81b721e0f785..ea8b4d940920 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -6,6 +6,7 @@ import torch +from sglang.multimodal_gen import envs from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.layers.attention.selector import ( component_attn_backend_context_manager, @@ -429,12 +430,30 @@ def load_customized( "--minimax-h3-adaln-online and --minimax-h3-adaln-cache-path " "are mutually exclusive" ) + if dit_config.arch_config.checkpoint_uses_diffusers_layout: + # The rebuild reads native tensor names straight from the + # shards; on a Diffusers-layout checkpoint it would KeyError + # on the first request instead of failing here. + raise ValueError( + "--minimax-h3-adaln-online requires the native-layout " + "MiniMax H3 checkpoint (FL2VA/transformer or " + "Ref2VA/transformer), not the Diffusers-layout one" + ) # Keep the weights off-device; the model rebuilds the AdaLN # outputs from the checkpoint for each request's timestep plan. init_params["adaln_weight_files"] = safetensors_list init_params["adaln_plan_width"] = ( component_server_args.minimax_h3_adaln_plan_width ) + init_params["adaln_max_plans"] = ( + envs.SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS + ) + init_params["adaln_host_cache_bytes"] = int( + component_server_args.minimax_h3_adaln_host_cache_gb * 1e9 + ) + init_params["adaln_precision"] = ( + "fp32" if envs.SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32 else "match" + ) checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter runtime_quant_config = init_params["quant_config"] diff --git a/python/sglang/multimodal_gen/runtime/models/dits/base.py b/python/sglang/multimodal_gen/runtime/models/dits/base.py index d58900174d8d..60108e366376 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/base.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/base.py @@ -109,6 +109,34 @@ def prepare_lora_adapter( """Apply model-specific LoRA transforms after names are normalized.""" return adapter + def validate_weight_update_source(self, *, weights_path: str | None) -> None: + """Reject a weight update this model cannot stay coherent under. + + Called before any weight is written, so a rejection leaves the served + model untouched. ``weights_path`` is the new on-disk source, or None + for in-memory (tensor RPC) updates. Default no-op; models that derive + served values from their weights override this. + """ + return None + + def validate_lora_layers(self, layer_names: list[str]) -> None: + """Reject LoRA layers this model cannot apply. + + Called before any LoRA weight is written. Default no-op; models that + prune or replace layers a LoRA may target override this so the update + fails instead of silently skipping those layers. + """ + return None + + def refresh_weight_derived_caches(self, *, weights_path: str | None) -> None: + """Invalidate caches derived from weights after a weight update. + + ``weights_path`` is the new on-disk source, or None for in-memory + (tensor RPC) updates. Default no-op; models that precompute values + from their weights override this. + """ + return None + @property def supported_attention_backends(self) -> set[AttentionBackendEnum]: return self._supported_attention_backends diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 50f572b3d47c..3a9e30cea29f 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -9,15 +9,12 @@ import math import os -import struct from collections import defaultdict from collections.abc import Iterable, Iterator -from contextlib import ExitStack from typing import Any, Callable import torch import torch.nn as nn -from safetensors.torch import safe_open from torch.distributed.tensor import DTensor from sglang.kernels.ops.activation.activation import ( @@ -45,7 +42,6 @@ ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_ring_ctx, - get_tp_rank, get_ulysses_ctx, ) from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( @@ -73,6 +69,16 @@ is_layerwise_offloaded_module, ) from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import ( + MINIMAX_H3_ADALN_MAX_PLAN_WIDTH, + MiniMaxH3AdalnCache, +) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import ( + _plan_key as _adaln_plan_key, +) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import ( + native_adaln_weight_files, +) from sglang.multimodal_gen.runtime.models.parameter import BlockQuantScaleParameter from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, @@ -189,6 +195,7 @@ def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any: "img_position_ids", "rope_cache", "unique_timesteps", + "adaln_cache_slot", "inverse_indices", "update_mask", "update_audio_mask", @@ -1226,303 +1233,6 @@ def forward(self, adaln_input: torch.Tensor) -> tuple[torch.Tensor, ...]: return self.split_output(x) -# A ref2va request carrying both a visual and an audio reference reaches four -# distinct timesteps in one step: video, audio, the imgvid condition and the -# audio reference. That is the widest case, so it is the default; a deployment -# serving only narrower tasks (t2va reaches 2, fl2va 3) can shrink the slab -# proportionally via --minimax-h3-adaln-plan-width. -MINIMAX_H3_ADALN_MAX_PLAN_WIDTH = 4 - - -def _plan_key(timesteps: torch.Tensor) -> tuple[int, ...]: - """One denoise step's unique timesteps as their exact fp32 bit patterns.""" - return tuple( - struct.unpack(" None: - super().__init__() - if (path is None) == (weight_files is None): - raise ValueError( - "MiniMax H3 AdaLN cache takes exactly one of path (prebuilt " - "sidecar) or weight_files (rebuild from the checkpoint)" - ) - if max_plans < 1: - raise ValueError("MiniMax H3 AdaLN cache max_plans must be positive") - if max_plan_width < 1: - raise ValueError( - "MiniMax H3 AdaLN cache max_plan_width must be positive; " - "set --minimax-h3-adaln-plan-width to at least 1" - ) - self.path = path - self.model_variant = model_variant - self.weight_files = weight_files - self.max_plans = max_plans - self.max_plan_width = max_plan_width - self.num_layers = arch.num_layers - self.hidden_size = arch.hidden_size - self.block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * arch.hidden_size - self.final_width = 2 * arch.hidden_size - # Rebuild path only: plan bit pattern -> slot, tracked on the host. - self._slots: dict[tuple[int, ...], int] = {} - self.rebuilds = 0 - - def load(self, device: torch.device) -> None: - if self.path is None: - self._allocate(device) - return - if not os.path.isfile(self.path): - raise ValueError(f"MiniMax H3 AdaLN cache does not exist: {self.path}") - - with safe_open(self.path, framework="pt", device="cpu") as cache_file: - metadata = cache_file.metadata() or {} - if metadata.get("format_version") != self._FORMAT_VERSION: - raise ValueError( - "MiniMax H3 AdaLN cache has an unsupported or missing format_version" - ) - cache_variant = metadata.get("model_variant") - if self.model_variant is not None and cache_variant != self.model_variant: - raise ValueError( - "MiniMax H3 AdaLN cache model_variant does not match the loaded " - f"variant ({cache_variant!r} != {self.model_variant!r})" - ) - plan_timesteps = cache_file.get_tensor("plan_timesteps") - plan_lengths = cache_file.get_tensor("plan_lengths") - block_params = cache_file.get_tensor("block_params") - final_params = cache_file.get_tensor("final_params") - - expected_block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * self.hidden_size - expected_final_width = 2 * self.hidden_size - if ( - plan_timesteps.dtype != _FP32_DTYPE - or plan_timesteps.ndim != 2 - or plan_lengths.dtype != torch.int64 - or plan_lengths.shape != (plan_timesteps.shape[0],) - or (plan_lengths < 1).any() - or (plan_lengths > plan_timesteps.shape[1]).any() - ): - raise ValueError("MiniMax H3 AdaLN cache has invalid timestep plans") - if block_params.dtype != _BF16_DTYPE or block_params.shape != ( - plan_timesteps.shape[0], - plan_timesteps.shape[1], - self.num_layers, - expected_block_width, - ): - raise ValueError("MiniMax H3 AdaLN cache has invalid block_params") - if final_params.dtype != _BF16_DTYPE or final_params.shape != ( - plan_timesteps.shape[0], - plan_timesteps.shape[1], - expected_final_width, - ): - raise ValueError("MiniMax H3 AdaLN cache has invalid final_params") - - self.register_buffer("plan_timesteps", plan_timesteps.to(device)) - self.register_buffer("plan_lengths", plan_lengths.to(device)) - self.register_buffer("block_params", block_params.to(device)) - self.register_buffer("final_params", final_params.to(device)) - - def _allocate(self, device: torch.device) -> None: - """Empty slab for the rebuild path; its pointers must never move. - - ``plan_lengths`` starts at zero and that is what keeps unused slots out - of ``lookup``: a real plan always has at least one timestep, so a zero - length can never match. Breakable CUDA graph keys its replay signature - on tensor pointers, so this is allocated once and only written in place. - """ - width = self.max_plan_width - self.register_buffer( - "plan_timesteps", - torch.zeros((self.max_plans, width), dtype=_FP32_DTYPE, device=device), - ) - self.register_buffer( - "plan_lengths", - torch.zeros((self.max_plans,), dtype=torch.int64, device=device), - ) - self.register_buffer( - "block_params", - torch.zeros( - (self.max_plans, width, self.num_layers, self.block_width), - dtype=_BF16_DTYPE, - device=device, - ), - ) - self.register_buffer( - "final_params", - torch.zeros( - (self.max_plans, width, self.final_width), - dtype=_BF16_DTYPE, - device=device, - ), - ) - logger.info( - "MiniMax H3 AdaLN rebuild slab: %d plans x %d timesteps = %.2f GiB", - self.max_plans, - width, - self.block_params.numel() * 2 / 2**30, - ) - - def build( - self, - step_timesteps: list[torch.Tensor], - *, - embed: Callable[[torch.Tensor], torch.Tensor], - ) -> None: - """Fill every plan this request will look up, in one streaming pass. - - Each plan keeps its own timestep count as the GEMM batch size, because - cuBLAS selects kernels by shape and the selection is not monotonic in M: - against the runtime's M == 2, results at M == 4/8/16/64/96 are - bit-identical while M == 32 differs in 11760 of 96768 elements and - M == 1 (the GEMV path the first denoise step takes) differs in 69. - Rebuilding a plan at any other batch size silently perturbs the output. - - The pass reads all 50 adaln_proj layers regardless of how many plans are - missing, so a request builds everything it needs before denoising rather - than filling in step by step. - """ - wanted: dict[tuple[int, ...], torch.Tensor] = {} - for timesteps in step_timesteps: - wanted.setdefault(_plan_key(timesteps), timesteps) - missing = {k: v for k, v in wanted.items() if k not in self._slots} - if not missing: - return - if len(wanted) > self.max_plans: - raise ValueError( - f"MiniMax H3 AdaLN rebuild needs {len(wanted)} plans but " - f"max_plans is {self.max_plans}" - ) - widest = max(timesteps.numel() for timesteps in wanted.values()) - if widest > self.max_plan_width: - raise ValueError( - f"MiniMax H3 AdaLN rebuild hit a {widest}-timestep plan but the " - f"slab was allocated for {self.max_plan_width}; raise " - "--minimax-h3-adaln-plan-width (t2va needs 2, fl2va 3, ref2va 4)" - ) - - reset = len(self._slots) + len(missing) > self.max_plans - # A reset also evicts this request's cache hits, so rebuild its complete - # plan set rather than only the plans that were initially missing. - plans_to_build = wanted if reset else missing - if reset: - self._slots.clear() - self.plan_lengths.zero_() - - device = self.block_params.device - slots = [] - pending_slots: dict[tuple[int, ...], int] = {} - for offset, (key, timesteps) in enumerate(plans_to_build.items()): - slot = len(self._slots) + offset - pending_slots[key] = slot - slots.append((slot, timesteps.numel(), embed(timesteps.to(device)))) - self.plan_timesteps[slot, : timesteps.numel()] = timesteps.to(device) - - # adaln_proj is a ColumnParallelLinear: each rank owns a slice of the - # output features and all-gathers afterwards. The rebuild has to do the - # same rather than read the full width in one go -- a sharded GEMM has a - # different N, so cuBLAS picks a different kernel and the outputs stop - # matching. It also cuts per-rank checkpoint reads to 1/tp. - tp_size = get_tp_world_size() - tp_rank = get_tp_rank() if tp_size > 1 else 0 - - with ExitStack() as stack: - handles = [ - stack.enter_context(safe_open(f, framework="pt", device=str(device))) - for f in self.weight_files - ] - index = {name: h for h in handles for name in h.keys()} - - def read_shard(name: str, out_features: int) -> torch.Tensor: - if tp_size == 1: - return index[name].get_tensor(name) - shard = out_features // tp_size - start = tp_rank * shard - return index[name].get_slice(name)[start : start + shard] - - def project(adaln_input: torch.Tensor, weight, bias) -> torch.Tensor: - out = nn.functional.linear(adaln_input, weight, bias) - return tensor_model_parallel_all_gather(out) if tp_size > 1 else out - - for layer in range(self.num_layers): - prefix = f"blocks.{layer}.adaln_proj.linear" - weight = read_shard(f"{prefix}.weight", self.block_width) - bias = read_shard(f"{prefix}.bias", self.block_width) - for slot, length, adaln_input in slots: - self.block_params[slot, :length, layer] = project( - adaln_input, weight, bias - ) - del weight, bias - prefix = "final_layer.adaln_proj.linear" - weight = read_shard(f"{prefix}.weight", self.final_width) - bias = read_shard(f"{prefix}.bias", self.final_width) - for slot, length, adaln_input in slots: - self.final_params[slot, :length] = project(adaln_input, weight, bias) - del weight, bias - - for slot, length, _ in slots: - self.plan_lengths[slot] = length - # Commit host metadata only after every layer has been written. If a - # checkpoint read or projection raises, the zero-length slots remain - # invisible and a later request can retry the rebuild. - self._slots.update(pending_slots) - self.rebuilds += 1 - logger.info( - "MiniMax H3 AdaLN: rebuilt %d plan(s), %d/%d resident, pass #%d", - len(plans_to_build), - len(self._slots), - self.max_plans, - self.rebuilds, - ) - - def lookup(self, unique_timesteps: torch.Tensor) -> torch.Tensor: - num_timesteps = unique_timesteps.shape[0] - matches = self.plan_lengths.eq(num_timesteps) & self.plan_timesteps[ - :, :num_timesteps - ].eq(unique_timesteps).all(dim=-1) - if not bool(matches.any()): - raise ValueError( - "MiniMax H3 AdaLN cache does not cover the request timestep plan" - ) - return matches.to(torch.int64).argmax() - - def block( - self, - index: int, - cache_plan_index: torch.Tensor, - num_timesteps: int, - ) -> tuple[torch.Tensor, ...]: - params = self.block_params[cache_plan_index, :num_timesteps, index] - params = params.reshape(-1, 6, self.hidden_size) - return tuple(params.unbind(dim=1)) - - def final( - self, - cache_plan_index: torch.Tensor, - num_timesteps: int, - ) -> tuple[torch.Tensor, ...]: - params = self.final_params[cache_plan_index, :num_timesteps] - return tuple(params.reshape(-1, 2, self.hidden_size).unbind(dim=1)) - - class MiniMaxH3TokenRefinerBlock(nn.Module): """Standard pre-norm transformer block without AdaLN or RoPE.""" @@ -1817,6 +1527,24 @@ def forward( return video, audio +def _reject_adaln_lora(names: list[str]) -> None: + """Reject LoRA names touching adaln_proj; callers gate on cache mode. + + Cache modes prune the adaln_proj modules, so these deltas have nothing to + attach to: they would be dropped without a trace while the rebuild keeps + reading base weights from the checkpoint. + """ + adaln_names = sorted(name for name in names if "adaln_proj" in name) + if not adaln_names: + return + raise ValueError( + "MiniMax H3 AdaLN cache modes (--minimax-h3-adaln-online / " + "--minimax-h3-adaln-cache-path) cannot apply LoRA deltas on " + f"adaln_proj ({len(adaln_names)} name(s), e.g. {adaln_names[0]!r}); " + "serve this adapter with resident AdaLN weights" + ) + + class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): _aliases = [ "MiniMaxH3Transformer3DModel", @@ -1839,6 +1567,8 @@ def prepare_lora_adapter( ) -> dict[str, torch.Tensor]: """Project released-checkpoint AdaLN LoRAs onto pruned coordinates.""" _reject_non_lora_delta_tensors(adapter) + if self._adaln_precomputed: + _reject_adaln_lora(list(adapter)) full_width = self.arch.adaln_affine_input_dim if full_width is None: return adapter @@ -1894,20 +1624,94 @@ def prepare_lora_adapter( ) return projected - def prepare_adaln_plans(self, step_timesteps: list[torch.Tensor]) -> None: + def prepare_adaln_plans( + self, step_timesteps: list[torch.Tensor] + ) -> torch.Tensor | None: """Fill the AdaLN cache for this request before denoising starts. - No-op for a prebuilt sidecar; the rebuild path needs the model's own - timestep embedding so a filled plan is bit-identical to what resident - adaln_proj weights would have produced. + Returns the per-step slab slots as a device tensor (None without a + cache); forward consumes one scalar view per step so it never has to + match timesteps on device. A prebuilt sidecar only resolves slots; the + rebuild path needs the model's own timestep embedding so a filled plan + is bit-identical to what resident adaln_proj weights would have + produced. """ - if self.adaln_cache is None or self.adaln_cache.weight_files is None: + cache = self.adaln_cache + if cache is None: + return None + # Keying costs one D2H sync per plan; compute the keys once and share + # them between build and resolve. + keys = [_adaln_plan_key(timesteps) for timesteps in step_timesteps] + if cache.weight_files is not None: + + def embed(timesteps: torch.Tensor) -> torch.Tensor: + out = nn.functional.silu(self.time_embedder(timesteps)) + # 'match' replicates forward's bf16 cast bit-exactly; 'fp32' + # keeps the embedding in fp32 for the one-time projection. + if cache.precision == "match": + out = out.to(_BF16_DTYPE) + return out + + cache.build(step_timesteps, embed=embed, keys=keys) + return cache.resolve_slots(step_timesteps, keys=keys) + + def validate_lora_layers(self, layer_names: list[str]) -> None: + if self._adaln_precomputed: + _reject_adaln_lora(layer_names) + + def validate_weight_update_source(self, *, weights_path: str | None) -> None: + """Reject a weight update the AdaLN cache cannot follow. + + Runs before any weight is written: cached AdaLN outputs are derived + from adaln_proj, so an update this cache cannot follow would pair new + transformer weights with the previous checkpoint's conditioning. + """ + cache = self.adaln_cache + if cache is None: return + if cache.weight_files is None: + raise ValueError( + "MiniMax H3 was started with a prebuilt AdaLN sidecar " + "(--minimax-h3-adaln-cache-path), which is built offline from " + "the startup checkpoint and cannot be regenerated online; " + "rebuild the sidecar against the new weights and restart, or " + "serve with --minimax-h3-adaln-online" + ) + if weights_path is None: + raise ValueError( + "MiniMax H3 --minimax-h3-adaln-online rebuilds AdaLN outputs " + "by streaming adaln_proj from a checkpoint directory, and a " + "tensor weight update carries no such directory (its " + "adaln_proj tensors have no resident modules to land in); " + "use update_weights_from_disk instead" + ) + if not ( + os.path.isdir(weights_path) and native_adaln_weight_files(weights_path) + ): + # The rebuild streams native tensor names; a Diffusers-layout or + # quantized export would defer a KeyError to the next request. + raise ValueError( + "MiniMax H3 --minimax-h3-adaln-online cannot retarget its " + f"AdaLN rebuild at {weights_path!r} (no native adaln_proj " + "safetensors there), and rebuilding from the original " + "checkpoint would serve stale conditioning" + ) - def embed(timesteps: torch.Tensor) -> torch.Tensor: - return nn.functional.silu(self.time_embedder(timesteps)).to(_BF16_DTYPE) + def refresh_weight_derived_caches(self, *, weights_path: str | None) -> None: + """Drop cached AdaLN plans after a weight swap; retarget the rebuild. - self.adaln_cache.build(step_timesteps, embed=embed) + Cached plans are weight-derived values; keeping them after an update + silently serves the previous checkpoint's conditioning. + """ + cache = self.adaln_cache + if cache is None: + return + # validate_weight_update_source ran before the weights were written and + # rejected every source this cannot follow; anything else is a broken + # call order rather than a deployment the cache can degrade through. + self.validate_weight_update_source(weights_path=weights_path) + cache.weight_files = native_adaln_weight_files(weights_path) + cache.invalidate() def _can_batch_block_adaln(self) -> bool: return ( @@ -1992,6 +1796,9 @@ def __init__( adaln_cache_model_variant: str | None = None, adaln_weight_files: list[str] | None = None, adaln_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH, + adaln_max_plans: int = 64, + adaln_host_cache_bytes: int = 0, + adaln_precision: str = "match", ) -> None: super().__init__(config=config, hf_config=hf_config) arch = self.config @@ -2125,7 +1932,10 @@ def __init__( path=adaln_cache_path, model_variant=adaln_cache_model_variant, weight_files=adaln_weight_files, + max_plans=adaln_max_plans, max_plan_width=adaln_plan_width, + host_cache_bytes=adaln_host_cache_bytes, + precision=adaln_precision, ) if self._adaln_precomputed else None @@ -2688,16 +2498,16 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: block_adaln_params = None adaln_cache_plan_index = None if self.adaln_cache is not None: - adaln_cache_plan_index = self.adaln_cache.lookup( - unique_timesteps.view(-1).to(device) - ) - block_adaln_params = tuple( - self.adaln_cache.block( - index, - adaln_cache_plan_index, - adaln_input.shape[0], + # prepare_adaln_plans resolved the slot on the host; the device + # lookup remains for callers that drive forward() directly. + adaln_cache_plan_index = kwargs.get("adaln_cache_slot") + if adaln_cache_plan_index is None: + adaln_cache_plan_index = self.adaln_cache.lookup( + unique_timesteps.view(-1).to(device) ) - for index in range(len(self.blocks)) + block_adaln_params = self.adaln_cache.block_all( + cache_plan_index=adaln_cache_plan_index, + num_timesteps=adaln_input.shape[0], ) elif self._can_batch_block_adaln(): local_adaln = torch.stack( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_adaln_cache.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_adaln_cache.py new file mode 100644 index 000000000000..a0e09efb8cf2 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_adaln_cache.py @@ -0,0 +1,853 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Precomputed AdaLN plan cache for the MiniMax H3 DiT.""" + +from __future__ import annotations + +import os +import struct +from collections import OrderedDict +from contextlib import ExitStack, contextmanager, nullcontext +from typing import Callable + +import msgspec +import torch +import torch.nn as nn +from safetensors.torch import safe_open + +from sglang.multimodal_gen.configs.models.dits.minimax_h3 import ( + MINIMAX_H3_ADALN_MODALITY_NUM, + MiniMaxH3DiTArchConfig, +) +from sglang.multimodal_gen.runtime.distributed import ( + get_tp_world_size, + tensor_model_parallel_all_gather, +) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_tp_rank, + get_world_group, + world_group_is_initialized, +) +from sglang.multimodal_gen.runtime.managers.memory_managers.host_memory_budget import ( + HOST_RESERVE_FRACTION, + MIN_HOST_RESERVE_BYTES, + host_memory_available_bytes, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +_BF16_DTYPE = torch.bfloat16 +_FP32_DTYPE = torch.float32 + +# The native adaln_proj tensor names the online rebuild streams; a checkpoint +# without them (Diffusers layout, quantized export) cannot serve as a rebuild +# source. +_NATIVE_ADALN_PROBE_KEY = "blocks.0.adaln_proj.linear.weight" + + +# A ref2va request carrying both a visual and an audio reference reaches four +# distinct timesteps in one step: video, audio, the imgvid condition and the +# audio reference. That is the widest case, so it is the default; a deployment +# serving only narrower tasks (t2va reaches 2, fl2va 3) can shrink the slab +# proportionally via --minimax-h3-adaln-plan-width. +MINIMAX_H3_ADALN_MAX_PLAN_WIDTH = 4 + + +def _plan_key(timesteps: torch.Tensor) -> tuple[int, ...]: + """One denoise step's unique timesteps as their exact fp32 bit patterns.""" + return tuple( + struct.unpack(" list[str]: + """Safetensors under ``weights_path`` usable as an online rebuild source. + + Returns [] when the directory holds no native-layout adaln tensors, so a + caller retargeting the rebuild source can fail closed instead of deferring + a KeyError to the next request. + """ + import glob as _glob + + files = sorted(_glob.glob(os.path.join(weights_path, "*.safetensors"))) + for file in files: + with safe_open(file, framework="pt", device="cpu") as handle: + if _NATIVE_ADALN_PROBE_KEY in handle.keys(): + return files + return [] + + +@contextmanager +def _fp32_gemm_guard(*, enabled: bool, device: torch.device): + """Neighboring stages flip the process-global TF32 switches; without this + an 'fp32' build silently becomes tf32-once.""" + if not enabled or device.type != "cuda": + yield + return + allow_tf32 = torch.backends.cuda.matmul.allow_tf32 + matmul_precision = torch.get_float32_matmul_precision() + torch.backends.cuda.matmul.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.set_float32_matmul_precision(matmul_precision) + + +class MiniMaxH3AdalnCacheStats(msgspec.Struct): + gpu_hit_plans: int = 0 + host_hit_plans: int = 0 + built_plans: int = 0 + host_evicted_groups: int = 0 + host_pressure_skips: int = 0 + + +class _HostPlan(msgspec.Struct): + pages: list[int] + refcount: int = 0 + + +class MiniMaxH3AdalnHostTier: + """Bounded pinned-host store of built AdaLN plans, group-LRU evicted. + + One page holds one (plan, timestep) entry: the 50-layer block params and + the final-layer params packed back to back (~9.25 MiB at full size). A + plan spans ``length`` pages that need not be contiguous; H2D swap-in is + two copies per page. Groups (one request's full plan set) are the LRU and + refcount unit because a rebuild pass streams every adaln_proj layer no + matter how few plans it misses -- a partial host hit saves nothing, so + admission and eviction are all-or-nothing per group. + + Rank determinism: capacity is MIN-reduced across ranks at construction and + every state transition happens synchronously inside prepare, as a pure + function of the (replica-broadcast) request stream. Any timing-based + state flip here would desynchronize build()'s collectives across ranks. + """ + + def __init__( + self, + *, + num_layers: int, + block_width: int, + final_width: int, + capacity_bytes: int, + device: torch.device, + ) -> None: + self.block_numel = num_layers * block_width + self.num_layers = num_layers + self.block_width = block_width + self.page_numel = self.block_numel + final_width + page_bytes = self.page_numel * 2 + budget = _host_cache_budget_bytes() + num_pages = max(0, min(capacity_bytes, budget)) // page_bytes + self._slab, self.pinned, num_pages = _allocate_host_slab( + num_pages=num_pages, page_numel=self.page_numel + ) + num_pages = _all_ranks_min(num_pages, device=device) + self.num_pages = num_pages + self._free_pages = list(range(num_pages)) + self._plans: dict[tuple[int, ...], _HostPlan] = {} + self._groups: OrderedDict[tuple[tuple[int, ...], ...], tuple] = OrderedDict() + self._stream = ( + torch.cuda.Stream(device=device) + if device.type == "cuda" and self.pinned + else None + ) + self.stats: MiniMaxH3AdalnCacheStats | None = None + logger.info( + "MiniMax H3 AdaLN host tier: %d pages x %.2f MiB (%s)", + num_pages, + page_bytes / 2**20, + "pinned" if self.pinned else "pageable", + ) + + def fence_prepare_start(self) -> None: + if self._stream is None: + return + current = torch.cuda.current_stream() + self._stream.wait_stream(current) + current.wait_stream(self._stream) + + def fence_gpu_reads(self) -> None: + if self._stream is None: + return + torch.cuda.current_stream().wait_stream(self._stream) + + def synchronize(self) -> None: + if self._stream is not None: + self._stream.synchronize() + + def has_all(self, keys) -> bool: + return all(key in self._plans for key in keys) + + def register_group(self, *, group_key, keys) -> None: + """Record (or refresh) a group whose plans are all resident already.""" + if group_key in self._groups: + self._groups.move_to_end(group_key) + return + keys = tuple(keys) + if not self.has_all(keys): + # Some plans were never stored (host_pressure skip); nothing to pin. + return + for key in keys: + self._plans[key].refcount += 1 + self._groups[group_key] = keys + + def copy_to_gpu( + self, + assignments: dict[tuple[int, ...], int], + *, + block_params: torch.Tensor, + final_params: torch.Tensor, + ) -> None: + """H2D-copy each assigned plan's pages into its reserved slab slot. + + Runs on the tier's copy stream; the caller fences the compute stream + afterwards (fence_gpu_reads) and only then flips plan_lengths. + """ + with self._stream_ctx(): + for key, slot in assignments.items(): + for index, page in enumerate(self._plans[key].pages): + row = self._slab[page] + block_params[slot, index].copy_( + row[: self.block_numel].view(self.num_layers, self.block_width), + non_blocking=True, + ) + final_params[slot, index].copy_( + row[self.block_numel :], non_blocking=True + ) + + def store_group( + self, + *, + group_key, + entries: dict[tuple[int, ...], tuple[int, int]], + block_params: torch.Tensor, + final_params: torch.Tensor, + ) -> None: + """D2H-copy a freshly built group into host pages and commit it. + + ``entries`` maps each plan key to its (GPU slot, timestep count). + Commits synchronously (the copies are waited on before the group + becomes visible) so the resident set stays a deterministic function of + the request stream on every rank. Over capacity the group is simply + not cached; the next occurrence recomputes. + """ + if group_key in self._groups: + self._groups.move_to_end(group_key) + return + new_plans = { + key: value for key, value in entries.items() if key not in self._plans + } + needed = sum(length for _, length in new_plans.values()) + if needed > self.num_pages: + self._skip_for_pressure(group_key) + return + # Pin the shared plans first so evicting other groups cannot free them + # (and then double-count them as new). + for key in entries: + if key in self._plans: + self._plans[key].refcount += 1 + while needed > len(self._free_pages) and self._groups: + self._evict_oldest_group() + if needed > len(self._free_pages): + for key in entries: + if key in self._plans: + self._unpin(key) + self._skip_for_pressure(group_key) + return + + if self._stream is not None: + # Order the D2H reads after the build's slab writes. + self._stream.wait_stream(torch.cuda.current_stream()) + staged: dict[tuple[int, ...], _HostPlan] = {} + with self._stream_ctx(): + for key, (slot, length) in new_plans.items(): + pages = [self._free_pages.pop() for _ in range(length)] + for index, page in enumerate(pages): + row = self._slab[page] + row[: self.block_numel].view( + self.num_layers, self.block_width + ).copy_(block_params[slot, index], non_blocking=True) + row[self.block_numel :].copy_( + final_params[slot, index], non_blocking=True + ) + staged[key] = _HostPlan(pages=pages, refcount=1) + self.synchronize() + self._plans.update(staged) + self._groups[group_key] = tuple(entries) + + def clear(self) -> None: + self._plans.clear() + self._groups.clear() + self._free_pages = list(range(self.num_pages)) + + def _stream_ctx(self): + if self._stream is None: + return nullcontext() + return torch.cuda.stream(self._stream) + + def _unpin(self, key) -> None: + plan = self._plans[key] + plan.refcount -= 1 + if plan.refcount == 0: + self._free_pages.extend(plan.pages) + del self._plans[key] + + def _evict_oldest_group(self) -> None: + _, keys = self._groups.popitem(last=False) + for key in keys: + self._unpin(key) + if self.stats is not None: + self.stats.host_evicted_groups += 1 + + def _skip_for_pressure(self, group_key) -> None: + if self.stats is not None: + self.stats.host_pressure_skips += 1 + logger.info( + "MiniMax H3 AdaLN host tier: group of %d plan(s) skipped under " + "memory pressure; it will recompute on its next occurrence", + len(group_key), + ) + + +def _host_cache_budget_bytes() -> int: + """Per-rank share of the host memory actually available right now. + + Co-located ranks each see the same free memory; without the split every + rank would size its tier against all of it (the HiCache lesson), and the + reserve mirrors HostPinBudget so this pinner honors the same headroom as + every other one. Assumes single-node deployments when LOCAL_WORLD_SIZE is + unset. + """ + ranks = int( + os.environ.get( + "LOCAL_WORLD_SIZE", + ( + torch.distributed.get_world_size() + if torch.distributed.is_initialized() + else 1 + ), + ) + ) + available = host_memory_available_bytes() + reserve = max(int(available * HOST_RESERVE_FRACTION), MIN_HOST_RESERVE_BYTES) + return max(0, available - reserve) // max(1, ranks) + + +def _allocate_host_slab( + *, num_pages: int, page_numel: int +) -> tuple[torch.Tensor, bool, int]: + """Pinned slab, halving the page count on failure; pageable as last resort.""" + pages = num_pages + while pages > 0: + try: + slab = torch.empty((pages, page_numel), dtype=_BF16_DTYPE, pin_memory=True) + if pages < num_pages: + logger.warning( + "MiniMax H3 AdaLN host tier shrank to %d of %d pages " + "(pinned allocation failures)", + pages, + num_pages, + ) + return slab, True, pages + except RuntimeError: + pages //= 2 + if num_pages > 0: + logger.warning( + "MiniMax H3 AdaLN host tier could not pin any memory; falling " + "back to pageable host memory (H2D copies run synchronously)" + ) + return torch.empty((num_pages, page_numel), dtype=_BF16_DTYPE), False, num_pages + return torch.empty((0, page_numel), dtype=_BF16_DTYPE), False, 0 + + +def _all_ranks_min(value: int, *, device: torch.device) -> int: + """Every rank must run the same tier capacity or build()'s collectives + desynchronize; MIN over the world group is the conservative agreement.""" + if not torch.distributed.is_initialized() or not world_group_is_initialized(): + return value + if torch.distributed.get_world_size() == 1: + return value + probe = torch.tensor( + [value], + dtype=torch.int64, + device=device if device.type == "cuda" else "cpu", + ) + probe = get_world_group().all_reduce(probe, op=torch.distributed.ReduceOp.MIN) + return int(probe.item()) + + +class MiniMaxH3AdalnCache(nn.Module): + """Precomputed AdaLN outputs for fixed FP32 timestep plans.""" + + _FORMAT_VERSION = "2" + plan_timesteps: torch.Tensor + plan_lengths: torch.Tensor + block_params: torch.Tensor + final_params: torch.Tensor + + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + *, + path: str | None = None, + model_variant: str | None = None, + weight_files: list[str] | None = None, + max_plans: int = 64, + max_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH, + host_cache_bytes: int = 0, + precision: str = "match", + ) -> None: + super().__init__() + if (path is None) == (weight_files is None): + raise ValueError( + "MiniMax H3 AdaLN cache takes exactly one of path (prebuilt " + "sidecar) or weight_files (rebuild from the checkpoint)" + ) + if max_plans < 1: + raise ValueError("MiniMax H3 AdaLN cache max_plans must be positive") + if max_plan_width < 1: + raise ValueError( + "MiniMax H3 AdaLN cache max_plan_width must be positive; " + "set --minimax-h3-adaln-plan-width to at least 1" + ) + if precision not in ("match", "fp32"): + raise ValueError( + "MiniMax H3 AdaLN cache precision must be 'match' (bit-exact " + f"with resident adaln_proj weights) or 'fp32', got {precision!r}" + ) + self.path = path + self.model_variant = model_variant + self.weight_files = weight_files + self.max_plans = max_plans + self.max_plan_width = max_plan_width + self.num_layers = arch.num_layers + self.hidden_size = arch.hidden_size + self.block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * arch.hidden_size + self.final_width = 2 * arch.hidden_size + # Plan bit pattern -> slot, tracked on the host in LRU order (oldest + # first). The rebuild path evicts per plan; a sidecar never evicts. + self._slots: OrderedDict[tuple[int, ...], int] = OrderedDict() + self._free_slots: list[int] = list(range(max_plans)) + self.host_cache_bytes = host_cache_bytes + self.precision = precision + # Constructed in load(): needs the device and the distributed runtime. + self._host_tier: MiniMaxH3AdalnHostTier | None = None + self.stats = MiniMaxH3AdalnCacheStats() + self.rebuilds = 0 + + def load(self, device: torch.device) -> None: + if self.path is None: + self._allocate(device) + if self.host_cache_bytes > 0: + self._host_tier = MiniMaxH3AdalnHostTier( + num_layers=self.num_layers, + block_width=self.block_width, + final_width=self.final_width, + capacity_bytes=self.host_cache_bytes, + device=device, + ) + self._host_tier.stats = self.stats + return + if not os.path.isfile(self.path): + raise ValueError(f"MiniMax H3 AdaLN cache does not exist: {self.path}") + + with safe_open(self.path, framework="pt", device="cpu") as cache_file: + metadata = cache_file.metadata() or {} + if metadata.get("format_version") != self._FORMAT_VERSION: + raise ValueError( + "MiniMax H3 AdaLN cache has an unsupported or missing format_version" + ) + cache_variant = metadata.get("model_variant") + if self.model_variant is not None and cache_variant != self.model_variant: + raise ValueError( + "MiniMax H3 AdaLN cache model_variant does not match the loaded " + f"variant ({cache_variant!r} != {self.model_variant!r})" + ) + plan_timesteps = cache_file.get_tensor("plan_timesteps") + plan_lengths = cache_file.get_tensor("plan_lengths") + block_params = cache_file.get_tensor("block_params") + final_params = cache_file.get_tensor("final_params") + + if ( + plan_timesteps.dtype != _FP32_DTYPE + or plan_timesteps.ndim != 2 + or plan_lengths.dtype != torch.int64 + or plan_lengths.shape != (plan_timesteps.shape[0],) + or (plan_lengths < 1).any() + or (plan_lengths > plan_timesteps.shape[1]).any() + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid timestep plans") + if block_params.dtype != _BF16_DTYPE or block_params.shape != ( + plan_timesteps.shape[0], + plan_timesteps.shape[1], + self.num_layers, + self.block_width, + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid block_params") + if final_params.dtype != _BF16_DTYPE or final_params.shape != ( + plan_timesteps.shape[0], + plan_timesteps.shape[1], + self.final_width, + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid final_params") + + for slot in range(plan_timesteps.shape[0]): + length = int(plan_lengths[slot]) + self._slots[_plan_key(plan_timesteps[slot, :length])] = slot + + # The 0.9-2.3 GiB slabs are derived data; keep them out of state_dict. + self.register_buffer( + "plan_timesteps", plan_timesteps.to(device), persistent=False + ) + self.register_buffer("plan_lengths", plan_lengths.to(device), persistent=False) + self.register_buffer("block_params", block_params.to(device), persistent=False) + self.register_buffer("final_params", final_params.to(device), persistent=False) + + def _allocate(self, device: torch.device) -> None: + """Empty slab for the rebuild path; its pointers must never move. + + ``plan_lengths`` starts at zero and that is what keeps unused slots out + of ``lookup``: a real plan always has at least one timestep, so a zero + length can never match. Breakable CUDA graph keys its replay signature + on tensor pointers, so this is allocated once and only written in place. + """ + width = self.max_plan_width + self.register_buffer( + "plan_timesteps", + torch.zeros((self.max_plans, width), dtype=_FP32_DTYPE, device=device), + persistent=False, + ) + self.register_buffer( + "plan_lengths", + torch.zeros((self.max_plans,), dtype=torch.int64, device=device), + persistent=False, + ) + self.register_buffer( + "block_params", + torch.zeros( + (self.max_plans, width, self.num_layers, self.block_width), + dtype=_BF16_DTYPE, + device=device, + ), + persistent=False, + ) + self.register_buffer( + "final_params", + torch.zeros( + (self.max_plans, width, self.final_width), + dtype=_BF16_DTYPE, + device=device, + ), + persistent=False, + ) + logger.info( + "MiniMax H3 AdaLN rebuild slab: %d plans x %d timesteps = %.2f GiB", + self.max_plans, + width, + self.block_params.numel() * 2 / 2**30, + ) + + def build( + self, + step_timesteps: list[torch.Tensor], + *, + embed: Callable[[torch.Tensor], torch.Tensor], + keys: list[tuple[int, ...]] | None = None, + ) -> None: + """Fill every plan this request will look up, in one streaming pass. + + Each plan keeps its own timestep count as the GEMM batch size, because + cuBLAS selects kernels by shape and the selection is not monotonic in M: + against the runtime's M == 2, results at M == 4/8/16/64/96 are + bit-identical while M == 32 differs in 11760 of 96768 elements and + M == 1 (the GEMV path the first denoise step takes) differs in 69. + Rebuilding a plan at any other batch size silently perturbs the output. + + The pass reads all 50 adaln_proj layers regardless of how many plans are + missing, so a request builds everything it needs before denoising rather + than filling in step by step. + """ + if keys is None: + keys = [_plan_key(timesteps) for timesteps in step_timesteps] + wanted: dict[tuple[int, ...], torch.Tensor] = {} + for key, timesteps in zip(keys, step_timesteps): + wanted.setdefault(key, timesteps) + missing = {k: v for k, v in wanted.items() if k not in self._slots} + group_key = tuple(wanted) + if not missing: + # A pure hit is still a use: without the touch, a hot schedule + # keeps its build-time LRU stamp and gets evicted first. + for key in wanted: + self._slots.move_to_end(key) + self.stats.gpu_hit_plans += len(wanted) + if self._host_tier is not None: + self._host_tier.register_group(group_key=group_key, keys=wanted) + return + if len(wanted) > self.max_plans: + raise ValueError( + f"MiniMax H3 AdaLN rebuild needs {len(wanted)} plans but the " + "slab holds " + f"{self.max_plans}; raise SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS" + ) + widest = max(timesteps.numel() for timesteps in wanted.values()) + if widest > self.max_plan_width: + raise ValueError( + f"MiniMax H3 AdaLN rebuild hit a {widest}-timestep plan but the " + f"slab was allocated for {self.max_plan_width}; raise " + "--minimax-h3-adaln-plan-width (t2va needs 2, fl2va 3, ref2va 4)" + ) + self.stats.gpu_hit_plans += len(wanted) - len(missing) + + if self._host_tier is not None: + # Order the tier's copy stream against everything the compute + # stream still has queued (a failed request never drains it) and + # vice versa, once per request, before any slot is touched. + self._host_tier.fence_prepare_start() + if self._host_tier.has_all(missing): + self._swap_in_from_host(wanted=wanted, missing=missing) + self._host_tier.register_group(group_key=group_key, keys=wanted) + return + + device = self.block_params.device + pending_slots = self._allocate_slots(wanted=wanted, missing=missing) + slots = [] + for key, timesteps in missing.items(): + slot = pending_slots[key] + adaln_input = embed(timesteps.to(device)) + if self.precision == "fp32": + # Compute the projections in fp32 once; the slab stays bf16. + adaln_input = adaln_input.float() + slots.append((slot, timesteps.numel(), adaln_input)) + self.plan_timesteps[slot, : timesteps.numel()] = timesteps.to(device) + + # adaln_proj is a ColumnParallelLinear: each rank owns a slice of the + # output features and all-gathers afterwards. The rebuild has to do the + # same rather than read the full width in one go -- a sharded GEMM has a + # different N, so cuBLAS picks a different kernel and the outputs stop + # matching. It also cuts per-rank checkpoint reads to 1/tp. + tp_size = get_tp_world_size() + tp_rank = get_tp_rank() if tp_size > 1 else 0 + + try: + with _fp32_gemm_guard(enabled=self.precision == "fp32", device=device): + self._project_plans_from_checkpoint( + slots, tp_size=tp_size, tp_rank=tp_rank, device=device + ) + except BaseException: + # The pending slots never became visible (their lengths stayed + # zero); hand them back so a retry does not leak capacity. + self._free_slots.extend(pending_slots.values()) + raise + + for slot, length, _ in slots: + self.plan_lengths[slot] = length + # Commit host metadata only after every layer has been written. If a + # checkpoint read or projection raises, the zero-length slots remain + # invisible and a later request can retry the rebuild. + self._slots.update(pending_slots) + self.rebuilds += 1 + self.stats.built_plans += len(missing) + logger.info( + "MiniMax H3 AdaLN: rebuilt %d plan(s), %d/%d resident, pass #%d", + len(missing), + len(self._slots), + self.max_plans, + self.rebuilds, + ) + if self._host_tier is not None: + self._host_tier.store_group( + group_key=group_key, + entries={ + key: (self._slots[key], wanted[key].numel()) for key in wanted + }, + block_params=self.block_params, + final_params=self.final_params, + ) + + def _allocate_slots( + self, + *, + wanted: dict[tuple[int, ...], torch.Tensor], + missing: dict[tuple[int, ...], torch.Tensor], + ) -> dict[tuple[int, ...], int]: + """Reserve one slab slot per missing plan, evicting LRU plans if needed. + + Reserved slots stay invisible (length 0) until the caller commits them + into ``self._slots``. + """ + for key in wanted: + if key in self._slots: + self._slots.move_to_end(key) + shortfall = len(missing) - len(self._free_slots) + if shortfall > 0: + # The wanted-count check in build() guarantees enough resident + # plans outside this request to cover the shortfall, in LRU order. + victims = [key for key in self._slots if key not in wanted][:shortfall] + for key in victims: + slot = self._slots.pop(key) + # Zero the length first: the slot must be invisible to lookup + # and resolve_slots before its contents are overwritten. + self.plan_lengths[slot] = 0 + self._free_slots.append(slot) + return {key: self._free_slots.pop() for key in missing} + + def _swap_in_from_host( + self, + *, + wanted: dict[tuple[int, ...], torch.Tensor], + missing: dict[tuple[int, ...], torch.Tensor], + ) -> None: + """Fill the missing GPU slots from pinned host pages, no weight pass.""" + assert self._host_tier is not None + device = self.block_params.device + assignments = self._allocate_slots(wanted=wanted, missing=missing) + try: + for key, timesteps in missing.items(): + slot = assignments[key] + self.plan_timesteps[slot, : timesteps.numel()] = timesteps.to(device) + self._host_tier.copy_to_gpu( + assignments, + block_params=self.block_params, + final_params=self.final_params, + ) + except BaseException: + # Lengths never flipped, so the slots stayed invisible; hand them + # back or the free-list invariant breaks and a later admission- + # accepted request pops from an empty list. + self._free_slots.extend(assignments.values()) + raise + # The compute stream must observe the H2D copies before any forward + # reads the slots; lengths flip last so half-filled slots stay hidden. + self._host_tier.fence_gpu_reads() + for key, timesteps in missing.items(): + self.plan_lengths[assignments[key]] = timesteps.numel() + self._slots.update(assignments) + self.stats.host_hit_plans += len(missing) + logger.info( + "MiniMax H3 AdaLN: %d plan(s) from the host cache, %d/%d resident", + len(missing), + len(self._slots), + self.max_plans, + ) + + def invalidate(self) -> None: + """Drop every cached plan; call between requests after a weight swap.""" + if self._host_tier is not None: + self._host_tier.synchronize() + self._host_tier.clear() + self._slots.clear() + self._free_slots = list(range(self.max_plans)) + self.plan_lengths.zero_() + + def _project_plans_from_checkpoint( + self, + slots: list[tuple[int, int, torch.Tensor]], + *, + tp_size: int, + tp_rank: int, + device: torch.device, + ) -> None: + with ExitStack() as stack: + handles = [ + stack.enter_context(safe_open(f, framework="pt", device=str(device))) + for f in self.weight_files + ] + index = {name: h for h in handles for name in h.keys()} + + def read_shard(name: str, out_features: int) -> torch.Tensor: + if tp_size == 1: + tensor = index[name].get_tensor(name) + else: + shard = out_features // tp_size + start = tp_rank * shard + tensor = index[name].get_slice(name)[start : start + shard] + if self.precision == "fp32": + tensor = tensor.float() + return tensor + + def project(adaln_input: torch.Tensor, weight, bias) -> torch.Tensor: + out = nn.functional.linear(adaln_input, weight, bias) + return tensor_model_parallel_all_gather(out) if tp_size > 1 else out + + for layer in range(self.num_layers): + prefix = f"blocks.{layer}.adaln_proj.linear" + weight = read_shard(f"{prefix}.weight", self.block_width) + bias = read_shard(f"{prefix}.bias", self.block_width) + for slot, length, adaln_input in slots: + self.block_params[slot, :length, layer] = project( + adaln_input, weight, bias + ) + del weight, bias + prefix = "final_layer.adaln_proj.linear" + weight = read_shard(f"{prefix}.weight", self.final_width) + bias = read_shard(f"{prefix}.bias", self.final_width) + for slot, length, adaln_input in slots: + self.final_params[slot, :length] = project(adaln_input, weight, bias) + del weight, bias + + def resolve_slots( + self, + step_timesteps: list[torch.Tensor], + *, + keys: list[tuple[int, ...]] | None = None, + ) -> torch.Tensor: + """Per-step slab slots as one device tensor, resolved on the host. + + Forward receives one scalar view per step. Keeping it a device tensor + matters twice over: a Python int would enter the breakable-CUDA-graph + replay signature (one graph per slot value), and an int baked into a + captured gather would read the wrong slab row after slot reuse. + """ + if keys is None: + keys = [_plan_key(timesteps) for timesteps in step_timesteps] + slots = [] + for key in keys: + slot = self._slots.get(key) + if slot is None: + raise ValueError( + "MiniMax H3 AdaLN cache does not cover the request timestep plan" + ) + slots.append(slot) + return torch.tensor(slots, dtype=torch.int64, device=self.block_params.device) + + def lookup(self, unique_timesteps: torch.Tensor) -> torch.Tensor: + num_timesteps = unique_timesteps.shape[0] + if num_timesteps > self.plan_timesteps.shape[1]: + # Without this the slice below silently clamps and the comparison + # dies on a shape mismatch instead of the real reason. + raise ValueError( + "MiniMax H3 AdaLN cache does not cover the request timestep plan" + ) + matches = self.plan_lengths.eq(num_timesteps) & self.plan_timesteps[ + :, :num_timesteps + ].eq(unique_timesteps).all(dim=-1) + if not bool(matches.any()): + raise ValueError( + "MiniMax H3 AdaLN cache does not cover the request timestep plan" + ) + return matches.to(torch.int64).argmax() + + def block_all( + self, + *, + cache_plan_index: torch.Tensor, + num_timesteps: int, + ) -> tuple[tuple[torch.Tensor, ...], ...]: + """Every block's AdaLN tuple via one layer-major slab gather.""" + stacked = self.block_params.permute(2, 0, 1, 3)[ + :, cache_plan_index, :num_timesteps + ] + stacked = stacked.reshape(self.num_layers, -1, 6, self.hidden_size) + return tuple(tuple(layer.unbind(dim=1)) for layer in stacked) + + def final( + self, + cache_plan_index: torch.Tensor, + num_timesteps: int, + ) -> tuple[torch.Tensor, ...]: + params = self.final_params[cache_plan_index, :num_timesteps] + return tuple(params.reshape(-1, 2, self.hidden_size).unbind(dim=1)) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py index b04643e61ec5..77f389b08ffe 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py @@ -319,6 +319,7 @@ def forward_kwargs( video_rows: torch.Tensor, audio_rows: torch.Tensor, step_timesteps: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + adaln_slot: torch.Tensor | None = None, ) -> dict[str, Any]: x = self.x_buffer audio_x = self.audio_x_buffer @@ -341,7 +342,7 @@ def forward_kwargs( 0, self.audio_target_seq_idx, audio_rows[self.audio_target_slice] ) unique_timesteps, inverse_indices, block_combined_indices = step_timesteps - return { + kwargs = { **self.static_kwargs, "x": x, "audio_x": audio_x, @@ -349,6 +350,11 @@ def forward_kwargs( "inverse_indices": inverse_indices, "block_combined_indices": block_combined_indices, } + if adaln_slot is not None: + # Device scalar, never a Python int: an int would key one breakable + # CUDA graph per slot value and go stale when LRU reuses the slot. + kwargs["adaln_cache_slot"] = adaln_slot + return kwargs def _expand_step_timesteps( self, @@ -537,7 +543,7 @@ def minimax_h3_denoise_loop( # Every step's timesteps are settled by now. Rebuilding AdaLN reads all # 24.2 GiB of adaln_proj whatever is missing, so fill the whole request in # one pass here instead of topping up step by step inside the loop. - model.prepare_adaln_plans([entry[0] for entry in timestep_plan]) + adaln_plan_slots = model.prepare_adaln_plans([entry[0] for entry in timestep_plan]) # match the scheduler's device-fp32 math once, then reuse one denoised # scratch per modality instead of allocating intermediates every step @@ -561,6 +567,9 @@ def minimax_h3_denoise_loop( video_rows=video_rows, audio_rows=audio_rows, step_timesteps=timestep_plan[step], + adaln_slot=( + None if adaln_plan_slots is None else adaln_plan_slots[step] + ), ) if attn_metadata is not None: attn_metadata.current_timestep = step diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py index bd8d0e7d0042..e1b06508e4d7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from typing import Any, Mapping +from sglang.multimodal_gen import envs from sglang.multimodal_gen.configs.sample.sampling_params import QUALITY_LEVELS from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage @@ -152,6 +153,19 @@ def forward(self, batch: Req, server_args: ServerArgs) -> Req: "MiniMax H3 requires num_inference_steps >= 2 because its " "video/audio sigma schedules include both interval endpoints" ) + gpu_plans = envs.SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS + if ( + server_args.minimax_h3_adaln_online + and batch.num_inference_steps - 1 > gpu_plans + ): + # Fail here, before the encode stages spend GPU time on a request + # whose AdaLN rebuild is guaranteed to overflow the slab. + raise ValueError( + f"num_inference_steps={batch.num_inference_steps} needs up to " + f"{batch.num_inference_steps - 1} AdaLN plans but the online " + f"slab holds {gpu_plans}; raise " + "SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS" + ) quality = getattr(batch.sampling_params, "quality", "lossless") if quality not in QUALITY_LEVELS: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py index 0056949d602b..2111beb2ccd6 100644 --- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py @@ -57,6 +57,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( is_layerwise_offloaded_module, ) +from sglang.multimodal_gen.runtime.models.dits.base import BaseDiT from sglang.multimodal_gen.runtime.pipelines.diffusers_pipeline import DiffusersPipeline from sglang.multimodal_gen.runtime.pipelines_core.lora.pipeline import ( LoRAPipeline, @@ -397,6 +398,16 @@ def update_weights_from_disk( logger.error(error_msg) return False, error_msg + try: + for module_name, module in modules_to_update: + if isinstance(module, BaseDiT): + module.validate_weight_update_source( + weights_path=weights_map[module_name] + ) + except ValueError as e: + logger.error(str(e)) + return False, str(e) + logger.info( f"Updating {len(weights_map)} modules: " + ", ".join(f"{n} <- {p}" for n, p in weights_map.items()) @@ -409,6 +420,14 @@ def update_weights_from_disk( self._module_weight_dirs[module_name] = weights_map[module_name] if target_modules is None: self.pipeline.model_path = local_model_path + for module_name, module in modules_to_update: + if isinstance(module, BaseDiT): + # Weight-derived caches must not survive a weight swap + # (regardless of flush_cache, which only governs + # optimization caches like TeaCache). + module.refresh_weight_derived_caches( + weights_path=weights_map[module_name] + ) gc.collect() torch.cuda.empty_cache() @@ -543,6 +562,14 @@ def update_weights_from_tensor( logger.error(str(e)) return False, str(e) + try: + for _module_name, module in modules_to_update: + if isinstance(module, BaseDiT): + module.validate_weight_update_source(weights_path=None) + except ValueError as e: + logger.error(str(e)) + return False, str(e) + updated_modules: list[str] = [] for module_name, module in modules_to_update: try: @@ -559,6 +586,13 @@ def update_weights_from_tensor( logger.error(error_msg, exc_info=True) return False, error_msg + for module_name, module in modules_to_update: + if isinstance(module, BaseDiT): + # Same invariant as the disk path. Any model whose derived + # state needs an on-disk source rejected this update above, so + # what reaches here only has caches to drop. + module.refresh_weight_derived_caches(weights_path=None) + gc.collect() torch.cuda.empty_cache() names = ", ".join(updated_modules) @@ -620,6 +654,19 @@ def _update_lora_from_tensor( if not pairs: return False, "No LoRA A/B tensor pairs found in payload" + dit_module = dict(modules_to_update).get(target_module) + if dit_module is None: + return False, f"No DiT module found for LoRA IPC target {target_module!r}" + if isinstance(dit_module, BaseDiT): + # Before convert_to_lora_layers() wraps anything: a layer the model + # cannot host is skipped as "unknown" further down, which would + # publish success for a partially applied adapter. + try: + dit_module.validate_lora_layers(list(pairs)) + except ValueError as e: + logger.error(str(e)) + return False, str(e) + lora_pipeline: LoRAPipeline = self.pipeline if not lora_pipeline.lora_initialized: convert_target = ( @@ -645,10 +692,6 @@ def _update_lora_from_tensor( logger.error(str(e)) return False, str(e) - dit_module = dict(modules_to_update).get(target_module) - if dit_module is None: - return False, f"No DiT module found for LoRA IPC target {target_module!r}" - updated = 0 skipped = 0 unknown_layers: list[str] = [] diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index bd0bc7f731f1..f7469416a692 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -352,6 +352,11 @@ class ServerArgs(DisaggServerArgsMixin): # Widest timestep plan the rebuild slab is sized for; see # MINIMAX_H3_ADALN_MAX_PLAN_WIDTH. minimax_h3_adaln_plan_width: int = 4 + # Pinned-host cache for built AdaLN plans (decimal GB, 0 disables). Plans + # evicted from the GPU slab swap back in from here instead of re-reading + # the 24.2 GiB checkpoint. Expert knobs (GPU slot count, fp32 rebuild) + # live in envs.py as SGLANG_DIFFUSION_MINIMAX_H3_ADALN_*. + minimax_h3_adaln_host_cache_gb: float = 8.0 # Explicit quantization method override (e.g. "mxfp8", "fp8", "modelslim"). # When set, the transformer loader uses it instead of auto-detection. quantization: str | None = None @@ -615,8 +620,21 @@ def _validate_parameters(self): self._validate_cfg_parallel() self._validate_batching() self._validate_breakable_cuda_graph() + self._validate_minimax_h3_adaln() self.pipeline_config.validate_server_args(self) + def _validate_minimax_h3_adaln(self) -> None: + # Warn, not raise: config-file and from_kwargs construction mark every + # provided key as explicit, so a shared base config pinning the + # default (or 0) must not fail non-online launches. + if self.minimax_h3_adaln_online: + return + if self.is_arg_explicitly_set("minimax_h3_adaln_host_cache_gb"): + logger.warning( + "--minimax-h3-adaln-host-cache-gb only takes effect with " + "--minimax-h3-adaln-online; ignoring it" + ) + def _validate_scheduler_rpc_timeout(self) -> None: timeout = self.scheduler_rpc_timeout if timeout is None: @@ -1925,6 +1943,20 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "A request exceeding it is rejected rather than truncated." ), ) + parser.add_argument( + "--minimax-h3-adaln-host-cache-gb", + type=float, + default=ServerArgs.minimax_h3_adaln_host_cache_gb, + help=( + "Pinned host memory (decimal GB, per rank) caching AdaLN plans " + "built by --minimax-h3-adaln-online, so a plan set evicted " + "from the GPU slab swaps back in over PCIe instead of " + "re-reading the 24.2 GiB checkpoint (measured 5.8-6.7 s). One " + "50-step schedule needs ~0.9 (t2va) / 1.33 (fl2va) / 1.77 " + "(ref2va) GB; the default 8 holds several. Groups are evicted " + "LRU and over-cap groups just recompute. 0 disables the tier." + ), + ) parser.add_argument( "--minimax-h3-adaln-cache-path", type=str, diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py index e7b162bd3bde..d6a27e01efbb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py @@ -14,7 +14,7 @@ maybe_init_distributed_environment_and_model_parallel, model_parallel_is_initialized, ) -from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( +from sglang.multimodal_gen.runtime.models.dits.minimax_h3_adaln_cache import ( MiniMaxH3AdalnCache, ) from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( @@ -37,20 +37,30 @@ def _ensure_single_process_parallel_runtime() -> None: maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) +def _weights_fill(*shape: int, scale: float) -> torch.Tensor: + values = torch.arange(int(torch.tensor(shape).prod()), dtype=torch.float32) + return ((values % 7) * 0.01 * scale).reshape(shape) + + def _write_online_weights( path: Path, *, omit: str | None = None, + scale: float = 0.0, ) -> None: - # These cache-state tests need checkpoint-compatible shapes, not values. + # State-machine tests only need checkpoint-compatible shapes (scale 0); + # value-equality tests pass a nonzero scale for distinguishable outputs. + def _fill(*shape: int) -> torch.Tensor: + return _weights_fill(*shape, scale=scale) + tensors: dict[str, torch.Tensor] = {} for layer in range(_ARCH.num_layers): prefix = f"blocks.{layer}.adaln_proj.linear" - tensors[f"{prefix}.weight"] = torch.zeros(_BLOCK_WIDTH, _ARCH.time_embed_dim) - tensors[f"{prefix}.bias"] = torch.zeros(_BLOCK_WIDTH) + tensors[f"{prefix}.weight"] = _fill(_BLOCK_WIDTH, _ARCH.time_embed_dim) + tensors[f"{prefix}.bias"] = _fill(_BLOCK_WIDTH) prefix = "final_layer.adaln_proj.linear" - tensors[f"{prefix}.weight"] = torch.zeros(_FINAL_WIDTH, _ARCH.time_embed_dim) - tensors[f"{prefix}.bias"] = torch.zeros(_FINAL_WIDTH) + tensors[f"{prefix}.weight"] = _fill(_FINAL_WIDTH, _ARCH.time_embed_dim) + tensors[f"{prefix}.bias"] = _fill(_FINAL_WIDTH) if omit is not None: tensors.pop(omit) save_file(tensors, path) @@ -62,20 +72,34 @@ def _online_cache( max_plans: int = 2, max_plan_width: int = 2, omit: str | None = None, + host_cache_bytes: int = 0, + scale: float = 0.0, ) -> MiniMaxH3AdalnCache: _ensure_single_process_parallel_runtime() weight_path = tmp_path / "model.safetensors" - _write_online_weights(weight_path, omit=omit) + _write_online_weights(weight_path, omit=omit, scale=scale) cache = MiniMaxH3AdalnCache( _ARCH, weight_files=[str(weight_path)], max_plans=max_plans, max_plan_width=max_plan_width, + host_cache_bytes=host_cache_bytes, ) cache.load(torch.device("cpu")) return cache +# One host-tier page for the tiny arch, in bytes (see MiniMaxH3AdalnHostTier). +_PAGE_BYTES = (_ARCH.num_layers * _BLOCK_WIDTH + _FINAL_WIDTH) * 2 + + +def _reference_block(cache, index, plan, num_timesteps): + # Local oracle for block_all: explicit slab indexing, kept out of the + # production class so a layout bug cannot "fix itself" in both sides. + params = cache.block_params[plan, :num_timesteps, index] + return tuple(params.reshape(-1, 6, _ARCH.hidden_size).unbind(dim=1)) + + def _embed(timesteps: torch.Tensor) -> torch.Tensor: return timesteps[:, None].expand(-1, _ARCH.time_embed_dim) @@ -113,7 +137,7 @@ def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path): cache.load(torch.device("cpu")) cache_plan_index = cache.lookup(plan_timesteps[1]) - block = cache.block(1, cache_plan_index, 2) + block = _reference_block(cache, 1, cache_plan_index, 2) final = cache.final(cache_plan_index, 2) # block() hands the forward pass six [num_timesteps * modality, hidden] @@ -127,8 +151,227 @@ def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path): assert torch.equal(torch.cat(final, dim=-1), final_params[1]) -def test_online_cache_reset_rebuilds_previously_resident_request_plans(tmp_path): - """A capacity reset must not drop plans reused by the current request.""" +def test_sidecar_resolve_slots_and_block_all_match_per_step_paths(tmp_path): + """Host-resolved slots and the batched gather must mirror lookup/block.""" + cache_path = tmp_path / "adaln.safetensors" + plan_timesteps = torch.tensor([[0.5, 0.0], [1.0, 2.0]]) + plan_lengths = torch.tensor([1, 2], dtype=torch.int64) + block_params = ( + torch.arange(2 * 2 * 2 * _BLOCK_WIDTH, dtype=torch.float32) + .reshape(2, 2, 2, _BLOCK_WIDTH) + .bfloat16() + ) + final_params = torch.zeros(2, 2, _FINAL_WIDTH, dtype=torch.bfloat16) + save_file( + { + "plan_timesteps": plan_timesteps, + "plan_lengths": plan_lengths, + "block_params": block_params, + "final_params": final_params, + }, + cache_path, + metadata={"format_version": "2", "model_variant": "fl2va"}, + ) + cache = MiniMaxH3AdalnCache(_ARCH, path=str(cache_path), model_variant="fl2va") + cache.load(torch.device("cpu")) + + slots = cache.resolve_slots([torch.tensor([0.5]), torch.tensor([1.0, 2.0])]) + assert slots.dtype == torch.int64 + assert slots.tolist() == [0, 1] + assert int(cache.lookup(torch.tensor([1.0, 2.0]))) == int(slots[1]) + + stacked = cache.block_all(cache_plan_index=slots[1], num_timesteps=2) + assert len(stacked) == _ARCH.num_layers + for index in range(_ARCH.num_layers): + expected = _reference_block(cache, index, slots[1], 2) + for got, want in zip(stacked[index], expected): + assert torch.equal(got, want) + assert got.stride() == want.stride() + + with pytest.raises(ValueError, match="does not cover"): + cache.resolve_slots([torch.tensor([9.0])]) + + +def test_online_cache_resolve_slots_after_build(tmp_path): + cache = _online_cache(tmp_path, max_plan_width=2) + plan_a = torch.tensor([1.0]) + plan_b = torch.tensor([2.0, 3.0]) + + cache.build([plan_a, plan_b, plan_a], embed=_embed) + slots = cache.resolve_slots([plan_a, plan_b, plan_a]) + assert slots.tolist()[0] == slots.tolist()[2] + assert int(cache.lookup(plan_b)) == int(slots[1]) + + +def test_slab_buffers_stay_out_of_state_dict(tmp_path): + cache = _online_cache(tmp_path) + assert not any("params" in key or "plan" in key for key in cache.state_dict()) + + +def test_host_tier_swap_in_restores_evicted_plans_bit_exactly(tmp_path): + """A GPU-evicted plan set must return from the host tier byte-identical, + without another checkpoint pass.""" + cache = _online_cache( + tmp_path, + max_plans=2, + max_plan_width=1, + host_cache_bytes=64 * _PAGE_BYTES, + scale=1.0, + ) + set_a = [torch.tensor([1.0]), torch.tensor([2.0])] + set_b = [torch.tensor([3.0]), torch.tensor([4.0])] + + cache.build(set_a, embed=_embed) + slots_a = cache.resolve_slots(set_a) + snapshot = [ + ( + [t.clone() for t in _reference_block(cache, 0, slots_a[i], 1)], + [t.clone() for t in cache.final(slots_a[i], 1)], + ) + for i in range(2) + ] + assert cache.stats.built_plans == 2 + + cache.build(set_b, embed=_embed) # evicts set_a from the GPU slab + passes = cache.rebuilds + cache.build(set_a, embed=_embed) # swaps back in from the host tier + assert cache.rebuilds == passes + assert cache.stats.host_hit_plans == 2 + + slots_a = cache.resolve_slots(set_a) + for i in range(2): + blocks, finals = snapshot[i] + for got, want in zip(_reference_block(cache, 0, slots_a[i], 1), blocks): + assert torch.equal(got, want) + for got, want in zip(cache.final(slots_a[i], 1), finals): + assert torch.equal(got, want) + assert float(cache.plan_timesteps[slots_a[i], 0]) == float(set_a[i][0]) + + +def test_host_tier_over_capacity_group_recomputes(tmp_path): + """A group that cannot fit is skipped (never raises) and rebuilt later.""" + cache = _online_cache( + tmp_path, + max_plans=2, + max_plan_width=1, + host_cache_bytes=1 * _PAGE_BYTES, # one page: a 2-plan group never fits + ) + set_a = [torch.tensor([1.0]), torch.tensor([2.0])] + set_b = [torch.tensor([3.0]), torch.tensor([4.0])] + + cache.build(set_a, embed=_embed) + assert cache.stats.host_pressure_skips == 1 + cache.build(set_b, embed=_embed) + passes = cache.rebuilds + cache.build(set_a, embed=_embed) # host tier empty: full rebuild again + assert cache.rebuilds == passes + 1 + + +def test_host_tier_lru_eviction_and_shared_plan_refcount(tmp_path): + cache = _online_cache( + tmp_path, + max_plans=4, + max_plan_width=1, + host_cache_bytes=3 * _PAGE_BYTES, + ) + plan_a = torch.tensor([1.0]) + plan_b = torch.tensor([2.0]) + plan_c = torch.tensor([3.0]) + + cache.build([plan_a, plan_b], embed=_embed) # group 1 {a, b}: 2 pages + cache.build([plan_a, plan_c], embed=_embed) # group 2 {a, c}: +1 page (a shared) + tier = cache._host_tier + assert tier is not None + assert len(tier._plans) == 3 and len(tier._free_pages) == 0 + + # A third group needs a page; group 1 is LRU. Its shared plan_a must + # survive because group 2 still references it. + plan_d = torch.tensor([4.0]) + cache.build([plan_a, plan_d], embed=_embed) + assert cache.stats.host_evicted_groups == 1 + import struct + + keys = { + tuple(struct.unpack(" torch.Tensor: + # The production embed stays fp32 in this mode; a bf16 input here + # proves the cache upcasts before the projection (a 'match'-style + # bf16 x fp32 GEMM would fail on dtype mismatch). + return _embed(timesteps).bfloat16() + + cache.build([plan], embed=bf16_embed) + slot = cache.resolve_slots([plan])[0] + + adaln_input = bf16_embed(plan).float() + weight = _weights_fill(_BLOCK_WIDTH, _ARCH.time_embed_dim, scale=1.0) + bias = _weights_fill(_BLOCK_WIDTH, scale=1.0) + for layer in range(_ARCH.num_layers): + expected = torch.nn.functional.linear(adaln_input, weight, bias).bfloat16() + got = torch.cat(_reference_block(cache, layer, slot, 1), dim=-1).reshape( + 1, _BLOCK_WIDTH + ) + assert torch.equal(got, expected) + + with pytest.raises(ValueError, match="precision"): + MiniMaxH3AdalnCache(_ARCH, weight_files=[str(weight_path)], precision="fp64") + + +def test_lora_guard_rejects_adaln_keys_in_cache_mode(): + from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( + MiniMaxH3DiTModel, + ) + + model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel) + torch.nn.Module.__init__(model) + model._adaln_precomputed = True + adapter = {"blocks.0.adaln_proj.linear.lora_A": torch.zeros(1)} + with pytest.raises(ValueError, match="adaln_proj"): + MiniMaxH3DiTModel.prepare_lora_adapter(model, adapter) + + +def test_invalidate_drops_all_tiers_and_allows_rebuild(tmp_path): + cache = _online_cache( + tmp_path, + max_plans=2, + max_plan_width=1, + host_cache_bytes=64 * _PAGE_BYTES, + ) + plan_a = torch.tensor([1.0]) + cache.build([plan_a], embed=_embed) + + cache.invalidate() + with pytest.raises(ValueError, match="does not cover"): + cache.lookup(plan_a) + with pytest.raises(ValueError, match="does not cover"): + cache.resolve_slots([plan_a]) + + passes = cache.rebuilds + cache.build([plan_a], embed=_embed) # host tier was cleared too + assert cache.rebuilds == passes + 1 + cache.lookup(plan_a) + + +def test_online_cache_eviction_preserves_in_flight_request_plans(tmp_path): + """Capacity eviction must never drop plans reused by the current request.""" cache = _online_cache(tmp_path, max_plan_width=1) plan_a = torch.tensor([1.0]) plan_b = torch.tensor([2.0]) @@ -139,6 +382,49 @@ def test_online_cache_reset_rebuilds_previously_resident_request_plans(tmp_path) cache.lookup(plan_a) cache.lookup(plan_c) + with pytest.raises(ValueError, match="does not cover"): + cache.lookup(plan_b) + + +def test_online_cache_lru_keeps_alternating_plan_sets_resident(tmp_path): + """Two alternating schedules must both stay resident once built. + + The pre-LRU slab did a full reset whenever slots overflowed, so two + alternating plan sets re-read the whole checkpoint on every request. + """ + cache = _online_cache(tmp_path, max_plans=4, max_plan_width=1) + set_a = [torch.tensor([1.0]), torch.tensor([2.0])] + set_b = [torch.tensor([3.0]), torch.tensor([4.0])] + + cache.build(set_a, embed=_embed) + cache.build(set_b, embed=_embed) + passes = cache.rebuilds + cache.build(set_a, embed=_embed) + cache.build(set_b, embed=_embed) + assert cache.rebuilds == passes + + slots_a = cache.resolve_slots(set_a) + slots_b = cache.resolve_slots(set_b) + assert sorted(slots_a.tolist() + slots_b.tolist()) == [0, 1, 2, 3] + + +def test_online_cache_evicts_least_recently_used_plan_first(tmp_path): + cache = _online_cache(tmp_path, max_plans=2, max_plan_width=1) + plan_a = torch.tensor([1.0]) + plan_b = torch.tensor([2.0]) + plan_c = torch.tensor([3.0]) + + cache.build([plan_a], embed=_embed) + cache.build([plan_b], embed=_embed) + # Touch plan_a so plan_b becomes the LRU entry, then overflow with plan_c. + cache.build([plan_a, plan_c], embed=_embed) + + cache.lookup(plan_a) + cache.lookup(plan_c) + with pytest.raises(ValueError, match="does not cover"): + cache.lookup(plan_b) + with pytest.raises(ValueError, match="does not cover"): + cache.resolve_slots([plan_b]) def test_online_cache_failed_rebuild_can_be_retried(tmp_path): @@ -168,3 +454,113 @@ def test_online_cache_width_rejection_preserves_resident_plans(tmp_path): cache.lookup(plan_a) cache.lookup(plan_b) + + +def _sidecar_cache(tmp_path: Path) -> MiniMaxH3AdalnCache: + # The weight-update guards only read which tier the cache was built as, so + # the sidecar never has to be loaded here. + return MiniMaxH3AdalnCache(_ARCH, path=str(tmp_path / "adaln.safetensors")) + + +def _cache_mode_model(cache: MiniMaxH3AdalnCache | None): + from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( + MiniMaxH3DiTModel, + ) + + model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel) + torch.nn.Module.__init__(model) + model._adaln_precomputed = True + model.adaln_cache = cache + return model + + +def test_sidecar_mode_rejects_weight_updates(tmp_path): + """A sidecar is built offline; no update can keep it in step.""" + model = _cache_mode_model(_sidecar_cache(tmp_path)) + for weights_path in (str(tmp_path), None): + with pytest.raises(ValueError, match="sidecar"): + model.validate_weight_update_source(weights_path=weights_path) + + +def test_online_cache_rejects_tensor_weight_updates(tmp_path): + """Tensor RPC carries no directory the rebuild could stream adaln from.""" + model = _cache_mode_model(_online_cache(tmp_path)) + with pytest.raises(ValueError, match="update_weights_from_disk"): + model.validate_weight_update_source(weights_path=None) + + +def test_online_cache_rejects_update_source_without_native_adaln(tmp_path): + model = _cache_mode_model(_online_cache(tmp_path)) + diffusers_layout = tmp_path / "diffusers" + diffusers_layout.mkdir() + save_file({"unrelated": torch.zeros(1)}, diffusers_layout / "model.safetensors") + + for weights_path in (str(diffusers_layout), str(tmp_path / "absent")): + with pytest.raises(ValueError, match="no native adaln_proj"): + model.validate_weight_update_source(weights_path=weights_path) + + +def test_disk_update_retargets_rebuild_source_and_drops_plans(tmp_path): + cache = _online_cache(tmp_path, max_plan_width=1) + model = _cache_mode_model(cache) + plan = torch.tensor([1.0]) + cache.build([plan], embed=_embed) + updated = tmp_path / "updated" + updated.mkdir() + _write_online_weights(updated / "model.safetensors") + + model.validate_weight_update_source(weights_path=str(updated)) + model.refresh_weight_derived_caches(weights_path=str(updated)) + + assert cache.weight_files == [str(updated / "model.safetensors")] + with pytest.raises(ValueError, match="does not cover"): + cache.lookup(plan) + + +def test_lora_ipc_layer_guard_rejects_adaln_in_cache_mode(): + """IPC passes module prefixes, not the '.lora_A' keys the disk path sees.""" + model = _cache_mode_model(None) + model.validate_lora_layers(["blocks.0.attn.qkv_proj"]) + with pytest.raises(ValueError, match="adaln_proj"): + model.validate_lora_layers(["blocks.0.adaln_proj.linear"]) + + +def _updater_for(model, model_path: str): + from types import SimpleNamespace + + from sglang.multimodal_gen.runtime.post_training.weights_updater import ( + WeightsUpdater, + ) + + model.register_parameter("probe", torch.nn.Parameter(torch.zeros(2))) + pipeline = SimpleNamespace(modules={"transformer": model}, model_path=model_path) + return WeightsUpdater(pipeline), pipeline + + +def test_weights_updater_rejects_sidecar_update_before_writing_weights(tmp_path): + model = _cache_mode_model(_sidecar_cache(tmp_path)) + updater, pipeline = _updater_for(model, str(tmp_path)) + new_checkpoint = tmp_path / "new" + (new_checkpoint / "transformer").mkdir(parents=True) + save_file( + {"probe": torch.ones(2)}, new_checkpoint / "transformer" / "model.safetensors" + ) + + ok, message = updater.update_weights_from_disk(str(new_checkpoint)) + + assert not ok + assert "sidecar" in message + # The rejection has to land before _apply_weights touches anything. + assert torch.equal(model.probe, torch.zeros(2)) + assert pipeline.model_path == str(tmp_path) + + +def test_weights_updater_rejects_tensor_update_in_online_cache_mode(tmp_path): + model = _cache_mode_model(_online_cache(tmp_path)) + updater, _ = _updater_for(model, str(tmp_path)) + + ok, message = updater.update_weights_from_tensor([("probe", torch.ones(2))]) + + assert not ok + assert "update_weights_from_disk" in message + assert torch.equal(model.probe, torch.zeros(2)) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py index ef9a18ab4e64..571917e0caeb 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py @@ -3,6 +3,7 @@ from __future__ import annotations +import os from types import SimpleNamespace from unittest.mock import patch @@ -297,6 +298,7 @@ def _quality_server_args(): enable_breakable_cuda_graph=False, enable_torch_compile=False, is_dit_layerwise_offload_selected=False, + minimax_h3_adaln_online=False, performance_mode="speed", quantization=None, transformer_weights_path=None, @@ -352,6 +354,36 @@ def test_high_quality_request_warns_when_bcg_suppresses_cache_dit(): ) +def test_admission_rejects_steps_exceeding_online_adaln_gpu_plans(): + metadata = MiniMaxH3ReleaseMetadata.from_model_index( + { + "_minimax_h3": { + "schema_version": 1, + "partition": "fl2va", + "tasks": ["t2va", "fl2va"], + "task_aliases": {}, + "sigma_shift_scales": {"video": 12.0, "audio": 3.0}, + } + } + ) + stage = MiniMaxH3PartitionAdmissionStage(metadata) + server_args = _quality_server_args() + server_args.minimax_h3_adaln_online = True + batch = SimpleNamespace( + sampling_params=SimpleNamespace(task="t2va", quality="lossless"), + num_inference_steps=50, + is_warmup=False, + ) + with patch.dict(os.environ, {"SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS": "8"}): + with pytest.raises( + ValueError, match="SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS" + ): + stage.forward(batch, server_args) + + batch.num_inference_steps = 9 + assert stage.forward(batch, server_args) is batch + + def test_extra_high_quality_does_not_enable_h3_cache_dit(): stage = MiniMaxH3DenoisingStage.__new__(MiniMaxH3DenoisingStage) stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=False) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py index 54a1b8f0fe4e..c99ce719c94e 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py @@ -65,6 +65,7 @@ def test_pruned_adaln_lora_projection_preserves_affine_term(): arch=SimpleNamespace(adaln_affine_input_dim=3, time_embed_dim=2), adaln_basis=torch.tensor([[1.0, 0.0, 2.0], [0.0, 1.0, -1.0]]), adaln_mean=torch.tensor([1.0, 2.0, 3.0]), + _adaln_precomputed=False, ) prefix = "blocks.0.adaln_proj.linear." a = torch.tensor([[2.0, 3.0, 4.0]]) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py index 269b8c81e5a6..cf79081ccc35 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py @@ -92,7 +92,10 @@ def test_fasth3_pipeline_config_gates_and_rejections() -> None: def test_fasth3_lora_bundle_is_rejected_loudly() -> None: - model = SimpleNamespace(arch=SimpleNamespace(adaln_affine_input_dim=None)) + model = SimpleNamespace( + arch=SimpleNamespace(adaln_affine_input_dim=None), + _adaln_precomputed=False, + ) plain = { "blocks.0.attn.qkv_proj.lora_A": torch.zeros(3, 64, 8), "blocks.0.attn.qkv_proj.lora_B": torch.zeros(3, 8, 64),