[diffusion] MiniMax-H3 Spectrum skip-step + fused RMSNorm/AdaLN

docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx

@@ -273,6 +273,19 @@ baseline; everything else stays identical. GPU peak stays about 18 GB
 either way because streaming offload is set by the offload buffers and VAE
 decode, not the weight dtype.
 
+To skip DiT steps on 4090 while keeping `fa`, add Spectrum to the generate
+command above. Do not combine it with `quality="high"` or Cache-DiT.
+On this card, 50-step BF16+FA (**1008.6 s**) drops to **285.3 s (3.53×)**
+with exact attention, up to **219.0 s (4.61×)** with Sage→Sol hybrid +
+Spectrum. See the 4090 table in Benchmarks.
+
+```bash
+  --enable-spectrum \
+  --spectrum-warmup-steps 11 \
+  --spectrum-window-size 5 \
+  --spectrum-flex-window 1.0
+```
+
 The first launch resolves every selected source through the normal Hub path. If
 a repository requires authentication, export a Hugging Face token in the
 server environment; no manual pre-download is required.
@@ -1417,6 +1430,7 @@ fold decision is not node-boundary aware:
 - `--use-fsdp-inference true` shards only the DiT. MiniMax-H3 preserves the original FP32 dtype of its patch, time, and output projections during FSDP all-gather, so this path does not trade numerical correctness for memory. On 4×H100, prefer TP2 + Ulysses2 for speed; use FSDP as an explicit capacity policy rather than assuming it is faster.
 - `speed` keeps model components resident, while `auto` applies the model-aware 120 GiB residency threshold. `memory` prioritizes avoiding OOM and includes the executable VAE decoder in its default layerwise set. A measured recipe with sufficient headroom can opt into `--component-residency vae=resident`; the 2×H100 CI recipe does this because the VAE's 4.8 GiB/GPU cost avoids repeated decoder transfers during tiled decode. DiT residency and prefetch knobs remain scoped to the DiT. Use `speed` only after confirming that the complete target workload fits.
 - Breakable CUDA graph execution is an explicit opt-in, not part of the recommended `speed` preset. It requires `--enable-breakable-cuda-graph`, every served size in `--warmup-resolutions`, and `--bcg-text-buckets` that cover the live H3 condition sequence. The validated 1344×768 Ref2VA recipe uses 5504; other task profiles and reference sets may need a different value. It preserves eager output for matching captured signatures, but graph capture consumes additional GPU memory and may provide little latency benefit when Ulysses attention and collectives dominate, so benchmark it on the target topology before enabling it.
+- `--enable-spectrum` skips selected DiT steps. On H3 it cannot combine with Cache-DiT / `quality="high"` or SP > 1. 50-step 4090 recipe is in §3.
 
 ## 11. Benchmarks
 
@@ -1436,7 +1450,7 @@ the configurations with collected measurements:
 | Ascend NPU | 8 NPUs, TP2 + SP4, Laser Attention | 4 NPUs, TP2 + SP2, Laser Attention |
 | MI300X / MI355X | 8× Ulysses8 resident | 1×, 2×, and 4× scaling runs |
 | RTX 5090 | 1× layerwise offload, measured on a physical desktop | 2× TP2 + layerwise offload |
-| RTX 4090 24 GB | 1× layerwise offload + `kitchen_int8` | Approximate attention backends are opt-in |
+| RTX 4090 24 GB | 1× layerwise offload + `kitchen_int8` | Spectrum and approximate attention backends are opt-in |
 
 ### GB300 single-host
 
@@ -2086,6 +2100,18 @@ and pixel fidelity rank in opposite orders there. Default remains
 `kitchen_int8` + `fa`. Cache-DiT can share this layerwise recipe; omit
 `quality` and see the quality-level section.
 
