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

2026-02-04 09:03:37 +08:00
parent ae004e15c9
commit da758ed601
2 changed files with 138 additions and 41 deletions
--- a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py
+++ b/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})"
+    )
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+++ b/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:
-                logger.warning(
-                    "num_inference_steps changed from %d to %d after cache-dit was enabled. "
-                    "Continuing with initial configuration (steps=%d).",
-                    self._cached_num_steps,
+            scm_preset = envs.SGLANG_CACHE_DIT_SCM_PRESET
+            scm_preset = None if scm_preset == "none" else scm_preset
+            if isinstance(num_inference_steps, tuple):
+                refresh_context_on_dual_transformer(
+                    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,11 +243,23 @@ 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,
+            )
+
        # build config for primary transformer (high-noise expert)
        primary_config = CacheDitConfig(
            enabled=True,
@@ -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)",
-                num_inference_steps,
+                "cache-dit enabled on dual transformers (steps=%d, %d)",
+                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
-        # initialize/refresh its context correctly.
-        cache_dit_num_inference_steps = batch.extra.get(
-            "cache_dit_num_inference_steps", batch.num_inference_steps
-        )
+
+        boundary_timestep = self._handle_boundary_ratio(server_args, batch)
+        # Get timesteps and calculate warmup steps
+        timesteps = batch.timesteps
+        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