diff --git a/Matrix-Game-3/tests/test_ulysses_attention.py b/Matrix-Game-3/tests/test_ulysses_attention.py new file mode 100644 index 0000000..c4f427d --- /dev/null +++ b/Matrix-Game-3/tests/test_ulysses_attention.py @@ -0,0 +1,85 @@ +"""Regression tests for the Matrix-Game-3 Ulysses attention adapter.""" + +import importlib.util +import sys +import types +from pathlib import Path + +import torch + + +ROOT = Path(__file__).parents[1] +ULYSSES_PATH = ROOT / "wan" / "distributed" / "ulysses.py" + + +def _load_ulysses_with_attention_stub(attention): + """Load the adapter with lightweight relative-import stubs.""" + + package_names = ( + "matrix_game_test", + "matrix_game_test.wan", + "matrix_game_test.wan.distributed", + "matrix_game_test.wan.modules", + ) + for name in package_names: + package = types.ModuleType(name) + package.__path__ = [] + sys.modules[name] = package + + attention_module = types.ModuleType("matrix_game_test.wan.modules.attention") + attention_module.flash_attention = lambda *args, **kwargs: None + attention_module.attention = attention + sys.modules[attention_module.__name__] = attention_module + + util_module = types.ModuleType("matrix_game_test.wan.distributed.util") + util_module.all_to_all = lambda value, **kwargs: value + sys.modules[util_module.__name__] = util_module + + spec = importlib.util.spec_from_file_location( + "matrix_game_test.wan.distributed.ulysses", ULYSSES_PATH + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_distributed_attention_forwards_version_keyword(): + calls = [] + + def attention(q, k, v, *, k_lens, window_size, version): + calls.append( + { + "k_lens": k_lens, + "window_size": window_size, + "version": version, + } + ) + return q + + module = _load_ulysses_with_attention_stub(attention) + + class SingleProcessDist: + @staticmethod + def is_initialized(): + return True + + @staticmethod + def get_world_size(): + return 1 + + module.dist = SingleProcessDist + q = torch.zeros(1, 2, 4, 3) + k = torch.zeros_like(q) + v = torch.zeros_like(q) + seq_lens = torch.tensor([2]) + + output = module.distributed_attention( + q, k, v, seq_lens, window_size=(3, 5), fa_version="3" + ) + + assert output is q + assert calls == [ + {"k_lens": seq_lens, "window_size": (3, 5), "version": "3"} + ] diff --git a/Matrix-Game-3/wan/distributed/ulysses.py b/Matrix-Game-3/wan/distributed/ulysses.py index e9cd669..c7ee3b4 100644 --- a/Matrix-Game-3/wan/distributed/ulysses.py +++ b/Matrix-Game-3/wan/distributed/ulysses.py @@ -66,7 +66,7 @@ def distributed_attention( v, k_lens=seq_lens, window_size=window_size, - fa_version=fa_version, + version=fa_version, ) # scatter q/k/v sequence