[Perf] diffusion: Remove per-step and per-request host syncs in Cosmos3 and Qwen-Image-2.1

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py

@@ -1238,7 +1238,7 @@ def _run_transformer(
             )
 
     @staticmethod
-    def _cfg_active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool:
+    def _cfg_active_at(t: float, interval: tuple[float, float] | None) -> bool:
         """Return True iff CFG should be applied at timestep ``t``.
 
         T2I uses a CFG window (e.g. ``[400, 1000]``) to skip guidance at low
@@ -1247,9 +1247,8 @@ def _cfg_active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> boo
         """
         if interval is None:
             return True
-        t_scalar = float(t.item()) if torch.is_tensor(t) else float(t)
         lo, hi = interval
-        return lo <= t_scalar <= hi
+        return lo <= t <= hi
 
     def _normalize_transfer_video(
         self, video: torch.Tensor, device: torch.device, dtype: torch.dtype
@@ -1585,6 +1584,9 @@ def _denoise_once(
             desc="Denoising",
             batch=batch,
         )
+        # host copy for the per-step CFG window checks: a device read per step
+        # would stall the host on every step's kernels
+        timestep_values = timesteps.tolist()
 
         for i, t in progress_bar:
             # Precision is chosen once per step, before any transformer call,
@@ -1595,7 +1597,9 @@ def _denoise_once(
             # Outside the CFG window the effective scale collapses to 1.0,
             # which reduces CFG to the cond branch (cfg-parallel safe).
             effective_scale = (
-                guidance_scale if self._cfg_active_at(t, guidance_interval) else 1.0
+                guidance_scale
+                if self._cfg_active_at(timestep_values[i], guidance_interval)
+                else 1.0
             )
             # Transfer control-CFG: active only when a control video is present,
             # ``control_guidance != 1.0``, and the step is inside the (optional)
@@ -1604,7 +1608,7 @@ def _denoise_once(
             control_cfg_active = (
                 control_latents is not None
                 and control_guidance != 1.0
-                and self._cfg_active_at(t, control_guidance_interval)
+                and self._cfg_active_at(timestep_values[i], control_guidance_interval)
             )
 
             if control_cfg_active:

python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/qwen_image21.py

@@ -80,6 +80,9 @@ def __init__(self, text_encoder, processor, vae, scheduler):
         )
         self.text_encoder.model.visual.fp32_position_interpolation = False
         self.text_encoder.model.visual.rotary_pos_emb.recompute_on_device_change = True
+        # (dtype, device) -> latent mean/std, uploaded once: a per-request upload
+        # blocks the host on the VAE encode it should overlap with prompt prep
+        self.latent_stats = {}
         self.image_token_id = processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")
         system_message = [
             {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]}
@@ -195,8 +198,13 @@ def forward(self, batch, server_args):
                     )
                     pixels = (2 * pixels - 1).to(device=device, dtype=torch.bfloat16)
                     latent = vae.encode(pixels).mode()
-                    mean = latent.new_tensor(ac.latents_mean).view(1, ac.z_dim, 1, 1, 1)
-                    std = latent.new_tensor(ac.latents_std).view(1, ac.z_dim, 1, 1, 1)
+                    stats_key = (latent.dtype, latent.device)
+                    if stats_key not in self.latent_stats:
+                        self.latent_stats[stats_key] = tuple(
+                            latent.new_tensor(values).view(1, ac.z_dim, 1, 1, 1)
+                            for values in (ac.latents_mean, ac.latents_std)
+                        )
+                    mean, std = self.latent_stats[stats_key]
                     conditions.append(
                         ((latent - mean) / std).flatten(2).transpose(1, 2)
                     )

python/sglang/multimodal_gen/test/unit/test_cosmos3_rollout.py

@@ -147,6 +147,61 @@ def test_serving_request_steps_unipc_without_rollout_kwargs(self):
             stage.forward(batch, server_args)
         torch.testing.assert_close(batch.latents, torch.ones_like(batch.latents))
 
+    def test_cfg_window_check_reads_host_timesteps(self):
+        """The per-step CFG window check reads a host copy of the schedule: a
+        device read per step blocks the host on that step's kernels."""
+
+        class DeviceReadForbidden(torch.Tensor):
+            def item(self):
+                raise AssertionError("device read inside the denoising loop")
+
+            # float(t) and `lo <= t <= hi` read the value without calling item()
+            __float__ = __bool__ = item
+
+        stage = cosmos3.Cosmos3DenoisingStage.__new__(cosmos3.Cosmos3DenoisingStage)
+        stage.scheduler = mock.Mock()
+        stage.scheduler.step.side_effect = lambda noise_pred, t, latents, **_: (
+            latents,
+        )
+        stage.transformer = mock.Mock()
+        stage._logged_parallel_config = True
+        stage.log_info = mock.Mock()
+        stage.progress_bar = lambda iterable, **kwargs: iterable
+        stage.use_declared_component = lambda **kwargs: nullcontext()
+        forward_batch_sizes = []
+
+        def run_transformer(**kwargs):
+            forward_batch_sizes.append(kwargs["latents"].shape[0])
+            return torch.zeros_like(kwargs["latents"])
+
+        stage._run_transformer = run_transformer
+        text = torch.ones(1, 2, dtype=torch.long)
+        sampling_params = SamplingParams(prompt="test", guidance_scale=6.0)
+        sampling_params.guidance_interval = (400.0, 1000.0)
+        batch = Req(
+            sampling_params=sampling_params,
+            latents=torch.ones(1, 1, 1, 1, 1),
+            timesteps=torch.tensor([999, 700, 400, 399, 100]).as_subclass(
+                DeviceReadForbidden
+            ),
+            extra={
+                "cond_text_ids": text,
+                "cond_text_mask": text,
+                "uncond_text_ids": text,
+                "uncond_text_mask": text,
+                "cond_text_seq_len": 2,
+                "uncond_text_seq_len": 2,
+                "video_shape": (1, 1, 1),
+            },
+        )
+        server_args = types.SimpleNamespace(
+            pipeline_config=None, enable_cfg_parallel=False
+        )
+        with mock.patch.object(cosmos3, "get_sp_world_size", return_value=1):
+            stage.forward(batch, server_args)
+        # batched CFG inside the [400, 1000] window, one cond forward outside it
+        self.assertEqual(forward_batch_sizes, [2, 2, 2, 1, 1])
+
 
 class _FusedParamModule(torch.nn.Module):
     """Diffusers-style q/k/v weights that map into one fused to_qkv param."""

python/sglang/multimodal_gen/test/unit/test_qwen_image21.py

@@ -778,6 +778,48 @@ def test_condition_pixels_match_reference_preprocessing(monkeypatch, tiling):
     assert actual.stride() == expected.stride()
 
 
+def test_condition_latents_reuse_uploaded_latent_stats(monkeypatch):
+    module = "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.qwen_image21"
+    monkeypatch.setattr(f"{module}.get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(f"{module}.set_forward_context", lambda **kwargs: nullcontext())
+    image = Image.frombytes("RGBA", (32, 32), bytes(range(256)) * 16)
+    latent = torch.randn(
+        1, 64, 1, 2, 2, generator=torch.Generator().manual_seed(0)
+    ).bfloat16()
+    vae = Mock()
+    vae.encode.return_value.mode.return_value = latent
+    processor = Mock()
+    processor.apply_chat_template.return_value = [[1]]
+    stage = QwenImage21EncodingStage(Mock(), processor, vae, SimpleNamespace(config={}))
+    stage.use_declared_component = Mock(return_value=nullcontext(vae))
+    stage.encode_prompt = Mock(
+        return_value=(torch.zeros(3, 8), torch.tensor([False, True, False]))
+    )
+    server_args = SimpleNamespace(pipeline_config=QwenImage21PipelineConfig())
+    ac = server_args.pipeline_config.vae_config.arch_config
+    mean = latent.new_tensor(ac.latents_mean).view(1, ac.z_dim, 1, 1, 1)
+    std = latent.new_tensor(ac.latents_std).view(1, ac.z_dim, 1, 1, 1)
+    expected = ((latent - mean) / std).flatten(2).transpose(1, 2)
+    stats = []
+    for _ in range(2):
+        batch = SimpleNamespace(
+            height=32,
+            width=32,
+            condition_image=image,
+            prompt="edit",
+            negative_prompt=None,
+            num_outputs_per_prompt=1,
+            do_classifier_free_guidance=False,
+            extra={},
+        )
+        stage.forward(batch, server_args)
+        actual = batch.extra["qwen21_positive"]["condition_latents"]
+        torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        stats.append(stage.latent_stats[(latent.dtype, latent.device)])
+    # the second request normalizes with the first request's device tensors
+    assert stats[0] is stats[1]
+
+
 def test_condition_image_loading_preserves_alpha(tmp_path):
     path = tmp_path / "condition.png"
     Image.new("RGBA", (32, 32), (12, 34, 56, 78)).save(path)