query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)  # [B, 32, T, 128]
key_states   = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)  # [B, 8, T, 128]
value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)  # [B, 8, T, 128]
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
if past_key_values is not None:
    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
attn_output, attn_weights = attention_interface(self, query_states, key_states, value_states,
                                                attention_mask, scaling=self.scaling)
attn_output = self.o_proj(attn_output.reshape(*input_shape, -1))