Got output to match Phi2 exactly

attention.py

@@ -370,11 +370,29 @@ class MHA(nn.Module):
                 dtype=kv.dtype,
                 device=kv.device,
             )
+
+        batch_start = kv_cache.batch_size_offset
+        batch_end = batch_start + kv.shape[0]
+        sequence_start = kv_cache.seqlen_offset
+        sequence_end = sequence_start + kv.shape[1]
+
+        # TODO: figure out why they're doing this
+        if sequence_end >= kv_cache.max_seqlen:
+            kv_cache.kv_block_map[block_n] = torch.concatenate(
+                (kv_cache.kv_block_map[block_n], kv),
+                dim=1,
+            )
         kv_cache.kv_block_map[block_n][
-            kv_cache.batch_size_offset: kv_cache.batch_size_offset + kv.shape[0],
-            kv_cache.seqlen_offset: kv_cache.seqlen_offset + kv.shape[1],
+            batch_start:batch_end,
+            sequence_start:sequence_end,
             ...
         ] = kv
+        kv = kv_cache.kv_block_map[block_n][
+            batch_start:batch_end,
+            :sequence_end,
+            ...
+        ]
+        return kv
 
     def _forward_cross_attn(
         self,
@@ -396,8 +414,9 @@ class MHA(nn.Module):
             ],
             dim=2,
         )
-        self._update_kv_cache(kv, kv_cache, self.block_n)
-        causal = False  # turning off causal mask for cross attention
+        kv = self._update_kv_cache(kv, kv_cache, self.block_n)
+
+        causal = (kv_cache.seqlen_offset == 0)
 
         if self.using_flash_attn and unpad_input and pad_input:  # not touching flash attention code
             batch_size, seqlen_q = q.shape[0], q.shape[1]
@@ -528,4 +547,5 @@ class ParallelAttentionBlock(nn.Module):
             key_padding_mask=key_padding_mask,
         )
         mlp_outputs = self.mlp(x)
-        return self.dropout(attn_outputs + mlp_outputs) + residual
+        x = self.dropout(attn_outputs + mlp_outputs) + residual
+        return x