diff --git a/examples/inference/basic/basic_wan2_2.py b/examples/inference/basic/basic_wan2_2.py
index f1b324c73..bd2207b19 100644
--- a/examples/inference/basic/basic_wan2_2.py
+++ b/examples/inference/basic/basic_wan2_2.py
@@ -11,7 +11,7 @@ def main():
     generator = VideoGenerator.from_pretrained(
         "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
         # FastVideo will automatically handle distributed setup
-        num_gpus=1,
+        num_gpus=2,
         use_fsdp_inference=False, # set to True if GPU is out of memory 
         dit_cpu_offload=True, # DiT need to be offloaded for MoE
         vae_cpu_offload=False,
diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py
index f1177dadd..698c64e4f 100644
--- a/fastvideo/attention/layer.py
+++ b/fastvideo/attention/layer.py
@@ -66,11 +66,11 @@ class DistributedAttention(nn.Module):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
+        original_seq_len: int | None = None,
         replicated_q: torch.Tensor | None = None,
         replicated_k: torch.Tensor | None = None,
         replicated_v: torch.Tensor | None = None,
         freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
-        attention_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Forward pass for distributed attention.
         
@@ -78,10 +78,10 @@ class DistributedAttention(nn.Module):
             q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
             k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
             v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
+            original_seq_len (int): Original (unpadded) full sequence length
             replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
             replicated_k (Optional[torch.Tensor]): Replicated key tensor
             replicated_v (Optional[torch.Tensor]): Replicated value tensor
-            attention_mask (Optional[torch.Tensor]): Attention mask [batch_size, seq_len]
             
         Returns:
             Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
@@ -91,7 +91,7 @@ class DistributedAttention(nn.Module):
         # Check input shapes
         assert q.dim() == 4 and k.dim() == 4 and v.dim(
         ) == 4, "Expected 4D tensors"
-        batch_size, seq_len, num_heads, head_dim = q.shape
+        batch_size, _, num_heads, _ = q.shape
         local_rank = get_sp_parallel_rank()
         world_size = get_sp_world_size()
 
@@ -107,15 +107,11 @@ class DistributedAttention(nn.Module):
                                                     scatter_dim=2,
                                                     gather_dim=1)
 
-        # After all-to-all, each rank has the full sequence but only a subset of heads
-        # The attention mask should now apply to the full sequence length
-        # Since mask is [batch, full_seq_len], it's already in the correct format
-
-        # LOAY TODO, instead of slicing repeatedly maintain an original qkv and rewrite into that
-        valid_seq_len = None
-        if attention_mask is not None:
-            valid_seq_len = (attention_mask[0] == 1).sum().item()
-            qkv = qkv[:, :valid_seq_len, :, :]
+        # After all-to-all, each rank has the full sequence but only a subset of heads.
+        # Trim away SP padding for attention compute, then pad back before returning.
+        original_seq_len = original_seq_len or qkv.shape[1]
+        pad_seq_len = qkv.shape[1] - original_seq_len
+        qkv = qkv[:, :original_seq_len, :, :]
 
         if freqs_cis is not None:
             cos, sin = freqs_cis
@@ -145,7 +141,7 @@ class DistributedAttention(nn.Module):
         # Redistribute back if using sequence parallelism
         replicated_output = None
         if replicated_q is not None:
-            split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len
+            split_idx = original_seq_len
             replicated_output = output[:, split_idx:]
             output = output[:, :split_idx]
             # TODO: make this asynchronous
@@ -154,9 +150,7 @@ class DistributedAttention(nn.Module):
         # Apply backend-specific postprocess_output
         output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
 
-        if attention_mask is not None:
-            pad_len = (attention_mask[0] == 0).sum().item()
-            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
+        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
 
         output = sequence_model_parallel_all_to_all_4D(output,
                                                        scatter_dim=1,
@@ -175,12 +169,12 @@ class DistributedAttention_VSA(DistributedAttention):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
+        original_seq_len: int,
         replicated_q: torch.Tensor | None = None,
         replicated_k: torch.Tensor | None = None,
         replicated_v: torch.Tensor | None = None,
         gate_compress: torch.Tensor | None = None,
         freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
-        attention_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Forward pass for distributed attention.
         
@@ -188,11 +182,11 @@ class DistributedAttention_VSA(DistributedAttention):
             q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
             k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
             v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
+            original_seq_len (int): Original (unpadded) full sequence length
             gate_compress (torch.Tensor): Gate compress tensor [batch_size, seq_len, num_heads, head_dim]
             replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
             replicated_k (Optional[torch.Tensor]): Replicated key tensor
             replicated_v (Optional[torch.Tensor]): Replicated value tensor
-            attention_mask (Optional[torch.Tensor]): Attention mask [batch_size, seq_len]
             
         Returns:
             Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
@@ -221,11 +215,8 @@ class DistributedAttention_VSA(DistributedAttention):
                                                      gather_dim=1)
 
         # After all-to-all, each rank has the full sequence but only a subset of heads
-        # The attention mask should now apply to the full sequence length
-
-        if attention_mask is not None:
-            valid_seq_len = (attention_mask[0] == 1).sum().item()
-            qkvg = qkvg[:, :valid_seq_len, :, :]
+        pad_seq_len = qkvg.shape[1] - original_seq_len
+        qkvg = qkvg[:, :original_seq_len, :, :]
 
         if freqs_cis is not None:
             cos, sin = freqs_cis
@@ -246,9 +237,7 @@ class DistributedAttention_VSA(DistributedAttention):
         # Apply backend-specific postprocess_output
         output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
 
