[diffusion] fix: make MiniMax-H3 AdaLN cache rebuild transactional (#34993)

Co-authored-by: Mick <mickjagger19@icloud.com>

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,

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)