Add torch.compile for all small ops (#432)

Author: Edenzzzz

fastvideo/v1/layers/activation.py

@@ -25,6 +25,7 @@ class SiluAndMul(CustomOp):
     def __init__(self) -> None:
         super().__init__()
 
+    @torch.compile(dynamic=True)
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         d = x.shape[-1] // 2
@@ -48,6 +49,7 @@ def __init__(self, approximate: str = "none"):
         if approximate not in ("none", "tanh"):
             raise ValueError(f"Unknown approximate mode: {approximate}")
 
+    @torch.compile(dynamic=True)
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         d = x.shape[-1] // 2
@@ -63,6 +65,7 @@ class NewGELU(CustomOp):
     def __init__(self):
         super().__init__()
 
+    @torch.compile(dynamic=True)
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         c = math.sqrt(2.0 / math.pi)
@@ -76,6 +79,7 @@ class QuickGELU(CustomOp):
     def __init__(self):
         super().__init__()
 
+    @torch.compile(dynamic=True)
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
         return x * torch.sigmoid(1.702 * x)

fastvideo/v1/layers/layernorm.py

@@ -37,6 +37,7 @@ def __init__(
         if self.has_weight:
             self.weight = nn.Parameter(self.weight)
 
+    @torch.compile(dynamic=True)
     def forward_native(
         self,
         x: torch.Tensor,
@@ -89,6 +90,7 @@ class ScaleResidual(nn.Module):
     def __init__(self, prefix: str = ""):
         super().__init__()
 
+    @torch.compile(dynamic=True)
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
                 gate: torch.Tensor) -> torch.Tensor:
         """Apply gated residual connection."""
@@ -128,6 +130,7 @@ def __init__(
         else:
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
+    @torch.compile(dynamic=True)
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
                 gate: torch.Tensor, shift: torch.Tensor,
                 scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -178,6 +181,7 @@ def __init__(
         else:
             raise NotImplementedError(f"Norm type {norm_type} not implemented")
 
+    @torch.compile(dynamic=True)
     def forward(self, x: torch.Tensor, shift: torch.Tensor,
                 scale: torch.Tensor) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""

fastvideo/v1/layers/mlp.py

@@ -39,6 +39,7 @@ def __init__(
                                        bias=bias,
                                        params_dtype=dtype)
 
+    @torch.compile(dynamic=True)
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x, _ = self.fc_in(x)
         x = self.act(x)

fastvideo/v1/layers/rotary_embedding.py

@@ -44,6 +44,7 @@ def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
     return x.flatten(-2)
 
 
+# @torch.compile(dynamic=True)
 def _apply_rotary_emb(
     x: torch.Tensor,
     cos: torch.Tensor,

fastvideo/v1/layers/visual_embedding.py

@@ -53,6 +53,7 @@ def __init__(self,
                               dtype=dtype)
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
+    @torch.compile(dynamic=True)
     def forward(self, x):
         x = self.proj(x)
         if self.flatten:
@@ -98,6 +99,7 @@ def forward(self, t: torch.Tensor) -> torch.Tensor:
         return t_emb
 
 
+@torch.compile(dynamic=True)
 def timestep_embedding(t: torch.Tensor,
                        dim: int,
                        max_period: int = 10000,
@@ -145,6 +147,7 @@ def __init__(
                                        params_dtype=dtype)
         self.act = get_act_fn(act_layer)
 
+    @torch.compile(dynamic=True)
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.act(x)
         x, _ = self.linear(x)