-        if attention_mask is not None:
-            pad_len = (attention_mask[0] == 0).sum().item()
-            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
+        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
 
         output = sequence_model_parallel_all_to_all_4D(output,
                                                        scatter_dim=1,
diff --git a/fastvideo/configs/models/encoders/qwen2_5.py b/fastvideo/configs/models/encoders/qwen2_5.py
index 3af6a6687..ea2d4f8f3 100644
--- a/fastvideo/configs/models/encoders/qwen2_5.py
+++ b/fastvideo/configs/models/encoders/qwen2_5.py
@@ -78,7 +78,6 @@ class Qwen2_5_VLArchConfig(TextEncoderArchConfig):
             "add_generation_prompt": True,
             "tokenize": True,
             "return_dict": True,
-            "padding": "max_length",
             "max_length": 1000 + 108,
             "truncation": True,
             "return_tensors": "pt",
diff --git a/fastvideo/configs/models/encoders/t5.py b/fastvideo/configs/models/encoders/t5.py
index 3fbe75cb0..bab60f3a3 100644
--- a/fastvideo/configs/models/encoders/t5.py
+++ b/fastvideo/configs/models/encoders/t5.py
@@ -63,7 +63,6 @@ class T5ArchConfig(TextEncoderArchConfig):
             self.dense_act_fn = "gelu_new"
 
         self.tokenizer_kwargs = {
-            "padding": "max_length",
             "truncation": True,
             "max_length": self.text_len,
             "add_special_tokens": True,
diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py
index cb7b64f86..36536c201 100644
--- a/fastvideo/configs/sample/base.py
+++ b/fastvideo/configs/sample/base.py
@@ -58,7 +58,6 @@ class SamplingParam:
 
     # Original dimensions (before VAE scaling)
     num_frames: int = 125
-    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
     height: int = 720
     width: int = 1280
     height_sr: int = 1072
diff --git a/fastvideo/configs/sample/hunyuan15.py b/fastvideo/configs/sample/hunyuan15.py
index d50adb186..ded72eaaa 100644
--- a/fastvideo/configs/sample/hunyuan15.py
+++ b/fastvideo/configs/sample/hunyuan15.py
@@ -15,8 +15,6 @@ class Hunyuan15_480P_SamplingParam(SamplingParam):
     fps: int = 24
 
     guidance_scale: float = 6.0
-    prompt_attention_mask: list = field(default_factory=list)
-    negative_attention_mask: list = field(default_factory=list)
     sigmas: list[float] | None = field(
         default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
 
diff --git a/fastvideo/distributed/utils.py b/fastvideo/distributed/utils.py
index 41c6f689b..d484e5245 100644
--- a/fastvideo/distributed/utils.py
+++ b/fastvideo/distributed/utils.py
@@ -81,43 +81,6 @@ def compute_padding_for_sp(seq_len: int, sp_world_size: int) -> tuple[int, int]:
     return padded_seq_len, padding_amount
 
 
-def create_attention_mask_for_padding(
-    seq_len: int,
-    padded_seq_len: int,
-    batch_size: int,
-    device: torch.device,
-    dtype: torch.dtype = torch.bool,
-) -> torch.Tensor | None:
-    """
-    Create attention mask to ignore padded tokens.
-    
-    Args:
-        seq_len: Original sequence length (before padding)
-        padded_seq_len: Padded sequence length
-        batch_size: Batch size
-        device: Device to create mask on
-        dtype: Data type for the mask (default: bool)
-        
-    Returns:
-        Tensor: Boolean mask [B, padded_seq_len] where True = valid token,
-                or None if no padding is needed
-    """
-    if seq_len == padded_seq_len:
-        return None
-
-    # Create mask: True for valid tokens, False for padding
-    attention_mask = torch.ones(
-        (batch_size, padded_seq_len),
-        dtype=dtype,
-        device=device,
-    )
-
-    # Mask out padding tokens
-    attention_mask[:, seq_len:] = 0
-
-    return attention_mask
-
-
 def pad_sequence_tensor(
     tensor: torch.Tensor,
     target_seq_len: int,
diff --git a/fastvideo/models/dits/hunyuanvideo.py b/fastvideo/models/dits/hunyuanvideo.py
index 7e16f88a0..1d25c1ad4 100644
--- a/fastvideo/models/dits/hunyuanvideo.py
+++ b/fastvideo/models/dits/hunyuanvideo.py
@@ -196,6 +196,7 @@ class MMDoubleStreamBlock(nn.Module):
         txt: torch.Tensor,
         vec: torch.Tensor,
         freqs_cis: tuple,
+        original_seq_len: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Process modulation vectors
         img_mod_outputs = self.img_mod(vec)
@@ -253,7 +254,7 @@ class MMDoubleStreamBlock(nn.Module):
         txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
 
         # Run distributed attention
-        img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)
+        img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)
         img_attn_out, _ = self.img_attn_proj(
             img_attn.view(batch_size, image_seq_len, -1))
         # Use fused operation for residual connection, normalization, and modulation
@@ -355,6 +356,7 @@ class MMSingleStreamBlock(nn.Module):
         vec: torch.Tensor,
         txt_len: int,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        original_seq_len: int | None = None,
     ) -> torch.Tensor:
         # Process modulation
         mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
@@ -386,7 +388,7 @@ class MMSingleStreamBlock(nn.Module):
 
 
         # Run distributed attention
-        img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, txt_q,
+        img_attn_output, txt_attn_output = self.attn(img_q, img_k, img_v, original_seq_len, txt_q,
                                                      txt_k, txt_v, freqs_cis = freqs_cis)
         attn_output = torch.cat((img_attn_output, txt_attn_output),
                                 dim=1).view(batch_size, seq_len, -1)
@@ -594,7 +596,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
 
         # Process through double stream blocks
         for index, block in enumerate(self.double_blocks):
-            double_block_args = [img, txt, vec, freqs_cis]
+            double_block_args = [img, txt, vec, freqs_cis, original_seq_len]
             img, txt = block(*double_block_args)
         # Merge txt and img to pass through single stream blocks
         x = torch.cat((img, txt), 1)
@@ -607,6 +609,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT):
                     vec,
                     txt_seq_len,
                     freqs_cis,
+                    original_seq_len,
                 ]
                 x = block(*single_block_args)
 
diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py
index 849c0fa95..b8e5562eb 100644
--- a/fastvideo/models/dits/hunyuanvideo15.py
+++ b/fastvideo/models/dits/hunyuanvideo15.py
@@ -16,9 +16,9 @@ from typing import Any, Dict, Optional, List
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from fastvideo.attention import DistributedAttention, LocalAttention
+from fastvideo.forward_context import set_forward_context
 from fastvideo.distributed.communication_op import (
     sequence_model_parallel_all_gather_with_unpad,
     sequence_model_parallel_shard)
@@ -34,11 +34,8 @@ from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.logger import init_logger
-from fastvideo.forward_context import set_forward_context
-from fastvideo.attention.backends.abstract import AttentionMetadata
 
 from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.distributed.utils import create_attention_mask_for_padding
 
 logger = init_logger(__name__)
 
@@ -290,10 +287,9 @@ class MMDoubleStreamBlock(nn.Module):
         self,
         img: torch.Tensor,
         txt: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
         vec: torch.Tensor,
         freqs_cis: tuple,
-        seq_attention_mask: torch.Tensor,
+        original_seq_len: int,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         # Process modulation vectors
         img_mod_outputs = self.img_mod(vec)
