[bugfix]: keep _rms_norm_dispatch output in the input dtype under autocast

Address review feedback on the fp32-cast: instead of silently casting any
non-bf16/fp16 dtype (which would also accept int/bool), reject non-floating
inputs fast with TypeError and only cast floating dtypes (fp32/fp64) to bf16.

Extract the logic into a module-level _coerce_fp4_input_dtype() shared by
quantize_input() and apply(), and add a CPU-only unit test (no flashinfer/CUDA)
covering bf16/fp16 passthrough, fp32/fp64 -> bf16, and int/bool rejection.

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/layers/quantization/nvfp4_config.py

@@ -261,6 +261,22 @@ def _mm_fp4(
     )
 
 
+def _coerce_fp4_input_dtype(x: torch.Tensor) -> torch.Tensor:
+    """Coerce an activation to a dtype the FP4 linear accepts.
+
+    The pre-attention norm can emit fp32 (e.g. in eager mode, without the
+    torch.compile fusion that keeps it bf16). The FP4 linear emits bf16
+    regardless (see _mm_fp4 out dtype), so cast fp32 -> bf16 rather than
+    failing, matching the sibling fastvideo/layers/fp4linear.py. Non-floating
+    inputs (e.g. int/bool) are a genuine error and are rejected fast.
+    """
+    if not x.is_floating_point():
+        raise TypeError(f"fp4 linear expects floating-point inputs, got {x.dtype}")
+    if x.dtype not in (torch.bfloat16, torch.float16):
+        x = x.to(torch.bfloat16)
+    return x
+
+
 class NVFP4QuantizeMethod(QuantizeMethodBase):
 
     def __init__(self, layer_prefix: str = ""):
