Spaces:
Running on L40S
Running on L40S
File size: 4,871 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""
Imaginaire4 Attention Subpackage:
Unified implementation for all Attention implementations.
Environment-related utilities.
"""
import os
import torch
from cosmos_framework.utils import log
# Controls all regions guarded against torch compile
# Logs, and certain assertions cause graph breaks.
def is_torch_compiling() -> bool:
try:
return torch.compiler.is_compiling()
except Exception as e:
log.exception(f"Exception occurred checking whether in torch compiled region: {e}")
# Assume too old to support torch compile
return False
def parse_backend_filter(env_var_value: str | None, default_backends: list[str]) -> list[str]:
"""
Parse backend filter from environment variable value.
Logic:
- If env_var_value is None or empty: return default_backends (no filtering)
- If all items start with "-": start with default_backends and remove specified backends (ban-list)
- Otherwise: only include specified backends (allow-list)
- Mixing bans and allows raises a ValueError
- For ban-list: invalid backend names only issue a warning (may not be available on this GPU)
- For allow-list: invalid backend names raise a ValueError (explicit request for unavailable backend)
Parameters:
env_var_value (str | None): The environment variable value (comma-separated list)
default_backends (list[str]): The default list of backends
Returns:
list[str]: Filtered list of backends
Raises:
ValueError: If the list mixes bans and allows, or if any backend name in allow-list is invalid
"""
if not env_var_value:
return default_backends
# Parse comma-separated list
items = [item.strip() for item in env_var_value.split(",") if item.strip()]
if not items:
return default_backends
# Check if items are bans or allows
bans = [item for item in items if item.startswith("-")]
allows = [item for item in items if not item.startswith("-")]
# Validate: cannot mix bans and allows
if bans and allows:
raise ValueError(
f"Cannot mix ban-list (items starting with '-') and allow-list (items without '-') in backend filter. "
f"Got bans: {bans}, allows: {allows}. "
f"Either specify all bans (e.g., '-flash2,-cudnn') or all allows (e.g., 'natten,flash2')."
)
default_backends_set = set(default_backends)
if bans:
# Ban-list mode: start with defaults and remove specified backends
ban_backends = {item[1:] for item in bans} # Remove "-" prefix
# Warn about banned backends that don't exist in default list
# (they may not be available on this GPU, which is fine)
invalid_bans = ban_backends - default_backends_set
if invalid_bans:
log.warning(
f"Attempting to ban backend(s) that are not in the available list: {sorted(invalid_bans)}. "
f"Available backends are: {default_backends}. "
f"This may be expected if these backends are not supported on this GPU."
)
filtered = [b for b in default_backends if b not in ban_backends]
log.debug(f"Backend filter (ban-list): removing {ban_backends} from {default_backends}, result: {filtered}")
return filtered
else:
# Allow-list mode: only include specified backends
allow_backends = set(allows)
# Validate: all allowed backends must exist in default list
invalid_allows = allow_backends - default_backends_set
if invalid_allows:
raise ValueError(
f"Invalid backend(s) in allow-list: {sorted(invalid_allows)}. "
f"Available backends are: {default_backends}. "
f"Check for typos in the backend names."
)
# Preserve order from default_backends for consistency
filtered = [b for b in default_backends if b in allow_backends]
log.debug(f"Backend filter (allow-list): allowing {allow_backends} from {default_backends}, result: {filtered}")
return filtered
def filter_attention_backends(default_backends: list[str]) -> list[str]:
env_var_value = os.environ.get("I4_ATTN_BACKENDS")
return parse_backend_filter(env_var_value, default_backends)
def filter_multi_dim_attention_backends(default_backends: list[str]) -> list[str]:
env_var_value = os.environ.get("I4_ATTN_BACKENDS_MULTIDIM")
return parse_backend_filter(env_var_value, default_backends)
def filter_attention_merge_backends(default_backends: list[str]) -> list[str]:
env_var_value = os.environ.get("I4_ATTN_BACKENDS_MERGE")
return parse_backend_filter(env_var_value, default_backends)
|