@@ -348,21 +344,8 @@ class MMDoubleStreamBlock(nn.Module):
         txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
         txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
 
-        # seq_len = txt_q.shape[1] + img_q.shape[1]
-        # attention_mask = F.pad(encoder_attention_mask, (seq_len - encoder_attention_mask.shape[1], 0), value=True)
-        # attention_mask = attention_mask.bool()
-        # self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
-        # self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
-        # attention_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
 
-        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
-        attn_metadata = FlashAttnMetadataBuilder().build(
-            current_timestep=0,
-            attn_mask=encoder_attention_mask,
-        )
-        # Run distributed attention
-        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
-            img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis, attention_mask=seq_attention_mask)
+        img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis)
         
         img_attn_out, _ = self.img_attn_proj(
             img_attn.view(batch_size, image_seq_len, -1))
@@ -469,15 +452,12 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
         encoder_hidden_states: List[torch.Tensor],
         timestep: torch.LongTensor,
         encoder_hidden_states_image: List[torch.Tensor],
-        encoder_attention_mask: List[torch.Tensor],
         guidance: Optional[torch.Tensor] = None,
         timestep_r: Optional[torch.LongTensor] = None,
         attention_kwargs: Optional[Dict[str, Any]] = None,
     ):
         encoder_hidden_states_image = encoder_hidden_states_image[0]
         encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
-        encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask
-
         batch_size, num_channels, num_frames, height, width = hidden_states.shape
         p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
         post_patch_num_frames = num_frames // p_t
@@ -503,18 +483,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
         sp_world_size = get_sp_world_size()
         padded_seq_len = current_seq_len * sp_world_size
         
-        if padded_seq_len > original_seq_len:
-            seq_attention_mask = create_attention_mask_for_padding(
-                seq_len=original_seq_len,
-                padded_seq_len=padded_seq_len,
-                batch_size=batch_size,
-                device=hidden_states.device,
-            )
-        else:
-            seq_attention_mask = None
+                
 
-        # qwen text embedding
-        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep)
 
         encoder_hidden_states_cond_emb = self.cond_type_embed(
             torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
@@ -532,19 +503,6 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
         # image embed
         encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
         is_t2v = torch.all(encoder_hidden_states_image == 0)
-        if is_t2v:
-            encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
-            encoder_attention_mask_3 = torch.zeros(
-                (batch_size, encoder_hidden_states_3.shape[1]),
-                dtype=encoder_attention_mask.dtype,
-                device=encoder_attention_mask.device,
-            )
-        else:
-            encoder_attention_mask_3 = torch.ones(
-                (batch_size, encoder_hidden_states_3.shape[1]),
-                dtype=encoder_attention_mask.dtype,
-                device=encoder_attention_mask.device,
-            )
         encoder_hidden_states_3_cond_emb = self.cond_type_embed(
             2
             * torch.ones_like(
@@ -554,52 +512,7 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
         )
         encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
 
-        # reorder and combine text tokens: combine valid tokens first, then padding
-        encoder_attention_mask = encoder_attention_mask.bool()
-        encoder_attention_mask_2 = encoder_attention_mask_2.bool()
-        encoder_attention_mask_3 = encoder_attention_mask_3.bool()
-        new_encoder_hidden_states = []
-        new_encoder_attention_mask = []
-
-        for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
-            encoder_hidden_states,
-            encoder_attention_mask,
-            encoder_hidden_states_2,
-            encoder_attention_mask_2,
-            encoder_hidden_states_3,
-            encoder_attention_mask_3,
-        ):
-            # Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm]
-            new_encoder_hidden_states.append(
-                torch.cat(
-                    [
-                        image[image_mask],  # valid image
-                        text_2[text_mask_2],  # valid byt5
-                        text[text_mask],  # valid mllm
-                        image[~image_mask],  # invalid image (zeroed)
-                        torch.zeros_like(text_2[~text_mask_2]),  # invalid byt5 (zeroed)
-                        torch.zeros_like(text[~text_mask]),  # invalid mllm (zeroed)
-                    ],
-                    dim=0,
-                )
-            )
-            # Apply same reordering to attention masks
-            new_encoder_attention_mask.append(
-                torch.cat(
-                    [
-                        image_mask[image_mask],
-                        text_mask_2[text_mask_2],
-                        text_mask[text_mask],
-                        image_mask[~image_mask],
-                        text_mask_2[~text_mask_2],
-                        text_mask[~text_mask],
-                    ],
-                    dim=0,
-                )
-            )
-
-        encoder_hidden_states = torch.stack(new_encoder_hidden_states)
-        encoder_attention_mask = torch.stack(new_encoder_attention_mask)
+        encoder_hidden_states = torch.cat([encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat([encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1)
 
         # 4. Transformer blocks
         if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -608,10 +521,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
                     block,
                     hidden_states,
                     encoder_hidden_states,
-                    encoder_attention_mask,
                     temb,
                     freqs_cis,
-                    seq_attention_mask
+                    original_seq_len
                 )
 
         else:
@@ -619,10 +531,9 @@ class HunyuanVideo15Transformer3DModel(BaseDiT):
                 hidden_states, encoder_hidden_states = block(
                     hidden_states,
                     encoder_hidden_states,
-                    encoder_attention_mask,
                     temb,
                     freqs_cis,
-                    seq_attention_mask
+                    original_seq_len
                 )
 
         # Final layer processing
@@ -691,7 +602,6 @@ class SingleTokenRefiner(nn.Module):
         # Get timestep embeddings
         timestep_aware_representations = self.t_embedder(t)
 
-        # Get context-aware representations
         original_dtype = x.dtype
         if mask is None:
             context_aware_representations = x.mean(dim=1)
@@ -706,7 +616,7 @@ class SingleTokenRefiner(nn.Module):
         x = self.input_embedder(x)
         # Process through refiner blocks
         for block in self.refiner_blocks:
-            x = block(x, c, mask)
+            x = block(x, c, mask=mask)
         return x
 
 class IndividualTokenRefinerBlock(nn.Module):
@@ -776,9 +686,6 @@ class IndividualTokenRefinerBlock(nn.Module):
         )
 
     def forward(self, x, c, mask=None):
-        if mask is not None:
-            mask = mask.clone().bool()
-            mask[:, 0] = True  # Prevent attention weights from becoming NaN
 
         # Get modulation parameters
         gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1)
@@ -789,18 +696,19 @@ class IndividualTokenRefinerBlock(nn.Module):
         batch_size, seq_len = qkv.shape[0], qkv.shape[1]
         qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1)
         q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