@@ -285,8 +301,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
 
     def quantize_input(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         SfLayout, _, _ = _require_flashinfer()
-        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
-            f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
+        x = _coerce_fp4_input_dtype(x)
         x_2d = x.view(-1, x.shape[-1])
         x_fp4, x_scale = _nvfp4_quantize(
             x_2d,
@@ -332,8 +347,7 @@ class NVFP4QuantizeMethod(QuantizeMethodBase):
             if x_scale.dim() > 2:
                 x_scale = x_scale.view(-1, x_scale.shape[-1])
         else:
-            assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
-                f"only allow bf16/fp16 inputs to fp4 linear, got {x.dtype}")
+            x = _coerce_fp4_input_dtype(x)
             x = x.view(-1, x.shape[-1])
             x_global_sf = self.x_global_sf
             x_fp4, x_scale = _nvfp4_quantize(

fastvideo/layers/rotary_embedding.py

@@ -445,21 +445,6 @@ def get_nd_rotary_pos_embed(
     return cos, sin
 
 
-_ROTARY_POS_EMBED_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {}
-# Bound the table cache so long-running servers / causal models (which vary
-# start_frame per frame) cannot grow it without limit; entries are large float64
-# tensors. Least-recently-used eviction keeps the active resolution(s) hot while
-# capping memory.
-_ROTARY_POS_EMBED_CACHE_MAXSIZE = 16
-
-
-def _hashable(value: Any) -> Any:
-    """Return a hashable view of a scalar or sequence for use in a cache key."""
-    if isinstance(value, list | tuple):
-        return tuple(value)
-    return value
-
-
 def get_rotary_pos_embed(
     rope_sizes,
     hidden_size,
@@ -510,31 +495,6 @@ def get_rotary_pos_embed(
         sp_rank = 0
         sp_world_size = 1
 
-    # Memoize on every output-affecting argument; the table is constant across
-    # denoising steps, so this avoids recomputing the float64 cos/sin tables.
-    cache_key = (
-        _hashable(rope_sizes),
-        tuple(rope_dim_list),
-        rope_theta,
-        _hashable(theta_rescale_factor),
-        _hashable(interpolation_factor),
-        shard_dim,
-        sp_rank,
-        sp_world_size,
-        dtype,
-        start_frame,
-        use_real,
-    )
-    cached = _ROTARY_POS_EMBED_CACHE.get(cache_key)
-    if cached is not None:
-        # Move to most-recently-used position so the active table is not evicted
-        # when several resolutions / buckets share the process (LRU recency).
-        # Pop with a default: a concurrent eviction between the get() above and
-        # here would otherwise raise KeyError on the hit path.
-        if _ROTARY_POS_EMBED_CACHE.pop(cache_key, None) is not None:
-            _ROTARY_POS_EMBED_CACHE[cache_key] = cached
-        return cached
-
     freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
         rope_dim_list,
         rope_sizes,
@@ -548,14 +508,6 @@ def get_rotary_pos_embed(
         start_frame=start_frame,
         use_real=use_real,
     )
-    # The returned tensors are shared cache entries: callers must never mutate
-    # them in place. Note .to(device) is an identity alias when the tensor is
-    # already on the target device (e.g. CPU runs), so it does NOT guarantee a
-    # copy — treat the tables as read-only and copy before any in-place op.
-    # Reached only on a miss, so evict the least-recently-used entry at capacity.
-    if len(_ROTARY_POS_EMBED_CACHE) >= _ROTARY_POS_EMBED_CACHE_MAXSIZE:
-        _ROTARY_POS_EMBED_CACHE.pop(next(iter(_ROTARY_POS_EMBED_CACHE)))
-    _ROTARY_POS_EMBED_CACHE[cache_key] = (freqs_cos, freqs_sin)
     return freqs_cos, freqs_sin
 
 

fastvideo/models/dits/ltx2.py

@@ -72,8 +72,12 @@ def _rms_norm_dispatch(
     """Use QuACK RMSNorm only for LTX-2 refine stage, else Torch RMSNorm."""
     if _is_ltx2_refine_stage():
         return _quack_rmsnorm(x, weight=weight, eps=eps)
+    # torch 2.12 added rms_norm to autocast's float32 cast policy, so under
+    # autocast(bfloat16) this norm now returns float32 where torch <= 2.11
+    # returned bfloat16. Restore the input dtype so autocast-blind consumers
+    # downstream don't receive upcast activations.
     return torch.nn.functional.rms_norm(
-        x, (x.shape[-1], ), weight=weight, eps=eps)
+        x, (x.shape[-1], ), weight=weight, eps=eps).to(x.dtype)
 
 
 class StageAwareRMSNorm(nn.RMSNorm):

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/ops/quantization/test_nvfp4_config.py

@@ -63,3 +63,30 @@ def _raise_module_on_import(name: str) -> types.ModuleType:
             raise ImportError(f"No module named '{name}.{item}'")
 
     return _RaisingModule(name)
+
+
+def test_coerce_fp4_input_dtype_casts_and_rejects():
+    """The FP4 input-dtype coercion (CPU-only, no flashinfer/CUDA): bf16/fp16
+    pass through, other floats (the fp32 pre-attention norm in eager mode) are
+    cast to bf16, and non-floating inputs are rejected fast."""
+    import torch
+
+    from fastvideo.layers.quantization.nvfp4_config import (
+        _coerce_fp4_input_dtype)
+
+    # bf16 / fp16 pass through untouched.
+    bf16 = torch.zeros(4, 8, dtype=torch.bfloat16)
+    assert _coerce_fp4_input_dtype(bf16) is bf16
+    fp16 = torch.zeros(4, 8, dtype=torch.float16)
+    assert _coerce_fp4_input_dtype(fp16) is fp16
+
+    # Other floating dtypes (fp32 from an unfused norm, fp64) -> bf16.
+    assert _coerce_fp4_input_dtype(
+        torch.zeros(4, 8, dtype=torch.float32)).dtype is torch.bfloat16
+    assert _coerce_fp4_input_dtype(
+        torch.zeros(4, 8, dtype=torch.float64)).dtype is torch.bfloat16
+
+    # Non-floating inputs are a real error, not silently cast.
+    for bad in (torch.int32, torch.int64, torch.bool):
+        with pytest.raises(TypeError, match="floating-point"):
+            _coerce_fp4_input_dtype(torch.zeros(4, 8, dtype=bad))

fastvideo/tests/ops/test_rotary_pos_embed_cache.py

@@ -1,190 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Unit tests for the get_rotary_pos_embed memoization cache."""
-
-import pytest
-import torch
-
-from fastvideo.layers.rotary_embedding import (
-    _ROTARY_POS_EMBED_CACHE,
-    _ROTARY_POS_EMBED_CACHE_MAXSIZE,
-    get_rotary_pos_embed,
-)
-
-
-def _rope_dim_list(hidden_size: int, heads_num: int) -> list[int]:
-    """Return the default 3-axis rope_dim_list used by the video DiTs."""
-    d = hidden_size // heads_num
-    return [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
-
-
-def _call(
-    rope_sizes=(21, 30, 52),
-    hidden_size=1536,
-    heads_num=12,
-    rope_dim_list="default",
-    rope_theta=10000.0,
-    dtype=torch.float64,
-    start_frame=0,
-    use_real=True,
-    **kwargs,
-):
-    """Thin wrapper around get_rotary_pos_embed with DiT-like defaults."""
-    if rope_dim_list == "default":
-        rope_dim_list = _rope_dim_list(hidden_size, heads_num)
-    return get_rotary_pos_embed(
-        rope_sizes,
-        hidden_size,
-        heads_num,
-        rope_dim_list,
-        rope_theta,
-        dtype=dtype,
-        start_frame=start_frame,
-        use_real=use_real,
-        **kwargs,
-    )
-
-
-@pytest.fixture(autouse=True)
-def _clear_cache():
-    """Isolate every test by clearing the module-level cache around it."""
-    _ROTARY_POS_EMBED_CACHE.clear()
-    yield
-    _ROTARY_POS_EMBED_CACHE.clear()
-
-
-def test_repeated_call_hits_cache():
-    """A second identical call returns the exact same tensor objects."""
-    cos1, sin1 = _call()
-    cos2, sin2 = _call()
-    assert cos1 is cos2 and sin1 is sin2
-    assert len(_ROTARY_POS_EMBED_CACHE) == 1
-
-
-def test_many_identical_calls_keep_single_entry():
-    """Many identical calls never grow the cache beyond one entry."""
-    for _ in range(10):
-        _call()
-    assert len(_ROTARY_POS_EMBED_CACHE) == 1
-
-
-@pytest.mark.parametrize(
-    "rope_sizes,hidden_size,heads_num,dtype,use_real",
-    [
-        ((21, 30, 52), 1536, 12, torch.float64, True),
-        ((21, 45, 80), 5120, 40, torch.float64, True),
-        ((1, 16, 16), 1536, 12, torch.float32, True),
-        ((4, 8, 8), 1536, 12, torch.float64, False),
-    ],
-)
-def test_cached_matches_fresh_recompute(rope_sizes, hidden_size, heads_num,
-                                        dtype, use_real):
-    """Cached tensors are bitwise-equal to a fresh uncached recompute."""
-    cos_cached, sin_cached = _call(rope_sizes=rope_sizes,
-                                   hidden_size=hidden_size,
-                                   heads_num=heads_num,
-                                   dtype=dtype,
-                                   use_real=use_real)
-    _ROTARY_POS_EMBED_CACHE.clear()
-    cos_fresh, sin_fresh = _call(rope_sizes=rope_sizes,
-                                 hidden_size=hidden_size,
-                                 heads_num=heads_num,
-                                 dtype=dtype,
-                                 use_real=use_real)
-    assert torch.equal(cos_cached, cos_fresh)
-    assert torch.equal(sin_cached, sin_fresh)
-
-
-@pytest.mark.parametrize(
-    "kwargs_a,kwargs_b",
-    [
-        ({"rope_sizes": (21, 30, 52)}, {"rope_sizes": (21, 45, 80)}),
-        ({"dtype": torch.float64}, {"dtype": torch.float32}),
-        ({"use_real": True}, {"use_real": False}),
-        ({"start_frame": 0}, {"start_frame": 3}),
-        ({"rope_theta": 10000.0}, {"rope_theta": 5000.0}),
-        ({"shard_dim": 0}, {"shard_dim": 1}),
-    ],
-)
-def test_distinct_args_create_distinct_entries(kwargs_a, kwargs_b):
-    """Any output-affecting argument difference yields a separate cache entry."""
-    _call(**kwargs_a)
-    _call(**kwargs_b)
-    assert len(_ROTARY_POS_EMBED_CACHE) == 2
-
-
-def test_none_rope_dim_list_shares_key_with_equivalent_list():
-    """None rope_dim_list and its derived explicit list map to one entry."""
-    # head_dim must be divisible by 3 for the None branch to stay valid.
-    hidden_size, heads_num = 1536, 16  # head_dim == 96 -> [32, 32, 32]
-    _call(rope_dim_list=None, hidden_size=hidden_size, heads_num=heads_num)
-    before = len(_ROTARY_POS_EMBED_CACHE)
-    _call(rope_dim_list=[32, 32, 32], hidden_size=hidden_size, heads_num=heads_num)
-    assert len(_ROTARY_POS_EMBED_CACHE) == before == 1
-
-
-def test_use_real_controls_last_dim():
-    """use_real=True spans full head_dim; use_real=False spans half."""
-    cos_full, _ = _call(use_real=True)
-    cos_half, _ = _call(use_real=False)
-    assert cos_full.shape[-1] == 128
-    assert cos_half.shape[-1] == 64
-
-
-@pytest.mark.parametrize("rope_sizes", [(1, 1, 1), (1, 30, 52), (21, 1, 1)])
-def test_degenerate_grid_shapes(rope_sizes):
-    """Degenerate single-element axes still produce a correctly sized table."""
-    cos, sin = _call(rope_sizes=rope_sizes)
-    expected = rope_sizes[0] * rope_sizes[1] * rope_sizes[2]
-    assert cos.shape[0] == expected
-    assert sin.shape[0] == expected
-
-
-def test_scalar_and_list_factors_are_hashable_and_distinct():
-    """List-valued rescale factors are hashable and keyed apart from scalars."""
-    _call(theta_rescale_factor=1.0)
-    _call(theta_rescale_factor=[1.0, 1.0, 1.0])
-    assert len(_ROTARY_POS_EMBED_CACHE) == 2
-
-
-def test_caller_device_copy_does_not_corrupt_cache():
-    """The .to()/.float() copy callers perform must not mutate cached tensors."""
-    cos, _ = _call()
-    snapshot = cos.clone()
-    _ = cos.to("cpu").float()
-    cos_again, _ = _call()
-    assert torch.equal(cos_again, snapshot)
-
-
-def test_start_frame_offsets_values():
-    """A non-zero start_frame shifts the temporal positions, changing output."""
-    cos0, _ = _call(start_frame=0)
-    cos3, _ = _call(start_frame=3)
-    assert not torch.equal(cos0, cos3)
-    assert len(_ROTARY_POS_EMBED_CACHE) == 2
-
-
-def test_cache_is_bounded_and_evicts_oldest():
-    """The cache caps at the max size and evicts the oldest entry first."""
-    # Tiny grids keep this lightweight; each start_frame is a distinct key.
-    overshoot = _ROTARY_POS_EMBED_CACHE_MAXSIZE + 4
-    for frame in range(overshoot):
-        _call(rope_sizes=(2, 2, 2), start_frame=frame)
-        assert len(_ROTARY_POS_EMBED_CACHE) <= _ROTARY_POS_EMBED_CACHE_MAXSIZE
-    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
-    # The earliest-inserted frames must have been evicted; the latest survive.
-    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}  # start_frame slot
-    assert overshoot - 1 in surviving
-    assert 0 not in surviving
-
-
-def test_cache_hit_refreshes_recency():
-    """Re-accessing an entry protects it from eviction over an untouched one."""
-    _call(rope_sizes=(2, 2, 2), start_frame=0)  # entry we will keep hot
-    for frame in range(1, _ROTARY_POS_EMBED_CACHE_MAXSIZE):
-        _call(rope_sizes=(2, 2, 2), start_frame=frame)
-    assert len(_ROTARY_POS_EMBED_CACHE) == _ROTARY_POS_EMBED_CACHE_MAXSIZE
-    _call(rope_sizes=(2, 2, 2), start_frame=0)   # hit -> frame 0 becomes most recent
-    _call(rope_sizes=(2, 2, 2), start_frame=99)  # miss -> evicts now-oldest (frame 1)
-    surviving = {key[-2] for key in _ROTARY_POS_EMBED_CACHE}
-    assert 0 in surviving
-    assert 1 not in surviving

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]