diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md index a2e245526..1e656c22d 100644 --- a/examples/offline_inference/sensenova_u1/README.md +++ b/examples/offline_inference/sensenova_u1/README.md @@ -38,9 +38,9 @@ python end2end.py --modality text2text \ python end2end.py --prompt "A futuristic cityscape at sunset" \ --width 2048 --height 1024 --think -# Text-to-image with Cache-DiT acceleration +# Text-to-image with cache acceleration (cache_dit or tea_cache) python end2end.py --prompt "A cute cat" \ - --cache-backend cache_dit \ + --cache-backend tea_cache ``` > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU. @@ -139,7 +139,7 @@ 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 | +| `--cache-backend` | None | Set to `cache_dit` or `tea_cache` for image generation acceleration | | `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation | ## Reducing GPU Memory Usage @@ -211,4 +211,4 @@ The img2img counterpart lives at ## Online Serving For OpenAI-compatible API serving, see -[`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/). +[`examples/online_serving/sensenova_u1/`](../../online_serving/sensenova_u1/README.md). diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py index 05f037dd7..7c727449e 100644 --- a/examples/offline_inference/sensenova_u1/end2end.py +++ b/examples/offline_inference/sensenova_u1/end2end.py @@ -164,7 +164,7 @@ def parse_args(): "--cache-backend", type=str, default=None, - choices=["cache_dit"], + choices=["cache_dit", "tea_cache"], help="Cache backend to use for image generation acceleration.", ) parser.add_argument( diff --git a/examples/online_serving/sensenova_u1/README.md b/examples/online_serving/sensenova_u1/README.md index 86eadb1d6..d66a3506a 100644 --- a/examples/online_serving/sensenova_u1/README.md +++ b/examples/online_serving/sensenova_u1/README.md @@ -6,11 +6,11 @@ vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 ``` -With Cache-DiT acceleration: +With cache acceleration (`cache_dit` or `tea_cache`): ```bash vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \ - --cache-backend cache_dit + --cache-backend tea_cache ``` Or use the convenience script: @@ -19,8 +19,8 @@ Or use the convenience script: cd examples/online_serving/sensenova_u1 bash run_server.sh -# Cache-DiT -CACHE_BACKEND=cache_dit bash run_server.sh +# Cache acceleration (cache_dit or tea_cache) +CACHE_BACKEND=tea_cache bash run_server.sh ``` ### Tensor Parallelism (TP) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 772dec789..f6f9db88d 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -57,6 +57,22 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: ) +def enable_sensenova_u1_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """Enable TeaCache for SenseNova-U1 denoising forwards.""" + teacache_config = TeaCacheConfig( + transformer_type="SenseNovaU1ForCausalLM", + rel_l1_thresh=config.rel_l1_thresh, + coefficients=config.coefficients, + ) + transformer = pipeline.denoising_transformer + apply_teacache_hook(transformer, teacache_config) + + logger.info( + f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " + f"transformer_class={teacache_config.transformer_type}" + ) + + def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: """ Enable TeaCache for Flux2 Klein model. @@ -80,6 +96,7 @@ CUSTOM_TEACACHE_ENABLERS = { "BagelPipeline": enable_bagel_teacache, "Flux2KleinPipeline": enable_flux2_klein_teacache, "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache, + "SenseNovaU1Pipeline": enable_sensenova_u1_teacache, } @@ -180,6 +197,8 @@ class TeaCacheBackend(CacheBackend): # Extract transformer from pipeline transformer = pipeline.transformer + if not hasattr(transformer, "_hook_registry") and hasattr(pipeline, "denoising_transformer"): + transformer = pipeline.denoising_transformer if hasattr(transformer, "_hook_registry"): hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME) diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 7efdd418e..bcf2d3fe4 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -37,6 +37,14 @@ _MODEL_COEFFICIENTS = { # Bagel transformer coefficients # Using Qwen's coefficients as reasonable default given shared architecture "Bagel": [1.33313129e06, -1.68644226e05, 7.95050740e03, -1.63747873e02, 1.26352397e00], + # SenseNova-U1 transformer coefficients + "SenseNovaU1ForCausalLM": [ + 9.07281930e04, + -2.17699186e04, + 1.83940990e03, + -6.30339273e01, + 7.61309272e-01, + ], # Z-Image transformer coefficients # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future "ZImageTransformer2DModel": [ diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 1f0c38aee..b0bd58bc8 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1187,6 +1187,68 @@ def extract_flux_context( ) +def extract_sensenova_u1_context( + module: nn.Module, + input_ids: torch.Tensor | None = None, + indexes: torch.Tensor | None = None, + attention_mask: dict[str, torch.Tensor] | torch.Tensor | None = None, + past_key_values: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + embed_only: bool = False, + compute_logits: bool = True, + **kwargs: Any, +) -> CacheContext: + """Extract cache context for SenseNovaU1ForCausalLM denoising forwards.""" + from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import ( + SenseNovaU1CausalLMOutput, + ) + + layer_kwargs = dict(kwargs) + layer_kwargs.pop("cache_dit_skip", None) + image_gen_indicators = layer_kwargs.pop("image_gen_indicators", None) + exist_und = (~image_gen_indicators).any().item() + exist_gen = image_gen_indicators.any().item() + causal_mask_mapping = attention_mask + + first_layer = module.model.layers[0] + modulated_input = first_layer.input_layernorm_mot_gen(inputs_embeds) + + def run_transformer_blocks(): + h = inputs_embeds + for layer in module.model.layers: + h = layer( + h, + image_gen_indicators=image_gen_indicators, + exist_und=exist_und, + exist_gen=exist_gen, + indexes=indexes, + attention_mask=causal_mask_mapping, + past_key_values=past_key_values, + **layer_kwargs, + ) + return (h,) + + def postprocess(h: torch.Tensor) -> SenseNovaU1CausalLMOutput: + h = module.model.norm_mot_gen(h) + + logits = module.logits_processor(module.lm_head, h) if compute_logits else None + return SenseNovaU1CausalLMOutput( + logits=logits, + past_key_values=past_key_values if use_cache else None, + hidden_states=h, + ) + + return CacheContext( + modulated_input=modulated_input, + hidden_states=inputs_embeds, + encoder_hidden_states=None, + temb=modulated_input, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + ) + + # Registry for model-specific extractors # Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -1202,6 +1264,7 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = { "Flux2Transformer2DModel": extract_flux2_context, "LongCatImageTransformer2DModel": extract_longcat_context, "FluxTransformer2DModel": extract_flux_context, + "SenseNovaU1ForCausalLM": extract_sensenova_u1_context, # Future models: # "CogVideoXTransformer3DModel": extract_cogvideox_context, } 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 73951fb35..9c4a052ae 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -490,6 +490,30 @@ def _optimized_scale(positive_flat, negative_flat): # --------------------------------------------------------------------------- +class SenseNovaU1DenoisingAdapter(nn.Module): + """Denoising-only entry point used by cache backends.""" + + def __init__(self, language_model: SenseNovaU1ForCausalLM): + super().__init__() + object.__setattr__(self, "language_model", language_model) + self.do_true_cfg = True + + @property + def model(self): + return self.language_model.model + + @property + def lm_head(self): + return self.language_model.lm_head + + @property + def logits_processor(self): + return self.language_model.logits_processor + + def forward(self, *args, **kwargs): + return self.language_model(*args, **kwargs) + + class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin): """SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni. @@ -536,8 +560,12 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli self.llm_cfg, prefix="language_model", ) - # We define this for Cache DiT compatibility + # Cache-DiT hooks pipeline.transformer(.blocks), so it must point at the + # real decoder module (exposes .blocks and real parameters). self.transformer = self.language_model.model + # TeaCache intercepts the ForCausalLM-level denoising forward; route it + # through a dedicated adapter so it does not collide with Cache-DiT. + self.denoising_transformer = SenseNovaU1DenoisingAdapter(self.language_model) # Vision model (understanding branch) self.vision_model = NEOVisionModel(self.vis_cfg) @@ -704,7 +732,8 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli **_kw, ): B, L = z.shape[0], z.shape[1] - outputs = self.language_model( + denoising_model = self.language_model if cache_dit_skip else self.denoising_transformer + outputs = denoising_model( inputs_embeds=input_embeds, image_gen_indicators=torch.ones( (input_embeds.shape[0], input_embeds.shape[1]), dtype=torch.bool, device=input_embeds.device