Update modeling_auristream.py

modeling_auristream.py

@@ -495,6 +495,8 @@ class CausalSelfAttention(nn.Module):
         q = q.view(B, T, self.n_head, self.head_dim)
         v = v.view(B, T, self.n_head, self.head_dim)
 
+        k_orig = k.clone()
+
         if self.rotary is not None:
             cos, sin = self.rotary(q)
             q = apply_rotary_emb(q, cos, sin)
@@ -529,7 +531,7 @@ class CausalSelfAttention(nn.Module):
 
         # return key and value caches if requested
         if return_kv:
-            return y, k, v
+            return y, k_orig, v
 
         return y