+Same card, 50-step run. Speedup is vs BF16+FA. PSNR is vs same-tree INT8+FA (Spectrum off).
+
+| Config | e2e | vs BF16+FA | PSNR mean / min |
+| --- | ---: | ---: | ---: |
+| BF16 + FlashAttention | 1008.6 s | 1.00× | — |
+| `kitchen_int8` + FA | 718.5 s | 1.40× | — |
+| `kitchen_int8` + FA + Spectrum `11 / 5 / 1.0` | **285.3 s** | **3.53×** | **27.4 / 25.3** |
+| `kitchen_int8` + `sage_attn` + Spectrum `15 / 5 / 0` | 245.3 s | 4.11× | 23.6 / 21.0 |
+| `kitchen_int8` + Sage→Sol hybrid + Spectrum `15 / 5 / 0.5` | **219.0 s** | **4.61×** | 20.6 / 17.8 |
+
+Prefer FA + `11 / 5 / 1.0` when the bar is FA-like. Sage / hybrid + Spectrum is faster; quality then follows the attention backend, not the skip schedule.
+
 ### AMD Instinct task and scaling runs
 
 The AMD recipes keep the released BF16/FP32 precision policy and use AITER

docs/docs/sglang-diffusion/spectrum.mdx

@@ -20,15 +20,19 @@ sglang generate \
 
 ## Scope and constraints
 
-- Available only on native FLUX.1, Wan, HunyuanVideo, and SD3 implementation
-  paths. It is not a `--backend diffusers` feature and does not currently
-  cover FLUX.2.
+- Available on native FLUX.1, Wan, HunyuanVideo, SD3, and MiniMax-H3
+  implementation paths. It is not a `--backend diffusers` feature and does
+  not currently cover FLUX.2.
 - The request control is available through `sglang generate` and Python
   sampling parameters. It is not a `sglang serve` or OpenAI-server request
   option yet.
-- Spectrum and `--enable-teacache` are mutually exclusive.
-- Start with the defaults. `--debug` adds shadow-prediction validation work, so
-  its latency is not representative of normal Spectrum execution.
+- Spectrum and `--enable-teacache` are mutually exclusive. On MiniMax-H3
+  it also cannot run with Cache-DiT / `quality="high"`, or with SP > 1.
+- MiniMax-H3 50-step recipe: `--spectrum-warmup-steps 11
+  --spectrum-window-size 5 --spectrum-flex-window 1.0`. Defaults are the
+  20-step schedule. See the [MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3).
+- Start with the defaults. `--debug` adds shadow-prediction validation work,
+  so its latency is not representative of normal Spectrum execution.
 
 ## Advanced controls
 

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md

@@ -168,7 +168,7 @@ These options **trade output quality** for speed or VRAM savings. Results will d
 | **Cache-DiT** | Native: per-request `--enable-cache-dit true\|false` + `--cache-dit-params <json>` (sampling params; also via API `extra_body`). `SGLANG_CACHE_DIT_ENABLED` / `SGLANG_CACHE_DIT_*` env vars are the server-wide defaults for requests that leave them unset. Diffusers backend: `--backend diffusers --cache-dit-config <yaml-or-json>` | Caches intermediate residuals across denoising steps and skips redundant computations via DBCache, TaylorSeer, and optional SCM. | ~1.5-2x on supported models | Quality depends on cache policy. Compatible with `--dit-layerwise-offload`: skipped blocks are not streamed, and the first layer after a skip may sync-load. Models that touch every layer before the block loop (for example a full-stack AdaLN prepass) must keep that prepass off while caching. Do not pass `--cache-dit-config` for native SGLang tuning unless you are intentionally using the diffusers backend flow. |
 | **CFG Gating** | Per-request `--cfg-gate-step 0.5` (sampling param; also via API `extra_body`). `SGLANG_DIFFUSION_CFG_GATE_STEP` is the server-wide default (1.0 = off). | After the given fraction of denoising steps, reuses the cached cond-uncond residual instead of running the unconditional branch each step. | Up to ~2x on the gated tail of CFG models (skips one of two branches) | Lossy; no-op without classifier-free guidance or with `--enable-cfg-parallel`. Lower fractions gate earlier and drift more. |
 | **TeaCache** | `--enable-teacache` (uses model sampling presets) | Reuses residuals when adjacent denoising steps are sufficiently similar. | Model- and threshold-dependent | Approximate and model-specific. Mutually exclusive with Spectrum. Fix prompt/seed/shape/steps and validate temporal consistency, not only single frames. |
