From 4ab1b4ff9cfc2ca233cccb5a4ba1f845288278fc Mon Sep 17 00:00:00 2001 From: Kristoffer Date: Tue, 11 Aug 2026 09:31:02 +0200 Subject: [PATCH] [Perf][Flux2][HunyuanVideo1.5] Skip attention-mask to avoid varlen path (#4645) Signed-off-by: kTorp Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../parallelism/sequence_parallel.md | 2 ++ .../models/flux2/flux2_transformer.py | 27 ++++++++++++++++++- .../hunyuan_video_15_transformer.py | 26 +++++++++++++++--- 3 files changed, 51 insertions(+), 4 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism/sequence_parallel.md b/docs/user_guide/diffusion/parallelism/sequence_parallel.md index e69b541f2..9bb5b5d7d 100644 --- a/docs/user_guide/diffusion/parallelism/sequence_parallel.md +++ b/docs/user_guide/diffusion/parallelism/sequence_parallel.md @@ -164,10 +164,12 @@ In `DiffusionParallelConfig`: | `ulysses_degree` | int | 1 | Number of GPUs for Ulysses-SP. Uses all-to-all communication. | | `ring_degree` | int | 1 | Number of GPUs for Ring-Attention. Uses P2P ring communication. | | `ulysses_mode` | str | `"default"` | Ulysses attention mode. Set to `"advanced_uaa"` to handle arbitrary sequence lengths and head counts without padding. | +| `mask_sp_padding` | bool | `False` | When the sequence length is not divisible by the SP size, tokens are auto-padded with zeros. Set to `True` to mask those padding tokens (strict, but uses the slower varlen attention path); the default `False` leaves them unmasked, keeping the fast path with negligible numerical impact. | **Notes:** - Total sequence parallel size equals to `ulysses_degree × ring_degree` - Degrees must evenly divide the sequence length for optimal performance (or use `ulysses_mode="advanced_uaa"` for Ulysses-SP) +- `mask_sp_padding` is an experimental feature, currently only supported by `Wan2.2`, `Wan2.2 Vace`, `Qwen-Image`, `Flux 2`, and `HunyuanVideo 1.5` ## Best Practices diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index f05632efd..9715cae0e 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -15,6 +15,7 @@ from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous from diffusers.utils import is_torch_npu_available from torch import nn +from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -38,6 +39,8 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +logger = init_logger(__name__) + def _join_prefix(prefix: str, name: str) -> str: return f"{prefix}.{name}" if prefix else name @@ -946,7 +949,13 @@ class Flux2Transformer2DModel(nn.Module): ctx = get_forward_context() else: ctx = None - if ctx is not None and ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + if ( + ctx is not None + and self.parallel_config.sequence_parallel_size > 1 + and self.parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): batch_size = hidden_states.shape[0] img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size hidden_states_mask = torch.ones( @@ -958,6 +967,22 @@ class Flux2Transformer2DModel(nn.Module): hidden_states_mask[:, ctx.sp_original_seq_len :] = False if hidden_states_mask.all(): hidden_states_mask = None + elif ( + ctx is not None + and self.parallel_config.sequence_parallel_size > 1 + and not self.parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + logger.warning_once( + "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). " + "Padding tokens are not masked from attention (mask_sp_padding=False), " + "which avoids the varlen attention path but may produce minor numerical differences. " + "Set parallel_config.mask_sp_padding=True to restore strict masking.", + ctx.sp_padding_size, + ctx.sp_original_seq_len, + self.parallel_config.sequence_parallel_size, + ) if hidden_states_mask is not None: joint_attention_kwargs["hidden_states_mask"] = hidden_states_mask 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 04862fe0a..0e670d731 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 @@ -814,7 +814,13 @@ class HunyuanVideo15Transformer3DModel(nn.Module): ctx = get_forward_context() hidden_states_mask = None - if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + if ( + self.parallel_config is not None + and self.parallel_config.sequence_parallel_size > 1 + and self.parallel_config.mask_sp_padding + and 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, @@ -823,10 +829,24 @@ class HunyuanVideo15Transformer3DModel(nn.Module): 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 + elif ( + self.parallel_config is not None + and self.parallel_config.sequence_parallel_size > 1 + and not self.parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + logger.warning_once( + "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). " + "Padding tokens are not masked from attention (mask_sp_padding=False), " + "which avoids the varlen attention path but may produce minor numerical differences. " + "Set parallel_config.mask_sp_padding=True to restore strict masking.", + ctx.sp_padding_size, + ctx.sp_original_seq_len, + self.parallel_config.sequence_parallel_size, + ) for block in self.transformer_blocks: hidden_states, encoder_hidden_states = block(