214 lines
7.1 KiB
Python
214 lines
7.1 KiB
Python
# 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
|