def get_qkv(
        self,
        x: torch.Tensor,
        prev_layer_key_value: Optional[tuple[torch.Tensor,
                                             torch.Tensor]] = None,
        key_value_states: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Computes and returns the query, key, and value tensors.

        Args:
            x (torch.Tensor): The input query tensor.
            prev_layer_key_value  (Optional[Tuple[torch.Tensor, torch.Tensor]]): The key value of the previous layer.
            key_value_states (Optional[torch.Tensor]): The input tensor for keys and values.

        Returns:
            query (torch.Tensor): The query tensor.
            key (torch.Tensor): The key tensor.
            value (torch.Tensor): The value tensor.
        """
        if self.reuse_kv_layer_idx is not None:
            if prev_layer_key_value is None:
                raise ValueError(
                    'prev_layer_key_value is None, cannot reuse_prev_layer_kv.',
                )
            key, value = prev_layer_key_value
            if self.attn_impl == 'torch':
                key = rearrange(key, 'b h d s -> b s (h d)')
                value = rearrange(value, 'b h s d -> b s (h d)')

            query = self.Wq(x)
            if self.clip_qkv:
                query = query.clamp(min=-self.clip_qkv, max=self.clip_qkv)

            if self.qk_ln or self.qk_gn:
                # Applying layernorm to qk
                q_shape = query.shape
                if self.qk_gn:
                    b, s = query.shape[:2]
                    query = query.view(b, s, self.n_heads, -1)
                dtype = query.dtype
                query = self.q_ln(query).to(dtype).view(q_shape)
            return query, key, value

        if self.fused_qkv:
            if key_value_states is not None:
                raise ValueError(
                    'Cannot use separate hidden and key_value states when fused_qkv = True.',
                )
            qkv = self.Wqkv(x)

            if self.clip_qkv:
                qkv = qkv.clamp(min=-self.clip_qkv, max=self.clip_qkv)

            query, key, value = qkv.split(
                [
                    self.d_model,
                    self.kv_n_heads * self.head_dim,
                    self.kv_n_heads * self.head_dim,
                ],
                dim=2,
            )
        else:
            query = self.Wq(x)
            if key_value_states is not None:
                key = self.Wk(key_value_states)
                value = self.Wv(key_value_states)
            else:
                key = self.Wk(x)
                value = self.Wv(x)

            if self.clip_qkv:
                query = query.clamp(min=-self.clip_qkv, max=self.clip_qkv)
                key = key.clamp(min=-self.clip_qkv, max=self.clip_qkv)
                value = value.clamp(min=-self.clip_qkv, max=self.clip_qkv)

        if self.qk_ln or self.qk_gn:
            # Applying layernorm to qk
            q_shape, k_shape = query.shape, key.shape
            if self.qk_gn:
                b, s = query.shape[:2]
                query = query.view(b, s, self.n_heads, -1)
                key = key.view(b, s, self.kv_n_heads, -1)
            dtype = query.dtype
            query = self.q_ln(query).to(dtype).view(q_shape)
            key = self.k_ln(key).to(dtype).view(k_shape)

        return query, key, value
Computes and returns the query, key, and value tensors.

        Args:
            x (torch.Tensor): The input query tensor.
            prev_layer_key_value  (Optional[Tuple[torch.Tensor, torch.Tensor]]): The key value of the previous layer.
            key_value_states (Optional[torch.Tensor]): The input tensor for keys and values.

        Returns:
            query (torch.Tensor): The query tensor.
            key (torch.Tensor): The key tensor.
            value (torch.Tensor): The value tensor.
        
llmfoundry/models/layers/attention.py