diff --git a/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py b/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py index 986138384..726f27321 100644 --- a/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py +++ b/python/sglang/kernels/ops/diffusion/triton/indexed_modulation.py @@ -53,6 +53,7 @@ def _indexed_gate_bf16_kernel( other_ptr, indices_ptr, hidden_size, + stride_output_row, stride_x_row, stride_gate_row, stride_other_row, @@ -76,7 +77,7 @@ def _indexed_gate_bf16_kernel( gated = round_bf16_to_fp32(gate * other) tl.store( - output_ptr + row * stride_x_row + columns, + output_ptr + row * stride_output_row + columns, x + gated, mask=mask, ) @@ -109,7 +110,8 @@ def indexed_scale_shift_bf16_( return x -def indexed_gate_bf16_( +def _indexed_gate_bf16( + output: torch.Tensor, x: torch.Tensor, gate: torch.Tensor, other: torch.Tensor, @@ -117,15 +119,16 @@ def indexed_gate_bf16_( ) -> torch.Tensor: rows, hidden_size = x.shape if rows == 0: - return x + return output block_n = triton.next_power_of_2(hidden_size) _indexed_gate_bf16_kernel[(rows,)]( - x, + output, x, gate, other, indices, hidden_size, + output.stride(0), x.stride(0), gate.stride(0), other.stride(0), @@ -133,4 +136,22 @@ def indexed_gate_bf16_( BLOCK_N=block_n, num_warps=8, ) - return x + return output + + +def indexed_gate_bf16_( + x: torch.Tensor, + gate: torch.Tensor, + other: torch.Tensor, + indices: torch.Tensor, +) -> torch.Tensor: + return _indexed_gate_bf16(x, x, gate, other, indices) + + +def indexed_gate_bf16( + x: torch.Tensor, + gate: torch.Tensor, + other: torch.Tensor, + indices: torch.Tensor, +) -> torch.Tensor: + return _indexed_gate_bf16(torch.empty_like(x), x, gate, other, indices) diff --git a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py index 819b8b9da..2616c9a37 100644 --- a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py +++ b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py @@ -273,16 +273,26 @@ DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS: dict[str, DualTransformerBlockAdapterSpec] } -# Custom BlockAdapter for DiT models absent from cache-dit's BlockAdapterRegister. -# Value: (blocks attr, forward_pattern). forward_pattern must -# match the block's forward signature (see cache_dit.ForwardPattern; e.g., ERNIE -# uses Pattern_3). has_separate_cfg follows the run (passed by -# enable_cache_on_transformer); cache-dit auto-resolves the remaining -# fields. -_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern]] = { - "ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3), - "Krea2Transformer2DModel": ("transformer_blocks", ForwardPattern.Pattern_3), - "MiniMaxH3DiTModel": ("blocks", ForwardPattern.Pattern_3), +@dataclass(frozen=True) +class CustomBlockAdapterSpec: + blocks_attr: str + forward_pattern: ForwardPattern + + +# Custom BlockAdapter metadata for models absent from cache-dit's registry. +_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, CustomBlockAdapterSpec] = { + "ErnieImageTransformer2DModel": CustomBlockAdapterSpec( + blocks_attr="layers", + forward_pattern=ForwardPattern.Pattern_3, + ), + "Krea2Transformer2DModel": CustomBlockAdapterSpec( + blocks_attr="transformer_blocks", + forward_pattern=ForwardPattern.Pattern_3, + ), + "MiniMaxH3DiTModel": CustomBlockAdapterSpec( + blocks_attr="blocks", + forward_pattern=ForwardPattern.Pattern_3, + ), } @@ -295,17 +305,16 @@ def _build_custom_block_adapter( spec = _CUSTOM_BLOCK_ADAPTER_SPECS.get(transformer.__class__.__name__) if spec is None: return None - blocks_attr, forward_pattern = spec - blocks = getattr(transformer, blocks_attr, None) + blocks = getattr(transformer, spec.blocks_attr, None) if blocks is None: raise ValueError( f"Transformer {transformer.__class__.__name__} has no attribute " - f"{blocks_attr!r} for cache-dit blocks." + f"{spec.blocks_attr!r} for cache-dit blocks." ) return BlockAdapter( transformer=transformer, blocks=blocks, - forward_pattern=forward_pattern, + forward_pattern=spec.forward_pattern, has_separate_cfg=has_separate_cfg, ) 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 37fff83a6..26d9f3dc5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -21,6 +21,7 @@ from sglang.kernels.ops.diffusion.qknorm_rope import ( fused_inplace_qknorm_rope, ) from sglang.kernels.ops.diffusion.triton.indexed_modulation import ( + indexed_gate_bf16, indexed_gate_bf16_, indexed_scale_shift_bf16_, ) @@ -240,8 +241,9 @@ def _modulate_gate( indices: torch.Tensor, *, dtype: torch.dtype, + allow_inplace: bool = True, ) -> torch.Tensor: - """Apply indexed gated residual, reusing disposable CUDA BF16 input.""" + """Apply an indexed gated residual, optionally reusing the input buffer.""" # Apply the per-index gated residual: x + gate[idx] * other. if ( x.is_cuda @@ -252,7 +254,9 @@ def _modulate_gate( and x.is_contiguous() and other.is_contiguous() ): - return indexed_gate_bf16_(x, gate, other, indices) + if allow_inplace: + return indexed_gate_bf16_(x, gate, other, indices) + return indexed_gate_bf16(x, gate, other, indices) return (x + gate.index_select(0, indices) * other).to(dtype) @@ -904,6 +908,7 @@ class MiniMaxH3DiTBlock(nn.Module): expand_ratio=6, modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, ) + self.preserve_input_for_cache_dit = False def forward( self, @@ -929,7 +934,9 @@ class MiniMaxH3DiTBlock(nn.Module): if adaln_params is None: adaln_params = self.adaln_proj(adaln_input) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln_params - + # Cache-DiT retains the inputs to its Fn and Mn block ranges. Only the + # 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( @@ -944,7 +951,14 @@ class MiniMaxH3DiTBlock(nn.Module): ulysses_active=ulysses_active, ring_active=ring_active, ) - x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) + x = _modulate_gate( + residual, + gate_msa, + h, + combined_indices, + dtype=_BF16_DTYPE, + allow_inplace=not self.preserve_input_for_cache_dit, + ) residual = x h = self.norm2(x) @@ -952,8 +966,14 @@ class MiniMaxH3DiTBlock(nn.Module): h, 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 + # even while Cache-DiT is attached. return _modulate_gate( - residual, gate_mlp, h, combined_indices, dtype=_BF16_DTYPE + residual, + gate_mlp, + h, + combined_indices, + dtype=_BF16_DTYPE, ) @@ -1190,6 +1210,22 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): self._resolved_attention_backend: AttentionBackendEnum | None = None self._mark_missing_params_required() + def set_cache_dit_input_preservation(self, enabled: bool) -> None: + """Stop the blocks from overwriting the input Cache-DiT holds by reference. + + Cache-DiT snapshots the block-stack input to measure its residuals, so a + block that rewrites its own input in place makes that residual read as + zero. Only the first gated residual of a block writes the block input; + the second one operates on a buffer this block just allocated, so it is + left on the in-place fused path either way. + + The caller owns the lifecycle. It has to be on before Cache-DiT mounts, + because mounting replaces `blocks` with a wrapper and the real blocks + stop being reachable by iterating it. + """ + for block in self.blocks: + block.preserve_input_for_cache_dit = enabled + def _resolve_attention_backend_once(self) -> None: if self._resolved_attention_backend is not None: return 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 568e82a0c..fdcc37ca2 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 @@ -417,16 +417,94 @@ class MiniMaxH3DenoisingStage(DenoisingStage): # a time. Combined with `quality` in the dynamic-batch signature, this # makes the process-wide hook transition safe at this batch boundary. if self._cache_dit_enabled and current_mode != desired_mode: + # Unmount first: the blocks must stay preserved for as long as + # Cache-DiT still holds references to their inputs. Settle the state + # fields before restoring the in-place path, so a failure there + # costs throughput rather than leaving the stage inconsistent. self.transformer = disable_cache_on_transformer(self.transformer) self._cache_dit_enabled = False self._cached_num_steps = None self._minimax_h3_cache_mode = None + self._set_cache_dit_input_preservation(False) if desired_mode is None: return - super()._maybe_enable_cache_dit(num_inference_steps, batch) + + # Arm before delegating whenever this H3 stage requests caching, + # without predicting whether the parent will accept the mount. + # cache_dit.enable_cache swaps `blocks` for a single CachedBlocks + # wrapper, so the real blocks are only reachable beforehand -- and a + # wrong prediction would hand Cache-DiT unpreserved blocks, whose + # residuals read as zero, which is silent. Arming is one boolean per + # block and nothing runs before the parent decides, so guessing + # conservatively costs approximately nothing. + was_enabled = self._cache_dit_enabled + if not was_enabled: + self._set_cache_dit_input_preservation(True) + try: + super()._maybe_enable_cache_dit(num_inference_steps, batch) + except Exception: + if not was_enabled: + self._disarm_after_failed_mount() + raise + if self._cache_dit_enabled: self._minimax_h3_cache_mode = desired_mode + elif not was_enabled: + # The parent declined to mount, for example because breakable + # CUDA graphs are enabled or this is an ordinary warmup. Nothing + # holds the block inputs, so go back to the in-place path. + self._set_cache_dit_input_preservation(False) + + def _disarm_after_failed_mount(self) -> None: + """Restore the in-place path only once Cache-DiT is confirmed gone. + + cache_dit.enable_cache swaps the block list before the rest of the mount + runs, so a later failure can leave caching attached. Disarming then + would hand Cache-DiT unpreserved blocks and silently reproduce the + zero-residual bug, so if the unmount does not succeed we stay armed and + pay throughput instead. + """ + try: + self.transformer = disable_cache_on_transformer(self.transformer) + except Exception: + logger.warning( + "Could not unmount Cache-DiT after a failed mount; leaving " + "MiniMax-H3 input preservation on", + exc_info=True, + ) + return + # The parent may have flipped these before failing; leaving them set + # would send the next request down the refresh path with nothing + # mounted. + self._cache_dit_enabled = False + self._cached_num_steps = None + self._minimax_h3_cache_mode = None + self._set_cache_dit_input_preservation(False) + + def _set_cache_dit_input_preservation(self, enabled: bool) -> None: + """Flip input preservation on the H3 model, failing closed. + + Skipping this silently would put us back where this stage started: cache + mounted, residuals reading as zero, no hits and nothing logged. If the + model cannot be reached, that is a wiring bug and should surface here. + """ + model = self.transformer + for _ in range(4): # unwrap compile/parallel wrappers, if any + if hasattr(model, "set_cache_dit_input_preservation"): + break + inner = getattr(model, "_orig_mod", None) or getattr(model, "module", None) + if inner is None: + break + model = inner + setter = getattr(model, "set_cache_dit_input_preservation", None) + if not callable(setter): + raise TypeError( + "MiniMax-H3 Cache-DiT requires set_cache_dit_input_preservation() " + f"on the transformer, but {type(self.transformer).__name__} does " + "not expose it" + ) + setter(enabled) def _cache_dit_scm_masks( self, primary_num_steps: int, secondary_num_steps: int | None = None 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 4a722ef48..4592a4e8d 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 @@ -25,8 +25,10 @@ from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( MINIMAX_H3_FP32_BUFFER_NAMES, MINIMAX_H3_FP32_PARAM_NAMES, + MiniMaxH3DiTBlock, MiniMaxH3DiTModel, _copy_grouped_qkv_tp_shard, + _modulate_gate, _reorder_grouped_qkv_to_qkv, ) from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import ( @@ -127,6 +129,97 @@ def test_native_weight_names_and_grouped_qkv_reorder(): ) +class _KwargIdentity(torch.nn.Module): + def forward(self, x, **_kwargs): + return x + + +def test_cache_dit_preservation_only_makes_first_gate_out_of_place(): + block = MiniMaxH3DiTBlock.__new__(MiniMaxH3DiTBlock) + torch.nn.Module.__init__(block) + block.norm1 = torch.nn.Identity() + block.norm2 = torch.nn.Identity() + block.attn = _KwargIdentity() + block.mlp = torch.nn.Identity() + gate_modes = [] + + def fake_gate(residual, _gate, _other, _indices, *, dtype, allow_inplace=True): + gate_modes.append(allow_inplace) + return residual.to(dtype) + + def run(preserve): + block.preserve_input_for_cache_dit = preserve + gate_modes.clear() + block( + torch.zeros(2, 4), + adaln_input=torch.zeros(1, 4), + combined_indices=torch.zeros(2, dtype=torch.long), + rope_cache=None, + cu_seqlens=torch.tensor([0, 2], dtype=torch.int32), + max_seqlen=2, + adaln_params=tuple(torch.zeros(1, 4) for _ in range(6)), + ) + return list(gate_modes) + + with ( + patch( + "sglang.multimodal_gen.runtime.models.dits.minimax_h3._modulate_scale_shift", + side_effect=lambda value, *_args, **_kwargs: value, + ), + patch( + "sglang.multimodal_gen.runtime.models.dits.minimax_h3._modulate_gate", + side_effect=fake_gate, + ), + ): + assert run(preserve=False) == [True, True] + # Only the first gated residual can alias the block input Cache-DiT + # holds by reference, so only it goes out-of-place. The second works on + # a block-local buffer and keeps the fused in-place kernel. + assert run(preserve=True) == [False, True] + + +def test_cache_dit_input_preservation_toggles_every_block(): + model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel) + torch.nn.Module.__init__(model) + model.blocks = torch.nn.ModuleList([torch.nn.Identity() for _ in range(5)]) + + model.set_cache_dit_input_preservation(True) + assert all(block.preserve_input_for_cache_dit for block in model.blocks) + + model.set_cache_dit_input_preservation(False) + assert not any(block.preserve_input_for_cache_dit for block in model.blocks) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cache_dit_out_of_place_gate_preserves_cuda_input(): + x = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) + original = x.clone() + gate = torch.randn(2, 16, device="cuda", dtype=torch.bfloat16) + other = torch.randn_like(x) + indices = torch.tensor([0, 1, 0, 1], device="cuda", dtype=torch.long) + expected = _modulate_gate( + x.clone(), + gate, + other, + indices, + dtype=torch.bfloat16, + allow_inplace=True, + ) + + output = _modulate_gate( + x, + gate, + other, + indices, + dtype=torch.bfloat16, + allow_inplace=False, + ) + + assert output.data_ptr() != x.data_ptr() + torch.testing.assert_close(x, original, rtol=0, atol=0) + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + def test_tp_and_ulysses_admission_uses_tp_local_shapes(): arch = MiniMaxH3DiTArchConfig() model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel)