fix: residual is kept in kwargs

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

mha.py

@@ -649,7 +649,8 @@ class MHA(nn.Module):
             if not self.return_residual:
                 qkv = self.Wqkv(x, **lora_kwargs)
             else:
-                lora_kwargs['residual'] = True
+                if lora_kwargs:
+                    lora_kwargs['residual'] = True
                 qkv, x = self.Wqkv(x, **lora_kwargs)
 
             if self.dwconv:
@@ -737,5 +738,6 @@ class MHA(nn.Module):
             else:
                 context = self._apply_rotary_update_kvcache_attention(q, kv, inference_params)
 
+        lora_kwargs.pop('residual', None)
         out = self.out_proj(rearrange(context, "... h d -> ... (h d)"), **lora_kwargs)
         return out if not self.return_residual else (out, x)