code: def get_kv_block_size(head_size: int) -> int:
      """
      Return preferred granulatity for blocked KV-cache implementation.
      """
      cc_major, cc_minor = torch.cuda.get_device_capability(
          get_accelerator().current_device()
      )  # ignore-cuda
  
      if cc_major < 8:
          raise RuntimeError(
              "Blocked attention requires CUDA compute capability >= 8.0"
          )
  
      if head_size <= 64:
          return 128
      elif head_size != 160 or cc_minor != 0:
          return 64
      else:
          return 32
  
file_path: deepspeed/deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.py
imports: import torch
  from deepspeed.accelerator import get_accelerator
  from ....inference_utils import DtypeEnum
  from deepspeed.ops.op_builder import RaggedOpsBuilder
  from ...None import DSKernelBase
label_desc: Use image node_id 2 for calling a global function with example usage: get_kv_block_size(head_size) and returns: int, int, int