Skip to content

selector

Classes

Functions

fastvideo.attention.selector.backend_name_to_enum

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.

Source code in fastvideo/attention/selector.py
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

fastvideo.attention.selector.get_env_variable_attn_backend

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
Source code in fastvideo/attention/selector.py
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))

fastvideo.attention.selector.get_global_forced_attn_backend

get_global_forced_attn_backend() -> AttentionBackendEnum | None

Get the currently-forced choice of attention backend, or None if auto-selection is currently enabled.

Source code in fastvideo/attention/selector.py
def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the currently-forced choice of attention backend,
    or None if auto-selection is currently enabled.
    '''
    return forced_attn_backend

fastvideo.attention.selector.global_force_attn_backend

global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None

Force all attention operations to use a specified backend.

Passing None for the argument re-enables automatic backend selection.,

Arguments:

  • attn_backend: backend selection (None to revert to auto)
Source code in fastvideo/attention/selector.py
def global_force_attn_backend(
        attn_backend: AttentionBackendEnum | None) -> None:
    '''
    Force all attention operations to use a specified backend.

    Passing `None` for the argument re-enables automatic
    backend selection.,

    Arguments:

    * attn_backend: backend selection (None to revert to auto)
    '''
    global forced_attn_backend
    forced_attn_backend = attn_backend

fastvideo.attention.selector.global_force_attn_backend_context_manager

global_force_attn_backend_context_manager(attn_backend: AttentionBackendEnum) -> Generator[None, None, None]

Globally force a FastVideo attention backend override within a context manager, reverting the global attention backend override to its prior state upon exiting the context manager.

Arguments:

  • attn_backend: attention backend to force

Returns:

  • Generator
Source code in fastvideo/attention/selector.py
@contextmanager
def global_force_attn_backend_context_manager(
        attn_backend: AttentionBackendEnum) -> Generator[None, None, None]:
    '''
    Globally force a FastVideo attention backend override within a
    context manager, reverting the global attention backend
    override to its prior state upon exiting the context
    manager.

    Arguments:

    * attn_backend: attention backend to force

    Returns:

    * Generator
    '''

    # Save the current state of the global backend override (if any)
    original_value = get_global_forced_attn_backend()

    # Globally force the new backend override
    global_force_attn_backend(attn_backend)

    # Yield control back to the enclosed code block
    try:
        yield
    finally:
        # Revert the original global backend override, if any
        global_force_attn_backend(original_value)