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 22f78948d..392107ed8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -849,6 +849,13 @@ class MiniMaxH3AdalnCache(nn.Module): "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 @@ -981,18 +988,12 @@ class MiniMaxH3AdalnCache(nn.Module): missing = {k: v for k, v in wanted.items() if k not in self._slots} if not missing: return - if len(self._slots) + len(missing) > self.max_plans: - # Every plan a request looks up has to stay resident for the whole - # denoise loop, so an overflow means the capacity is too small -- - # evicting part of it would only move the failure into lookup(). - self._slots.clear() - self.plan_lengths.zero_() - if len(missing) > self.max_plans: + if len(wanted) > self.max_plans: raise ValueError( - f"MiniMax H3 AdaLN rebuild needs {len(missing)} plans but " + f"MiniMax H3 AdaLN rebuild needs {len(wanted)} plans but " f"max_plans is {self.max_plans}" ) - widest = max(timesteps.numel() for timesteps in missing.values()) + 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 " @@ -1000,11 +1001,20 @@ class MiniMaxH3AdalnCache(nn.Module): "--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 = [] - for key, timesteps in missing.items(): - slot = len(self._slots) - self._slots[key] = slot + 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) @@ -1052,10 +1062,14 @@ class MiniMaxH3AdalnCache(nn.Module): 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(missing), + len(plans_to_build), len(self._slots), self.max_plans, self.rebuilds, 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 c4dbc8296..e7b162bd3 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 @@ -1,31 +1,99 @@ # SPDX-License-Identifier: Apache-2.0 +from pathlib import Path + +import pytest import torch from safetensors.torch import save_file from sglang.multimodal_gen.configs.models.dits.minimax_h3 import ( + MINIMAX_H3_ADALN_MODALITY_NUM, MiniMaxH3DiTArchConfig, ) +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + maybe_init_distributed_environment_and_model_parallel, + model_parallel_is_initialized, +) from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( MiniMaxH3AdalnCache, ) +from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( + ensure_distributed_env_defaults, +) + +_ARCH = MiniMaxH3DiTArchConfig( + num_layers=2, + hidden_size=4, + time_embed_dim=3, +) +_BLOCK_WIDTH = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * _ARCH.hidden_size +_FINAL_WIDTH = 2 * _ARCH.hidden_size + + +def _ensure_single_process_parallel_runtime() -> None: + if model_parallel_is_initialized(): + return + ensure_distributed_env_defaults() + maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1) + + +def _write_online_weights( + path: Path, + *, + omit: str | None = None, +) -> None: + # These cache-state tests need checkpoint-compatible shapes, not values. + 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) + 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) + if omit is not None: + tensors.pop(omit) + save_file(tensors, path) + + +def _online_cache( + tmp_path: Path, + *, + max_plans: int = 2, + max_plan_width: int = 2, + omit: str | None = None, +) -> MiniMaxH3AdalnCache: + _ensure_single_process_parallel_runtime() + weight_path = tmp_path / "model.safetensors" + _write_online_weights(weight_path, omit=omit) + cache = MiniMaxH3AdalnCache( + _ARCH, + weight_files=[str(weight_path)], + max_plans=max_plans, + max_plan_width=max_plan_width, + ) + cache.load(torch.device("cpu")) + return cache + + +def _embed(timesteps: torch.Tensor) -> torch.Tensor: + return timesteps[:, None].expand(-1, _ARCH.time_embed_dim) def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path): - arch = MiniMaxH3DiTArchConfig( - num_layers=2, - hidden_size=4, - time_embed_dim=3, - ) cache_path = tmp_path / "adaln.safetensors" plan_timesteps = torch.tensor([[0.0, 0.0], [1.0, 2.0]]) plan_lengths = torch.tensor([1, 2], dtype=torch.int64) block_params = ( - torch.arange(2 * 2 * 2 * 72, dtype=torch.float32) - .reshape(2, 2, 2, 72) + torch.arange(2 * 2 * 2 * _BLOCK_WIDTH, dtype=torch.float32) + .reshape(2, 2, 2, _BLOCK_WIDTH) + .bfloat16() + ) + final_params = ( + torch.arange(2 * 2 * _FINAL_WIDTH, dtype=torch.float32) + .reshape(2, 2, _FINAL_WIDTH) .bfloat16() ) - final_params = torch.arange(32, dtype=torch.float32).reshape(2, 2, 8).bfloat16() save_file( { "plan_timesteps": plan_timesteps, @@ -38,7 +106,7 @@ def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path): ) cache = MiniMaxH3AdalnCache( - arch, + _ARCH, path=str(cache_path), model_variant="fl2va", ) @@ -57,3 +125,46 @@ def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path): block_params[1, :, 1], ) 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.""" + cache = _online_cache(tmp_path, 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, plan_b], embed=_embed) + cache.build([plan_a, plan_c], embed=_embed) + + cache.lookup(plan_a) + cache.lookup(plan_c) + + +def test_online_cache_failed_rebuild_can_be_retried(tmp_path): + """A failed rebuild must not publish a cache hit that blocks its retry.""" + missing_name = "final_layer.adaln_proj.linear.bias" + cache = _online_cache(tmp_path, omit=missing_name) + plan_a = torch.tensor([1.0]) + + with pytest.raises(KeyError, match=missing_name): + cache.build([plan_a], embed=_embed) + + _write_online_weights(tmp_path / "model.safetensors") + cache.build([plan_a], embed=_embed) + cache.lookup(plan_a) + + +def test_online_cache_width_rejection_preserves_resident_plans(tmp_path): + """Rejecting an over-width plan must not evict usable resident plans.""" + cache = _online_cache(tmp_path, max_plan_width=1) + plan_a = torch.tensor([1.0]) + plan_b = torch.tensor([2.0]) + wide_plan = torch.tensor([3.0, 4.0]) + + cache.build([plan_a, plan_b], embed=_embed) + with pytest.raises(ValueError, match="--minimax-h3-adaln-plan-width"): + cache.build([wide_plan], embed=_embed) + + cache.lookup(plan_a) + cache.lookup(plan_b)