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: Saves the generated key and value tensors to the end of the buffers in inference_params.
  Returns the full size keys and values from the provided inference_params, as well as
  adjusted rotary_pos_emb.
  
  Returns a tuple: (key, value, rotary_pos_emb)
docstr_tok: ["Saves","the","generated","key","and","value","tensors","to","the","end","of","the","buffers","in","inference_params",".","Returns","the","full","size","keys","and","values","from","the","provided","inference_params",",","as","well","as","adjusted","rotary_pos_emb",".","Returns","a","tuple",":","(","key",",","value",",","rotary_pos_emb",")"]
code: def _adjust_key_value_for_inference(
      self, inference_params, key, value, rotary_pos_emb
  ):
      """
      Saves the generated key and value tensors to the end of the buffers in inference_params.
      Returns the full size keys and values from the provided inference_params, as well as
      adjusted rotary_pos_emb.
  
      Returns a tuple: (key, value, rotary_pos_emb)
  
      """
      attn_mask_type = self.attn_mask_type
      if inference_params is None:
          return key, value, rotary_pos_emb, attn_mask_type
  
      # =================================================
      # Pre-allocate memory for key-values for inference.
      # =================================================
      is_first_step = False
      if (
          self.layer_number
          not in inference_params.key_value_memory_dict
      ):
          inf_max_seq_length = inference_params.max_sequence_length
          inf_max_batch_size = inference_params.max_batch_size
          inference_key_memory = self._allocate_memory(
              inf_max_seq_length, inf_max_batch_size, key.dtype
          )
          inference_value_memory = self._allocate_memory(
              inf_max_seq_length, inf_max_batch_size, value.dtype
          )
          inference_params.key_value_memory_dict[self.layer_number] = (
              inference_key_memory,
              inference_value_memory,
          )
          is_first_step = True
      else:
          # Get the pre-allocated buffers for this layer
          (
              inference_key_memory,
              inference_value_memory,
          ) = inference_params.key_value_memory_dict[self.layer_number]
          attn_mask_type = AttnMaskType.no_mask
  
      batch_start = inference_params.batch_size_offset
      batch_end = batch_start + key.size(1)
      assert batch_end <= inference_key_memory.size(1)
      sequence_start = inference_params.sequence_len_offset
      sequence_end = sequence_start + key.size(0)
      assert sequence_end <= inference_key_memory.size(0)
      # Copy key and values.
      inference_key_memory[
          sequence_start:sequence_end, batch_start:batch_end, ...
      ] = key
      inference_value_memory[
          sequence_start:sequence_end, batch_start:batch_end, ...
      ] = value
      key = inference_key_memory[
          :sequence_end, batch_start:batch_end, ...
      ]
      value = inference_value_memory[
          :sequence_end, batch_start:batch_end, ...
      ]
  
      # adjust the key rotary positional embedding
      if rotary_pos_emb is not None:
          q_pos_emb, k_pos_emb = rotary_pos_emb
          # need to cross check this condition during inference
          # if not set_inference_key_value_memory:
          if not is_first_step:
              # In inference, we compute one token at a time.
              # Select the correct positional embedding
              # (only the last token in the sequence)
              q_pos_emb = q_pos_emb[sequence_end - 1 : sequence_end]
          else:
              # In the first forward pass of inference,
              # we use the entire provided prefix.
              # q_pos_emb here has the rope embeddings of the entire
              # prefix + to-be-generated output so
              # we slice to just the prefix.
              q_pos_emb = q_pos_emb[:sequence_end, :, :, :]
          k_pos_emb = k_pos_emb[:sequence_end, :, :, :]
          rotary_pos_emb = (q_pos_emb, k_pos_emb)
  
      return key, value, rotary_pos_emb, attn_mask_type
  
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 4 for calling the Attention obj's underlying member method code with example usage: obj._adjust_key_value_for_inference(inference_params, key, value, rotary_pos_emb) and returns: key,  value,  rotary_pos_emb,  attn_mask_type, key,  value,  rotary_pos_emb,  attn_mask_type