-
-        # Run scaled dot product attention
-        from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
-        attn_metadata = FlashAttnMetadataBuilder().build(
-            current_timestep=0,
-            attn_mask=mask,
-        )
-        # Run distributed attention
-        with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
-            attn_output = self.attn(q, k, v)  # [B, L, H, D]
+        if mask is None:
+            attn_output = self.attn(q, k, v) 
+        else:
+            from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder
+            attn_metadata = FlashAttnMetadataBuilder().build(
+                current_timestep=0,
+                attn_mask=mask,
+            )
+            # Run distributed attention
+            with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
+                attn_output = self.attn(q, k, v)  # [B, L, H, D]
         attn_output = attn_output.reshape(batch_size, seq_len,
-                                          -1)  # [B, L, H*D]
+                                          -1)
 
         # Project and apply residual connection with gating
         attn_out, _ = self.self_attn_proj(attn_output)
diff --git a/fastvideo/models/dits/hyworld/hyworld.py b/fastvideo/models/dits/hyworld/hyworld.py
index 81a7f6cc2..7197fdf5b 100644
--- a/fastvideo/models/dits/hyworld/hyworld.py
+++ b/fastvideo/models/dits/hyworld/hyworld.py
@@ -35,7 +35,6 @@ from fastvideo.platforms import AttentionBackendEnum
 from fastvideo.logger import init_logger
 from fastvideo.forward_context import set_forward_context
 from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.distributed.utils import create_attention_mask_for_padding
 
 from .camera_rope import prope_qkv
 
@@ -88,7 +87,7 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
         vec: torch.Tensor,
         vec_txt: torch.Tensor,
         freqs_cis: tuple,
-        seq_attention_mask: torch.Tensor,
+        original_seq_len: int,
         viewmats: torch.Tensor,
         Ks: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -101,7 +100,7 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
             encoder_attention_mask: Text attention mask
             vec: Modulation vector
             freqs_cis: Rotary embedding frequencies
-            seq_attention_mask: Sequence attention mask
+            original_seq_len: Original (unpadded) image sequence length
             viewmats: Camera view matrices for ProPE [B, T, 4, 4]
             Ks: Camera intrinsics for ProPE [B, T, 3, 3]
 
@@ -181,7 +180,16 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
         )
         # Run distributed attention
         with set_forward_context(current_timestep=0, attn_metadata=attn_metadata):
-            img_attn, txt_attn = self.attn(img_q, img_k, img_v, txt_q, txt_k, txt_v, freqs_cis=freqs_cis, attention_mask=seq_attention_mask)
+            img_attn, txt_attn = self.attn(
+                img_q,
+                img_k,
+                img_v,
+                original_seq_len,
+                txt_q,
+                txt_k,
+                txt_v,
+                freqs_cis=freqs_cis,
+            )
         
         # begin hyworld
         # attention with prope
@@ -193,8 +201,14 @@ class HYWorldDoubleStreamBlock(MMDoubleStreamBlock):
         # NOTE: Do NOT pass freqs_cis to prope attention - HY-WorldPlay does not apply RoPE to prope
         with set_forward_context(current_timestep=0, attn_metadata=attn_metadata_prope):
             img_attn_prope, _ = self.attn(
-                img_q_prope, img_k_prope, img_v_prope, txt_q, txt_k, txt_v, 
-                freqs_cis=None, attention_mask=seq_attention_mask  # No RoPE for prope attention
+                img_q_prope,
+                img_k_prope,
+                img_v_prope,
+                original_seq_len,
+                txt_q,
+                txt_k,
+                txt_v,
+                freqs_cis=None,  # No RoPE for prope attention
             )
             img_attn_prope = img_attn_prope.reshape(batch_size, image_seq_len, -1)
             img_attn_prope = rearrange(
@@ -399,20 +413,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
 
         hidden_states = self.img_in(hidden_states)
         hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
-
-        current_seq_len = hidden_states.shape[1]
         sp_world_size = get_sp_world_size()
-        padded_seq_len = current_seq_len * sp_world_size
-
-        if padded_seq_len > original_seq_len:
-            seq_attention_mask = create_attention_mask_for_padding(
-                seq_len=original_seq_len,
-                padded_seq_len=padded_seq_len,
-                batch_size=batch_size,
-                device=hidden_states.device,
-            )
-        else:
-            seq_attention_mask = None
 
         viewmats_seq = repeat(
             viewmats, "B T M N->B (T H W) M N",
@@ -535,7 +536,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
                     temb,
                     temb_txt,
                     freqs_cis,
-                    seq_attention_mask,
+                    original_seq_len,
                     viewmats_seq,
                     Ks_seq,
                 )
@@ -548,7 +549,7 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
                     temb,
                     temb_txt,
                     freqs_cis,
-                    seq_attention_mask,
+                    original_seq_len,
                     viewmats_seq,
                     Ks_seq,
                 )
diff --git a/fastvideo/models/dits/lingbotworld/model.py b/fastvideo/models/dits/lingbotworld/model.py
index 4cc84cb78..c15014015 100644
--- a/fastvideo/models/dits/lingbotworld/model.py
+++ b/fastvideo/models/dits/lingbotworld/model.py
@@ -30,7 +30,6 @@ from fastvideo.models.dits.wanvideo import (
 from fastvideo.platforms import AttentionBackendEnum, current_platform
 
 from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.distributed.utils import create_attention_mask_for_padding
 
 logger = init_logger(__name__)
 
@@ -145,7 +144,7 @@ class LingBotWorldTransformerBlock(nn.Module):
         encoder_hidden_states: torch.Tensor,
         temb: torch.Tensor,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: torch.Tensor | None = None,
+        original_seq_len: int,
         c2ws_plucker_emb: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if hidden_states.dim() == 4:
@@ -189,7 +188,13 @@ class LingBotWorldTransformerBlock(nn.Module):
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
 
-        attn_output, _ = self.attn1(query, key, value, freqs_cis=freqs_cis, attention_mask=attention_mask)
+        attn_output, _ = self.attn1(
+            query,
+            key,
+            value,
+            original_seq_len,
+            freqs_cis=freqs_cis,
+        )
         attn_output = attn_output.flatten(2)
         attn_output, _ = self.to_out(attn_output)
         attn_output = attn_output.squeeze(1)
@@ -292,7 +297,6 @@ class LingBotWorldTransformer3DModel(BaseDiT):
             torch.randn(1, 2, inner_dim) / inner_dim**0.5)
 
         self.gradient_checkpointing = False
