[diffusion] Add CFG gating for denoising (#25848)

Co-authored-by: BBuf <bbuf@example.com>

python/sglang/multimodal_gen/envs.py

@@ -32,6 +32,7 @@ if TYPE_CHECKING:
     VERBOSE: bool = False
     SGLANG_DIFFUSION_SERVER_DEV_MODE: bool = False
     SGLANG_DIFFUSION_STAGE_LOGGING: bool = False
+    SGLANG_DIFFUSION_CFG_GATE_STEP: float = 1.0
     # cache-dit env vars (primary transformer)
     SGLANG_CACHE_DIT_ENABLED: bool = False
     SGLANG_CACHE_DIT_FN: int = 1
@@ -250,6 +251,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # If set, sgl_diffusion will enable stage logging, which will print the time
     # taken for each stage
     "SGLANG_DIFFUSION_STAGE_LOGGING": _lazy_bool("SGLANG_DIFFUSION_STAGE_LOGGING"),
+    # Fraction of denoising steps that run both CFG branches before reusing the
+    # last conditional-minus-unconditional residual. Keep 1.0 to disable.
+    "SGLANG_DIFFUSION_CFG_GATE_STEP": _lazy_float(
+        "SGLANG_DIFFUSION_CFG_GATE_STEP", 1.0
+    ),
     "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D": _lazy_str(
         "SGLANG_DIFFUSION_VAE_CHANNELS_LAST_3D", "auto"
     ),

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -804,6 +804,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         """Prepare scheduler state before entering the shared denoising loop."""
         self._reset_scheduler_loop_state(ctx.scheduler)
         ctx.scheduler.set_begin_index(0)
+        self._init_cfg_gate_state(ctx, batch, server_args)
 
     def _reset_scheduler_loop_state(self, scheduler) -> None:
         if hasattr(scheduler, "_step_index"):
@@ -824,6 +825,55 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             if hasattr(scheduler, "timestep_list"):
                 scheduler.timestep_list = [None] * solver_order
 
+    def _init_cfg_gate_state(
+        self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
+    ) -> None:
+        """Initialize optional CFG residual reuse for the current denoising loop."""
+        fraction = envs.SGLANG_DIFFUSION_CFG_GATE_STEP
+        if not 0.0 <= fraction <= 1.0:
+            raise ValueError(
+                "SGLANG_DIFFUSION_CFG_GATE_STEP must be between 0.0 and 1.0, "
+                f"got {fraction}."
+            )
+
+        num_steps = len(ctx.timesteps)
+        requested = fraction < 1.0 and batch.do_classifier_free_guidance
+        active = requested and not server_args.enable_cfg_parallel
+        gate_step = int(num_steps * fraction) if active else num_steps + 1
+        ctx.extra["cfg_gate_state"] = {
+            "fraction": fraction,
+            "requested": requested,
+            "active": active,
+            "gate_step": gate_step,
+            "delta": None,
+            "model_id": None,
+            "fresh_uncond": 0,
+            "reused": 0,
+            "invalidations": 0,
+        }
+
+        if ctx.is_warmup or get_world_group().local_rank != 0:
+            return
+
+        if active:
+            logger.info(
+                "CFG gating enabled: reuse unconditioned residual after step %d/%d "
+                "(fraction=%.3f).",
+                gate_step,
+                num_steps,
+                fraction,
+            )
+            if batch.guidance_rescale > 0:
+                logger.warning(
+                    "CFG gating is enabled with guidance_rescale=%s; benchmark image "
+                    "quality before using this setting in production.",
+                    batch.guidance_rescale,
+                )
+        elif requested:
+            logger.info(
+                "CFG gating requested but skipped because CFG parallel is enabled."
+            )
+
     def _get_transformer_attr(self, name: str) -> Any:
         seen: set[int] = set()
         stack = [self.transformer]
@@ -958,6 +1008,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             target_dtype=ctx.target_dtype,
             current_guidance_scale=step.current_guidance_scale,
             cfg_policy=ctx.cfg_policy,
+            cfg_gate_state=ctx.extra.get("cfg_gate_state"),
             server_args=server_args,
             guidance=ctx.guidance,
             latents=ctx.latents,
@@ -1000,6 +1051,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         self, ctx: DenoisingContext, batch: Req, server_args: ServerArgs
     ) -> None:
         """Finalize the shared loop by handing state to post-denoising processing."""
+        self._log_cfg_gate_summary(ctx, batch)
         self._post_denoising_loop(
             batch=batch,
             latents=ctx.latents,
@@ -1009,6 +1061,28 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
             is_warmup=ctx.is_warmup,
         )
 
