[diffusion] fix: fix server cache-dit bug under continuous dynamic requests (#17140)

python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py

@@ -446,11 +446,20 @@ def enable_cache_on_dual_transformer(
         compute_steps = sum(primary_config.steps_computation_mask)
         cache_steps = len(primary_config.steps_computation_mask) - compute_steps
         logger.info(
-            "  SCM enabled: %d compute steps, %d cache steps, policy=%s",
+            "  SCM enabled for primary transformer: %d compute steps, %d cache steps, policy=%s",
             compute_steps,
             cache_steps,
             primary_config.steps_computation_policy,
         )
+    if secondary_config.steps_computation_mask:
+        compute_steps = sum(secondary_config.steps_computation_mask)
+        cache_steps = len(secondary_config.steps_computation_mask) - compute_steps
+        logger.info(
+            "  SCM enabled for secondary transformer: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            secondary_config.steps_computation_policy,
+        )
 
     parallelism_config = _build_parallelism_config(sp_group, tp_group)
     if parallelism_config is not None:
@@ -508,3 +517,60 @@ def enable_cache_on_dual_transformer(
                 context_manager._sglang_tp_sp_group = tp_sp_group
 
     return transformer, transformer_2
+
+
+def refresh_context_on_transformer(
+    transformer: torch.nn.Module,
+    num_inference_steps: int,
+    scm_preset: str | None = None,
+    verbose: bool = False,
+) -> None:
+    """Refresh cache-dit context for transformer."""
+    cache_dit.refresh_context(
+        transformer,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_inference_steps,
+            steps_computation_mask=cache_dit.steps_mask(
+                mask_policy=scm_preset, total_steps=num_inference_steps
+            ),
+            steps_computation_policy=scm_preset,
+        ),
+        verbose=verbose,
+    )
+    logger.debug(f"cache-dit refreshed on transformer (steps={num_inference_steps})")
+
+
+def refresh_context_on_dual_transformer(
+    transformer: torch.nn.Module,
+    transformer_2: torch.nn.Module,
+    num_high_noise_steps: int,
+    num_low_noise_steps: int,
+    scm_preset: str | None = None,
+    verbose: bool = False,
+) -> None:
+    """Refresh cache-dit context for dual transformers."""
+    cache_dit.refresh_context(
+        transformer,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_high_noise_steps,
+            steps_computation_mask=cache_dit.steps_mask(
+                mask_policy=scm_preset, total_steps=num_high_noise_steps
+            ),
+            steps_computation_policy=scm_preset,
+        ),
+        verbose=verbose,
+    )
+    cache_dit.refresh_context(
+        transformer_2,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_low_noise_steps,
+            steps_computation_mask=cache_dit.steps_mask(
+                mask_policy=scm_preset, total_steps=num_low_noise_steps
+            ),
+            steps_computation_policy=scm_preset,
+        ),
+        verbose=verbose,
+    )
+    logger.debug(
+        f"cache-dit refreshed on dual transformers (steps={num_high_noise_steps}, {num_low_noise_steps})"
+    )

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

@@ -25,12 +25,23 @@ from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     Wan2_2_TI2V_5B_Config,
     WanI2V480PConfig,
 )
