375 lines
16 KiB
Python
375 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Golden tests for explicit per-component attention-backend resolution.
|
|
|
|
The selector's contract after the refactor:
|
|
* precedence is unchanged: global force > scoped/explicit request >
|
|
env var > layer default > platform auto — the scoped request occupies
|
|
exactly the position the process-global force held when the train
|
|
stack used it for per-role models;
|
|
* every selection input is part of the resolution cache key, so no
|
|
mutation (scope, force, env) ever needs a ``cache_clear()``;
|
|
* an active ``_component_attention_backend_scope`` suppresses the env var unless
|
|
``consult_env=True`` — ``_component_attention_backend_scope(None)`` therefore
|
|
means "automatic selection, ignore the process-wide request" (the
|
|
dense teacher/critic case);
|
|
* scopes are exception-safe and per-component (component identity and
|
|
the active device index are cache-key inputs, so two components — or
|
|
two devices with different capabilities — never share a resolution).
|
|
|
|
CPU-only: the platform is faked (same seam as the pre-existing
|
|
role-override tests); no kernels, no GPU.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import fastvideo.attention.selector as selector
|
|
import fastvideo.platforms as platforms
|
|
from fastvideo.platforms import AttentionBackendEnum
|
|
|
|
FLASH = AttentionBackendEnum.FLASH_ATTN
|
|
SDPA = AttentionBackendEnum.TORCH_SDPA
|
|
SAGE = AttentionBackendEnum.SAGE_ATTN
|
|
|
|
SUPPORTED = (FLASH, SDPA, SAGE)
|
|
KWARGS = {
|
|
"head_size": 64,
|
|
"dtype": torch.bfloat16,
|
|
"supported_attention_backends": SUPPORTED,
|
|
}
|
|
|
|
|
|
class _FakePlatform:
|
|
device_name = "fake"
|
|
|
|
@classmethod
|
|
def is_mps(cls) -> bool:
|
|
# Not decoration. The autouse fixture swaps this stub in process-wide,
|
|
# so anything that constructs FastVideoArgs under it reaches
|
|
# check_fastvideo_args, which calls current_platform.is_mps(). Without
|
|
# this, the two env-folding tests raise AttributeError before they
|
|
# assert anything.
|
|
return False
|
|
|
|
@classmethod
|
|
def get_attn_backend_cls(
|
|
cls,
|
|
selected_backend: AttentionBackendEnum | None,
|
|
head_size: int,
|
|
dtype: torch.dtype,
|
|
) -> str:
|
|
del head_size, dtype
|
|
# Auto-selection (None) resolves to FLASH on this fake platform so
|
|
# the oracle can distinguish "auto" from every explicit request.
|
|
return (selected_backend or FLASH).name
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _fake_platform(monkeypatch):
|
|
monkeypatch.setattr(platforms, "_current_platform", _FakePlatform())
|
|
monkeypatch.setattr(selector, "resolve_obj_by_qualname", lambda name: name)
|
|
monkeypatch.delenv("FASTVIDEO_ATTENTION_BACKEND", raising=False)
|
|
selector._cached_get_attn_backend.cache_clear()
|
|
yield
|
|
# Fake-platform resolutions must not outlive the test: mutations no
|
|
# longer flush the cache (that is the feature), so evict explicitly.
|
|
selector._cached_get_attn_backend.cache_clear()
|
|
|
|
|
|
def _oracle(requested, env, default):
|
|
"""The documented precedence, spelled out independently."""
|
|
selected = requested
|
|
if selected is None and env is not None:
|
|
selected = AttentionBackendEnum[env]
|
|
if selected is None and default is not None:
|
|
selected = default
|
|
if selected is not None and selected not in SUPPORTED:
|
|
selected = default if default in SUPPORTED else None
|
|
return (selected or FLASH).name
|
|
|
|
|
|
@pytest.mark.parametrize("scoped", [None, "unset", SAGE])
|
|
@pytest.mark.parametrize("env", [None, "TORCH_SDPA"])
|
|
@pytest.mark.parametrize("default", [None, SAGE])
|
|
def test_precedence_matrix_matches_previous_semantics(monkeypatch, scoped, env, default) -> None:
|
|
if env is not None:
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", env)
|
|
|
|
if scoped == "unset":
|
|
# No scope and no explicit request: env is consulted.
|
|
got = selector.get_attn_backend(default_backend=default, **KWARGS)
|
|
expected = _oracle(None, env, default)
|
|
else:
|
|
# Scope active: env suppressed, scoped request used.
|
|
with selector._component_attention_backend_scope(scoped, component="dit"):
|
|
got = selector.get_attn_backend(default_backend=default, **KWARGS)
|
|
expected = _oracle(scoped, None, default)
|
|
assert got == expected
|
|
|
|
|
|
@pytest.mark.parametrize("requested", [None, SAGE])
|
|
@pytest.mark.parametrize("env", [None, "TORCH_SDPA"])
|
|
@pytest.mark.parametrize("default", [None, SAGE])
|
|
def test_explicit_request_matches_the_same_oracle(monkeypatch, requested, env, default) -> None:
|
|
"""An explicit `requested=` resolves exactly as the same request would via
|
|
a scope, and never consults the environment behind it."""
|
|
if env is not None:
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", env)
|
|
|
|
got = selector.get_attn_backend(default_backend=default, requested=requested, **KWARGS)
|
|
assert got == _oracle(requested, None, default)
|
|
|
|
|
|
def test_scope_none_ignores_env_request(monkeypatch) -> None:
|
|
"""The dense teacher/critic case: auto-select, ignore the process-wide
|
|
request — previously implemented by popping the env var and flushing
|
|
the selector cache around the build."""
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
with selector._component_attention_backend_scope(None, component="transformer"):
|
|
assert selector.get_attn_backend(**KWARGS) == "FLASH_ATTN" # auto
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
|
|
|
|
def test_interleaved_component_scopes_need_no_cache_clear(monkeypatch) -> None:
|
|
"""Per-role/per-component builds interleave freely; identical
|
|
shape/dtype keys resolve per scope with zero cache management."""
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
|
|
for _ in range(2):
|
|
with selector._component_attention_backend_scope(SDPA, component="student"):
|
|
assert selector.get_attn_backend(**KWARGS) == "TORCH_SDPA"
|
|
with selector._component_attention_backend_scope(None, component="teacher"):
|
|
assert selector.get_attn_backend(**KWARGS) == "FLASH_ATTN"
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
|
|
|
|
def test_scope_is_exception_safe(monkeypatch) -> None:
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
with selector._component_attention_backend_scope(SDPA, component="dit"):
|
|
raise RuntimeError("boom")
|
|
assert selector._active_component_attention_backend_scope() is None
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
|
|
|
|
def test_explicit_none_means_auto_not_absence(monkeypatch) -> None:
|
|
"""`requested=None` is an answer ("this component resolved to automatic
|
|
selection"), so the env var must NOT be consulted behind it. Omitting the
|
|
argument is the absence, and still falls back to the env."""
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
|
|
# Explicit "automatic selection" -> the fake platform's auto answer.
|
|
assert selector.get_attn_backend(requested=None, **KWARGS) == "FLASH_ATTN"
|
|
# Counterfactual: same call, no opinion -> env still wins.
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
assert selector.get_attn_backend(requested=selector.NO_REQUEST, **KWARGS) == "SAGE_ATTN"
|
|
|
|
|
|
def test_explicit_request_outranks_the_construction_scope(monkeypatch) -> None:
|
|
"""A caller that knows its component's decision is not overridden by an
|
|
ambient scope that happens to be active."""
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
|
|
with selector._component_attention_backend_scope(SDPA, component="dit"):
|
|
assert selector.get_attn_backend(**KWARGS) == "TORCH_SDPA"
|
|
assert selector.get_attn_backend(requested=SAGE, **KWARGS) == "SAGE_ATTN"
|
|
|
|
|
|
def test_env_change_takes_effect_without_cache_clear(monkeypatch) -> None:
|
|
"""Selection inputs live in the cache key, so an env change is simply a
|
|
different key. (Previously a changed env var was silently ignored until
|
|
someone remembered to flush the cache.)"""
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
|
|
assert selector.get_attn_backend(**KWARGS) == "TORCH_SDPA"
|
|
|
|
|
|
def test_scope_typo_fails_fast() -> None:
|
|
with pytest.raises(ValueError, match="Unknown attention backend"):
|
|
with selector._component_attention_backend_scope("flash_atn", component="dit"):
|
|
pass
|
|
|
|
|
|
def test_consult_env_scope_keeps_env_visible(monkeypatch) -> None:
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
with selector._component_attention_backend_scope(None, component="dit", consult_env=True):
|
|
assert selector.get_attn_backend(**KWARGS) == "SAGE_ATTN"
|
|
|
|
|
|
def test_active_device_is_part_of_the_cache_key(monkeypatch) -> None:
|
|
"""Platform auto-selection probes the *current* device's capability
|
|
(AttnQatInferBackend, for one, resolves a different kernel per
|
|
capability set), so two devices must not share one cached resolution.
|
|
The fake platform here stands in for that probe."""
|
|
current = {"index": 0}
|
|
per_device = {0: FLASH, 1: SDPA}
|
|
|
|
class _PerDevicePlatform:
|
|
device_name = "fake"
|
|
|
|
@classmethod
|
|
def get_attn_backend_cls(cls, selected_backend, head_size, dtype) -> str:
|
|
del head_size, dtype
|
|
return (selected_backend or per_device[current["index"]]).name
|
|
|
|
monkeypatch.setattr(platforms, "_current_platform", _PerDevicePlatform())
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
|
monkeypatch.setattr(torch.cuda, "current_device", lambda: current["index"])
|
|
|
|
assert selector.get_attn_backend(**KWARGS) == "FLASH_ATTN"
|
|
current["index"] = 1
|
|
assert selector.get_attn_backend(**KWARGS) == "TORCH_SDPA"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The decision is carried on the component
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_env_is_folded_into_the_typed_request_once(monkeypatch):
|
|
"""``FastVideoArgs.attention_backend`` is the parse-once adapter."""
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "SAGE_ATTN")
|
|
assert FastVideoArgs(model_path="x").attention_backend == "SAGE_ATTN"
|
|
|
|
# An explicit request always wins over the environment.
|
|
assert FastVideoArgs(model_path="x", attention_backend="TORCH_SDPA").attention_backend == "TORCH_SDPA"
|
|
|
|
|
|
def test_unparseable_env_falls_through_instead_of_raising(monkeypatch):
|
|
"""The env var keeps its permissive parse; only explicit requests raise."""
|
|
from fastvideo.fastvideo_args import FastVideoArgs
|
|
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "flash_atn")
|
|
assert FastVideoArgs(model_path="x").attention_backend is None
|
|
|
|
with pytest.raises(ValueError, match="Unknown attention backend"):
|
|
FastVideoArgs(model_path="x", attention_backend="flash_atn")
|
|
|
|
|
|
def test_recorded_decision_is_readable_from_the_component():
|
|
"""After load, the component carries what it resolved."""
|
|
from fastvideo.configs.models.dits.base import DiTConfig
|
|
|
|
config = DiTConfig()
|
|
assert config._resolved_attention_backend is None
|
|
|
|
with selector._component_attention_backend_scope(SAGE, component="transformer"):
|
|
assert selector.record_resolved_attention_backend(config) == SAGE
|
|
assert config._resolved_attention_backend == SAGE
|
|
|
|
|
|
def test_recorded_decision_is_the_narrowed_one():
|
|
"""A loader that narrows the request records the narrowed value.
|
|
|
|
The DMD teacher/critic transformers build dense inside a nested scope
|
|
while the run as a whole requested a quantized backend; the component
|
|
must report what it actually resolved, not the run-wide request.
|
|
"""
|
|
from fastvideo.configs.models.dits.base import DiTConfig
|
|
|
|
student, teacher = DiTConfig(), DiTConfig()
|
|
with selector._component_attention_backend_scope(SAGE, component="transformer"):
|
|
selector.record_resolved_attention_backend(student)
|
|
with selector._component_attention_backend_scope(None, component="transformer"):
|
|
selector.record_resolved_attention_backend(teacher)
|
|
|
|
assert student._resolved_attention_backend == SAGE
|
|
assert teacher._resolved_attention_backend is None
|
|
|
|
|
|
def test_no_request_records_no_decision():
|
|
from fastvideo.configs.models.encoders.base import TextEncoderConfig
|
|
|
|
config = TextEncoderConfig()
|
|
assert selector.record_resolved_attention_backend(config) is None
|
|
assert config._resolved_attention_backend is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Reading the component's recorded decision back
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_component_attention_backend_reads_the_recorded_decision():
|
|
from fastvideo.configs.models.dits.base import DiTConfig
|
|
|
|
class _Component:
|
|
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
config = DiTConfig()
|
|
with selector._component_attention_backend_scope(SAGE, component="transformer"):
|
|
selector.record_resolved_attention_backend(config)
|
|
|
|
assert selector.component_attention_backend(_Component(config)) == SAGE
|
|
|
|
|
|
def test_component_with_no_concrete_decision_reports_no_request(monkeypatch):
|
|
""""Stamped with no request" and "never stamped" are the SAME stored value.
|
|
|
|
``ModelConfig._resolved_attention_backend`` is a field defaulting to None,
|
|
so the attribute always exists and ``record_resolved_attention_backend``
|
|
writes None whenever no scope is active. Neither state may suppress the env
|
|
var at a call site that previously honoured it, so both report NO_REQUEST.
|
|
"""
|
|
from fastvideo.configs.models.dits.base import DiTConfig
|
|
|
|
class _Component:
|
|
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
# (a) stamped outside any scope -> stored None -> NO_REQUEST
|
|
stamped = DiTConfig()
|
|
assert selector.record_resolved_attention_backend(stamped) is None
|
|
assert selector.component_attention_backend(_Component(stamped)) is selector.NO_REQUEST
|
|
|
|
# (b) never stamped at all -> same stored value, same answer
|
|
untouched = DiTConfig()
|
|
assert untouched._resolved_attention_backend is None # the field always exists
|
|
assert selector.component_attention_backend(_Component(untouched)) is selector.NO_REQUEST
|
|
|
|
# (c) and it does NOT suppress the env var, which is the whole point
|
|
monkeypatch.setenv("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
|
|
selector._cached_get_attn_backend.cache_clear()
|
|
passthrough = selector.get_attn_backend(requested=selector.component_attention_backend(_Component(untouched)),
|
|
**KWARGS)
|
|
assert passthrough == selector.get_attn_backend(**KWARGS)
|
|
|
|
|
|
def test_component_with_a_concrete_decision_reports_it():
|
|
"""A real recorded backend is an answer, and overrides the ambient fallback."""
|
|
from fastvideo.configs.models.dits.base import DiTConfig
|
|
|
|
class _Component:
|
|
|
|
def __init__(self, config):
|
|
self.config = config
|
|
|
|
config = DiTConfig()
|
|
with selector._component_attention_backend_scope(SAGE, component="transformer"):
|
|
assert selector.record_resolved_attention_backend(config) == SAGE
|
|
assert selector.component_attention_backend(_Component(config)) == SAGE
|
|
|
|
|
|
def test_component_without_a_recorded_decision_reports_no_request():
|
|
"""A transformer built outside a loader (or carrying a non-FastVideo
|
|
config) has no decision, so callers keep the previous fallback."""
|
|
|
|
class _Bare:
|
|
pass
|
|
|
|
class _HFStyle:
|
|
|
|
def __init__(self):
|
|
self.config = object()
|
|
|
|
assert selector.component_attention_backend(_Bare()) is selector.NO_REQUEST
|
|
assert selector.component_attention_backend(_HFStyle()) is selector.NO_REQUEST
|