Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import importlib | |
| from functools import lru_cache | |
| import torch | |
| import torch.nn.functional as F | |
| SUPPORTED_KERNELS = { | |
| "auto", | |
| "sdpa", | |
| "flash_attn2", | |
| "flash_attn3", | |
| "block_sparse", | |
| "triton_block_sparse", | |
| } | |
| def _optional_module(name: str): | |
| try: | |
| return importlib.import_module(name) | |
| except (ImportError, OSError): | |
| return None | |
| def attention_kernel_status() -> dict[str, bool]: | |
| return { | |
| "sdpa": True, | |
| "flash_attn2": _optional_module("flash_attn") is not None, | |
| "flash_attn3": _optional_module("flash_attn_interface") is not None, | |
| "block_sparse": _optional_module("block_sparse_attn") is not None, | |
| "triton_block_sparse": _optional_module("flashvsr_attn") is not None, | |
| } | |
| def _require_kernel(name: str): | |
| module_name = { | |
| "flash_attn2": "flash_attn", | |
| "flash_attn3": "flash_attn_interface", | |
| "block_sparse": "block_sparse_attn", | |
| }[name] | |
| module = _optional_module(module_name) | |
| if module is None: | |
| raise ImportError( | |
| f"attention kernel {name!r} requires the optional module {module_name!r}" | |
| ) | |
| return module | |
| def _run_flash_attention( | |
| q: torch.Tensor, | |
| k: torch.Tensor, | |
| v: torch.Tensor, | |
| kernel: str, | |
| ) -> torch.Tensor: | |
| module = _require_kernel(kernel) | |
| function = getattr(module, "flash_attn_func", None) | |
| if function is None: | |
| raise AttributeError(f"{module.__name__} does not expose flash_attn_func") | |
| # FlashAttention 2/3 both consume [B,S,H,D]. | |
| output = function( | |
| q.transpose(1, 2).contiguous(), | |
| k.transpose(1, 2).contiguous(), | |
| v.transpose(1, 2).contiguous(), | |
| ) | |
| if isinstance(output, tuple): | |
| output = output[0] | |
| return output.transpose(1, 2) | |
| def _expand_block_mask( | |
| block_mask: torch.Tensor, | |
| block_token_size: int, | |
| query_tokens: int, | |
| key_tokens: int, | |
| ) -> torch.Tensor: | |
| mask = block_mask.repeat_interleave(block_token_size, dim=-2) | |
| mask = mask.repeat_interleave(block_token_size, dim=-1) | |
| mask = mask[..., :query_tokens, :key_tokens].bool() | |
| empty = ~mask.any(dim=-1, keepdim=True) | |
| if empty.any(): | |
| fallback = torch.zeros_like(mask) | |
| diagonal = torch.arange(min(query_tokens, key_tokens), device=mask.device) | |
| fallback[..., diagonal, diagonal] = True | |
| mask = torch.where(empty, fallback, mask) | |
| return mask | |
| def _run_block_sparse( | |
| q: torch.Tensor, | |
| k: torch.Tensor, | |
| v: torch.Tensor, | |
| block_mask: torch.Tensor, | |
| ) -> torch.Tensor: | |
| module = _require_kernel("block_sparse") | |
| function = getattr(module, "block_sparse_attn_func", None) | |
| if function is None: | |
| raise AttributeError("block_sparse_attn does not expose block_sparse_attn_func") | |
| batch, heads, query_tokens, head_dim = q.shape | |
| key_tokens = k.shape[2] | |
| expected_query_blocks = (query_tokens + 127) // 128 | |
| expected_key_blocks = (key_tokens + 127) // 128 | |
| expected_shape = (batch, heads, expected_query_blocks, expected_key_blocks) | |
| if tuple(block_mask.shape) != expected_shape: | |
| raise ValueError( | |
| "block_sparse_attn requires a compact 128-token block mask with " | |
| f"shape {expected_shape}, got {tuple(block_mask.shape)}" | |
| ) | |
| q_flat = q.transpose(1, 2).reshape(batch * query_tokens, heads, head_dim).contiguous() | |
| k_flat = k.transpose(1, 2).reshape(batch * key_tokens, heads, head_dim).contiguous() | |
| v_flat = v.transpose(1, 2).reshape(batch * key_tokens, heads, head_dim).contiguous() | |
| cu_q = torch.arange(batch + 1, device=q.device, dtype=torch.int32) * query_tokens | |
| cu_k = torch.arange(batch + 1, device=q.device, dtype=torch.int32) * key_tokens | |
| head_mask_type = torch.ones(heads, device=q.device, dtype=torch.int32) | |
| output = function( | |
| q_flat, | |
| k_flat, | |
| v_flat, | |
| cu_q, | |
| cu_k, | |
| head_mask_type, | |
| None, | |
| block_mask.contiguous(), | |
| query_tokens, | |
| key_tokens, | |
| 0.0, | |
| deterministic=False, | |
| softmax_scale=None, | |
| is_causal=False, | |
| exact_streaming=False, | |
| return_attn_probs=False, | |
| ) | |
| if isinstance(output, tuple): | |
| output = output[0] | |
| return output.reshape(batch, query_tokens, heads, head_dim).transpose(1, 2) | |
| def run_attention( | |
| q: torch.Tensor, | |
| k: torch.Tensor, | |
| v: torch.Tensor, | |
| *, | |
| kernel: str = "auto", | |
| token_mask: torch.Tensor | None = None, | |
| block_mask: torch.Tensor | None = None, | |
| block_token_size: int | None = None, | |
| ) -> torch.Tensor: | |
| """Run one explicitly selected attention implementation. | |
| Inputs and output use [B,H,S,D]. FlashAttention cannot consume arbitrary | |
| masks; block masks are routed to block-sparse attention or an exact SDPA | |
| expansion. | |
| """ | |
| kernel = kernel.lower().replace("-", "_") | |
| aliases = {"fa2": "flash_attn2", "fa3": "flash_attn3"} | |
| aliases["triton_sparse"] = "triton_block_sparse" | |
| kernel = aliases.get(kernel, kernel) | |
| if kernel not in SUPPORTED_KERNELS: | |
| raise ValueError(f"Unsupported attention kernel: {kernel}") | |
| if token_mask is not None and block_mask is not None: | |
| raise ValueError("Provide token_mask or block_mask, not both") | |
| has_mask = token_mask is not None or block_mask is not None | |
| cuda_half = q.is_cuda and q.dtype in {torch.float16, torch.bfloat16} | |
| if kernel == "auto": | |
| if ( | |
| block_mask is not None | |
| and block_token_size == 128 | |
| and cuda_half | |
| and _optional_module("block_sparse_attn") is not None | |
| ): | |
| kernel = "block_sparse" | |
| elif has_mask: | |
| kernel = "sdpa" | |
| elif cuda_half and _optional_module("flash_attn_interface") is not None: | |
| kernel = "flash_attn3" | |
| elif cuda_half and _optional_module("flash_attn") is not None: | |
| kernel = "flash_attn2" | |
| else: | |
| kernel = "sdpa" | |
| if kernel in {"flash_attn2", "flash_attn3"}: | |
| if has_mask: | |
| raise ValueError( | |
| f"{kernel} supports only unmasked full attention in this API; " | |
| "use block_sparse or sdpa for causal/sparse masks" | |
| ) | |
| if not cuda_half: | |
| raise ValueError(f"{kernel} requires CUDA FP16/BF16 tensors") | |
| return _run_flash_attention(q, k, v, kernel) | |
| if kernel == "block_sparse": | |
| if block_mask is None: | |
| raise ValueError("block_sparse requires a block_mask") | |
| if not cuda_half: | |
| raise ValueError("block_sparse requires CUDA FP16/BF16 tensors") | |
| if block_token_size != 128: | |
| raise ValueError( | |
| "block_sparse_attn requires 128 tokens per block; " | |
| f"the configured 3D window produces {block_token_size}" | |
| ) | |
| # block_sparse_attn consumes the compact [B,H,Q_blocks,K_blocks] | |
| # layout directly. Each entry represents one 128-token window. | |
| return _run_block_sparse(q, k, v, block_mask) | |
| if kernel == "triton_block_sparse": | |
| raise ValueError( | |
| "triton_block_sparse accepts the unflattened stream window layout; " | |
| "call it through forward_stream()" | |
| ) | |
| if block_mask is not None: | |
| if block_token_size is None: | |
| raise ValueError("block_token_size is required when SDPA expands a block mask") | |
| mask_entries = q.shape[0] * q.shape[1] * q.shape[2] * k.shape[2] | |
| if mask_entries > 2_000_000_000: | |
| raise RuntimeError( | |
| "SDPA block-mask fallback would allocate too many mask entries; " | |
| "install block_sparse_attn or select a shorter sequence" | |
| ) | |
| token_mask = _expand_block_mask( | |
| block_mask, block_token_size, q.shape[2], k.shape[2] | |
| ) | |
| return F.scaled_dot_product_attention( | |
| q, k, v, attn_mask=token_mask, dropout_p=0.0 | |
| ) | |