diff --git a/python/sglang/multimodal_gen/envs.py b/python/sglang/multimodal_gen/envs.py index a5341a777..3dc73dcac 100644 --- a/python/sglang/multimodal_gen/envs.py +++ b/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" ), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 55fe2118b..761a09cfa 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_cfg_gating.py b/python/sglang/multimodal_gen/test/unit/test_cfg_gating.py new file mode 100644 index 000000000..37f66c9c7 --- /dev/null +++ b/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()