[diffusion] optimization: reuse minimax h3 prompt refinement across outputs (#36027)

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,

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"]