diff --git a/README.md b/README.md
index dc349e7eb..18df501b3 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@ Dev in progress and highly experimental.
 
 
 ## Change Log
-- ```2025/02/14```: Release the inference code and kernel for Sliding Tile Attention(https://hao-ai-lab.github.io/blogs/sta/).
+- ```2025/02/14```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
 - ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
 - ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
 - ```2024/12/17```: `FastVideo` v1.0 is released.
@@ -41,11 +41,16 @@ To try sliding tile attention (optional), please follow the instruction in [csrc
 ## 🚀 Inference
 
 ### Inference HunyuanVideo with Sliding Tile Attention
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
-# Inference with FA2 full attn
-sh scripts/inference/inference_hunyuan_STA.sh
+First, download the model:
 ```
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
+```
+We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
+```
+scripts/inference/inference_hunyuan_STA.sh
+```
+### Video Demos using STA + Teacache
+Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on a single H100.
 
 ### Inference FastHunyuan on single RTX4090
 We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
diff --git a/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py b/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py
index 9a2aa3c10..215cabc78 100644
--- a/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py
+++ b/fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py
@@ -869,13 +869,20 @@ class HunyuanVideoPipeline(DiffusionPipeline):
         num_warmup_steps = len(
             timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
-
+        def dict_to_3d_list(mask_strategy, t_max=50, l_max=60, h_max=24):
+            result = [[[None for _ in range(h_max)] 
+                    for _ in range(l_max)] 
+                    for _ in range(t_max)]
+            if mask_strategy is None:
+                return result
+            for key, value in mask_strategy.items():
+                t, l, h = map(int, key.split('_'))
+                result[t][l][h] = value
+            return result
+        mask_strategy = dict_to_3d_list(mask_strategy)
         # if is_progress_bar:
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
-                mask_param = [
-                    mask_strategy, i
-                ]  # if mask_strategy is None, STA will not be used
                 if self.interrupt:
                     continue
 
@@ -912,7 +919,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                         encoder_hidden_states,
                         t_expand,
                         prompt_mask,
-                        mask_param=mask_param,
+                        mask_strategy=mask_strategy[i],
                         guidance=guidance_expand,
                         return_dict=False,
                     )[0]
diff --git a/fastvideo/models/hunyuan/inference.py b/fastvideo/models/hunyuan/inference.py
index df58b94a4..064c9efdc 100644
--- a/fastvideo/models/hunyuan/inference.py
+++ b/fastvideo/models/hunyuan/inference.py
@@ -100,6 +100,8 @@ class Inference(object):
         )
         model = model.to(device)
         model = Inference.load_state_dict(args, model, pretrained_model_path)
+        if args.enable_torch_compile:
+            model = torch.compile(model)
         model.eval()
 
         # ============================= Build extra models ========================
diff --git a/fastvideo/models/hunyuan/modules/attenion.py b/fastvideo/models/hunyuan/modules/attenion.py
index aae630f1d..d1e2f5843 100644
--- a/fastvideo/models/hunyuan/modules/attenion.py
+++ b/fastvideo/models/hunyuan/modules/attenion.py
@@ -76,7 +76,7 @@ def parallel_attention(q,
                        img_q_len,
                        img_kv_len,
                        text_mask,
-                       mask_param=None):
+                       mask_strategy=None):
     query, encoder_query = q
     key, encoder_key = k
     value, encoder_value = v
