[bugfix]: bound WanCausalModel num_frames_per_block override to <= 3

The constructor override only checked >= 1, bypassing the <= 3 limit the
config path enforces in CausalWanTransformer3DModel. Apply the same upper
bound with a clear error message.

docker/Dockerfile

@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
 # flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
 # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
-ARG FA4_CUTE_REF=940cd9680f3315f2f06b43ab5bea2c2cf2d96806
+ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
 
 # Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
 # select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
@@ -161,7 +161,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
         uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
     fi
 
-# Overlay the cutlass-4.5-safe upstream FA4 cute (FA4_CUTE_REF) over the
+# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
 # wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
 # runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
 # the same deps the [dreamverse] extra already installs in CI; the installed torch

fastvideo/distributed/parallel_state.py

@@ -23,6 +23,7 @@ If you only need to use the distributed environment without model parallelism,
  you can skip the model parallel initialization and destruction steps.
 """
 import contextlib
+import gc
 import os
 import pickle
 import weakref
@@ -1006,6 +1007,13 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
     if shutdown_ray:
         import ray  # Lazy import Ray
         ray.shutdown()
+    # Actually free GPU memory, as the name promises. FSDP-wrapped modules sit
+    # in reference cycles (module <-> FSDP state <-> hooks), so dropping the
+    # last user reference does not free their parameters until a gc pass runs.
+    # Without this, back-to-back model-load tests in one pytest session OOM.
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
 
 
 def get_same_node_ranks(pg: ProcessGroup | StatelessProcessGroup, source_rank: int = 0) -> list[int]:

fastvideo/tests/api/test_compat_translation.py

@@ -17,46 +17,6 @@ class TestLegacyTorchCompileKwargsTranslation:
     first-class :class:`CompileConfig` fields and anything unknown falls
     into ``extras``."""
 
