content: #trellis\modules\sparse\attention\full_attn.py
  from typing import *
  import torch
  import math
  from .. import SparseTensor
  from trellis.backend_config import (
      get_attention_backend,
      get_debug_mode,
      get_available_backends
  )
  import logging
  
  logger = logging.getLogger(__name__)
  
  # Get configuration from central config
  ATTN = get_attention_backend()
  DEBUG = get_debug_mode()
  
  # Get available backends and import if active
  available_backends = get_available_backends()
  
  if ATTN not in ['xformers', 'flash_attn']:
      logger.warning(f"Attention backend {ATTN} not supported for sparse attention. Only 'xformers' and 'flash_attn' are available. Defaulting to 'flash_attn'")
      ATTN = 'flash_attn'
  
  if ATTN == 'xformers' and available_backends['xformers']:
      import xformers.ops as xops
  elif ATTN == 'flash_attn' and available_backends['flash_attn']:
      import flash_attn
  else:
      raise ImportError(f"Could not import {ATTN}. Please install either xformers or flash-attn for sparse attention support.")
  
  
  
  __all__ = [
      'sparse_scaled_dot_product_attention',
  ]
  
  
  @overload
  def sparse_scaled_dot_product_attention(qkv: SparseTensor) -> SparseTensor:
      """
      Apply scaled dot product attention to a sparse tensor.
  
      Args:
          qkv (SparseTensor): A [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
      """
      ...
  
  @overload
  def sparse_scaled_dot_product_attention(q: SparseTensor, kv: Union[SparseTensor, torch.Tensor]) -> SparseTensor:
      """
      Apply scaled dot product attention to a sparse tensor.
  
      Args:
          q (SparseTensor): A [N, *, H, C] sparse tensor containing Qs.
          kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor or a [N, L, 2, H, C] dense tensor containing Ks and Vs.
      """
      ...
  
  @overload
  def sparse_scaled_dot_product_attention(q: torch.Tensor, kv: SparseTensor) -> torch.Tensor:
      """
      Apply scaled dot product attention to a sparse tensor.
  
      Args:
          q (SparseTensor): A [N, L, H, C] dense tensor containing Qs.
          kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor containing Ks and Vs.
      """
      ...
  
  @overload
  def sparse_scaled_dot_product_attention(q: SparseTensor, k: SparseTensor, v: SparseTensor) -> SparseTensor:
      """
      Apply scaled dot product attention to a sparse tensor.
  
      Args:
          q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
          k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
          v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
  
      Note:
          k and v are assumed to have the same coordinate map.
      """
      ...
  
  @overload
  def sparse_scaled_dot_product_attention(q: SparseTensor, k: torch.Tensor, v: torch.Tensor) -> SparseTensor:
      """
      Apply scaled dot product attention to a sparse tensor.
  
      Args:
          q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
          k (torch.Tensor): A [N, L, H, Ci] dense tensor containing Ks.
          v (torch.Tensor): A [N, L, H, Co] dense tensor containing Vs.
      """
      ...
  
  @overload
  def sparse_scaled_dot_product_attention(q: torch.Tensor, k: SparseTensor, v: SparseTensor) -> torch.Tensor:
      """
      Apply scaled dot product attention to a sparse tensor.
  
      Args:
          q (torch.Tensor): A [N, L, H, Ci] dense tensor containing Qs.
          k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
          v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
      """
      ...
  
  def sparse_scaled_dot_product_attention(*args, **kwargs):
      arg_names_dict = {
          1: ['qkv'],
          2: ['q', 'kv'],
          3: ['q', 'k', 'v']
      }
      num_all_args = len(args) + len(kwargs)
      assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
      for key in arg_names_dict[num_all_args][len(args):]:
          assert key in kwargs, f"Missing argument {key}"
  
      if num_all_args == 1:
          qkv = args[0] if len(args) > 0 else kwargs['qkv']
          assert isinstance(qkv, SparseTensor), f"qkv must be a SparseTensor, got {type(qkv)}"
          assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
          device = qkv.device
  
          s = qkv
          q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
          kv_seqlen = q_seqlen
          qkv = qkv.feats     # [T, 3, H, C]
  
      elif num_all_args == 2:
          q = args[0] if len(args) > 0 else kwargs['q']
          kv = args[1] if len(args) > 1 else kwargs['kv']
          assert isinstance(q, SparseTensor) and isinstance(kv, (SparseTensor, torch.Tensor)) or \
                 isinstance(q, torch.Tensor) and isinstance(kv, SparseTensor), \
                 f"Invalid types, got {type(q)} and {type(kv)}"
          assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
          device = q.device
  
          if isinstance(q, SparseTensor):
              assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
              s = q
              q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
              q = q.feats     # [T_Q, H, C]
          else:
              assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
              s = None
              N, L, H, C = q.shape
              q_seqlen = [L] * N
              q = q.reshape(N * L, H, C)   # [T_Q, H, C]
  
          if isinstance(kv, SparseTensor):
              assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
              kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
              kv = kv.feats     # [T_KV, 2, H, C]
          else:
              assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
              N, L, _, H, C = kv.shape
              kv_seqlen = [L] * N
              kv = kv.reshape(N * L, 2, H, C)   # [T_KV, 2, H, C]
  
      elif num_all_args == 3:
          q = args[0] if len(args) > 0 else kwargs['q']
          k = args[1] if len(args) > 1 else kwargs['k']
          v = args[2] if len(args) > 2 else kwargs['v']
          assert isinstance(q, SparseTensor) and isinstance(k, (SparseTensor, torch.Tensor)) and type(k) == type(v) or \
                 isinstance(q, torch.Tensor) and isinstance(k, SparseTensor) and isinstance(v, SparseTensor), \
                 f"Invalid types, got {type(q)}, {type(k)}, and {type(v)}"
          assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
          device = q.device
  
          if isinstance(q, SparseTensor):
              assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, Ci]"
              s = q
              q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
              q = q.feats     # [T_Q, H, Ci]
          else:
              assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
              s = None
              N, L, H, CI = q.shape
              q_seqlen = [L] * N
              q = q.reshape(N * L, H, CI)  # [T_Q, H, Ci]
  
          if isinstance(k, SparseTensor):
              assert len(k.shape) == 3, f"Invalid shape for k, got {k.shape}, expected [N, *, H, Ci]"
              assert len(v.shape) == 3, f"Invalid shape for v, got {v.shape}, expected [N, *, H, Co]"
              kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
              k = k.feats     # [T_KV, H, Ci]
              v = v.feats     # [T_KV, H, Co]
          else:
              assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
              assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
              N, L, H, CI, CO = *k.shape, v.shape[-1]
              kv_seqlen = [L] * N
              k = k.reshape(N * L, H, CI)     # [T_KV, H, Ci]
              v = v.reshape(N * L, H, CO)     # [T_KV, H, Co]
  
      if DEBUG:
          if s is not None:
              for i in range(s.shape[0]):
                  assert (s.coords[s.layout[i]] == i).all(), f"SparseScaledDotProductSelfAttention: batch index mismatch"
          if num_all_args in [2, 3]:
              assert q.shape[:2] == [1, sum(q_seqlen)], f"SparseScaledDotProductSelfAttention: q shape mismatch"
          if num_all_args == 3:
              assert k.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: k shape mismatch"
              assert v.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: v shape mismatch"
  
      if ATTN == 'xformers':
          if num_all_args == 1:
              q, k, v = qkv.unbind(dim=1)
          elif num_all_args == 2:
              k, v = kv.unbind(dim=1)
          q = q.unsqueeze(0)
          k = k.unsqueeze(0)
          v = v.unsqueeze(0)
          mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
          out = xops.memory_efficient_attention(q, k, v, mask)[0]
      elif ATTN == 'flash_attn':
          cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
          if num_all_args in [2, 3]:
              cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
          if num_all_args == 1:
              out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
          elif num_all_args == 2:
              out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
          elif num_all_args == 3:
              out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
      else:
          raise ValueError(f"Unknown attention module: {ATTN}")
      
      if s is not None:
          return s.replace(out)
      else:
          return out.reshape(N, L, H, -1)