minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: mickqian <mickqian@users.noreply.github.com>
This commit is contained in:
triple-mu
2026-09-03 22:09:21 +08:00
committed by GitHub
co-authored by mickqian
parent 4e37882a93
commit bf71035d39
15 changed files with 1621 additions and 335 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/envs.py
+15
View File
@@ -37,6 +37,8 @@ if TYPE_CHECKING:
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 @@ environment_variables: dict[str, Callable[[], Any]] = {
# 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
+19
View File
@@ -6,6 +6,7 @@ from typing import Any
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 @@ class TransformerLoader(OnlineQuantizationComponentLoader):
"--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
+28
View File
@@ -109,6 +109,34 @@ class BaseDiT(nn.Module, ABC):
"""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
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+128 -318
View File
@@ -9,15 +9,12 @@ from __future__ import annotations
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 import (
)
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 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
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,
@@ -192,6 +198,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset(
"img_position_ids",
"rope_cache",
"unique_timesteps",
"adaln_cache_slot",
"inverse_indices",
"update_mask",
"update_audio_mask",
@@ -1229,303 +1236,6 @@ class MiniMaxH3AdalnProj(nn.Module):
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("<I", struct.pack("<f", float(value)))[0]
for value in timesteps.tolist()
)
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,
) -> 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."""
@@ -1820,6 +1530,24 @@ class MiniMaxH3FinalLayer(nn.Module):
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",
@@ -1842,6 +1570,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
) -> 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
@@ -1897,20 +1627,94 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
)
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 (
@@ -1994,6 +1798,9 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
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
@@ -2126,7 +1933,10 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
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
@@ -2689,16 +2499,16 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
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(
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_adaln_cache.py
+853
View File
@@ -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("<I", struct.pack("<f", float(value)))[0]
for value in timesteps.tolist()
)
def native_adaln_weight_files(weights_path: str) -> 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))
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py
+11 -2
View File
@@ -319,6 +319,7 @@ class MiniMaxH3DenoiseBranch:
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 @@ class MiniMaxH3DenoiseBranch:
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 @@ class MiniMaxH3DenoiseBranch:
"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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py
+14
View File
@@ -7,6 +7,7 @@ import math
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 @@ class MiniMaxH3PartitionAdmissionStage(PipelineStage):
"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(
python/sglang/multimodal_gen/runtime/post_training/weights_updater.py
+47 -4
View File
@@ -57,6 +57,7 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import (
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 @@ class WeightsUpdater:
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 @@ class WeightsUpdater:
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 @@ class WeightsUpdater:
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 @@ class WeightsUpdater:
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 @@ class WeightsUpdater:
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 @@ class WeightsUpdater:
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] = []
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+32
View File
@@ -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 @@ class ServerArgs(DisaggServerArgsMixin):
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 @@ class ServerArgs(DisaggServerArgsMixin):
"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,
python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py
+406 -10
View File
@@ -14,7 +14,7 @@ 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 (
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("<f", struct.pack("<I", bits))[0] for bits in key)
for key in tier._plans
}
assert keys == {(1.0,), (3.0,), (4.0,)}
def test_precision_fp32_projects_in_fp32_then_stores_bf16(tmp_path):
_ensure_single_process_parallel_runtime()
weight_path = tmp_path / "model.safetensors"
_write_online_weights(weight_path, scale=1.0)
cache = MiniMaxH3AdalnCache(
_ARCH,
weight_files=[str(weight_path)],
max_plans=2,
max_plan_width=1,
precision="fp32",
)
cache.load(torch.device("cpu"))
plan = torch.tensor([0.31])
def bf16_embed(timesteps: torch.Tensor) -> 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))
python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
+32
View File
@@ -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)
python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
+1
View File
@@ -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]])
python/sglang/multimodal_gen/test/unit/test_minimax_h3_fasth3.py
+4 -1
View File
@@ -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),