-    def test_all_typed_keys_promoted(self) -> None:
-        config = legacy_from_pretrained_to_config(
-            "/models/ltx2",
-            {
-                "enable_torch_compile": True,
-                "torch_compile_kwargs": {
-                    "backend": "inductor",
-                    "fullgraph": True,
-                    "mode": "max-autotune-no-cudagraphs",
-                    "dynamic": False,
-                },
-            },
-        )
-        compile_config = config.engine.compile
-        assert compile_config.enabled is True
-        assert compile_config.backend == "inductor"
-        assert compile_config.fullgraph is True
-        assert compile_config.mode == "max-autotune-no-cudagraphs"
-        assert compile_config.dynamic is False
-        assert compile_config.extras == {}
-
-    def test_unknown_keys_land_in_extras(self) -> None:
-        config = legacy_from_pretrained_to_config(
-            "/models/ltx2",
-            {
-                "enable_torch_compile": True,
-                "torch_compile_kwargs": {
-                    "backend": "inductor",
-                    "options": {"triton.cudagraphs": False},
-                    "disable": False,
-                },
-            },
-        )
-        compile_config = config.engine.compile
-        assert compile_config.backend == "inductor"
-        assert compile_config.extras == {
-            "options": {"triton.cudagraphs": False},
-            "disable": False,
-        }
-
     def test_empty_kwargs_produces_empty_extras(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",
@@ -117,13 +77,6 @@ class TestLegacyLtx2VaeTilingTranslation:
     ``generator.pipeline.vae_tiling``; reverse direction emits the
     legacy name back to FastVideoArgs."""
 
-    def test_forward_routes_to_pipeline_vae_tiling(self) -> None:
-        config = legacy_from_pretrained_to_config(
-            "/models/ltx2",
-            {"ltx2_vae_tiling": False},
-        )
-        assert config.pipeline.vae_tiling is False
-
     def test_true_round_trips(self) -> None:
         config = legacy_from_pretrained_to_config(
             "/models/ltx2",

fastvideo/tests/contract/test_ci_test_collection.py

@@ -35,7 +35,6 @@ ALLOWLIST = {
     "hooks": "no lane yet — run manually",
     "layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
     "nightly": "by design: nightly cadence, not per-PR",
-    "ops": "no lane yet — GPU op tests, run manually",
     "modal": "CI infrastructure itself, not a test suite",
 }
 

fastvideo/tests/modal/pr_test.py

@@ -187,6 +187,8 @@ def run_transformer_tests():
 
 
 @app.function(gpu="L40S:4",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -201,6 +203,8 @@ def run_training_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=900,
               secrets=[
@@ -276,7 +280,7 @@ def run_self_forcing_tests():
 @app.function(gpu="L40S:1", image=image, timeout=900)
 def run_unit_test():
     run_test(
-        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
+        "pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ ./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ ./fastvideo/tests/stages/ ./fastvideo/tests/ops/ --ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py --ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs"
     )
 
 
@@ -318,6 +322,8 @@ def run_dreamverse_app_tests():
 
 
 @app.function(gpu="L40S:1",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[
@@ -400,6 +406,8 @@ def run_lora_extraction_tests():
 
 
 @app.function(gpu="L40S:2",
+              cpu=8.0,
+              memory=32768,
               image=image,
               timeout=1800,
               secrets=[

fastvideo/tests/ops/quantization/test_absmax_fp8.py

@@ -34,6 +34,9 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
             method._convert_scale(scale)
 
     def test_create_weights_rejects_invalid_dtype(self):
+        # float64 is outside the supported set (bf16/fp16/fp32). The test
+        # originally passed float32, which create_weights has accepted since
+        # PR #981 — it never ran in CI, so the mismatch went unnoticed.
         method = AbsMaxFP8LinearMethod()
         layer = nn.Module()
         with self.assertRaisesRegex(AssertionError, "only supports"):
@@ -43,7 +46,7 @@ class TestAbsMaxFP8LinearMethod(unittest.TestCase):
                 output_partition_sizes=[3],
                 input_size=2,
                 output_size=3,
-                params_dtype=torch.float32,
+                params_dtype=torch.float64,
             )
 
     def test_absmax_fp8_parameter_weight_loader(self):

fastvideo/tests/transformers/test_hunyuanvideo.py

@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import glob
-import json
-import os
-import pytest
-import torch
-
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.utils import maybe_download_model
-from fastvideo.forward_context import set_forward_context
-from fastvideo.configs.models.dits import HunyuanVideoConfig
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
-
-LOCAL_RANK = 0
-RANK = 0
-WORLD_SIZE = 1
-
-# Latent generated on commit c021e8a27cf437ac22827f2bc58b7f006561317f with 1 x L40S
-REFERENCE_LATENT = 93.23263549804688
-
-
-@pytest.mark.usefixtures("distributed_setup")
-@pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
-def test_hunyuanvideo_distributed():
-    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
-
-    torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
-
-    # Get tensor parallel info
-    sp_rank = get_sp_parallel_rank()
-    sp_world_size = get_sp_world_size()
-
-    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
-
-    config = json.load(open(CONFIG_PATH))
-    # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
-    # TODO: write normalize config function
-    config.pop("_class_name")
-    config.pop("_diffusers_version")
-
-    precision_str = "bf16"
-    args = FastVideoArgs(
-        model_path=TRANSFORMER_PATH,
-        dit_cpu_offload=True,
-        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
-    )
-    args.device = torch.device(f"cuda:{LOCAL_RANK}")
-
-    loader = TransformerLoader()
-    model = loader.load(TRANSFORMER_PATH, args)
-
-    model.eval()
-
-    # Create random inputs for testing
-    batch_size = 1
-    seq_len = 3
-    device = torch.device(f"cuda:{LOCAL_RANK}")
-
-    # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
-    chunk_per_rank = hidden_states.shape[2] // sp_world_size
-    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
-
-    # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
-
-    # Timestep
-    timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
-    forward_batch = ForwardBatch(
-        data_type="dummy",
-    )
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Run inference on model
-        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16):
-            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
-                output = model(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    timestep=timestep,
-                )
-
-    latent = output.double().sum().item()
-
-    # Check if latents are similar
-    diff_output_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"

fastvideo/tests/transformers/test_matrixgame2.py

@@ -1,133 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import pytest
-import torch
-
-from fastvideo.configs.models.dits.matrixgame2 import MatrixGame2WanVideoConfig
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.distributed.parallel_state import (
-    get_sp_parallel_rank,
-    get_sp_world_size,
-)
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import TransformerLoader
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29505"
-os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
-
-MODEL_PATH = maybe_download_model("FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers")
-TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
-REFERENCE_LATENT = -111835.58895772696
-
-
-@pytest.mark.skip(reason="Not reliably reproducible")
-@pytest.mark.usefixtures("distributed_setup")
-def test_matrixgame2_transformer():
-    transformer_path = TRANSFORMER_PATH
-
-    sp_rank = get_sp_parallel_rank()
-    sp_world_size = get_sp_world_size()
-
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-
-    args = FastVideoArgs(
-        model_path=transformer_path,
-        dit_cpu_offload=False,
-        use_fsdp_inference=False,
-        pipeline_config=PipelineConfig(
-            dit_config=MatrixGame2WanVideoConfig(), dit_precision=precision_str
-        ),
-    )
-    args.device = device
-
-    loader = TransformerLoader()
-    model = loader.load(transformer_path, args).to(device, dtype=precision)
-    model.eval()
-
-    total_params = sum(p.numel() for p in model.parameters())
-    weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters())
-    weight_mean = weight_sum / total_params
-    logger.info("Total parameters: %s", total_params)
-    logger.info("Weight sum: %s", weight_sum)
-    logger.info("Weight mean: %s", weight_mean)
-
-    torch.manual_seed(42)
-
-    batch_size = 1
-    latent_frames = 15
-    latent_height = 44
-    latent_width = 80
-    image_dim = 1280
-    image_seq_len = 257
-
-    # Set num_frame_per_block to match actual latent_frames for action_model
-    model.num_frame_per_block = latent_frames
-    # Force full attention mode
-    model.local_attn_size = -1
-    # Clear cached block masks so they regenerate with correct settings
-    model.block_mask = None
-    model.block_mask_mouse = None
-    model.block_mask_keyboard = None
-
-    # Video latents [B, C, T, H, W]
-    x = torch.randn(batch_size, 16, latent_frames, latent_height, latent_width,
-                    device=device, dtype=precision)
-    cond_concat = torch.randn(batch_size, 20, latent_frames, latent_height, latent_width,
-                              device=device, dtype=precision)
-    hidden_states = torch.cat([x, cond_concat], dim=1)
-
-    if sp_world_size > 1:
-        chunk_per_rank = hidden_states.shape[2] // sp_world_size
-        hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank]
-
-    # Image embeddings
-    encoder_hidden_states_image = [
-        torch.randn(batch_size, image_seq_len, image_dim, device=device, dtype=precision)
-    ]
-
-    # Timestep
-    timestep = torch.full((batch_size, latent_frames), 500, device=device, dtype=precision)
-
-    # Action conditions
-    num_pixel_frames = (latent_frames - 1) * 4 + 1
-    mouse_cond = torch.randn(batch_size, num_pixel_frames, 2, device=device, dtype=precision)
-    keyboard_cond = torch.randn(batch_size, num_pixel_frames, 4, device=device, dtype=precision)
-
-    # Text embeddings (empty for i2v)
-    encoder_hidden_states = torch.zeros(batch_size, 0, model.config.hidden_size,
-                                        device=device, dtype=precision)
-
-    forward_batch = ForwardBatch(data_type="dummy")
-    model.num_frame_per_block = latent_frames
-
-    with torch.no_grad():
-        with torch.amp.autocast('cuda', dtype=precision):
-            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
-                output = model(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_hidden_states_image=encoder_hidden_states_image,
-                    timestep=timestep,
-                    mouse_cond=mouse_cond,
-                    keyboard_cond=keyboard_cond,
-                )
-
-    latent = output.double().sum().item()
-
-    diff = abs(REFERENCE_LATENT - latent)
-    relative_diff = diff / abs(REFERENCE_LATENT)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%")
-
-    # The 0.19% diff now
-    assert relative_diff < 0.005, f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)"

fastvideo/tests/vaes/test_hunyuan_vae.py

@@ -1,82 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import json
-import os
-
-import pytest
-import torch
-from safetensors.torch import load_file
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.logger import init_logger
-
-# from fastvideo.models.vaes.hunyuanvae import (
-#     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import HunyuanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
-
-# Latent generated on commit d71a4ebffc2034922fc379568b6a6aa722f3744c with 1 x A40
-# torch 2.7.1
-A40_REFERENCE_LATENT = -106.22467041015625
-
-# Latent generated on commit 2b54068960c41d42221e8b8719a374b499855029 with 1 x L40S
-L40S_REFERENCE_LATENT = -157.09130859375
-
-
-@pytest.mark.usefixtures("distributed_setup")
-@pytest.mark.skip(reason="Skipping test_hunyuan_vae")
-def test_hunyuan_vae():
-    device = torch.device("cuda:0")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(
-        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
-    )
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model = loader.load(VAE_PATH, args)
-
-    model.enable_tiling(
-        tile_sample_min_height=32,
-        tile_sample_min_width=32,
-        tile_sample_min_num_frames=8,
-        tile_sample_stride_height=16,
-        tile_sample_stride_width=16,
-        tile_sample_stride_num_frames=4,
-    )
-
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        latent = model.encode(input_tensor).mean.double().sum().item()
-
-    # Check if latents are similar
-    device_name = torch.cuda.get_device_name()
-    if "A40" in device_name:
-        REFERENCE_LATENT = A40_REFERENCE_LATENT
-    elif "L40S" in device_name:
-        REFERENCE_LATENT = L40S_REFERENCE_LATENT
-    else:
-        raise ValueError(f"Unknown device: {device_name}")
-
-    diff_encoded_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
-    assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"

fastvideo/tests/vaes/test_official_wan_vae.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-
-import numpy as np
-import pytest
-import torch
-from diffusers import AutoencoderKLWan
-
-from fastvideo.configs.pipelines import PipelineConfig
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.loader.component_loader import VAELoader
-from fastvideo.configs.models.vaes import WanVAEConfig
-from fastvideo.utils import maybe_download_model
-
-try:
-    from wan.modules.vae2_2 import Wan2_2_VAE
-except ImportError:
-    Wan2_2_VAE = None
-
-logger = init_logger(__name__)
-
-os.environ["MASTER_ADDR"] = "localhost"
-os.environ["MASTER_PORT"] = "29503"
-
-BASE_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH)
-VAE_PATH = os.path.join(MODEL_PATH, "vae")
-
-
-@pytest.mark.skip(reason="disable test")
-@pytest.mark.usefixtures("distributed_setup")
-def test_official_wan_vae():
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    precision = torch.bfloat16
-    precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
-    args.device = device
-    args.vae_cpu_offload = False
-
-    loader = VAELoader()
-    model2 = loader.load(VAE_PATH, args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
-
-    model1 = Wan2_2_VAE(vae_pth="/mnt/weka/home/hao.zhang/wei/Wan2.2_VAE.pth", device=device, dtype=precision)
-
-    # Create identical inputs for both models
-    batch_size = 1
-
-    # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               81,
-                               32,
-                               32,
-                               device=device,
-                               dtype=precision)
-    # latent_tensor = torch.randn(batch_size,
-    #                             16,
-    #                             21,
-    #                             32,
-    #                             32,
-    #                             device=device,
-    #                             dtype=precision)
-
-    # Disable gradients for inference
-    with torch.no_grad():
-        # Test encoding
-        logger.info("Testing encoding...")
-        latent1 = model1.encode([input_tensor.squeeze(0)])[0].unsqueeze(0)
-        print("--------------------------------")
-        latent2 = model2.encode(input_tensor)
-        # Check if latents have the same shape
-        assert latent1.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
-        # Check if latents are similar
-        max_diff_encode = torch.max(torch.abs(latent1 - latent2.mean))
-        mean_diff_encode = torch.mean(torch.abs(latent1 - latent2.mean))
-        logger.info("Maximum difference between encoded latents: %s",
-                    max_diff_encode.item())
-        logger.info("Mean difference between encoded latents: %s",
-                    mean_diff_encode.item())
-        assert max_diff_encode < 1e-5, f"Encoded latents differ significantly: max diff = {mean_diff_encode.item()}"
-        # Test decoding
-        logger.info("Testing decoding...")
-        # latent1_tensor = latent1.mode()
-        # mean1 = (torch.tensor(model1.config.latents_mean).view(
-        #     1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
-        #                                         input_tensor.dtype))
-        # std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
-        #     1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
-        #                                         input_tensor.dtype)
-        # latent1_tensor = latent1_tensor / std1 + mean1
-        # output1 = model1.decode(latent1_tensor).sample
-
-        output1 = model1.decode([latent1.squeeze(0)])
-        output1 = output1[0].unsqueeze(0)
-
-        mean2 = model2.config.arch_config.shift_factor.to(input_tensor.device, input_tensor.dtype)
-        std2 = model2.config.arch_config.scaling_factor.to(input_tensor.device, input_tensor.dtype)
-        latent2_tensor = latent2.mode()
-        latent2_tensor = latent2_tensor / std2 + mean2
-        output2 = model2.decode(latent2_tensor)
-        # Check if outputs have the same shape
-        assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"
-
-        # Check if outputs are similar
-        max_diff_decode = torch.max(torch.abs(output1.float() - output2.float()))
-        mean_diff_decode = torch.mean(torch.abs(output1.float() - output2.float()))
-        logger.info("Maximum difference between decoded outputs: %s",
-                    max_diff_decode.item())
-        logger.info("Mean difference between decoded outputs: %s",
-                    mean_diff_decode.item())
-        assert max_diff_decode < 1e-5, f"Decoded outputs differ significantly: max diff = {max_diff_decode.item()}"

pyproject.toml

@@ -114,7 +114,7 @@ fastvideo-kernel = [
 imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
 # FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
 # comes from FastVideo's own custom_op wrappers.
-flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "940cd9680f3315f2f06b43ab5bea2c2cf2d96806", subdirectory = "flash_attn/cute" }
+flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
 
 [project.optional-dependencies]
 

tests/local_tests/test_nvfp4_fa4.py

@@ -10,15 +10,6 @@ import pytest
 
 os.environ["CUTE_DSL_ENABLE_TVM_FFI"] = "1"
 
-# Warm up flashinfer JIT before any FA4 imports
-from flashinfer.quantization import nvfp4_quantize, SfLayout
-_warmup = nvfp4_quantize(
-    torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
-    torch.ones(1, device="cuda", dtype=torch.float32),
-    sfLayout=SfLayout.layout_128x4, do_shuffle=False,
-)
-del _warmup
-
 
 def _make_impl(nheads, headdim, nvfp4=False):
     from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
@@ -45,7 +36,8 @@ def _cuda_timer(fn, warmup=5, iters=20):
 
 
 @pytest.mark.skipif(
-    torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
+    not torch.cuda.is_available()
+    or torch.cuda.get_device_capability() not in [(10, 0), (10, 3)],
     reason="Requires Blackwell GPU (sm100a or sm103a)"
 )
 class TestNVFP4FA4:
@@ -55,6 +47,17 @@ class TestNVFP4FA4:
     MODEL_HEADDIM = 128
     MODEL_SEQLEN = 32760  # 480x832 video, 81 frames
 
+    @pytest.fixture(scope="class", autouse=True)
+    def _flashinfer_warmup(self):
+        # Warm up flashinfer JIT before any FA4 imports
+        from flashinfer.quantization import SfLayout, nvfp4_quantize
+        _warmup = nvfp4_quantize(
+            torch.randn(4, 128, device="cuda", dtype=torch.bfloat16),
+            torch.ones(1, device="cuda", dtype=torch.float32),
+            sfLayout=SfLayout.layout_128x4, do_shuffle=False,
+        )
+        del _warmup
+
     def test_quantize_fn_shapes(self):
         """_nvfp4_quantize_for_fa4 produces correct shapes and strides."""
         from fastvideo.attention.backends.flash_attn import _nvfp4_quantize_for_fa4