code: def _get_block_size_n(device, head_dim, is_dropout, is_causal):
      # This should match the block sizes in the CUDA kernel
      assert head_dim <= 256
      major, minor = torch.cuda.get_device_capability(device)
      is_sm8x = (
          major == 8 and minor > 0
      )  # Only include sm86 and sm89, exclude sm80 (A100)
      is_sm80 = major == 8 and minor == 0
      is_sm90 = major == 9 and minor == 0
      if head_dim <= 32:
          return 128
      if head_dim <= 64:
          return 128 if not is_dropout else 64
      elif head_dim <= 96:
          return 64
      elif head_dim <= 128:
          if is_sm8x:
              return 64 if (not is_dropout and is_causal) else 32
          else:
              return 64 if not is_dropout else 32
      elif head_dim <= 160:
          if is_sm8x:
              return 64
          else:
              return 32
      elif head_dim <= 192:
          return 64
      elif head_dim <= 224:
          return 64
      elif head_dim <= 256:
          return 64
  
imports: from typing import Optional, Union
  import torch
  import torch.nn
  import flash_attn_2_cuda
label_desc: Use image node_id 1 for calling a global function with example usage: _get_block_size_n(device, head_dim, is_dropout, is_causal) and returns: int, unknown, int, unknown, unknown, int, int, int, int, int