@@ -101,7 +101,7 @@ def parallel_attention(q,
     sequence_length = query.size(1)
     encoder_sequence_length = encoder_query.size(1)
 
-    if mask_param[0] is not None:
+    if mask_strategy is not None:
         query = torch.cat([tile(query, nccl_info.sp_size), encoder_query],
                           dim=1).transpose(1, 2)
         key = torch.cat([tile(key, nccl_info.sp_size), encoder_key],
@@ -110,9 +110,8 @@ def parallel_attention(q,
                           dim=1).transpose(1, 2)
 
         head_num = query.size(1)
-        mask_strategy, time_step, layer_idx = mask_param
         windows = [
-            mask_strategy[f"{time_step}_{layer_idx}_{head_idx}"]
+            mask_strategy[head_idx]
             for head_idx in range(head_num)
         ]
 
@@ -135,7 +134,7 @@ def parallel_attention(q,
     hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
         (sequence_length, encoder_sequence_length), dim=1)
 
-    if mask_param[0] is not None:
+    if mask_strategy is not None:
         hidden_states = untile(hidden_states, nccl_info.sp_size)
 
     if get_sequence_parallel_state():
diff --git a/fastvideo/models/hunyuan/modules/models.py b/fastvideo/models/hunyuan/modules/models.py
index cddb54b82..b4e741e2d 100644
--- a/fastvideo/models/hunyuan/modules/models.py
+++ b/fastvideo/models/hunyuan/modules/models.py
@@ -138,7 +138,7 @@ class MMDoubleStreamBlock(nn.Module):
         vec: torch.Tensor,
         freqs_cis: tuple = None,
         text_mask: torch.Tensor = None,
-        mask_param=None,
+        mask_strategy=None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         (
             img_mod1_shift,
@@ -215,7 +215,7 @@ class MMDoubleStreamBlock(nn.Module):
             img_q_len=img_q.shape[1],
             img_kv_len=img_k.shape[1],
             text_mask=text_mask,
-            mask_param=mask_param,
+            mask_strategy=mask_strategy,
         )
 
         # attention computation end
@@ -319,7 +319,7 @@ class MMSingleStreamBlock(nn.Module):
         txt_len: int,
         freqs_cis: Tuple[torch.Tensor, torch.Tensor] = None,
         text_mask: torch.Tensor = None,
-        mask_param=None,
+        mask_strategy=None,
     ) -> torch.Tensor:
         mod_shift, mod_scale, mod_gate = self.modulation(vec).chunk(3, dim=-1)
         x_mod = modulate(self.pre_norm(x), shift=mod_shift, scale=mod_scale)
@@ -365,7 +365,7 @@ class MMSingleStreamBlock(nn.Module):
             img_q_len=img_q.shape[1],
             img_kv_len=img_k.shape[1],
             text_mask=text_mask,
-            mask_param=mask_param,
+            mask_strategy=mask_strategy,
         )
 
         # attention computation end
@@ -597,7 +597,7 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         encoder_attention_mask: torch.Tensor,
-        mask_param=None,
+        mask_strategy=None,
         output_features=False,
         output_features_stride=8,
         attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -656,28 +656,24 @@ class HYVideoDiffusionTransformer(ModelMixin, ConfigMixin):
         # --------------------- Pass through DiT blocks ------------------------
 
         for index, block in enumerate(self.double_blocks):
-            mask_param.append(index)
             double_block_args = [
-                img, txt, vec, freqs_cis, text_mask, mask_param
+                img, txt, vec, freqs_cis, text_mask, mask_strategy[index]
             ]
             img, txt = block(*double_block_args)
-            mask_param = mask_param[:-1]
         # Merge txt and img to pass through single stream blocks.
         x = torch.cat((img, txt), 1)
         if output_features:
             features_list = []
         if len(self.single_blocks) > 0:
             for index, block in enumerate(self.single_blocks):
-                mask_param.append(index + len(self.double_blocks))
                 single_block_args = [
                     x,
                     vec,
                     txt_seq_len,
                     (freqs_cos, freqs_sin),
                     text_mask,
-                    mask_param,
+                    mask_strategy[index + len(self.double_blocks)],
                 ]
-                mask_param = mask_param[:-1]
                 x = block(*single_block_args)
                 if output_features and _ % output_features_stride == 0:
                     features_list.append(x[:, :img_seq_len, ...])