+from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
+    CacheDitConfig,
+    enable_cache_on_dual_transformer,
+    enable_cache_on_transformer,
+    get_scm_mask,
+    refresh_context_on_dual_transformer,
+    refresh_context_on_transformer,
+)
 from sglang.multimodal_gen.runtime.distributed import (
     cfg_model_parallel_all_reduce,
     get_local_torch_device,
+    get_sp_group,
     get_sp_parallel_rank,
     get_sp_world_size,
+    get_tp_group,
     get_world_group,
+    get_world_size,
 )
 from sglang.multimodal_gen.runtime.distributed.communication_op import (
     sequence_model_parallel_all_gather,
@@ -136,7 +147,9 @@ class DenoisingStage(PipelineStage):
         # TODO(triple-mu): support customized fullgraph and dynamic in the future
         module.compile(mode=mode, fullgraph=False, dynamic=None)
 
-    def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None:
+    def _maybe_enable_cache_dit(
+        self, num_inference_steps: int | tuple[int, int], batch: Req
+    ) -> None:
         """Enable cache-dit on the transformers if configured (idempotent).
 
         This method should be called after the transformer is fully loaded
@@ -146,32 +159,33 @@ class DenoisingStage(PipelineStage):
         transformers with (potentially) different configurations.
 
         """
+        if isinstance(num_inference_steps, tuple):
+            num_high_noise_steps, num_low_noise_steps = num_inference_steps
+
+        # NOTE: When a new request arrives, we need to refresh the cache-dit context.
         if self._cache_dit_enabled:
-            if self._cached_num_steps != num_inference_steps:
+            scm_preset = envs.SGLANG_CACHE_DIT_SCM_PRESET
-                logger.warning(
+            scm_preset = None if scm_preset == "none" else scm_preset
-                    "num_inference_steps changed from %d to %d after cache-dit was enabled. "
+            if isinstance(num_inference_steps, tuple):
-                    "Continuing with initial configuration (steps=%d).",
+                refresh_context_on_dual_transformer(
-                    self._cached_num_steps,
+                    self.transformer,
+                    self.transformer_2,
+                    num_high_noise_steps,
+                    num_low_noise_steps,
+                    scm_preset=scm_preset,
+                )
+            else:
+                refresh_context_on_transformer(
+                    self.transformer,
                     num_inference_steps,
-                    self._cached_num_steps,
+                    scm_preset=scm_preset,
                 )
             return
+
         # check if cache-dit is enabled in config
         if not envs.SGLANG_CACHE_DIT_ENABLED or batch.is_warmup:
             return
 
-        from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
-            CacheDitConfig,
-            enable_cache_on_dual_transformer,
-            enable_cache_on_transformer,
-            get_scm_mask,
-        )
-        from sglang.multimodal_gen.runtime.distributed import (
-            get_sp_group,
-            get_tp_group,
-            get_world_size,
-        )
-
         world_size = get_world_size()
         parallelized = world_size > 1
 
@@ -229,7 +243,19 @@ class DenoisingStage(PipelineStage):
         # cache-dit handles step count validation and scaling internally
         steps_computation_mask = get_scm_mask(
             preset=scm_preset,
-            num_inference_steps=num_inference_steps,
+            num_inference_steps=(
+                num_inference_steps
+                if isinstance(num_inference_steps, int)
+                else num_high_noise_steps
+            ),
+            compute_bins=scm_compute_bins,
+            cache_bins=scm_cache_bins,
+        )
+
+        if isinstance(num_inference_steps, tuple):
+            steps_computation_mask_2 = get_scm_mask(
+                preset=scm_preset,
+                num_inference_steps=num_low_noise_steps,
                 compute_bins=scm_compute_bins,
                 cache_bins=scm_cache_bins,
             )
@@ -244,7 +270,11 @@ class DenoisingStage(PipelineStage):
             max_continuous_cached_steps=envs.SGLANG_CACHE_DIT_MC,
             enable_taylorseer=envs.SGLANG_CACHE_DIT_TAYLORSEER,
             taylorseer_order=envs.SGLANG_CACHE_DIT_TS_ORDER,
-            num_inference_steps=num_inference_steps,
+            num_inference_steps=(
+                num_inference_steps
+                if isinstance(num_inference_steps, int)
+                else num_high_noise_steps
+            ),
             # SCM fields
             steps_computation_mask=steps_computation_mask,
             steps_computation_policy=scm_policy,
@@ -263,9 +293,9 @@ class DenoisingStage(PipelineStage):
                 max_continuous_cached_steps=envs.SGLANG_CACHE_DIT_SECONDARY_MC,
                 enable_taylorseer=envs.SGLANG_CACHE_DIT_SECONDARY_TAYLORSEER,
                 taylorseer_order=envs.SGLANG_CACHE_DIT_SECONDARY_TS_ORDER,
-                num_inference_steps=num_inference_steps,
+                num_inference_steps=num_low_noise_steps,
                 # SCM fields - shared with primary
-                steps_computation_mask=steps_computation_mask,
+                steps_computation_mask=steps_computation_mask_2,
                 steps_computation_policy=scm_policy,
             )
 
@@ -281,8 +311,9 @@ class DenoisingStage(PipelineStage):
                 tp_group=tp_group,
             )
             logger.info(
-                "cache-dit enabled on dual transformers (steps=%d)",
+                "cache-dit enabled on dual transformers (steps=%d, %d)",
-                num_inference_steps,
+                num_high_noise_steps,
+                num_low_noise_steps,
             )
         else:
             # single transformer
@@ -482,12 +513,21 @@ class DenoisingStage(PipelineStage):
         """
         assert self.transformer is not None
         pipeline = self.pipeline() if self.pipeline else None
-        # NOTE: In warmup requests we may override req.num_inference_steps (e.g. set to 1)
+
-        # for latency amortization, but cache-dit needs the *original* total steps to
+        boundary_timestep = self._handle_boundary_ratio(server_args, batch)
-        # initialize/refresh its context correctly.
+        # Get timesteps and calculate warmup steps
-        cache_dit_num_inference_steps = batch.extra.get(
+        timesteps = batch.timesteps
-            "cache_dit_num_inference_steps", batch.num_inference_steps
+        num_inference_steps = batch.num_inference_steps
-        )
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        if self.transformer_2 is not None:
+            assert boundary_timestep is not None, "boundary_timestep must be provided"
+            num_high_noise_steps = (timesteps >= boundary_timestep).sum().item()
+            num_low_noise_steps = num_inference_steps - num_high_noise_steps
+            cache_dit_num_inference_steps = (num_high_noise_steps, num_low_noise_steps)
+        else:
+            cache_dit_num_inference_steps = num_inference_steps
+
         if not server_args.model_loaded["transformer"]:
             # FIXME: reuse more code
             loader = TransformerLoader()
@@ -515,13 +555,6 @@ class DenoisingStage(PipelineStage):
             target_dtype != torch.float32
         ) and not server_args.disable_autocast
 
-        # Get timesteps and calculate warmup steps
-        timesteps = batch.timesteps
-        if timesteps is None:
-            raise ValueError("Timesteps must be provided")
-        num_inference_steps = batch.num_inference_steps
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
-
         # Prepare image latents and embeddings for I2V generation
         image_embeds = batch.image_embeds
         if len(image_embeds) > 0:
@@ -543,8 +576,6 @@ class DenoisingStage(PipelineStage):
             assert neg_prompt_embeds is not None
             # Removed Tensor truthiness assert to avoid GPU sync
 
-        boundary_timestep = self._handle_boundary_ratio(server_args, batch)
-
         # specifically for Wan2_2_TI2V_5B_Config, not applicable for FastWan2_2_TI2V_5B_Config
         should_preprocess_for_wan_ti2v = (
             server_args.pipeline_config.task_type == ModelTaskType.TI2V