diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..27eec68 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +testpaths = tests +python_files = test_*.py diff --git a/tests/test_backend_fallbacks.py b/tests/test_backend_fallbacks.py index bc4cbfe..df8d604 100644 --- a/tests/test_backend_fallbacks.py +++ b/tests/test_backend_fallbacks.py @@ -10,6 +10,25 @@ import pytest ROOT = Path(__file__).resolve().parents[1] +@pytest.mark.skipif(platform.system() != "Darwin", reason="macOS backend default") +def test_dense_attention_defaults_to_sdpa_on_macos(): + env = os.environ.copy() + env.pop("ATTN_BACKEND", None) + completed = subprocess.run( + [ + sys.executable, + "-c", + "from trellis2.modules.attention import config; print('BACKEND=' + config.BACKEND)", + ], + cwd=ROOT, + env=env, + text=True, + capture_output=True, + check=True, + ) + assert "BACKEND=sdpa" in completed.stdout + + def test_disable_metal_selects_controlled_pytorch_fallback(): env = os.environ.copy() env.update( @@ -49,22 +68,40 @@ print('REPORT=' + json.dumps({ @pytest.mark.skipif(platform.system() != "Darwin", reason="Metal probe requires macOS") def test_metal_sparse_attention_has_its_own_numerical_probe(): - import torch - - if not torch.backends.mps.is_available(): + # Pipeline imports performed by other test modules may already have made a + # controlled fallback choice. Exercise auto-detection in a fresh process, + # which also mirrors a real CLI invocation. + code = r""" +import json +import torch +if not torch.backends.mps.is_available(): + raise SystemExit(77) +from trellis2.modules.sparse import config +from trellis2.backends import probe_metal_backends +print('REPORT=' + json.dumps({ + 'probe': config.probe_flex_gemm_sparse_attention_on_mps(), + 'conv': config.CONV, + 'attention': config.ATTN, + 'metal': probe_metal_backends(), +}, sort_keys=True)) +""" + completed = subprocess.run( + [sys.executable, "-c", code], + cwd=ROOT, + text=True, + capture_output=True, + check=False, + ) + if completed.returncode == 77: pytest.skip("MPS is unavailable") - - from trellis2.modules.sparse import config - - assert config.probe_flex_gemm_sparse_attention_on_mps() - assert config.CONV == "flex_gemm" - assert config.ATTN == "flex_gemm_sparse_attn" - - from trellis2.backends import probe_metal_backends - - probes = probe_metal_backends() - assert probes["rasterizer"]["ok"], probes["rasterizer"] - assert probes["mesh_bvh"]["ok"], probes["mesh_bvh"] + assert completed.returncode == 0, completed.stderr + line = next(line for line in completed.stdout.splitlines() if line.startswith("REPORT=")) + report = json.loads(line.removeprefix("REPORT=")) + assert report["probe"] + assert report["conv"] == "flex_gemm" + assert report["attention"] == "flex_gemm_sparse_attn" + assert report["metal"]["rasterizer"]["ok"], report["metal"]["rasterizer"] + assert report["metal"]["mesh_bvh"]["ok"], report["metal"]["mesh_bvh"] @pytest.mark.skipif(platform.system() != "Darwin", reason="Metal parity requires macOS") diff --git a/trellis2/modules/attention/config.py b/trellis2/modules/attention/config.py index a6d5180..e8e2798 100644 --- a/trellis2/modules/attention/config.py +++ b/trellis2/modules/attention/config.py @@ -1,6 +1,9 @@ -from typing import * +import platform +from typing import Literal -BACKEND = 'flash_attn' +# flash-attn is a CUDA-only optional dependency. Keep the upstream CUDA +# default, while making SDPA the source-native dense-attention path on macOS. +BACKEND = 'sdpa' if platform.system() == 'Darwin' else 'flash_attn' DEBUG = False def __from_env(): @@ -23,7 +26,9 @@ def __from_env(): __from_env() -def set_backend(backend: Literal['xformers', 'flash_attn']): +def set_backend( + backend: Literal['xformers', 'flash_attn', 'flash_attn_3', 'sdpa', 'naive'], +): global BACKEND BACKEND = backend