From 4ea2fa9478a312d910dff33cdcf11531ece4dd37 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Wed, 19 Aug 2026 13:36:09 +0000
Subject: [PATCH 1/8] [diffusion] Wire Spectrum skip-step into MiniMax-H3.

Forecast packed target audio/video so H3 can skip the DiT stack without Cache-DiT, while keeping layerwise offload and a bit-exact default path.
---
 .../configs/sample/minimax_h3.py              |  21 ++
 .../multimodal_gen/runtime/cache/spectrum.py  | 100 +++++++++
 .../runtime/models/dits/minimax_h3.py         | 199 ++++++++++++++----
 .../minimax_h3/stages/denoising.py            |   3 +
 .../test/unit/test_minimax_h3_spectrum.py     |  78 +++++++
 5 files changed, 361 insertions(+), 40 deletions(-)
 create mode 100644 python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py

diff --git a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
index 603500183441..3b3f2c480c39 100644
--- a/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
+++ b/python/sglang/multimodal_gen/configs/sample/minimax_h3.py
@@ -214,6 +214,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 "
diff --git a/python/sglang/multimodal_gen/runtime/cache/spectrum.py b/python/sglang/multimodal_gen/runtime/cache/spectrum.py
index 1bcb9b9de69e..4e39e45bfcc9 100644
--- a/python/sglang/multimodal_gen/runtime/cache/spectrum.py
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
index 4079ad8e8c19..9934c096ee38 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -36,6 +36,12 @@
     MiniMaxH3DiTConfig,
 )
 from sglang.multimodal_gen.configs.models.fsdp import is_block
+from sglang.multimodal_gen.runtime.cache.spectrum import (
+    SpectrumMixin,
+    blend_h3_spectrum_prediction,
+    local_target_hidden,
+    target_audio_video_span,
+)
 from sglang.multimodal_gen.runtime.distributed import (
     get_tp_world_size,
     tensor_model_parallel_all_gather,
@@ -58,6 +64,7 @@
     QuantizationConfig,
 )
 from sglang.multimodal_gen.runtime.layers.usp import _ring_attention_varlen
+from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
     is_layerwise_offloaded_module,
@@ -1546,7 +1553,7 @@ def forward(
         return video, audio
 
 
-class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
+class MiniMaxH3DiTModel(SpectrumMixin, BaseDiT, LayerwiseOffloadableModuleMixin):
     _fsdp_shard_conditions = [is_block]
     # refine_prompt_embeds drives a forward pass outside __call__.
     _fsdp_forward_methods = ("refine_prompt_embeds",)
@@ -1734,7 +1741,11 @@ def __init__(
                 for index in range(arch.num_layers)
             ]
         )
