[diffusion] MiniMax-H3: tiered AdaLN plan cache (pinned-host tier + per-plan LRU)

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

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.

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(

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"]

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