-        self._logged_attention_mask = False
         self.__post_init__()
 
     def forward(self,
@@ -350,26 +354,6 @@ class LingBotWorldTransformer3DModel(BaseDiT):
         if c2ws_plucker_emb is not None:
             c2ws_plucker_emb, _ = sequence_model_parallel_shard(c2ws_plucker_emb, dim=1)
         
-        # Create attention mask for padded tokens if padding was applied
-        current_seq_len = hidden_states.shape[1]
-        sp_world_size = get_sp_world_size()
-        padded_seq_len = current_seq_len * sp_world_size
-        
-        if padded_seq_len > original_seq_len:
-            if not self._logged_attention_mask:
-                logger.info(f"Padding applied, original seq len: {original_seq_len}, padded seq len: {padded_seq_len}")
-                self._logged_attention_mask = True
-            attention_mask = create_attention_mask_for_padding(
-                seq_len=original_seq_len,
-                padded_seq_len=padded_seq_len,
-                batch_size=batch_size,
-                device=hidden_states.device,
-            )
-        else:
-            if not self._logged_attention_mask:
-                logger.info(f"Padding not applied")
-                self._logged_attention_mask = True
-            attention_mask = None
 
         # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
         if timestep.dim() == 2:
@@ -406,11 +390,11 @@ class LingBotWorldTransformer3DModel(BaseDiT):
             for block in self.blocks:
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, encoder_hidden_states,
-                    timestep_proj, freqs_cis, attention_mask, c2ws_plucker_emb)
+                    timestep_proj, freqs_cis, original_seq_len, c2ws_plucker_emb)
         else:
             for block in self.blocks:
                 hidden_states = block(hidden_states, encoder_hidden_states,
-                                      timestep_proj, freqs_cis, attention_mask,
+                                      timestep_proj, freqs_cis, original_seq_len,
                                       c2ws_plucker_emb)
         # 5. Output norm, projection & unpatchify
         if temb.dim() == 3:
diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py
index 55abc4fd5..09e687ead 100644
--- a/fastvideo/models/dits/ltx2.py
+++ b/fastvideo/models/dits/ltx2.py
@@ -27,7 +27,6 @@ from fastvideo.distributed.communication_op import (
     sequence_model_parallel_shard,
 )
 from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
-from fastvideo.distributed.utils import create_attention_mask_for_padding
 from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
@@ -1020,11 +1019,11 @@ class LTXDistributedAttention(DistributedAttention):
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
+        original_seq_len: int,
         replicated_q: torch.Tensor | None = None,
         replicated_k: torch.Tensor | None = None,
         replicated_v: torch.Tensor | None = None,
         gate_compress: torch.Tensor | None = None,
-        attention_mask: torch.Tensor | None = None,
         ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         """Forward pass with LTX-2 style RoPE application.
@@ -1036,14 +1035,14 @@ class LTXDistributedAttention(DistributedAttention):
             replicated_q: Replicated query tensor for text tokens
             replicated_k: Replicated key tensor
             replicated_v: Replicated value tensor
-            attention_mask: Attention mask [batch_size, seq_len]
+            original_seq_len: Original (unpadded) full sequence length
             ltx_freqs_cis: LTX-2 style RoPE (cos, sin) with shape [B, H, T, D]
 
         Returns:
             Tuple of output tensor and optional replicated output
         """
         assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
-        batch_size, seq_len, num_heads, head_dim = q.shape
+        batch_size, _, num_heads, _ = q.shape
         local_rank = get_sp_parallel_rank()
         world_size = get_sp_world_size()
 
@@ -1067,19 +1066,16 @@ class LTXDistributedAttention(DistributedAttention):
             qkvg = sequence_model_parallel_all_to_all_4D(qkvg,
                                                          scatter_dim=2,
                                                          gather_dim=1)
-
-            valid_seq_len = None
-            if attention_mask is not None:
-                valid_seq_len = (attention_mask[0] == 1).sum().item()
-                qkvg = qkvg[:, :valid_seq_len, :, :]
+            pad_seq_len = qkvg.shape[1] - original_seq_len
+            qkvg = qkvg[:, :original_seq_len, :, :]
 
             if ltx_freqs_cis is not None:
                 cos, sin = ltx_freqs_cis
                 heads_per_rank = num_heads // world_size
                 head_start = local_rank * heads_per_rank
                 head_end = head_start + heads_per_rank
-                cos_local = cos[:, head_start:head_end, :valid_seq_len, :]
-                sin_local = sin[:, head_start:head_end, :valid_seq_len, :]
+                cos_local = cos[:, head_start:head_end, :original_seq_len, :]
+                sin_local = sin[:, head_start:head_end, :original_seq_len, :]
                 qk_part = qkvg[:batch_size * 2].transpose(1, 2)
                 qk_part = apply_ltx_rotary_emb_4d(
                     qk_part, (cos_local, sin_local), self.rope_type)
@@ -1090,10 +1086,7 @@ class LTXDistributedAttention(DistributedAttention):
             output = self.attn_impl.forward(  # type: ignore[call-arg]
                 q, k, v, gate_compress, ctx_attn_metadata)
             output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
-
-            if attention_mask is not None:
-                pad_len = (attention_mask[0] == 0).sum().item()
-                output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
+            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
 
             output = sequence_model_parallel_all_to_all_4D(output,
                                                            scatter_dim=1,
@@ -1107,10 +1100,8 @@ class LTXDistributedAttention(DistributedAttention):
         qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
 
         # After all-to-all, each rank has the full sequence but only a subset of heads
-        valid_seq_len = None
-        if attention_mask is not None:
-            valid_seq_len = (attention_mask[0] == 1).sum().item()
-            qkv = qkv[:, :valid_seq_len, :, :]
+        pad_seq_len = qkv.shape[1] - original_seq_len
+        qkv = qkv[:, :original_seq_len, :, :]
 
         # Apply LTX-2 style RoPE after all-to-all (when we have full sequence)
         if ltx_freqs_cis is not None:
@@ -1119,8 +1110,8 @@ class LTXDistributedAttention(DistributedAttention):
             head_start = local_rank * heads_per_rank
             head_end = head_start + heads_per_rank
             # Slice to this rank's heads: [B, H/SP, T, D]
-            cos_local = cos[:, head_start:head_end, :valid_seq_len, :]
-            sin_local = sin[:, head_start:head_end, :valid_seq_len, :]
+            cos_local = cos[:, head_start:head_end, :original_seq_len, :]
+            sin_local = sin[:, head_start:head_end, :original_seq_len, :]
 
             # Apply RoPE to Q and K together (first 2*batch_size in dim 0)
             qk_part = qkv[:batch_size * 2]
@@ -1152,7 +1143,7 @@ class LTXDistributedAttention(DistributedAttention):
         # Redistribute back if using sequence parallelism
         replicated_output = None
         if replicated_q is not None:
-            split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len
+            split_idx = original_seq_len
             replicated_output = output[:, split_idx:]
             output = output[:, :split_idx]
             replicated_output = sequence_model_parallel_all_gather(
@@ -1161,9 +1152,7 @@ class LTXDistributedAttention(DistributedAttention):
         # Apply backend-specific postprocess_output
         output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
 
-        if attention_mask is not None:
-            pad_len = (attention_mask[0] == 0).sum().item()
-            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
+        output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_seq_len))
 
         output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)
 