diff --git a/fastvideo/sample/sample_t2v_hunyuan_STA.py b/fastvideo/sample/sample_t2v_hunyuan_STA.py
index 67c8a3f42..7f218f078 100644
--- a/fastvideo/sample/sample_t2v_hunyuan_STA.py
+++ b/fastvideo/sample/sample_t2v_hunyuan_STA.py
@@ -23,7 +23,7 @@ def teacache_forward(
     encoder_hidden_states: torch.Tensor,
     timestep: torch.LongTensor,
     encoder_attention_mask: torch.Tensor,
-    mask_param=None,
+    mask_strategy=None,
     output_features=False,
     output_features_stride=8,
     attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -122,12 +122,10 @@ def teacache_forward(
             ori_img = img.clone()
             # --------------------- Pass through DiT blocks ------------------------
             for index, block in enumerate(self.double_blocks):
-                mask_param.append(index)
                 double_block_args = [
-                    img, txt, vec, freqs_cis, text_mask, mask_param
+                    img, txt, vec, freqs_cis, text_mask, mask_strategy[index]
                 ]
                 img, txt = block(*double_block_args)
-                mask_param = mask_param[:-1]
 
             # Merge txt and img to pass through single stream blocks.
             x = torch.cat((img, txt), 1)
@@ -135,16 +133,14 @@ def teacache_forward(
                 features_list = []
             if len(self.single_blocks) > 0:
                 for index, block in enumerate(self.single_blocks):
-                    mask_param.append(index + len(self.double_blocks))
                     single_block_args = [
                         x,
                         vec,
                         txt_seq_len,
                         (freqs_cos, freqs_sin),
                         text_mask,
-                        mask_param,
+                        mask_strategy[index + len(self.double_blocks)],
                     ]
-                    mask_param = mask_param[:-1]
                     x = block(*single_block_args)
                     if output_features and _ % output_features_stride == 0:
                         features_list.append(x[:, :img_seq_len, ...])
@@ -154,28 +150,24 @@ def teacache_forward(
     else:
         # --------------------- Pass through DiT blocks ------------------------
         for index, block in enumerate(self.double_blocks):
-            mask_param.append(index)
             double_block_args = [
-                img, txt, vec, freqs_cis, text_mask, mask_param
+                img, txt, vec, freqs_cis, text_mask, mask_strategy[index]
             ]
             img, txt = block(*double_block_args)
-            mask_param = mask_param[:-1]
         # Merge txt and img to pass through single stream blocks.
         x = torch.cat((img, txt), 1)
         if output_features:
             features_list = []
         if len(self.single_blocks) > 0:
             for index, block in enumerate(self.single_blocks):
-                mask_param.append(index + len(self.double_blocks))
                 single_block_args = [
                     x,
                     vec,
                     txt_seq_len,
                     (freqs_cos, freqs_sin),
                     text_mask,
-                    mask_param,
+                    mask_strategy[index + len(self.double_blocks)],
                 ]
-                mask_param = mask_param[:-1]
                 x = block(*single_block_args)
                 if output_features and _ % output_features_stride == 0:
                     features_list.append(x[:, :img_seq_len, ...])
@@ -214,7 +206,8 @@ def main(args):
     models_root_path = Path(args.model_path)
     if not models_root_path.exists():
         raise ValueError(f"`models_root` not exists: {models_root_path}")
-
+    if args.enable_teacache and args.enable_torch_compile:
+        parser.error("--enable_teacache and --enable_torch_compile cannot be used simultaneously. Please enable only one of these options.")
     # Create save folder to save the samples
     save_path = args.output_path
     os.makedirs(os.path.dirname(save_path), exist_ok=True)
@@ -450,6 +443,11 @@ if __name__ == "__main__":
         action="store_true",
         help="Use teacache for speeding up inference",
     )
+    parser.add_argument(
+        "--enable_torch_compile",
+        action="store_true",
+        help="Use torch.compile for speeding up STA inference without teacache",
+    )
     parser.add_argument("--mask_strategy_file_path",
                         type=str,
                         default="assets/mask_strategy.json")
diff --git a/scripts/inference/inference_hunyuan_STA.sh b/scripts/inference/inference_hunyuan_STA.sh
index 8a3512403..bc6793ae1 100644
--- a/scripts/inference/inference_hunyuan_STA.sh
+++ b/scripts/inference/inference_hunyuan_STA.sh
@@ -1,9 +1,10 @@
 #!/bin/bash
 
+# Inference with STA + Teacache
 num_gpus=1
 mask_strategy_file_path=assets/mask_strategy.json
 export MODEL_BASE=data/hunyuan
-rel_l1_thresh=0.2
+rel_l1_thresh=0.15
 CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
     fastvideo/sample/sample_t2v_hunyuan_STA.py \
     --height 768 \
@@ -22,4 +23,27 @@ CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_p
     --rel_l1_thresh $rel_l1_thresh \
     --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
     --vae-sp \
-    --enable_teacache 
\ No newline at end of file
+    --enable_teacache 
+
+# Inference with STA only
+num_gpus=1
+mask_strategy_file_path=assets/mask_strategy.json
+export MODEL_BASE=data/hunyuan
+CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29603 \
+    fastvideo/sample/sample_t2v_hunyuan_STA.py \
+    --height 768 \
+    --width 1280 \
+    --num_frames 117 \
+    --num_inference_steps 50 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 7 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 12345 \
+    --output_path outputs_video/hunyuan_STA/ \
+    --model_path $MODEL_BASE \
+    --mask_strategy_file_path $mask_strategy_file_path \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
+    --vae-sp \
+    --enable_torch_compile
\ No newline at end of file