def get_query_key_value_tensors(self, hidden_states, key_value_states):
        """
        This method needs to be implemented based on whether the derived class
        is "self-attn" or "cross-attn".
        """

        This method needs to be implemented based on whether the derived class
        is "self-attn" or "cross-attn".
        
training/DeepSpeed-Domino/megatron/core/transformer/attention.py