class_docstr: Attention layer abstract class.
  
  This layer only contains common modules required for the "self attn" and
  "cross attn" specializations.
class_docstr_tok: ["Attention","layer","abstract","class",".","This","layer","only","contains","common","modules","required","for","the","``","self","attn","''","and","''","cross","attn","''","specializations","."]
docstr: Allocate memory to store kv cache during inference.
docstr_tok: ["Allocate","memory","to","store","kv","cache","during","inference","."]
code: def _allocate_memory(
      self, inference_max_sequence_length, batch_size, dtype
  ):
      """Allocate memory to store kv cache during inference."""
  
      return torch.empty(
          inference_max_sequence_length,
          batch_size,
          self.num_query_groups_per_partition,
          self.hidden_size_per_attention_head,
          dtype=dtype,
          device=torch.cuda.current_device(),
      )
  
imports: from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import Union
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb
  from megatron.core.transformer.custom_layers.transformer_engine import SplitAlongDim
  from megatron.core.transformer.enums import AttnMaskType
  from megatron.core.transformer.identity_op import IdentityFuncOp, IdentityOp
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.spec_utils import ModuleSpec, build_module
  from megatron.core.transformer.transformer_config import TransformerConfig
  from megatron.core.utils import divide
  from .enums import AttnMaskType
  from .transformer_config import TransformerConfig
  from .utils import make_sharded_tensors_for_checkpoint
label_desc: Use image node_id 3 for calling the Attention obj's underlying member method code with example usage: obj._allocate_memory(inference_max_sequence_length, batch_size, dtype) and returns: torch