-        self.layer_names = ["token_refiner.blocks", "blocks"]
+        # Offload the 50-layer DiT stack before token_refiner. Each
+        # LayerwiseOffloadManager ends with model.to(device) for leftovers;
+        # putting the 2-layer refiner first would copy the full INT8 DiT onto
+        # a 24GB card and OOM.
+        self.layer_names = ["blocks", "token_refiner.blocks"]
         self.final_layer = MiniMaxH3FinalLayer(
             arch,
             quant_config,
@@ -1754,6 +1765,9 @@ def __init__(
         )
         self._resolved_attention_backend: AttentionBackendEnum | None = None
         self._mark_missing_params_required()
+        self._init_spectrum_state()
+        self._h3_spectrum_last_audio: torch.Tensor | None = None
+        self._h3_spectrum_last_video: torch.Tensor | None = None
 
     def set_cache_dit_input_preservation(self, enabled: bool) -> None:
         """Stop the blocks from overwriting the input Cache-DiT holds by reference.
@@ -1771,6 +1785,78 @@ def set_cache_dit_input_preservation(self, enabled: bool) -> None:
         for block in self.blocks:
             block.preserve_input_for_cache_dit = enabled
 
+    def _h3_spectrum_requested(self, *, sp_ws: int) -> bool:
+        """True when this forward should use Spectrum. Off path stays bit-exact."""
+        try:
+            forward_batch = get_forward_context().forward_batch
+        except AssertionError:
+            return False
+        if forward_batch is None or not getattr(
+            forward_batch, "enable_spectrum", False
+        ):
+            return False
+        if sp_ws > 1:
+            raise ValueError(
+                "MiniMax-H3 Spectrum requires sequence-parallel size 1 so "
+                "target audio/video rows stay on one rank. Disable "
+                "--enable-spectrum or run without SP / Ulysses / ring."
+            )
+        return True
+
+    def _h3_spectrum_record_targets(
+        self,
+        hidden: torch.Tensor,
+        audio_pos: torch.Tensor,
+        infer_out_pos: torch.Tensor,
+        row_start: int,
+        row_stop: int,
+    ) -> None:
+        audio_start, audio_stop, video_stop = target_audio_video_span(
+            audio_pos, infer_out_pos
+        )
+        sliced = local_target_hidden(
+            hidden, audio_start, video_stop, row_start, row_stop
+        )
+        if sliced is None:
+            return
+        local, _, _ = sliced
+        n_audio = audio_stop - audio_start
+        cpu_local = local.detach().to(device="cpu", dtype=local.dtype)
+        self._h3_spectrum_last_audio = cpu_local[:n_audio].contiguous()
+        self._h3_spectrum_last_video = cpu_local[n_audio:].contiguous()
+        self.spectrum_record_features(cpu_local)
+
+    def _h3_spectrum_predict_targets(
+        self,
+        hidden: torch.Tensor,
+        audio_pos: torch.Tensor,
+        infer_out_pos: torch.Tensor,
+        row_start: int,
+        row_stop: int,
+    ) -> torch.Tensor:
+        audio_start, audio_stop, video_stop = target_audio_video_span(
+            audio_pos, infer_out_pos
+        )
+        sliced = local_target_hidden(
+            hidden, audio_start, video_stop, row_start, row_stop
+        )
+        if sliced is None:
+            return hidden
+        local, local_start, local_stop = sliced
+        predicted = self.spectrum_predict_features(
+            local.detach().to(device="cpu", dtype=local.dtype)
+        )
+        predicted = blend_h3_spectrum_prediction(
+            predicted,
+            n_audio=audio_stop - audio_start,
+            last_audio=self._h3_spectrum_last_audio,
+            last_video=self._h3_spectrum_last_video,
+        )
+        hidden[local_start:local_stop].copy_(
+            predicted.to(device=hidden.device, dtype=hidden.dtype)
+        )
+        return hidden
+
     def _resolve_attention_backend_once(self) -> None:
         if self._resolved_attention_backend is not None:
             return
@@ -2243,48 +2329,81 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
 
         hidden = decoder_input
         cu_seqlens = cu_seqlens.to(device)
+        spectrum_on = self._h3_spectrum_requested(sp_ws=sp_ws)
+        run_transformer_blocks = (
+            self.begin_spectrum_step() if spectrum_on else True
+        )
+        if spectrum_on:
+            try:
+                timestep = get_forward_context().current_timestep
+            except AssertionError:
+                timestep = None
+            if timestep == 0:
+                self._h3_spectrum_last_audio = None
+                self._h3_spectrum_last_video = None
         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],
+        # Skip-step forecasts still run embed + final_layer. Do not touch
+        # per-block AdaLN or the 50-layer stack — that is what lets Spectrum
+        # coexist with DiT layerwise offload.
+        if run_transformer_blocks:
+            if self.adaln_cache is not None:
+                adaln_cache_plan_index = self.adaln_cache.lookup(
+                    unique_timesteps.view(-1).to(device)
                 )
-                for index in range(len(self.blocks))
-            )
-        elif self._can_batch_block_adaln():
-            local_adaln = torch.stack(
-                [block.adaln_proj.project_local(adaln_input) for block in self.blocks]
-            )
-            gathered_adaln = tensor_model_parallel_all_gather(local_adaln)
-            block_adaln_params = tuple(
-                block.adaln_proj.split_output(output)
-                for block, output in zip(self.blocks, gathered_adaln)
+                block_adaln_params = tuple(
+                    self.adaln_cache.block(
+                        index,
+                        adaln_cache_plan_index,
+                        adaln_input.shape[0],
+                    )
+                    for index in range(len(self.blocks))
+                )
+            elif self._can_batch_block_adaln():
+                local_adaln = torch.stack(
+                    [
+                        block.adaln_proj.project_local(adaln_input)
+                        for block in self.blocks
+                    ]
+                )
+                gathered_adaln = tensor_model_parallel_all_gather(local_adaln)
+                block_adaln_params = tuple(
+                    block.adaln_proj.split_output(output)
+                    for block, output in zip(self.blocks, gathered_adaln)
+                )
+            # With sequence parallelism, shard rows across the group for the
+            # block stack. Attention trades sequence for heads internally
+            # (Ulysses) and/or ring-rotates KV across ring ranks; everything
+            # else, including the final layer, is row-local. Only the narrow
+            # video/audio logits are gathered after the final layer.
+            for index, block in enumerate(self.blocks):
+                hidden = block(
+                    hidden,
+                    adaln_input=adaln_input,
+                    combined_indices=block_combined,
+                    rope_cache=rope_cache,
+                    cu_seqlens=cu_seqlens,
+                    cu_seqlens_host=cu_seqlens_host,
+                    max_seqlen=max_seqlen,
+                    ulysses_active=ulysses_ws > 1,
+                    ring_active=ring_ws > 1,
+                    adaln_params=(
+                        None
+                        if block_adaln_params is None
+                        else block_adaln_params[index]
+                    ),
+                )
+            if spectrum_on:
+                self._h3_spectrum_record_targets(
+                    hidden, audio_pos, infer_out_pos, row_start, row_stop
+                )
+        elif spectrum_on:
+            hidden = self._h3_spectrum_predict_targets(
+                hidden, audio_pos, infer_out_pos, row_start, row_stop
             )
-        # With sequence parallelism, shard rows across the group for the
-        # block stack. Attention trades sequence for heads internally
-        # (Ulysses) and/or ring-rotates KV across ring ranks; everything
-        # else, including the final layer, is row-local. Only the narrow
-        # video/audio logits are gathered after the final layer.
-        for index, block in enumerate(self.blocks):
-            hidden = block(
-                hidden,
-                adaln_input=adaln_input,
-                combined_indices=block_combined,
-                rope_cache=rope_cache,
-                cu_seqlens=cu_seqlens,
-                cu_seqlens_host=cu_seqlens_host,
-                max_seqlen=max_seqlen,
-                ulysses_active=ulysses_ws > 1,
-                ring_active=ring_ws > 1,
-                adaln_params=(
-                    None if block_adaln_params is None else block_adaln_params[index]
-                ),
+        if self.adaln_cache is not None and adaln_cache_plan_index is None:
+            adaln_cache_plan_index = self.adaln_cache.lookup(
+                unique_timesteps.view(-1).to(device)
             )
         self.materialize_mps_non_layer_weights("final_layer")
         video_logits, audio_logits = self.final_layer(
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
index 381d257e440f..f8fca9bb3e50 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
@@ -415,6 +415,9 @@ def _maybe_enable_cache_dit(
         if enable_override is False:
             # The per-request kill switch wins over quality="high".
             desired_mode = None
+        elif getattr(batch.sampling_params, "enable_spectrum", False):
+            # Spectrum skips the block stack; Cache-DiT wraps those blocks.
+            desired_mode = None
         elif quality == "high":
             desired_mode = "high"
         else:
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py
new file mode 100644
index 000000000000..ac2bd0e56c6e
--- /dev/null
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""H3 Spectrum packing helpers and sampling-param gates."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
+from sglang.multimodal_gen.runtime.cache.spectrum import (
+    H3_SPECTRUM_HISTORY_SIZE,
+    blend_h3_spectrum_prediction,
+    local_target_hidden,
+    target_audio_video_span,
+)
+
+
+def test_target_span_uses_contiguous_audio_before_video():
+    audio_pos = torch.tensor([4, 5, 6, 7, 8, 9])
+    infer_out_pos = torch.tensor([10, 11, 12, 13])
+    start, stop, video_stop = target_audio_video_span(audio_pos, infer_out_pos)
+    assert (start, stop, video_stop) == (4, 10, 14)
+
+
+def test_target_span_drops_earlier_ref_audio_gap():
+    audio_pos = torch.tensor([1, 2, 8, 9])
+    infer_out_pos = torch.tensor([10, 11])
+    start, stop, video_stop = target_audio_video_span(audio_pos, infer_out_pos)
+    assert (start, stop, video_stop) == (8, 10, 12)
+
+
+def test_local_target_hidden_is_rank_local_view():
+    hidden = torch.arange(20, dtype=torch.float32).unsqueeze(-1).expand(20, 3).clone()
+    sliced = local_target_hidden(hidden, audio_start=4, video_stop=10, row_start=0, row_stop=20)
+    assert sliced is not None
+    local, start, stop = sliced
+    assert start == 4 and stop == 10
+    assert local.shape == (6, 3)
+
+
+def test_blend_holds_audio_and_mixes_video():
+    predicted = torch.ones(6, 2)
+    last_audio = torch.full((2, 2), 3.0)
+    last_video = torch.full((4, 2), 5.0)
+    out = blend_h3_spectrum_prediction(
+        predicted.clone(),
+        n_audio=2,
+        last_audio=last_audio,
+        last_video=last_video,
+        audio_blend=0.0,
+        video_blend=0.5,
+    )
+    assert torch.equal(out[:2], last_audio)
+    assert torch.allclose(out[2:], torch.full((4, 2), 3.0))
+
+
+def test_h3_sampling_params_reject_quality_high_with_spectrum():
+    with pytest.raises(ValueError, match="enable_spectrum cannot be combined"):
+        MiniMaxH3SamplingParams(
+            prompt="cat",
+            enable_spectrum=True,
+            quality="high",
+        )
+
+
+def test_h3_sampling_params_default_history_is_comfy_window():
+    params = MiniMaxH3SamplingParams(prompt="cat", enable_spectrum=True)
+    assert params.spectrum_params is not None
+    assert params.spectrum_params.history_size == H3_SPECTRUM_HISTORY_SIZE
+
+
+def test_h3_sampling_params_keep_explicit_history_size():
+    params = MiniMaxH3SamplingParams(
+        prompt="cat",
+        enable_spectrum=True,
+        spectrum_params={"history_size": 6},
+    )
+    assert params.spectrum_params.history_size == 6

From 80c028c4d6a4519dc086e3805d7a9a1fa9c02611 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Thu, 20 Aug 2026 07:54:23 +0000
Subject: [PATCH 2/8] [diffusion] Fuse RMSNorm + indexed AdaLN on MiniMax-H3.

Skip materializing the normalized tensor on the H3 block/final-layer AdaLN path.
---
 .../sglang/kernels/ops/diffusion/__init__.py  |   9 ++
 .../modulate/indexed_modulation_triton.py     | 125 +++++++++++++++++-
 .../runtime/models/dits/minimax_h3.py         |  56 ++++++--
 3 files changed, 178 insertions(+), 12 deletions(-)

diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py
index 2759c5d251e0..1ca89859c5ec 100644
--- a/python/sglang/kernels/ops/diffusion/__init__.py
+++ b/python/sglang/kernels/ops/diffusion/__init__.py
@@ -186,6 +186,13 @@
         _CUDA,
         "LTX-2 nine-way adaLN value split.",
     ),
+    (
+        "diffusion.fused_indexed_rmsnorm_scale_shift",
+        KernelBackend.TRITON,
+        "modulate.indexed_modulation_triton:fused_indexed_rmsnorm_scale_shift",
+        _CUDA,
+        "Indexed RMSNorm + adaLN scale/shift (MiniMax-H3).",
+    ),
     (
         "diffusion.fused_inplace_qknorm_rope",
         KernelBackend.JIT,
@@ -364,6 +371,8 @@
     "can_use_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_triton",
     "zimage_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_triton",
     # adaLN modulation, gating and timestep conditioning
+    "can_use_fused_indexed_rmsnorm_scale_shift": "modulate.indexed_modulation_triton",
+    "fused_indexed_rmsnorm_scale_shift": "modulate.indexed_modulation_triton",
     "indexed_gate_bf16": "modulate.indexed_modulation_triton",
     "indexed_gate_bf16_": "modulate.indexed_modulation_triton",
     "indexed_scale_shift_bf16_": "modulate.indexed_modulation_triton",
diff --git a/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py b/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py
index f6aff00832ae..dac8ce37b13b 100644
--- a/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py
+++ b/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py
@@ -4,7 +4,11 @@
 import triton
 import triton.language as tl
 
-from sglang.kernels.ops.diffusion.common.numerics import round_bf16_to_fp32
+from sglang.kernels.ops.diffusion.common.numerics import (
+    cuda_rsqrtf,
+    mul_rn_f32,
+    round_bf16_to_fp32,
+)
 
 
 @triton.jit
@@ -155,3 +159,122 @@ def indexed_gate_bf16(
     indices: torch.Tensor,
 ) -> torch.Tensor:
     return _indexed_gate_bf16(torch.empty_like(x), x, gate, other, indices)
+
+
+@triton.jit
+def _fused_indexed_rmsnorm_scale_shift_kernel(
+    out_ptr,
+    x_ptr,
+    weight_ptr,
+    shift_ptr,
+    scale_ptr,
+    indices_ptr,
+    hidden_size,
+    stride_x_row,
+    stride_out_row,
+    stride_shift_row,
+    stride_scale_row,
+    stride_indices,
+    eps,
+    BLOCK_N: tl.constexpr,
+):
+    """RMSNorm(x) * (1 + scale[idx]) + shift[idx], one kernel, leave x intact."""
+    row = tl.program_id(0)
+    columns = tl.arange(0, BLOCK_N)
+    acc = 0.0
+    offs = 0
+    while offs < hidden_size:
+        cols = offs + columns
+        mask = cols < hidden_size
+        x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(
+            tl.float32
+        )
+        acc += tl.sum(mul_rn_f32(x, x), axis=0)
+        offs += BLOCK_N
+    rstd = cuda_rsqrtf(acc / hidden_size + eps)
+    index = tl.load(indices_ptr + row * stride_indices)
+    offs = 0
+    while offs < hidden_size:
+        cols = offs + columns
+        mask = cols < hidden_size
+        x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(
+            tl.float32
+        )
+        weight = tl.load(weight_ptr + cols, mask=mask, other=1.0).to(tl.float32)
+        shift = tl.load(
+            shift_ptr + index * stride_shift_row + cols, mask=mask, other=0.0
+        ).to(tl.float32)
+        scale = tl.load(
+            scale_ptr + index * stride_scale_row + cols, mask=mask, other=0.0
+        ).to(tl.float32)
+        y = round_bf16_to_fp32(x * rstd * weight)
+        one_plus = round_bf16_to_fp32(1.0 + scale)
+        tl.store(
+            out_ptr + row * stride_out_row + cols,
+            round_bf16_to_fp32(y * one_plus) + shift,
+            mask=mask,
+        )
+        offs += BLOCK_N
+
+
+def can_use_fused_indexed_rmsnorm_scale_shift(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    indices: torch.Tensor,
+) -> bool:
+    return (
+        x.is_cuda
+        and x.dim() == 2
+        and x.dtype == torch.bfloat16
+        and x.is_contiguous()
+        and weight.is_cuda
+        and weight.dtype == torch.bfloat16
+        and weight.is_contiguous()
+        and weight.shape == (x.shape[-1],)
+        and shift.is_cuda
+        and scale.is_cuda
+        and shift.dtype == scale.dtype == torch.bfloat16
+        and shift.shape == scale.shape
+        and shift.shape[-1] == x.shape[-1]
+        and shift.stride(-1) == 1
+        and scale.stride(-1) == 1
+        and indices.is_cuda
+        and indices.dtype in (torch.int32, torch.int64)
+        and indices.numel() == x.shape[0]
+    )
+
+
+def fused_indexed_rmsnorm_scale_shift(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    indices: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """``norm(x) * (1 + scale[idx]) + shift[idx]`` without materializing norm(x)."""
+    rows, hidden_size = x.shape
+    out = torch.empty_like(x)
+    if rows == 0:
+        return out
+    block_n = min(triton.next_power_of_2(hidden_size), 1024)
+    _fused_indexed_rmsnorm_scale_shift_kernel[(rows,)](
+        out,
+        x,
+        weight,
+        shift,
+        scale,
+        indices,
+        hidden_size,
+        x.stride(0),
+        out.stride(0),
+        shift.stride(0),
+        scale.stride(0),
+        indices.stride(0),
+        float(eps),
+        BLOCK_N=block_n,
+        num_warps=8,
+    )
+    return out
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
index 9934c096ee38..a361d6974a13 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -21,7 +21,9 @@
     silu_and_mul_with_activation_rounding_,
 )
 from sglang.kernels.ops.diffusion import (
+    can_use_fused_indexed_rmsnorm_scale_shift,
     can_use_fused_inplace_qknorm_rope,
+    fused_indexed_rmsnorm_scale_shift,
     fused_inplace_qknorm_rope,
     indexed_gate_bf16,
     indexed_gate_bf16_,
@@ -236,6 +238,29 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
     return torch.cat((-x2, x1), dim=-1)
 
 
+def _modulate_rmsnorm_scale_shift(
+    x: torch.Tensor,
+    norm: nn.RMSNorm,
+    shift: torch.Tensor,
+    scale: torch.Tensor,
+    indices: torch.Tensor,
+    *,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """RMSNorm + indexed AdaLN without materializing the normalized tensor."""
+    if (
+        not torch.compiler.is_compiling()
+        and dtype == _BF16_DTYPE
+        and can_use_fused_indexed_rmsnorm_scale_shift(
+            x, norm.weight, shift, scale, indices
+        )
+    ):
+        return fused_indexed_rmsnorm_scale_shift(
+            x, norm.weight, shift, scale, indices, float(norm.eps)
+        )
+    return _modulate_scale_shift(norm(x), shift, scale, indices, dtype=dtype)
+
+
 def _modulate_scale_shift(
     x: torch.Tensor,
     shift: torch.Tensor,
@@ -1405,9 +1430,13 @@ def forward(
         # first gated residual writes to that tensor; the second one operates on
         # a block-local buffer.
         residual = x
-        h = self.norm1(x)
-        h = _modulate_scale_shift(
-            h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE
+        h = _modulate_rmsnorm_scale_shift(
+            x,
+            self.norm1,
+            shift_msa,
+            scale_msa,
+            combined_indices,
+            dtype=_BF16_DTYPE,
         )
         h = self.attn(
             h,
@@ -1428,9 +1457,13 @@ def forward(
         )
 
         residual = x
-        h = self.norm2(x)
-        h = _modulate_scale_shift(
-            h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE
+        h = _modulate_rmsnorm_scale_shift(
+            x,
+            self.norm2,
+            shift_mlp,
+            scale_mlp,
+            combined_indices,
+            dtype=_BF16_DTYPE,
         )
         h = self.mlp(h)
         # `residual` is block-local here (see above), so this stays in-place
@@ -1516,9 +1549,9 @@ def forward(
             video = audio = None
             for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE):
                 stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0])
-                h = self.norm(x[start:stop])
-                h = _modulate_scale_shift(
-                    h,
+                h = _modulate_rmsnorm_scale_shift(
+                    x[start:stop],
+                    self.norm,
                     shift,
                     scale,
                     inverse_indices[start:stop],
@@ -1544,8 +1577,9 @@ def forward(
                 torch.mps.empty_cache()
             assert video is not None and audio is not None
             return video, audio
-        h = self.norm(x)
-        h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE)
+        h = _modulate_rmsnorm_scale_shift(
+            x, self.norm, shift, scale, inverse_indices, dtype=_BF16_DTYPE
+        )
         # Preserve full precision through both final output projections.
         h = h.to(_FP32_DTYPE)
         video, _ = self.video_out(h)

From c7c32fa649d23e30fdb63d5350e8becdc920f771 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Thu, 20 Aug 2026 07:54:23 +0000
Subject: [PATCH 3/8] [diffusion] Fix lint on MiniMax-H3 Spectrum files.

---
 .../sglang/multimodal_gen/runtime/models/dits/minimax_h3.py   | 4 +---
 .../multimodal_gen/test/unit/test_minimax_h3_spectrum.py      | 4 +++-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
index a361d6974a13..7a343cecab01 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -2364,9 +2364,7 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
         hidden = decoder_input
         cu_seqlens = cu_seqlens.to(device)
         spectrum_on = self._h3_spectrum_requested(sp_ws=sp_ws)
-        run_transformer_blocks = (
-            self.begin_spectrum_step() if spectrum_on else True
-        )
+        run_transformer_blocks = self.begin_spectrum_step() if spectrum_on else True
         if spectrum_on:
             try:
                 timestep = get_forward_context().current_timestep
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py
index ac2bd0e56c6e..e8fb6acc809a 100644
--- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_spectrum.py
@@ -31,7 +31,9 @@ def test_target_span_drops_earlier_ref_audio_gap():
 
 def test_local_target_hidden_is_rank_local_view():
     hidden = torch.arange(20, dtype=torch.float32).unsqueeze(-1).expand(20, 3).clone()
-    sliced = local_target_hidden(hidden, audio_start=4, video_stop=10, row_start=0, row_stop=20)
+    sliced = local_target_hidden(
+        hidden, audio_start=4, video_stop=10, row_start=0, row_stop=20
+    )
     assert sliced is not None
     local, start, stop = sliced
     assert start == 4 and stop == 10

From 2ee047dd45e5b17b22c4bde67258ad2503bdc7be Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Thu, 20 Aug 2026 09:20:20 +0000
Subject: [PATCH 4/8] [diffusion] docs: document MiniMax-H3 Spectrum on 4090.

---
 .../cookbook/diffusion/MiniMax/MiniMax-H3.mdx | 28 ++++++++++++++++++-
 .../sglang-diffusion/compatibility_matrix.mdx |  2 +-
 docs/docs/sglang-diffusion/spectrum.mdx       | 16 +++++++----
 .../sglang-diffusion-performance/SKILL.md     |  2 +-
 4 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx
index 164608275b32..73258f5a0441 100644
--- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx
+++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx
@@ -198,6 +198,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 downloads the model through the selected Hub. If the Hugging
 Face repository requires authentication, export a Hugging Face token in the
 server environment.
@@ -847,6 +860,7 @@ the strict `quality="high"` deployment contract.
 - `--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 &gt; 1. 50-step 4090 recipe is in §3.
 
 ## 9. Benchmarks
 
@@ -863,7 +877,7 @@ the configurations with collected measurements:
 | H100 | 4× TP2 + Ulysses2 resident | 4× TP4 + Ulysses1; 4× FSDP + Ulysses4 |
 | MI300X / MI355X | 8× Ulysses8 resident | 1×, 2×, and 4× scaling runs |
 | RTX 5090 | 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 |
 
 ### B300 precision and encoder placement
 
@@ -1077,6 +1091,18 @@ timed (seed 42); only the timed pass is reported. GPU peak stayed about
 and pixel fidelity rank in opposite orders there. Default remains
 `kitchen_int8` + `fa`.
 
+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
diff --git a/docs/docs/sglang-diffusion/compatibility_matrix.mdx b/docs/docs/sglang-diffusion/compatibility_matrix.mdx
index f07ef109e450..baa8ab1e79f3 100644
--- a/docs/docs/sglang-diffusion/compatibility_matrix.mdx
+++ b/docs/docs/sglang-diffusion/compatibility_matrix.mdx
@@ -135,7 +135,7 @@ Rows are grouped when a family shares the same runtime path or optimization supp
             <td>MiniMax-H3</td>
             <td><div className="sgd-id-list"><code>MiniMaxAI/MiniMax-H3</code></div></td>
             <td>T2VA / FL2VA / Ref2VA, 768p at 24 fps with synchronized audio</td>
-            <td><span className="sgd-chip">Cache-DiT</span><span className="sgd-chip">Sage</span><span className="sgd-chip">Online FP8</span></td>
+            <td><span className="sgd-chip">Cache-DiT</span><span className="sgd-chip">Spectrum</span><span className="sgd-chip">Sage</span><span className="sgd-chip">Online FP8</span></td>
           </tr>
           <tr>
             <td>Wan2.1 Fun</td>
diff --git a/docs/docs/sglang-diffusion/spectrum.mdx b/docs/docs/sglang-diffusion/spectrum.mdx
index 55ca7fc7cba8..cdc657517bc5 100644
--- a/docs/docs/sglang-diffusion/spectrum.mdx
+++ b/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 &gt; 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
 
diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md
index db4703345180..f427053c5a93 100644
--- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md
+++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md
@@ -61,7 +61,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. Incompatible with `--dit-layerwise-offload`. 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. |

From 7175caf9826b14fa3bd9c42191db6d55b6b18aa5 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Mon, 24 Aug 2026 08:04:34 +0000
Subject: [PATCH 5/8] [diffusion] Fix isort on MiniMax-H3 Spectrum imports.

---
 python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
index 3bc0d4a5d7ec..4a273df882dd 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -68,8 +68,8 @@
     QuantizationConfig,
 )
 from sglang.multimodal_gen.runtime.layers.usp import _ring_attention_varlen
-from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
 from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context
 from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import (
     LayerwiseOffloadableModuleMixin,
     is_layerwise_offloaded_module,

From 52a3e6ff67e4f3ae1a287324ce1ffcd8f5e6a4a6 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Sat, 29 Aug 2026 02:05:20 +0000
Subject: [PATCH 6/8] [diffusion] Fix MiniMax-H3 unit tests after Spectrum and
 fused AdaLN.

Keep the enable_spectrum gate stub complete, and patch the new RMSNorm+AdaLN entry so Cache-DiT preservation still isolates gate in-place behavior.
---
 .../multimodal_gen/test/unit/test_minimax_h3_admission.py       | 1 +
 .../multimodal_gen/test/unit/test_minimax_h3_dit_contract.py    | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
index 0daf8e023738..c09d479cdaae 100644
--- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
@@ -336,6 +336,7 @@ def test_high_quality_request_warns_when_bcg_suppresses_cache_dit():
             _explicit_fields={"quality"},
             enable_cache_dit=None,
             cache_dit_params=None,
+            enable_spectrum=False,
         )
     )
 
diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
index b4066892adab..122f753cc3f7 100644
--- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_dit_contract.py
@@ -252,7 +252,7 @@ def run(preserve):
 
     with (
         patch(
-            "sglang.multimodal_gen.runtime.models.dits.minimax_h3._modulate_scale_shift",
+            "sglang.multimodal_gen.runtime.models.dits.minimax_h3._modulate_rmsnorm_scale_shift",
             side_effect=lambda value, *_args, **_kwargs: value,
         ),
         patch(

From 360c1f7c9a644d815ac9c07d8f0c24cdada791b0 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Sat, 29 Aug 2026 11:21:50 +0000
Subject: [PATCH 7/8] [diffusion] Drop fused H3 RMSNorm+AdaLN to keep 2-GPU
 consistency GT.

The kernel is not bit-exact vs nn.RMSNorm on 5376-d packed rows and the ~0.3% e2e win is not worth the CI PSNR/audio drift.
---
 .../sglang/kernels/ops/diffusion/__init__.py  |   9 --
 .../modulate/indexed_modulation_triton.py     | 125 +-----------------
 .../runtime/models/dits/minimax_h3.py         |  18 +--
 3 files changed, 6 insertions(+), 146 deletions(-)

diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py
index 644bfc7ed29e..67793a51df2d 100644
--- a/python/sglang/kernels/ops/diffusion/__init__.py
+++ b/python/sglang/kernels/ops/diffusion/__init__.py
@@ -186,13 +186,6 @@
         _CUDA,
         "LTX-2 nine-way adaLN value split.",
     ),
-    (
-        "diffusion.fused_indexed_rmsnorm_scale_shift",
-        KernelBackend.TRITON,
-        "modulate.indexed_modulation_triton:fused_indexed_rmsnorm_scale_shift",
-        _CUDA,
-        "Indexed RMSNorm + adaLN scale/shift (MiniMax-H3).",
-    ),
     (
         "diffusion.fused_inplace_qknorm_rope",
         KernelBackend.JIT,
@@ -393,8 +386,6 @@
     "can_use_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_triton",
     "zimage_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_triton",
     # adaLN modulation, gating and timestep conditioning
-    "can_use_fused_indexed_rmsnorm_scale_shift": "modulate.indexed_modulation_triton",
-    "fused_indexed_rmsnorm_scale_shift": "modulate.indexed_modulation_triton",
     "indexed_gate_bf16": "modulate.indexed_modulation_triton",
     "indexed_gate_bf16_": "modulate.indexed_modulation_triton",
     "indexed_scale_shift_bf16_": "modulate.indexed_modulation_triton",
diff --git a/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py b/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py
index dac8ce37b13b..f6aff00832ae 100644
--- a/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py
+++ b/python/sglang/kernels/ops/diffusion/modulate/indexed_modulation_triton.py
@@ -4,11 +4,7 @@
 import triton
 import triton.language as tl
 
-from sglang.kernels.ops.diffusion.common.numerics import (
-    cuda_rsqrtf,
-    mul_rn_f32,
-    round_bf16_to_fp32,
-)
+from sglang.kernels.ops.diffusion.common.numerics import round_bf16_to_fp32
 
 
 @triton.jit
@@ -159,122 +155,3 @@ def indexed_gate_bf16(
     indices: torch.Tensor,
 ) -> torch.Tensor:
     return _indexed_gate_bf16(torch.empty_like(x), x, gate, other, indices)
-
-
-@triton.jit
-def _fused_indexed_rmsnorm_scale_shift_kernel(
-    out_ptr,
-    x_ptr,
-    weight_ptr,
-    shift_ptr,
-    scale_ptr,
-    indices_ptr,
-    hidden_size,
-    stride_x_row,
-    stride_out_row,
-    stride_shift_row,
-    stride_scale_row,
-    stride_indices,
-    eps,
-    BLOCK_N: tl.constexpr,
-):
-    """RMSNorm(x) * (1 + scale[idx]) + shift[idx], one kernel, leave x intact."""
-    row = tl.program_id(0)
-    columns = tl.arange(0, BLOCK_N)
-    acc = 0.0
-    offs = 0
-    while offs < hidden_size:
-        cols = offs + columns
-        mask = cols < hidden_size
-        x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(
-            tl.float32
-        )
-        acc += tl.sum(mul_rn_f32(x, x), axis=0)
-        offs += BLOCK_N
-    rstd = cuda_rsqrtf(acc / hidden_size + eps)
-    index = tl.load(indices_ptr + row * stride_indices)
-    offs = 0
-    while offs < hidden_size:
-        cols = offs + columns
-        mask = cols < hidden_size
-        x = tl.load(x_ptr + row * stride_x_row + cols, mask=mask, other=0.0).to(
-            tl.float32
-        )
-        weight = tl.load(weight_ptr + cols, mask=mask, other=1.0).to(tl.float32)
-        shift = tl.load(
-            shift_ptr + index * stride_shift_row + cols, mask=mask, other=0.0
-        ).to(tl.float32)
-        scale = tl.load(
-            scale_ptr + index * stride_scale_row + cols, mask=mask, other=0.0
-        ).to(tl.float32)
-        y = round_bf16_to_fp32(x * rstd * weight)
-        one_plus = round_bf16_to_fp32(1.0 + scale)
-        tl.store(
-            out_ptr + row * stride_out_row + cols,
-            round_bf16_to_fp32(y * one_plus) + shift,
-            mask=mask,
-        )
-        offs += BLOCK_N
-
-
-def can_use_fused_indexed_rmsnorm_scale_shift(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    shift: torch.Tensor,
-    scale: torch.Tensor,
-    indices: torch.Tensor,
-) -> bool:
-    return (
-        x.is_cuda
-        and x.dim() == 2
-        and x.dtype == torch.bfloat16
-        and x.is_contiguous()
-        and weight.is_cuda
-        and weight.dtype == torch.bfloat16
-        and weight.is_contiguous()
-        and weight.shape == (x.shape[-1],)
-        and shift.is_cuda
-        and scale.is_cuda
-        and shift.dtype == scale.dtype == torch.bfloat16
-        and shift.shape == scale.shape
-        and shift.shape[-1] == x.shape[-1]
-        and shift.stride(-1) == 1
-        and scale.stride(-1) == 1
-        and indices.is_cuda
-        and indices.dtype in (torch.int32, torch.int64)
-        and indices.numel() == x.shape[0]
-    )
-
-
-def fused_indexed_rmsnorm_scale_shift(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    shift: torch.Tensor,
-    scale: torch.Tensor,
-    indices: torch.Tensor,
-    eps: float,
-) -> torch.Tensor:
-    """``norm(x) * (1 + scale[idx]) + shift[idx]`` without materializing norm(x)."""
-    rows, hidden_size = x.shape
-    out = torch.empty_like(x)
-    if rows == 0:
-        return out
-    block_n = min(triton.next_power_of_2(hidden_size), 1024)
-    _fused_indexed_rmsnorm_scale_shift_kernel[(rows,)](
-        out,
-        x,
-        weight,
-        shift,
-        scale,
-        indices,
-        hidden_size,
-        x.stride(0),
-        out.stride(0),
-        shift.stride(0),
-        scale.stride(0),
-        indices.stride(0),
-        float(eps),
-        BLOCK_N=block_n,
-        num_warps=8,
-    )
-    return out
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
index 4245ec964995..42668823edb7 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
@@ -24,9 +24,7 @@
     silu_and_mul_with_activation_rounding_,
 )
 from sglang.kernels.ops.diffusion import (
-    can_use_fused_indexed_rmsnorm_scale_shift,
     can_use_fused_inplace_qknorm_rope,
-    fused_indexed_rmsnorm_scale_shift,
     fused_inplace_qknorm_rope,
     indexed_gate_bf16,
     indexed_gate_bf16_,
@@ -331,17 +329,11 @@ def _modulate_rmsnorm_scale_shift(
     *,
     dtype: torch.dtype,
 ) -> torch.Tensor:
-    """RMSNorm + indexed AdaLN without materializing the normalized tensor."""
-    if (
-        not torch.compiler.is_compiling()
-        and dtype == _BF16_DTYPE
-        and can_use_fused_indexed_rmsnorm_scale_shift(
-            x, norm.weight, shift, scale, indices
-        )
-    ):
-        return fused_indexed_rmsnorm_scale_shift(
-            x, norm.weight, shift, scale, indices, float(norm.eps)
-        )
+    """RMSNorm + indexed AdaLN.
+
+    Keep ``nn.RMSNorm``; a fused RMSNorm+AdaLN kernel drifted 2-GPU
+    consistency GT for ~0.3% e2e, so it was removed.
+    """
     return _modulate_scale_shift(norm(x), shift, scale, indices, dtype=dtype)
 
 

From 1a22c5486b2f9422bba58ed1dde157efd406e373 Mon Sep 17 00:00:00 2001
From: niehen6174 <1639206518@qq.com>
Date: Fri, 25 Sep 2026 13:47:21 +0000
Subject: [PATCH 8/8] [diffusion] Set enable_spectrum on the extra-high
 Cache-DiT stub

The merged admission test builds a SimpleNamespace, and the Spectrum path reads that field.
---
 .../sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
index 870eca022d7d..a3cfe95c267c 100644
--- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
+++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
@@ -466,6 +466,7 @@ def test_extra_high_quality_does_not_enable_h3_cache_dit():
             _explicit_fields={"quality"},
             enable_cache_dit=None,
             cache_dit_params=None,
+            enable_spectrum=False,
         )
     )