[misc]: harden the rotary cache hit path

fastvideo/layers/rotary_embedding.py

@@ -445,6 +445,21 @@ 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,
@@ -495,6 +510,31 @@ 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,
@@ -508,6 +548,14 @@ 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/tests/ops/test_rotary_pos_embed_cache.py

@@ -0,0 +1,190 @@
+# 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