code: def __init__(
      self,
      max_batch_size,
      max_seq_length,
      n_heads,
      head_dim,
      dtype=torch.bfloat16,
  ):
      super().__init__()
      cache_shape = (max_batch_size, n_heads, max_seq_length, head_dim)
      self.register_buffer(
          "k_cache", torch.zeros(cache_shape, dtype=dtype)
      )
      self.register_buffer(
          "v_cache", torch.zeros(cache_shape, dtype=dtype)
      )
  
imports: from dataclasses import dataclass
  from typing import Optional
  import torch
  import torch.nn
  from torch import Tensor
  from torch.nn import functional
label_desc: Use image node_id 1 to create a new KVCache object from inherited base classes: nn with example: obj = KVCache(max_batch_size, max_seq_length, n_heads, head_dim, dtype)