@@ -1407,7 +1396,7 @@ class LTXDistributedSelfAttention(nn.Module):
         context: torch.Tensor | None = None,
         pe: tuple[torch.Tensor, torch.Tensor] | None = None,
         k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
-        attention_mask: torch.Tensor | None = None,
+        original_seq_len: int | None = None,
     ) -> torch.Tensor:
         """Forward pass for distributed self-attention.
 
@@ -1417,7 +1406,7 @@ class LTXDistributedSelfAttention(nn.Module):
             pe: Rotary position embeddings for Q (cos, sin) with shape [B, H, T_full, D]
                 NOTE: For SP, this must be the FULL sequence RoPE, not sharded!
             k_pe: Rotary position embeddings for K (cos, sin), or None to use pe
-            attention_mask: Attention mask for padding [B, padded_seq_len]
+            original_seq_len: Original (unpadded) full sequence length
         """
         q = self.to_q(x)
         context = x if context is None else context
@@ -1442,10 +1431,14 @@ class LTXDistributedSelfAttention(nn.Module):
 
         # Pass full RoPE to distributed attention - it will apply after all-to-all
         # and slice to this rank's heads
+        if original_seq_len is None:
+            raise ValueError("original_seq_len must be provided for distributed attention")
         out, _ = self.attn(
-            q, k, v,
+            q,
+            k,
+            v,
+            original_seq_len,
             gate_compress=gate_compress,
-            attention_mask=attention_mask,
             ltx_freqs_cis=pe,
         )
 
@@ -1633,8 +1626,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        video_attention_mask: torch.Tensor | None = None,
-        audio_attention_mask: torch.Tensor | None = None,
+        video_original_seq_len: int | None = None,
+        audio_original_seq_len: int | None = None,
         skip_cross_modal_attn: bool = False,
         skip_video_self_attn: bool = False,
         skip_audio_self_attn: bool = False,
@@ -1644,8 +1637,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
         Args:
             video: Video transformer args
             audio: Audio transformer args
-            video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
-            audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
+            video_original_seq_len: Original unpadded video sequence length
+            audio_original_seq_len: Original unpadded audio sequence length
             skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
                 attention (used for the modality-isolated CFG pass).
             skip_video_self_attn: If True, skip video self-attention
@@ -1667,9 +1660,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
             )
             if not skip_video_self_attn:
                 norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
-                # Self-attention: pass SP attention mask for distributed attention
                 if self.use_distributed_attention:
-                    vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings, attention_mask=video_attention_mask) * vgate_msa
+                    vx = vx + self.attn1(
+                        norm_vx,
+                        pe=video.positional_embeddings,
+                        original_seq_len=video_original_seq_len,
+                    ) * vgate_msa
                 else:
                     vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa
             # Text cross-attention: no SP mask needed (text is replicated, uses local attention)
@@ -1685,9 +1681,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
             )
             if not skip_audio_self_attn:
                 norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
-                # Self-attention: pass SP attention mask for distributed attention
                 if self.use_distributed_attention:
-                    ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings, attention_mask=audio_attention_mask) * agate_msa
+                    ax = ax + self.audio_attn1(
+                        norm_ax,
+                        pe=audio.positional_embeddings,
+                        original_seq_len=audio_original_seq_len,
+                    ) * agate_msa
                 else:
                     ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa
             # Text cross-attention: no SP mask needed (text is replicated, uses local attention)
@@ -2116,8 +2115,8 @@ class LTXModel(torch.nn.Module):
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        video_attention_mask: torch.Tensor | None = None,
-        audio_attention_mask: torch.Tensor | None = None,
+        video_original_seq_len: int | None = None,
+        audio_original_seq_len: int | None = None,
         skip_cross_modal_attn: bool = False,
         skip_video_self_attn_blocks: list[int] | None = None,
         skip_audio_self_attn_blocks: list[int] | None = None,
@@ -2132,8 +2131,8 @@ class LTXModel(torch.nn.Module):
             video, audio = block(
                 video=video,
                 audio=audio,
-                video_attention_mask=video_attention_mask,
-                audio_attention_mask=audio_attention_mask,
+                video_original_seq_len=video_original_seq_len,
+                audio_original_seq_len=audio_original_seq_len,
                 skip_cross_modal_attn=skip_cross_modal_attn,
                 skip_video_self_attn=skip_v_sa,
                 skip_audio_self_attn=skip_a_sa,
@@ -2161,8 +2160,8 @@ class LTXModel(torch.nn.Module):
         self,
         video: Modality | None,
         audio: Modality | None,
-        video_attention_mask: torch.Tensor | None = None,
-        audio_attention_mask: torch.Tensor | None = None,
+        video_original_seq_len: int | None = None,
+        audio_original_seq_len: int | None = None,
         skip_cross_modal_attn: bool = False,
         skip_video_self_attn_blocks: list[int] | None = None,
         skip_audio_self_attn_blocks: list[int] | None = None,
@@ -2172,8 +2171,8 @@ class LTXModel(torch.nn.Module):
         Args:
             video: Video modality input
             audio: Audio modality input
-            video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
-            audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
+            video_original_seq_len: Original unpadded video sequence length
+            audio_original_seq_len: Original unpadded audio sequence length
             skip_cross_modal_attn: If True, skip A2V and V2A cross-modal
                 attention in all transformer blocks (modality-isolated
                 CFG pass).
@@ -2202,8 +2201,8 @@ class LTXModel(torch.nn.Module):
         video_out, audio_out = self._process_transformer_blocks(
             video_args,
             audio_args,
-            video_attention_mask=video_attention_mask,
-            audio_attention_mask=audio_attention_mask,
+            video_original_seq_len=video_original_seq_len,
+            audio_original_seq_len=audio_original_seq_len,
             skip_cross_modal_attn=skip_cross_modal_attn,
             skip_video_self_attn_blocks=skip_video_self_attn_blocks,
             skip_audio_self_attn_blocks=skip_audio_self_attn_blocks,
@@ -2310,7 +2309,6 @@ class LTX2Transformer3DModel(BaseDiT):
         self.hidden_size = arch.num_attention_heads * arch.attention_head_dim
         self.num_attention_heads = arch.num_attention_heads
         self.num_channels_latents = arch.num_channels_latents
-        self._logged_attention_mask = False
 
         if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1":
             detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "")
@@ -2388,11 +2386,8 @@ class LTX2Transformer3DModel(BaseDiT):
     ) -> torch.Tensor:
         if isinstance(encoder_hidden_states, list):
             encoder_hidden_states = encoder_hidden_states[0]
