diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 5bf31108b..436ea2966 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -1383,6 +1383,61 @@ def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[ return refresh_cache_context +def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for HunyuanVideo 1.5 pipeline. + + HunyuanVideo 1.5 uses a single transformer with has_separate_cfg=True + (separate conditional/unconditional forward passes). The _sp_plan scatter + is applied at the transformer input boundary (empty-string key), so CacheDiT + sees already-sharded hidden_states throughout transformer_blocks. + """ + 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}") + + logger.info( + f"Enabling cache-dit on HunyuanVideo15 transformer: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + cache_dit.enable_cache( + BlockAdapter( + transformer=pipeline.transformer, + blocks=pipeline.transformer.transformer_blocks, + forward_pattern=ForwardPattern.Pattern_0, + 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: + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + pipeline.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 + + # Register custom cache-dit enablers after function definitions CUSTOM_DIT_ENABLERS.update( { @@ -1406,6 +1461,8 @@ CUSTOM_DIT_ENABLERS.update( "GlmImagePipeline": enable_cache_for_glm_image, "Flux2Pipeline": enable_cache_for_flux2, "ErnieImagePipeline": enable_cache_for_ernie_image, + "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15, + "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15, } ) diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 6600b17d5..30ba62184 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py @@ -24,6 +24,10 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module +from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput +from vllm_omni.diffusion.distributed.sp_sharding import sp_shard_with_padding +from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward @@ -395,7 +399,8 @@ class HunyuanVideo15Attention(nn.Module): encoder_hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states_mask: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: qkv, _ = self.to_qkv(hidden_states) q_size = self.to_qkv.num_heads * self.head_dim kv_size = self.to_qkv.num_kv_heads * self.head_dim @@ -430,22 +435,40 @@ class HunyuanVideo15Attention(nn.Module): encoder_query = self.norm_added_q(encoder_query) encoder_key = self.norm_added_k(encoder_key) - query = torch.cat([query, encoder_query], dim=1) - key = torch.cat([key, encoder_key], dim=1) - value = torch.cat([value, encoder_value], dim=1) - attn_metadata = None - if attention_mask is not None: - seq_len = query.shape[1] - # Pad mask to full sequence length (video + encoder tokens) - # Keep mask 2D (batch_size, seq_len) — each attention backend - # handles reshaping internally (flash_attn uses unpadding, - # SDPA expands to 4D via _maybe_reshape_attn_mask). - attention_mask = F.pad(attention_mask, (seq_len - attention_mask.shape[1], 0), value=True) - attention_mask = attention_mask.bool() - attn_metadata = AttentionMetadata(attn_mask=attention_mask) + ctx = get_forward_context() + if ctx.sp_active and encoder_hidden_states is not None: + # Under Ulysses SP, encoder tokens are passed via joint_* + # metadata so they can be head-sliced separately from the + # all-to-all'd video tokens in UlyssesParallelAttention. + attn_metadata = AttentionMetadata( + joint_query=encoder_query, + joint_key=encoder_key, + joint_value=encoder_value, + joint_strategy="rear", + ) + if attention_mask is not None: + attn_metadata.joint_attn_mask = attention_mask.bool() + if hidden_states_mask is not None: + attn_metadata.attn_mask = hidden_states_mask + hidden_states = self.attn(query, key, value, attn_metadata) + else: + if encoder_hidden_states is not None: + query = torch.cat([query, encoder_query], dim=1) + key = torch.cat([key, encoder_key], dim=1) + value = torch.cat([value, encoder_value], dim=1) - hidden_states = self.attn(query, key, value, attn_metadata) + if attention_mask is not None: + seq_len = query.shape[1] + # Pad mask to full sequence length (video + encoder tokens) + # Keep mask 2D (batch_size, seq_len) - each attention backend + # handles reshaping internally (flash_attn uses unpadding, + # SDPA expands to 4D via _maybe_reshape_attn_mask). + attention_mask = F.pad(attention_mask, (seq_len - attention_mask.shape[1], 0), value=True) + attention_mask = attention_mask.bool() + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) @@ -499,7 +522,8 @@ class HunyuanVideo15TransformerBlock(nn.Module): temb: torch.Tensor, attention_mask: torch.Tensor | None = None, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states_mask: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb) norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context( encoder_hidden_states, emb=temb @@ -510,6 +534,7 @@ class HunyuanVideo15TransformerBlock(nn.Module): encoder_hidden_states=norm_encoder_hidden_states, attention_mask=attention_mask, image_rotary_emb=freqs_cis, + hidden_states_mask=hidden_states_mask, ) hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) @@ -547,6 +572,14 @@ class HunyuanVideo15Transformer3DModel(nn.Module): _hsdp_shard_conditions = [is_transformer_block_module] + _sp_plan = { + "rope": { + 0: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + 1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + }, + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, od_config: OmniDiffusionConfig, @@ -634,6 +667,18 @@ class HunyuanVideo15Transformer3DModel(nn.Module): hidden_states = self.x_embedder(hidden_states) + # Scatter hidden_states along seq dim for sequence parallelism. + # Done here (after x_embedder, before transformer_blocks) so that + # CacheDiT sees already-sharded hidden_states when saving + # original_hidden_states at the start of CachedBlocks.forward. + if get_sequence_parallel_world_size() > 1: + hidden_states, _pad_size = sp_shard_with_padding(hidden_states, dim=1) + if _pad_size > 0: + ctx = get_forward_context() + if ctx.sp_original_seq_len is None: + ctx.sp_padding_size = _pad_size + ctx.sp_original_seq_len = hidden_states.shape[1] * get_sequence_parallel_world_size() - _pad_size + encoder_hidden_states = self.context_embedder(encoder_hidden_states, timestep, encoder_attention_mask) encoder_hidden_states_cond_emb = self.cond_type_embed( @@ -717,6 +762,23 @@ class HunyuanVideo15Transformer3DModel(nn.Module): encoder_hidden_states = torch.stack(new_encoder_hidden_states) encoder_attention_mask = torch.stack(new_encoder_attention_mask) + # Create explicit attn_mask for image tokens when SP auto_pad is active. + ctx = get_forward_context() + hidden_states_mask = None + if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + hidden_states_mask = torch.ones( + batch_size, + padded_seq_len, + dtype=torch.bool, + device=hidden_states.device, + ) + hidden_states_mask[:, ctx.sp_original_seq_len :] = False + + # if mask is all true, set it to None + if hidden_states_mask.all(): + hidden_states_mask = None + for block in self.transformer_blocks: hidden_states, encoder_hidden_states = block( hidden_states, @@ -724,6 +786,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module): temb, encoder_attention_mask, image_rotary_emb, + hidden_states_mask=hidden_states_mask, ) hidden_states = self.norm_out(hidden_states, temb)