fix: use SDPA for dense attention on macOS
This commit is contained in:
parent
1488ba1505
commit
847bd59c25
3
pytest.ini
Normal file
3
pytest.ini
Normal file
@ -0,0 +1,3 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
python_files = test_*.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():
|
||||
# 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():
|
||||
pytest.skip("MPS is unavailable")
|
||||
|
||||
raise SystemExit(77)
|
||||
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"]
|
||||
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")
|
||||
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")
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user