-
         # Get SP parameters
         sp_world_size = get_sp_world_size()
-        sp_rank = get_sp_parallel_rank()
-        batch_size = hidden_states.shape[0]
 
         # Get fps for position computation
         fps = None
@@ -2413,29 +2408,14 @@ class LTX2Transformer3DModel(BaseDiT):
 
         # Shard video latents and timestep across SP ranks
         video_original_seq_len = latents.shape[1]
-        video_attention_mask = None
+        video_padded_seq_len = video_original_seq_len
         video_timestep = timestep
         if sp_world_size > 1:
             latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1)
             # Shard timestep along sequence dimension (timestep has shape [batch, seq_len])
             video_timestep, _ = sequence_model_parallel_shard(timestep, dim=1)
-            # Create attention mask for padded tokens
             current_seq_len = latents.shape[1]
-            padded_seq_len = current_seq_len * sp_world_size
-            if padded_seq_len > video_original_seq_len:
-                if not self._logged_attention_mask:
-                    logger.info(
-                        f"Video padding applied, original seq len: {video_original_seq_len}, "
-                        f"padded seq len: {padded_seq_len}"
-                    )
-                    self._logged_attention_mask = True
-                video_attention_mask = create_attention_mask_for_padding(
-                    seq_len=video_original_seq_len,
-                    padded_seq_len=padded_seq_len,
-                    batch_size=batch_size,
-                    device=hidden_states.device,
-                )
-
+            video_padded_seq_len = current_seq_len * sp_world_size
         # Compute RoPE positions for the FULL sequence (before sharding)
         # This is necessary because sequence sharding may not align to frame boundaries.
         # The full RoPE will be passed to DistributedAttention which applies it after
@@ -2451,8 +2431,8 @@ class LTX2Transformer3DModel(BaseDiT):
         ).to(hidden_states.dtype)
 
         # Pad positions to match padded sequence length for SP cross-attention
-        if sp_world_size > 1 and padded_seq_len > video_original_seq_len:
-            padding_needed = padded_seq_len - video_original_seq_len
+        if sp_world_size > 1 and video_padded_seq_len > video_original_seq_len:
+            padding_needed = video_padded_seq_len - video_original_seq_len
             last_pos = positions[:, :, -1:, :].expand(-1, -1, padding_needed, -1)
             positions = torch.cat([positions, last_pos], dim=2)
 
@@ -2482,7 +2462,6 @@ class LTX2Transformer3DModel(BaseDiT):
         audio_modality = None
         audio_shape = None
         audio_original_seq_len = 0
-        audio_attention_mask = None
 
         if audio_hidden_states is not None and audio_encoder_hidden_states is not None and audio_timestep is not None:
             audio_shape = AudioLatentShape.from_torch_shape(audio_hidden_states.shape)
@@ -2495,16 +2474,6 @@ class LTX2Transformer3DModel(BaseDiT):
                 audio_latents, audio_original_seq_len = sequence_model_parallel_shard(audio_latents, dim=1)
                 # Shard audio timestep along sequence dimension
                 sharded_audio_timestep, _ = sequence_model_parallel_shard(audio_timestep, dim=1)
-                # Create attention mask for padded tokens
-                audio_current_seq_len = audio_latents.shape[1]
-                audio_padded_seq_len = audio_current_seq_len * sp_world_size
-                if audio_padded_seq_len > audio_original_seq_len:
-                    audio_attention_mask = create_attention_mask_for_padding(
-                        seq_len=audio_original_seq_len,
-                        padded_seq_len=audio_padded_seq_len,
-                        batch_size=batch_size,
-                        device=audio_hidden_states.device,
-                    )
 
             # Compute audio RoPE positions for the FULL sequence (before sharding)
             # Same as video: full RoPE is applied after all-to-all in DistributedAttention
@@ -2534,12 +2503,12 @@ class LTX2Transformer3DModel(BaseDiT):
                     f"latent_checksum={audio_checksum:.6f}"
                 )
 
-        # Run transformer with attention masks
+        # Run transformer with original sequence lengths
         video_out, audio_out = self.model(
             video=video_modality,
             audio=audio_modality,
-            video_attention_mask=video_attention_mask,
-            audio_attention_mask=audio_attention_mask,
+            video_original_seq_len=video_original_seq_len,
+            audio_original_seq_len=audio_original_seq_len,
             skip_cross_modal_attn=skip_cross_modal_attn,
             skip_video_self_attn_blocks=skip_video_self_attn_blocks,
             skip_audio_self_attn_blocks=skip_audio_self_attn_blocks,
diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py
index 86aed7868..292026479 100644
--- a/fastvideo/models/dits/wanvideo.py
+++ b/fastvideo/models/dits/wanvideo.py
@@ -28,7 +28,6 @@ from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum, current_platform
 
 from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.distributed.utils import create_attention_mask_for_padding
 
 logger = init_logger(__name__)
 
@@ -324,7 +323,7 @@ class WanTransformerBlock(nn.Module):
         encoder_hidden_states: torch.Tensor,
         temb: torch.Tensor,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: torch.Tensor | None = None,
+        original_seq_len: int,
     ) -> torch.Tensor:
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
@@ -367,7 +366,13 @@ class WanTransformerBlock(nn.Module):
         key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
         value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
 
-        attn_output, _ = self.attn1(query, key, value, freqs_cis=freqs_cis, attention_mask=attention_mask)
+        attn_output, _ = self.attn1(
+            query,
+            key,
+            value,
+            original_seq_len,
+            freqs_cis=freqs_cis,
+        )
         attn_output = attn_output.flatten(2)
         attn_output, _ = self.to_out(attn_output)
         attn_output = attn_output.squeeze(1)
