diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 9400c828e..5d633ea22 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -129,7 +129,7 @@ The following tables show which models support each feature: | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **SenseNova-U1** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | | **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md index 8052c90d3..3943b8d4e 100644 --- a/examples/offline_inference/sensenova_u1/README.md +++ b/examples/offline_inference/sensenova_u1/README.md @@ -37,6 +37,10 @@ python end2end.py --modality text2text \ # Custom resolution (image generation) python end2end.py --prompt "A futuristic cityscape at sunset" \ --width 2048 --height 1024 --think + +# Text-to-image with Cache-DiT acceleration +python end2end.py --prompt "A cute cat" \ + --cache-backend cache_dit \ ``` > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU. @@ -135,6 +139,8 @@ python end2end.py \ | `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism | | `--enforce-eager` | False | Disable torch.compile | | `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM | +| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration | +| `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation | ## Reducing GPU Memory Usage diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py index 12f28da6f..ebb0d7f24 100644 --- a/examples/offline_inference/sensenova_u1/end2end.py +++ b/examples/offline_inference/sensenova_u1/end2end.py @@ -155,6 +155,18 @@ def parse_args(): "of extra CPU<->GPU transfers." ), ) + parser.add_argument( + "--cache-backend", + type=str, + default=None, + choices=["cache_dit"], + help="Cache backend to use for image generation acceleration.", + ) + parser.add_argument( + "--enable-cache-dit-summary", + action="store_true", + help="Enable cache-dit summary logging after diffusion forward passes.", + ) nullify_stage_engine_defaults(parser) return parser.parse_args() @@ -182,6 +194,8 @@ def main(): tensor_parallel_size=args.tensor_parallel_size, enforce_eager=args.enforce_eager, enable_cpu_offload=args.enable_cpu_offload, + cache_backend=args.cache_backend, + enable_cache_dit_summary=args.enable_cache_dit_summary, ) extra_args = { @@ -225,6 +239,7 @@ def main(): print(f" Seed : {args.seed}") print(f" Think mode : {args.think}") print(f" TP size : {args.tensor_parallel_size}") + print(f" Cache backend : {args.cache_backend or 'none'}") print(f"{'=' * 60}\n") # Build prompt dict diff --git a/examples/online_serving/sensenova_u1/README.md b/examples/online_serving/sensenova_u1/README.md index 161e5de3b..86eadb1d6 100644 --- a/examples/online_serving/sensenova_u1/README.md +++ b/examples/online_serving/sensenova_u1/README.md @@ -6,11 +6,21 @@ vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 ``` +With Cache-DiT acceleration: + +```bash +vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \ + --cache-backend cache_dit +``` + Or use the convenience script: ```bash cd examples/online_serving/sensenova_u1 bash run_server.sh + +# Cache-DiT +CACHE_BACKEND=cache_dit bash run_server.sh ``` ### Tensor Parallelism (TP) diff --git a/examples/online_serving/sensenova_u1/run_server.sh b/examples/online_serving/sensenova_u1/run_server.sh index 8edf54cc3..9dd8a9662 100644 --- a/examples/online_serving/sensenova_u1/run_server.sh +++ b/examples/online_serving/sensenova_u1/run_server.sh @@ -3,10 +3,18 @@ MODEL="${MODEL:-SenseNova/SenseNova-U1-8B-MoT}" PORT="${PORT:-8091}" +CACHE_BACKEND="${CACHE_BACKEND:-none}" echo "Starting SenseNova-U1 server..." echo "Model: $MODEL" echo "Port: $PORT" +echo "Cache backend: $CACHE_BACKEND" + +EXTRA_ARGS=() +if [ "$CACHE_BACKEND" != "none" ]; then + EXTRA_ARGS+=(--cache-backend "$CACHE_BACKEND") +fi vllm serve "$MODEL" --omni \ - --port "$PORT" + --port "$PORT" \ + "${EXTRA_ARGS[@]}" diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img.py index 10f5f05dc..fedd8e90e 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img.py @@ -167,3 +167,18 @@ def test_sensenova_u1_img2img(run_level): generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image) if run_level == "advanced_model": _validate_pixels(generated_image) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cache +@hardware_test(res={"cuda": "H100"}) +def test_sensenova_u1_img2img_cache_dit(): + """Test SenseNova-U1 img2img with Cache-DiT enabled.""" + input_image = _load_input_image() + with OmniRunner( + "SenseNova/SenseNova-U1-8B-MoT", + stage_configs_path=None, + cache_backend="cache_dit", + ) as runner: + _generate_sensenova_u1_img2img(runner.omni, input_image) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img.py index 334381067..fc9b5411e 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img.py @@ -156,3 +156,17 @@ def test_sensenova_u1_text2img(run_level): generated_image = _generate_sensenova_u1_image(runner.omni) if run_level == "advanced_model": _validate_pixels(generated_image) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cache +@hardware_test(res={"cuda": "H100"}) +def test_sensenova_u1_text2img_cache_dit(): + """Test SenseNova-U1 text2img with Cache-DiT enabled.""" + with OmniRunner( + "SenseNova/SenseNova-U1-8B-MoT", + stage_configs_path=None, + cache_backend="cache_dit", + ) as runner: + _generate_sensenova_u1_image(runner.omni) diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 7c27c31eb..005857642 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -18,6 +18,7 @@ from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifie from cache_dit.caching.block_adapters import FakeDiffusionPipeline from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 +from cache_dit.caching.cache_blocks.pattern_3_4_5 import CachedBlocks_Pattern_3_4_5 from cache_dit.caching.cache_contexts import BasicCacheConfig from cache_dit.caching.cache_contexts.cache_manager import CachedContextManager from vllm.logger import init_logger @@ -1170,6 +1171,162 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], return refresh_cache_context +class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5): + """ + Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden + states during denoising. + """ + + @classmethod + def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool: + if kwargs.get("cache_dit_skip", False): + return False + + # Prefix/text forwards either omit image_gen_indicators or update the + # DynamicCache. Denoising forwards are gen-only and use update_cache=False. + if kwargs.get("update_cache", True): + return False + + exist_gen = kwargs.get("exist_gen") + exist_und = kwargs.get("exist_und") + if exist_gen is None or exist_und is None: + image_gen_indicators = kwargs.get("image_gen_indicators") + if image_gen_indicators is None: + return False + exist_gen = image_gen_indicators.any().item() + exist_und = (~image_gen_indicators).any().item() + + return exist_gen and not exist_und + + @staticmethod + def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + kwargs = dict(kwargs) + kwargs.pop("cache_dit_skip", None) + return kwargs + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs): + block_kwargs = self._strip_cache_only_kwargs(kwargs) + if not self._is_denoising_call(kwargs): + hidden_states, new_encoder_hidden_states = self.call_blocks( + hidden_states, + *args, + **block_kwargs, + ) + return self._process_forward_outputs(hidden_states, new_encoder_hidden_states) + + return super().forward(hidden_states, *args, **block_kwargs) + + +class SensenovaCachedAdapter(CachedAdapter): + """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks.""" + + @classmethod + def collect_unified_blocks( + cls, + block_adapter: BlockAdapter, + contexts_kwargs: list[dict], + ) -> list[dict[str, torch.nn.ModuleList]]: + BlockAdapter.assert_normalized(block_adapter) + + total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = [] + assert hasattr(block_adapter.pipe, "_context_manager") + + for i in range(len(block_adapter.transformer)): + unified_blocks_bind_context = {} + for j in range(len(block_adapter.blocks[i])): + cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"] + unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList( + [ + SensenovaCachedBlocks( + # 0. Transformer blocks configuration + block_adapter.blocks[i][j], + transformer=block_adapter.transformer[i], + forward_pattern=block_adapter.forward_pattern[i][j], + check_forward_pattern=block_adapter.check_forward_pattern, + check_num_outputs=block_adapter.check_num_outputs, + # 1. Cache/Prune context configuration + cache_prefix=block_adapter.blocks_name[i][j], + cache_context=block_adapter.unique_blocks_name[i][j], + context_manager=block_adapter.pipe._context_manager, + cache_type=cache_config.cache_type, + ) + ] + ) + + total_cached_blocks.append(unified_blocks_bind_context) + + return total_cached_blocks + + +def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for SenseNova-U1 model. + + Args: + pipeline: The SenseNova-U1 pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + transformer = getattr(getattr(pipeline, "language_model", None), "model", None) + if transformer is None or not hasattr(transformer, "layers"): + raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.") + + # Build DBCacheConfig + db_cache_config = _build_db_cache_config(cache_config) + + calibrator_config = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + + modifier = ParamsModifier( + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + logger.info( + f"Enabling cache-dit on SenseNova-U1 decoder layers: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + SensenovaCachedAdapter.apply( + BlockAdapter( + transformer=transformer, + blocks=transformer.layers, + forward_pattern=ForwardPattern.Pattern_3, + params_modifiers=[modifier], + check_forward_pattern=True, + has_separate_cfg=True, + ), + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + transformer = pipeline.language_model.model + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]: """Enable cache-dit for Flux.2-dev pipeline. @@ -1587,6 +1744,7 @@ CUSTOM_DIT_ENABLERS.update( "LTX23Pipeline": enable_cache_for_ltx2, "LTX23ImageToVideoPipeline": enable_cache_for_ltx2, "BagelPipeline": enable_cache_for_bagel, + "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1, "GlmImagePipeline": enable_cache_for_glm_image, "Flux2Pipeline": enable_cache_for_flux2, "DreamIDOmniPipeline": enable_cache_for_dreamid_omni, diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index 906eac974..fe8695b07 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -707,7 +707,17 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli return out.past_key_values, out.hidden_states def _t2i_predict_v( - self, input_embeds, indexes_image, attn_mask, past_key_values, t, z, image_token_num, image_size=None, **_kw + self, + input_embeds, + indexes_image, + attn_mask, + past_key_values, + t, + z, + image_token_num, + image_size=None, + cache_dit_skip=False, + **_kw, ): B, L = z.shape[0], z.shape[1] outputs = self.language_model( @@ -719,6 +729,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli attention_mask=attn_mask, past_key_values=past_key_values, update_cache=False, + cache_dit_skip=cache_dit_skip, use_cache=True, compute_logits=False, ) @@ -1066,6 +1077,13 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli kv_cond = caches["cond"] idx_cond = caches["idx_cond"] mask_cond = caches["mask_cond"] + is_it2i = "img_cond" in caches + if is_it2i: + use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0 + needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1) + has_cached_partner = use_cfg and needs_cfg + else: + has_cached_partner = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1 out_cond = self._t2i_predict_v( image_embeds, @@ -1076,9 +1094,9 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli z, image_token_num=ns.token_h * ns.token_w, image_size=p.image_size, + cache_dit_skip=not has_cached_partner, ) - is_it2i = "img_cond" in caches if is_it2i: return self._denoise_step_it2i(out_cond, image_embeds, ns, t, z, caches, p, step_i) return self._denoise_step_t2i(out_cond, image_embeds, ns, t, z, caches, p, step_i) @@ -1153,6 +1171,7 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli caches["img_cond"], t, z, + cache_dit_skip=True, **predict_kw, ) out_uncond = self._t2i_predict_v( diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py index da43f55df..f9ff8aaf4 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py +++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py @@ -587,6 +587,12 @@ class SenseNovaU1DecoderLayer(nn.Module): past_key_values=None, **kwargs, ): + if isinstance(attention_mask, dict): + attention_mask = attention_mask.get( + self.attention_type, + attention_mask.get("full_attention"), + ) + if exist_und and not exist_gen: return self._forward_und(hidden_states, indexes, attention_mask, past_key_values, **kwargs) if not exist_und and exist_gen: @@ -663,7 +669,7 @@ class SenseNovaU1Model(nn.Module): exist_und=exist_und, exist_gen=exist_gen, indexes=indexes, - attention_mask=causal_mask_mapping.get(layer.attention_type, causal_mask_mapping.get("full_attention")), + attention_mask=causal_mask_mapping, past_key_values=past_key_values, use_cache=use_cache, **kwargs, diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index fe7d00c77..a381640b5 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -287,7 +287,6 @@ DiffusionModelRegistry = _ModelRegistry( _NO_CACHE_ACCELERATION = { # Pipelines that do not support cache acceleration (cache_dit / tea_cache). "NextStep11Pipeline", - "SenseNovaU1Pipeline", "AudioXPipeline", }