+    def _log_cfg_gate_summary(self, ctx: DenoisingContext, batch: Req) -> None:
+        state = ctx.extra.get("cfg_gate_state")
+        if (
+            not state
+            or not state["requested"]
+            or ctx.is_warmup
+            or get_world_group().local_rank != 0
+        ):
+            return
+
+        logger.info(
+            "CFG gating summary: fraction=%.3f, gate_step=%d/%d, "
+            "fresh_uncond=%d, reused=%d, invalidations=%d, guidance_rescale=%s.",
+            state["fraction"],
+            state["gate_step"],
+            len(ctx.timesteps),
+            state["fresh_uncond"],
+            state["reused"],
+            state["invalidations"],
+            batch.guidance_rescale,
+        )
+
     def _post_denoising_loop(
         self,
         batch: Req,
@@ -1398,6 +1472,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
         target_dtype,
         current_guidance_scale,
         cfg_policy: CFGPolicy,
+        cfg_gate_state: dict[str, Any] | None,
         server_args: ServerArgs,
         guidance: torch.Tensor,
         latents: torch.Tensor,
@@ -1429,6 +1504,48 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
                 )
             return _unwrap(pred_t)
 
+        if (
+            cfg_gate_state
+            and cfg_gate_state["active"]
+            and not server_args.enable_cfg_parallel
+            and type(cfg_policy) is CFGPolicy
+            and len(cfg_policy.branches) == 2
+            and cfg_policy.branches[0].is_conditional
+            and not cfg_policy.branches[1].is_conditional
+        ):
+            model_id = id(current_model)
+            if cfg_gate_state["model_id"] not in (None, model_id):
+                cfg_gate_state["delta"] = None
+                cfg_gate_state["invalidations"] += 1
+
+            pos_pred = predict_fn(cfg_policy.branches[0])
+            pos_t = _wrap(pos_pred)
+            delta_t = cfg_gate_state["delta"]
+            can_reuse = (
+                timestep_index >= cfg_gate_state["gate_step"]
+                and delta_t is not None
+                and len(pos_t) == len(delta_t)
+            )
+
+            if can_reuse:
+                neg_pred = _unwrap(tuple(p - d for p, d in zip(pos_t, delta_t)))
+                cfg_gate_state["reused"] += 1
+            else:
+                neg_pred = predict_fn(cfg_policy.branches[1])
+                neg_t = _wrap(neg_pred)
+                cfg_gate_state["delta"] = tuple(
+                    p.detach() - n.detach() for p, n in zip(pos_t, neg_t)
+                )
+                cfg_gate_state["model_id"] = model_id
+                cfg_gate_state["fresh_uncond"] += 1
+
+            return cfg_policy.combine(
+                [pos_pred, neg_pred],
+                batch,
+                cfg_scale,
+                server_args.pipeline_config,
+            )
+
         if server_args.enable_cfg_parallel:
             if (
                 len(cfg_policy.branches) == 2

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

@@ -0,0 +1,180 @@
+import os
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
+    DenoisingStage,
+)
+
+
+class _PipelineConfig:
+    def get_classifier_free_guidance_scale(self, batch, current_guidance_scale):
+        return current_guidance_scale
+
+    def slice_noise_pred(self, noise_pred, latents):
+        return noise_pred
+
+    def postprocess_cfg_noise(self, batch, noise_pred, noise_pred_cond):
+        return noise_pred
+
+
+class TestCFGGating(unittest.TestCase):
+    def _make_server_args(self, enable_cfg_parallel=False):
+        return SimpleNamespace(
+            enable_cfg_parallel=enable_cfg_parallel,
+            pipeline_config=_PipelineConfig(),
+        )
+
+    def _make_batch(self):
+        return SimpleNamespace(
+            cfg_normalization=0,
+            guidance_rescale=0,
+            do_classifier_free_guidance=True,
+            is_cfg_negative=False,
+        )
+
+    def _make_gate_state(self, gate_step=1, model_id=None, delta=None):
+        return {
+            "fraction": 0.5,
+            "requested": True,
+            "active": True,
+            "gate_step": gate_step,
+            "delta": delta,
+            "model_id": model_id,
+            "fresh_uncond": 0,
+            "reused": 0,
+            "invalidations": 0,
+        }
+
+    def test_reuses_unconditional_residual_after_gate_step(self):
+        stage = DenoisingStage.__new__(DenoisingStage)
+        batch = self._make_batch()
+        server_args = self._make_server_args()
+        policy = CFGPolicy().build(batch, {}, {}, {})
+        calls = []
+
+        def fake_predict_noise(**kwargs):
+            calls.append("neg" if batch.is_cfg_negative else "pos")
+            timestep = kwargs["timestep"]
+            timestep_value = float(timestep.item())
+            offset = 0.25 if batch.is_cfg_negative else 1.25
+            return torch.tensor([timestep_value + offset])
+
+        stage._predict_noise = fake_predict_noise
+        model = torch.nn.Identity()
+        latents = torch.zeros(1)
+        state = self._make_gate_state(gate_step=1)
+
+        first = stage._predict_noise_with_cfg(
+            current_model=model,
+            latent_model_input=latents,
+            timestep=torch.tensor(0),
+            batch=batch,
+            timestep_index=0,
+            attn_metadata=None,
+            target_dtype=torch.float32,
+            current_guidance_scale=4.0,
+            cfg_policy=policy,
+            cfg_gate_state=state,
+            server_args=server_args,
+            guidance=None,
+            latents=latents,
+        )
+        second = stage._predict_noise_with_cfg(
+            current_model=model,
+            latent_model_input=latents,
+            timestep=torch.tensor(1),
+            batch=batch,
+            timestep_index=1,
+            attn_metadata=None,
+            target_dtype=torch.float32,
+            current_guidance_scale=4.0,
+            cfg_policy=policy,
+            cfg_gate_state=state,
+            server_args=server_args,
+            guidance=None,
+            latents=latents,
+        )
+
+        self.assertTrue(torch.equal(first, torch.tensor([4.25])))
+        self.assertTrue(torch.equal(second, torch.tensor([5.25])))
+        self.assertEqual(calls, ["pos", "neg", "pos"])
+        self.assertEqual(state["fresh_uncond"], 1)
+        self.assertEqual(state["reused"], 1)
+        self.assertEqual(state["invalidations"], 0)
+
+    def test_model_switch_invalidates_cached_delta(self):
+        stage = DenoisingStage.__new__(DenoisingStage)
+        batch = self._make_batch()
+        server_args = self._make_server_args()
+        policy = CFGPolicy().build(batch, {}, {}, {})
+        calls = []
+
+        def fake_predict_noise(**kwargs):
+            calls.append("neg" if batch.is_cfg_negative else "pos")
+            value = 3.0 if batch.is_cfg_negative else 10.0
+            return torch.tensor([value])
+
+        stage._predict_noise = fake_predict_noise
+        old_model = torch.nn.Identity()
+        new_model = torch.nn.Identity()
+        latents = torch.zeros(1)
+        state = self._make_gate_state(
+            gate_step=0,
+            model_id=id(old_model),
+            delta=(torch.tensor([2.0]),),
+        )
+
+        output = stage._predict_noise_with_cfg(
+            current_model=new_model,
+            latent_model_input=latents,
+            timestep=torch.tensor(2),
+            batch=batch,
+            timestep_index=2,
+            attn_metadata=None,
+            target_dtype=torch.float32,
+            current_guidance_scale=2.0,
+            cfg_policy=policy,
+            cfg_gate_state=state,
+            server_args=server_args,
+            guidance=None,
+            latents=latents,
+        )
+
+        self.assertTrue(torch.equal(output, torch.tensor([17.0])))
+        self.assertEqual(calls, ["pos", "neg"])
+        self.assertEqual(state["model_id"], id(new_model))
+        self.assertEqual(state["fresh_uncond"], 1)
+        self.assertEqual(state["reused"], 0)
+        self.assertEqual(state["invalidations"], 1)
+
+    def test_cfg_parallel_disables_gate_state(self):
+        stage = DenoisingStage.__new__(DenoisingStage)
+        ctx = SimpleNamespace(timesteps=torch.arange(10), extra={}, is_warmup=True)
+        batch = self._make_batch()
+        server_args = self._make_server_args(enable_cfg_parallel=True)
+
+        with patch.dict(os.environ, {"SGLANG_DIFFUSION_CFG_GATE_STEP": "0.5"}):
+            stage._init_cfg_gate_state(ctx, batch, server_args)
+
+        self.assertTrue(ctx.extra["cfg_gate_state"]["requested"])
+        self.assertFalse(ctx.extra["cfg_gate_state"]["active"])
+        self.assertEqual(ctx.extra["cfg_gate_state"]["gate_step"], 11)
+
+    def test_rejects_invalid_gate_fraction(self):
+        stage = DenoisingStage.__new__(DenoisingStage)
+        ctx = SimpleNamespace(timesteps=torch.arange(10), extra={}, is_warmup=True)
+        batch = self._make_batch()
+        server_args = self._make_server_args()
+
+        with patch.dict(os.environ, {"SGLANG_DIFFUSION_CFG_GATE_STEP": "1.5"}):
+            with self.assertRaises(ValueError):
+                stage._init_cfg_gate_state(ctx, batch, server_args)
+
+
+if __name__ == "__main__":
+    unittest.main()