-| **Spectrum** | `--enable-spectrum` plus optional `--spectrum-*` controls | Forecasts DiT features and skips selected denoising steps. | Defaults target an accuracy/speed tradeoff; aggressive windows can be much faster | Native `sglang generate` only for FLUX.1, Wan, HunyuanVideo, and SD3; not FLUX.2 or server requests. Mutually exclusive with TeaCache. `--debug` adds shadow validation and is not representative latency. |
+| **Spectrum** | `--enable-spectrum` plus optional `--spectrum-*` controls | Forecasts DiT features and skips selected denoising steps. | Defaults target an accuracy/speed tradeoff; aggressive windows can be much faster | Native `sglang generate` for FLUX.1, Wan, HunyuanVideo, SD3, and MiniMax-H3; not FLUX.2 or server requests. Mutually exclusive with TeaCache. `--debug` is not representative latency. |
 | **Progressive Resolution** | `--progressive-mode dct_rewind --progressive-levels N --progressive-delta D` | Runs early denoising at lower latent resolution, then spectrally upsamples and switches to the target resolution. | Model- and schedule-dependent | Approximate and pipeline-specific. Keep the switch schedule fixed and compare detail, composition, and temporal stability. |
 | **Causal KV-Cache Quantization** | `--kv-cache-quant int4\|int2` plus optional `--kv-cache-quant-*` controls | Compresses completed causal KV-cache chunks with Quant-VideoGen PRQ while keeping the mutable/current chunk and recent chunks in BF16. | Primarily a long-session memory saving | Currently limited to LingBot World realtime causal serving; requires `quant-videogen`. INT4 is the starting point; INT2 saves more memory with more error. It quantizes cache state, not checkpoint weights. |
 | **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path <path>` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. |

python/sglang/multimodal_gen/configs/sample/minimax_h3.py

@@ -240,6 +240,27 @@ def _validate(self) -> None:
                 "MiniMax H3 does not support enable_teacache: its packed "
                 "video/audio denoise loop has no lossless TeaCache contract"
             )
+        if self.enable_spectrum:
+            if self.quality == "high":
+                raise ValueError(
+                    "MiniMax H3 enable_spectrum cannot be combined with "
+                    'quality="high" (Cache-DiT). Use one skip-step accelerator.'
+                )
+            if self.enable_cache_dit:
+                raise ValueError(
+                    "MiniMax H3 enable_spectrum cannot be combined with "
+                    "enable_cache_dit. Use one skip-step accelerator."
+                )
+            from sglang.multimodal_gen.configs.sample.spectrum import SpectrumParams
+            from sglang.multimodal_gen.runtime.cache.spectrum import (
+                apply_h3_spectrum_param_defaults,
+            )
+
+            if isinstance(self.spectrum_params, dict):
+                self.spectrum_params = SpectrumParams(**self.spectrum_params)
+            elif self.spectrum_params is None:
+                self.spectrum_params = SpectrumParams()
+            apply_h3_spectrum_param_defaults(self.spectrum_params)
         if self.rollout:
             raise ValueError(
                 "MiniMax H3 does not support rollout: its coupled video/audio "

python/sglang/multimodal_gen/runtime/cache/spectrum.py

@@ -644,3 +644,103 @@ def spectrum_predict_features(self, template: torch.Tensor) -> torch.Tensor:
         predicted = forecaster.predict(step_idx)
         self._spectrum_ctx = None
         return predicted.to(dtype=template.dtype, device=template.device)
+
+
+# MiniMax-H3: forecast packed target audio|video only (Comfy / Wan2GP).
+
+H3_SPECTRUM_HISTORY_SIZE = 8
+H3_AUDIO_BLEND_WEIGHT = 0.0
+H3_VIDEO_BLEND_WEIGHT = 0.5
+
+
+def apply_h3_spectrum_param_defaults(params: SpectrumParams) -> SpectrumParams:
+    """Shrink the generic 100-slot GPU buffer default to the H3 CPU window."""
+    if params.history_size == 100:
+        params.history_size = H3_SPECTRUM_HISTORY_SIZE
+    return params
+
+
+def target_audio_video_span(
+    audio_pos: torch.Tensor,
+    infer_out_pos: torch.Tensor,
+) -> tuple[int, int, int]:
+    """Return ``(audio_start, audio_stop, video_stop)`` in packed row indices.
+
+    Target video is ``infer_out_pos``. Target audio is the contiguous
+    ``audio_pos`` run that ends immediately before that video span — the
+    packed t2va / fl2va / ref2va suffix. Earlier reference-audio rows are
+    dropped, matching Wan2GP ``hidden[audio_start:]``.
+    """
+    audio_ids = audio_pos.detach().reshape(-1).to(dtype=torch.long, device="cpu")
+    video_ids = infer_out_pos.detach().reshape(-1).to(dtype=torch.long, device="cpu")
+    if audio_ids.numel() == 0 or video_ids.numel() == 0:
+        raise ValueError(
+            "MiniMax-H3 Spectrum needs non-empty target audio and video rows"
+        )
+    video_start = int(video_ids.min())
+    video_stop = int(video_ids.max()) + 1
+    if int(video_ids.numel()) != video_stop - video_start:
+        raise ValueError(
+            "MiniMax-H3 Spectrum target video rows must be a contiguous packed span"
+        )
+
+    preceding = audio_ids[audio_ids < video_start]
+    if preceding.numel() == 0 or int(preceding.max()) + 1 != video_start:
+        raise ValueError(
+            "MiniMax-H3 Spectrum expects packed target audio rows immediately "
+            f"before target video rows starting at {video_start}"
+        )
+    ordered = preceding.sort().values
+    expected = video_start - 1
+    cursor = int(ordered.numel()) - 1
+    while cursor >= 0 and int(ordered[cursor]) == expected:
+        expected -= 1
+        cursor -= 1
+    audio_start = expected + 1
+    return audio_start, video_start, video_stop
+
+
+def local_target_hidden(
+    hidden: torch.Tensor,
+    audio_start: int,
+    video_stop: int,
+    row_start: int,
+    row_stop: int,
+) -> tuple[torch.Tensor, int, int] | None:
+    """Slice rank-local ``hidden`` to the overlap with packed target media."""
+    start = max(audio_start, row_start)
+    stop = min(video_stop, row_stop)
+    if stop <= start:
+        return None
+    local_start = start - row_start
+    local_stop = stop - row_start
+    return hidden[local_start:local_stop], local_start, local_stop
+
+
+def blend_h3_spectrum_prediction(
+    predicted: torch.Tensor,
+    *,
+    n_audio: int,
+    last_audio: torch.Tensor | None,
+    last_video: torch.Tensor | None,
+    audio_blend: float = H3_AUDIO_BLEND_WEIGHT,
+    video_blend: float = H3_VIDEO_BLEND_WEIGHT,
+) -> torch.Tensor:
+    """Apply Comfy segment blends on a ``[audio | video]`` forecast."""
+    if n_audio < 0 or n_audio > int(predicted.shape[0]):
+        raise ValueError(
+            f"n_audio={n_audio} is outside predicted rows {predicted.shape[0]}"
+        )
+    out = predicted
+    audio = out[:n_audio]
+    video = out[n_audio:]
+    if last_audio is not None and audio.numel():
+        held = last_audio.to(device=audio.device, dtype=audio.dtype)
+        if audio_blend <= 0.0:
+            audio.copy_(held)
+        elif audio_blend < 1.0:
+            audio.copy_(torch.lerp(held, audio, audio_blend))
+    if last_video is not None and video.numel() and video_blend < 1.0:
+        held = last_video.to(device=video.device, dtype=video.dtype)
+        video.copy_(torch.lerp(held, video, video_blend))
+    return out