code: def scaled_dot_product_attention(
      q: torch.Tensor,
      k: torch.Tensor,
      v: torch.Tensor,
      att_mask: Optional[
          Union[AttentionMask, "SparseCS", torch.Tensor]
      ],
      dropout: Optional[torch.nn.Module] = None,
      block_size: int = 128,
  ) -> torch.Tensor:
      autocast_disabled = (
          _has_cpp_library
          and isinstance(att_mask, SparseCS)
          or (att_mask is not None and att_mask.is_sparse)
      )
      seq_len = q.shape[-2]
  
      # switch if:
      #   causal is required but mask is not sparse
      #   fp16 or under amp context
      #   sequence length is divisible by block size
      #   same seq len for K and Q
      switch_to_blocksparse = (
          _is_blocksparse_available
          and (att_mask is not None and not att_mask.is_sparse)
          and (
              isinstance(att_mask, AttentionMask) and att_mask.is_causal
          )
          and (q.dtype == torch.float16 or torch.is_autocast_enabled())
          and not seq_len % block_size
          and q.shape[-2] == k.shape[-2]
      )
  
      if switch_to_blocksparse:
          logger.info(
              "Switching causal attention to Triton blocksparse..."
          )
          return blocksparse_attention(q, k, v, dropout, block_size)
  
      with torch.cuda.amp.autocast(enabled=False) if autocast_disabled else nullcontext():  # type: ignore
          if autocast_disabled:
              q, k, v = q.float(), k.float(), v.float()
  
          att = scaled_query_key_softmax(q, k, att_mask=att_mask)
  
          #  Optional dropout, could be part of the masking in the future
          att = _apply_dropout(att, dropout)
  
          # Get to the predicted values, for all heads
          # y = att @ v  # (N, S, S) x (N, S, hs) -> (N, S, hs)
          y = bmm(att, v)
      return y
  
file_path: xformers/xformers/components/attention/core.py
imports: import logging
  import math
  from contextlib import nullcontext
  from functools import lru_cache
  from typing import Optional, Union
  import torch
  from xformers import _has_cpp_library, _is_triton_available
  from xformers.components.attention.attention_mask import AttentionMask
label_desc: Use image node_id 9 for calling a global function with example usage: scaled_dot_product_attention(q, k, v, att_mask, dropout, block_size) and returns: y, blocksparse_attention