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 3b7918b83..1337459f9 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 @@ -47,6 +47,7 @@ _REF2VA_VIDEO_CHAINS = { "video.reference_preserve", "video_audio.reference_preserve", } +_REFINED_PROMPT_EMBEDS_KEY = "_minimax_h3_refined_prompt_embeds" def minimax_h3_condition_noise_aug(sampling: Any) -> tuple[float, float]: @@ -339,25 +340,28 @@ def _precompute_refined_prompt_embeds( positive: Any, *, device: torch.device, + shared_conditioning: dict[str, Any] | None = None, ) -> bool: - """Move request-static text refinement out of the denoise hot loop.""" + """Refine text once for outputs that share the encoded presentation.""" refine = getattr(model, "refine_prompt_embeds", None) if not callable(refine): return False static_kwargs = positive.static_kwargs prompt_embeds = static_kwargs["prompt_embeds"] - refiner_params = static_kwargs["refiner_packed_seq_params"] - if isinstance(refiner_params, dict): - refiner_cu = refiner_params["cu_seqlens_q"] - else: - refiner_cu = refiner_params.cu_seqlens_q - with torch.inference_mode(): - refined = refine( - prompt_embeds, - refiner_cu, - device=device, - ) + refined = ( + shared_conditioning.get(_REFINED_PROMPT_EMBEDS_KEY) + if shared_conditioning is not None + else None + ) + if refined is None: + refiner_params = static_kwargs["refiner_packed_seq_params"] + if isinstance(refiner_params, dict): + refiner_cu = refiner_params["cu_seqlens_q"] + else: + refiner_cu = refiner_params.cu_seqlens_q + with torch.inference_mode(): + refined = refine(prompt_embeds, refiner_cu, device=device) if not torch.is_tensor(refined): raise TypeError("MiniMax H3 refine_prompt_embeds must return a torch.Tensor") if int(refined.shape[0]) != int(prompt_embeds.shape[0]): @@ -365,6 +369,8 @@ def _precompute_refined_prompt_embeds( "MiniMax H3 refined prompt row count changed: " f"{int(prompt_embeds.shape[0])} -> {int(refined.shape[0])}" ) + if shared_conditioning is not None: + shared_conditioning.setdefault(_REFINED_PROMPT_EMBEDS_KEY, refined) static_kwargs["prompt_embeds"] = refined static_kwargs["refined_prompt_embeds_length"] = int(refined.shape[0]) return True @@ -660,6 +666,7 @@ class MiniMaxH3DenoisingStage(DenoisingStage): model, positive, device=device, + shared_conditioning=emb, ) _precompute_rope_cache( model, diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py index 1c5e60f02..2a57c638d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_denoise_loop.py @@ -21,6 +21,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.m minimax_h3_packed_sequence, minimax_h3_packed_sequence_ref2va_blocks, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.stages.denoising import ( + _precompute_refined_prompt_embeds, +) def _branch( @@ -193,3 +196,30 @@ def test_rank_local_token_tags_match_reference_slice(): torch.testing.assert_close( branch.static_kwargs["block_token_tags"], expected, rtol=0, atol=0 ) + + +def test_grouped_outputs_share_prompt_refinement(): + class Refiner: + calls = 0 + + def refine_prompt_embeds(self, prompt_embeds, refiner_cu, *, device): + del refiner_cu + self.calls += 1 + return torch.ones( + prompt_embeds.shape[0], 5376, dtype=prompt_embeds.dtype, device=device + ) + + model = Refiner() + conditioning = {} + first, second = _branch("t2va"), _branch("t2va") + + for branch in (first, second): + assert _precompute_refined_prompt_embeds( + model, + branch, + device=torch.device("cpu"), + shared_conditioning=conditioning, + ) + + assert model.calls == 1 + assert first.static_kwargs["prompt_embeds"] is second.static_kwargs["prompt_embeds"]