Added Whisper TF Upgrade

Signed-off-by: Asmita Goswami <asmigosw@qti.qualcomm.com>

Author: asmigosw

QEfficient/transformers/models/whisper/modeling_whisper.py

@@ -63,18 +63,23 @@ def forward(
         is_cross_attention: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         """Input shape: Batch x Time x Channel"""
-        bsz, tgt_len, _ = hidden_states.size()
+        # determine input shapes
+        bsz, tgt_len = hidden_states.shape[:-1]
+        q_input_shape = (bsz, tgt_len, -1, self.head_dim)
 
-        # get query proj
-        query_states = self._shape(self.q_proj(hidden_states) * self.scaling, tgt_len, bsz)
+        query_states = self.q_proj(hidden_states) * self.scaling
+        query_states = query_states.view(*q_input_shape)
+        query_states = query_states.transpose(1, 2).contiguous()
 
         if self.is_decoder:
             if is_cross_attention and past_key_value:
                 # cross_attentions
                 key_states_old = past_key_value[self.layer_idx][0]
                 value_states_old = past_key_value[self.layer_idx][1]
-                key_states = self._shape(self.k_proj(key_value_states), -1, bsz)
-                value_states = self._shape(self.v_proj(key_value_states), -1, bsz)
+                key_states = self.k_proj(key_value_states).view(bsz, -1, self.num_heads, self.head_dim)
+                value_states = self.v_proj(key_value_states).view(bsz, -1, self.num_heads, self.head_dim)
+                key_states = key_states.transpose(1, 2).contiguous()
+                value_states = value_states.transpose(1, 2).contiguous()
                 indices = (torch.arange(bsz),)
                 key_states_new = torch.index_put(key_states_old, indices, key_states)
                 value_states_new = torch.index_put(value_states_old, indices, value_states)
@@ -85,21 +90,25 @@ def forward(
                     input_features.shape[2] == torch.tensor(1), value_states_old, value_states_new
                 )
 
-                past_key_value.key_cache[self.layer_idx] = key_states
-                past_key_value.value_cache[self.layer_idx] = value_states
+                past_key_value.layers[self.layer_idx].keys = key_states
+                past_key_value.layers[self.layer_idx].values = value_states
             else:
                 # self attention decoder
-                key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
-                value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+                key_states = self.k_proj(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
+                value_states = self.v_proj(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
+                key_states = key_states.transpose(1, 2).contiguous()
+                value_states = value_states.transpose(1, 2).contiguous()
                 if past_key_value is not None:
                     cache_kwargs = {"position_ids": position_ids_layer}
                     key_states, value_states = past_key_value.update(
                         key_states, value_states, self.layer_idx, cache_kwargs
                     )
         else:
             # self_attention Encoder
-            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
-            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+            key_states = self.k_proj(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
+            value_states = self.v_proj(hidden_states).view(bsz, -1, self.num_heads, self.head_dim)
+            key_states = key_states.transpose(1, 2).contiguous()
+            value_states = value_states.transpose(1, 2).contiguous()
 
         src_len = key_states.size(2)
 
@@ -150,7 +159,7 @@ def forward(
 
         attn_output = self.out_proj(attn_output)
 
-        return [attn_output, attn_weights, past_key_value]
+        return [attn_output, attn_weights]
 
 
 class QEffWhisperDecoderLayer(WhisperDecoderLayer):
@@ -203,7 +212,7 @@ def forward(
 
         # Self Attention
         self_attn_past_key_value = past_key_value.self_attention_cache if past_key_value is not None else None
-        hidden_states, self_attn_weights, self_attn_present_key_value = self.self_attn(
+        hidden_states, self_attn_weights = self.self_attn(
             hidden_states=hidden_states,
             past_key_value=self_attn_past_key_value,
             attention_mask=attention_mask,
@@ -217,13 +226,12 @@ def forward(
         hidden_states = residual + hidden_states
 
         # Cross-Attention Block
-        cross_attn_present_key_value = None
         cross_attn_weights = None
         if is_encoder_decoder:
             residual = hidden_states
             hidden_states = self.encoder_attn_layer_norm(hidden_states)
             cross_attn_past_key_value = past_key_value.cross_attention_cache if past_key_value is not None else None
-            hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
+            hidden_states, cross_attn_weights = self.encoder_attn(
                 hidden_states=hidden_states,
                 key_value_states=encoder_hidden_states,
                 attention_mask=encoder_attention_mask,
@@ -237,13 +245,6 @@ def forward(
             hidden_states = nn.functional.dropout(hidden_states, p=self.dropout)
             hidden_states = residual + hidden_states
 
-            # update the cached past_key_values accordingly
-            past_key_value.self_attention_cache = self_attn_present_key_value
-            past_key_value.cross_attention_cache = cross_attn_present_key_value
-        else:
-            # if no cross_attention, still need to update self_attn cache
-            past_key_value = self_attn_present_key_value
-
         # Fully Connected
         residual = hidden_states
         hidden_states = self.final_layer_norm(hidden_states)