From da758ed601270b21e1cfb404306ff0ca5c816a3f Mon Sep 17 00:00:00 2001 From: wxy <1908865287@qq.com> Date: Wed, 4 Feb 2026 09:03:37 +0800 Subject: [PATCH] [diffusion] fix: fix server cache-dit bug under continuous dynamic requests (#17140) --- .../runtime/cache/cache_dit_integration.py | 68 ++++++++++- .../pipelines_core/stages/denoising.py | 111 +++++++++++------- 2 files changed, 138 insertions(+), 41 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py index 7d7654ad4..e08124881 100644 --- 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})" + ) 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 57bc2b02d..13da5f65a 100644 --- 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