selector ¶

def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
    """
    Convert a string backend name to a _Backend enum value.

    Returns:
    * _Backend: enum value if backend_name is a valid in-tree type
    * None: otherwise it's an invalid in-tree type or an out-of-tree platform is
            loaded.
    """
    assert backend_name is not None
    return AttentionBackendEnum[backend_name] if backend_name in AttentionBackendEnum.__members__ else \
          None

def coerce_attn_backend(attn_backend: AttentionBackendEnum | str | None, ) -> AttentionBackendEnum | None:
    """Normalize an explicit backend selection.

    Environment-variable parsing remains permissive via
    :func:`backend_name_to_enum`, but typed/config-driven call sites should
    fail fast on typos instead of silently falling back to another backend.
    """
    if attn_backend is None or isinstance(attn_backend, AttentionBackendEnum):
        return attn_backend
    if not isinstance(attn_backend, str) or not attn_backend.strip():
        raise ValueError("attention backend must be a non-empty string, "
                         f"an AttentionBackendEnum, or None; got {attn_backend!r}")

    backend_name = attn_backend.strip().upper()
    backend = backend_name_to_enum(backend_name)
    if backend is None:
        raise ValueError(f"Unknown attention backend {attn_backend!r}. "
                         f"Expected one of {sorted(AttentionBackendEnum.__members__)}")
    return backend

def component_attention_backend(component: object) -> AttentionBackendEnum | _NoRequest:
    """Read back the decision :func:`record_resolved_attention_backend` wrote.

    Returns ``NO_REQUEST`` unless the component recorded a *concrete* backend,
    so a caller that passes this through only overrides the ambient fallback
    when there is a real decision to override it with.

    ``NO_REQUEST`` rather than ``None`` for the no-decision case is deliberate,
    and cannot be derived from the attribute's presence: ``ModelConfig`` declares
    ``_resolved_attention_backend`` as a field defaulting to ``None``, so the
    attribute always exists and a ``getattr`` default can never fire. Worse,
    ``record_resolved_attention_backend`` writes ``None`` whenever no scope is
    active, so "resolved to automatic selection" and "never recorded" are the
    same stored value. Neither state should suppress the environment variable at
    a call site that previously honoured it, and collapsing both to
    ``NO_REQUEST`` keeps that behavior identical.
    """
    resolved = getattr(getattr(component, "config", None), "_resolved_attention_backend", None)
    return NO_REQUEST if resolved is None else resolved

def get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
    supported_attention_backends: tuple[AttentionBackendEnum, ...]
    | None = None,
    default_backend: AttentionBackendEnum | None = None,
    *,
    requested: AttentionBackendEnum | None | _NoRequest = NO_REQUEST,
) -> type[AttentionBackend]:
    """Resolve the attention backend class for one call site.

    ``requested`` is the decision made for the component this call site belongs
    to — read from ``ModelConfig._resolved_attention_backend``. Passing it means
    the caller knows the answer, so nothing ambient is consulted:

    * ``requested=SOME_BACKEND`` — use it (subject to the layer's declared
      support), ignoring the construction scope and the environment;
    * ``requested=None`` — that component resolved to *automatic selection*.
      That is an answer, not an absence, so the environment is not consulted
      behind it;
    * ``requested`` omitted (``NO_REQUEST``) — the caller has no opinion; fall
      back to the construction scope, then the environment. This is the path
      every layer built inside a loader still takes.
    """
    # Resolve every selection-affecting input BEFORE the cache so all of them
    # live in the cache key: no mutation (env, scope) ever needs a cache_clear
    # again, and components with different requests get distinct cache entries
    # (per-component resolution).
    if not isinstance(requested, _NoRequest):
        # Explicit: the component's own decision outranks anything ambient.
        component = None
        env_backend = None
    else:
        scope = _SCOPE.get()
        if scope is not None:
            requested = scope.backend
            component = scope.component
            env_backend = envs.FASTVIDEO_ATTENTION_BACKEND if scope.consult_env else None
        else:
            requested = None
            component = None
            env_backend = envs.FASTVIDEO_ATTENTION_BACKEND
    # The active device is a real selection input, not bookkeeping: the
    # platform's backend resolution runs capability probes against the
    # *current* device (e.g. AttnQatInferBackend's per-arch capability sets
    # decide sm_12x CUTLASS vs sm_100/sm_103 FP4 FA4 vs FlashAttention
    # fallback). Keying on it means a resolution taken for one device is
    # never handed to another.
    device_index = torch.cuda.current_device() if torch.cuda.is_available() else None
    return _cached_get_attn_backend(
        head_size,
        dtype,
        supported_attention_backends,
        default_backend,
        requested=requested,
        env_backend=env_backend,
        component=component,
        device_index=device_index,
    )

def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the backend override specified by the FastVideo attention
    backend environment variable, if one is specified.

    Returns:

    * _Backend enum value if an override is specified
    * None otherwise
    '''
    backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
    return (None if backend_name is None else backend_name_to_enum(backend_name))

def record_resolved_attention_backend(config: object) -> AttentionBackendEnum | None:
    """Write the request governing construction *right now* onto a component config.

    Called by each loader while its component is being built, so the decision
    ends up on the object the component keeps and is readable from the component
    after load. The loader is the right place rather than the caller above it:
    a loader may narrow the request for one component (the DMD teacher/critic
    transformers build dense inside a nested scope), and only it knows that.
    """
    scope = _SCOPE.get()
    resolved = scope.backend if scope is not None else None
    config._resolved_attention_backend = resolved  # type: ignore[attr-defined]
    return resolved