@@ -479,7 +484,7 @@ class WanTransformerBlock_VSA(nn.Module):
         encoder_hidden_states: torch.Tensor,
         temb: torch.Tensor,
         freqs_cis: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: torch.Tensor | None = None,
+        original_seq_len: int,
     ) -> torch.Tensor:
         if hidden_states.dim() == 4:
             hidden_states = hidden_states.squeeze(1)
@@ -510,12 +515,14 @@ class WanTransformerBlock_VSA(nn.Module):
         gate_compress = gate_compress.squeeze(1).unflatten(
             2, (self.num_attention_heads, -1))
 
-        attn_output, _ = self.attn1(query,
-                                    key,
-                                    value,
-                                    freqs_cis = freqs_cis,
-                                    gate_compress=gate_compress,
-                                    attention_mask=attention_mask)
+        attn_output, _ = self.attn1(
+            query,
+            key,
+            value,
+            original_seq_len,
+            freqs_cis=freqs_cis,
+            gate_compress=gate_compress,
+        )
         attn_output = attn_output.flatten(2)
         attn_output, _ = self.to_out(attn_output)
         attn_output = attn_output.squeeze(1)
@@ -610,7 +617,6 @@ class WanTransformer3DModel(BaseDiT):
             torch.randn(1, 2, inner_dim) / inner_dim**0.5)
 
         self.gradient_checkpointing = False
-        self._logged_attention_mask = False
         self.__post_init__()
 
     def forward(self,
@@ -656,26 +662,10 @@ class WanTransformer3DModel(BaseDiT):
         # Shard with padding support - returns (sharded_tensor, original_seq_len)
         hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
         
-        # Create attention mask for padded tokens if padding was applied
         current_seq_len = hidden_states.shape[1]
         sp_world_size = get_sp_world_size()
         padded_seq_len = current_seq_len * sp_world_size
         
-        if padded_seq_len > original_seq_len:
-            if not self._logged_attention_mask:
-                logger.info(f"Padding applied, original seq len: {original_seq_len}, padded seq len: {padded_seq_len}")
-                self._logged_attention_mask = True
-            attention_mask = create_attention_mask_for_padding(
-                seq_len=original_seq_len,
-                padded_seq_len=padded_seq_len,
-                batch_size=batch_size,
-                device=hidden_states.device,
-            )
-        else:
-            if not self._logged_attention_mask:
-                logger.info(f"Padding not applied")
-                self._logged_attention_mask = True
-            attention_mask = None
 
         # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
         if timestep.dim() == 2:
@@ -712,11 +702,11 @@ class WanTransformer3DModel(BaseDiT):
             for block in self.blocks:
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, encoder_hidden_states,
-                    timestep_proj, freqs_cis, attention_mask)
+                    timestep_proj, freqs_cis, original_seq_len)
         else:
             for block in self.blocks:
                 hidden_states = block(hidden_states, encoder_hidden_states,
-                                      timestep_proj, freqs_cis, attention_mask)
+                                      timestep_proj, freqs_cis, original_seq_len)
         # 5. Output norm, projection & unpatchify
         if temb.dim() == 3:
             # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py
index 768816cfa..ee9086e82 100644
--- a/fastvideo/pipelines/pipeline_batch_info.py
+++ b/fastvideo/pipelines/pipeline_batch_info.py
@@ -150,7 +150,6 @@ class ForwardBatch:
     height_latents: list[int] | int | None = None
     width_latents: list[int] | int | None = None
     num_frames: list[int] | int = 1  # Default for image models
-    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
 
     # Original dimensions (before VAE scaling)
     height: list[int] | int | None = None
diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py
index 22891e7bb..3a1fbcd0a 100644
--- a/fastvideo/pipelines/stages/base.py
+++ b/fastvideo/pipelines/stages/base.py
@@ -203,10 +203,3 @@ class PipelineStage(ABC):
             The updated batch information after this stage's processing.
         """
         raise NotImplementedError
-
-    def backward(
-        self,
-        batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> ForwardBatch:
-        raise NotImplementedError
diff --git a/fastvideo/pipelines/stages/text_encoding.py b/fastvideo/pipelines/stages/text_encoding.py
index 51f5244c4..6d3362caa 100644
--- a/fastvideo/pipelines/stages/text_encoding.py
+++ b/fastvideo/pipelines/stages/text_encoding.py
@@ -230,11 +230,12 @@ class TextEncodingStage(PipelineStage):
                 if processed_text is not None:
                     processed_texts.append(processed_text)
                 else:
-                    # Assuming batch_size = 1
-                    prompt_embeds = torch.zeros((1, tok_kwargs["max_length"],
-                                                 encoder_config.hidden_size),
-                                                device=target_device)
-                    attention_mask = torch.zeros((1, tok_kwargs["max_length"]),
+                    # Assuming batch_size = 1, special case for hunyuanvideo1.5 where there is no glyph text
+                    assert len(texts) == 1
+                    prompt_embeds = torch.zeros(
+                        (1, 0, encoder_config.hidden_size),
+                        device=target_device)
+                    attention_mask = torch.zeros((1, 0),
                                                  device=target_device,
                                                  dtype=torch.int64)
                     embeds_list.append(prompt_embeds)
@@ -274,13 +275,11 @@ class TextEncodingStage(PipelineStage):
                 if dtype is not None:
                     audio_embed = audio_embed.to(dtype=dtype)
                 audio_embeds_list.append(audio_embed)
-
             if dtype is not None:
                 prompt_embeds = prompt_embeds.to(dtype=dtype)
             embeds_list.append(prompt_embeds)
             if return_attention_mask:
                 attn_masks_list.append(attention_mask)
-
         self._last_audio_embeds = audio_embeds_list if is_ltx2 else None
         return self.return_embeds(embeds_list, attn_masks_list, return_type,
                                   return_attention_mask, indices)
diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py
index 35a421fc2..417e900ad 100644
--- a/fastvideo/training/training_pipeline.py
+++ b/fastvideo/training/training_pipeline.py
@@ -304,6 +304,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 # Make sure that the timesteps are the same across all sp processes.
                 sp_group = get_sp_group()
                 sp_group.broadcast(timesteps, src=0)
+                sp_group.broadcast(noise, src=0)
             sigmas = get_sigmas(
                 self.noise_scheduler,
                 latents.device,