add kwargs to _scaled_dot_product_attention__tensorrt (#2332)

Author: AllentDan

mmdeploy/pytorch/functions/multi_head_attention_forward.py

@@ -45,12 +45,11 @@ def symbolic(g, q, k, v, mask):
 @FUNCTION_REWRITER.register_rewriter(
     func_name='torch.nn.functional._scaled_dot_product_attention',
     backend=Backend.TENSORRT.value)
-def _scaled_dot_product_attention__tensorrt(
-    q: Tensor,
-    k: Tensor,
-    v: Tensor,
-    attn_mask: Optional[Tensor] = None,
-    dropout_p: float = 0.0,
-) -> Tuple[Tensor, Tensor]:
+def _scaled_dot_product_attention__tensorrt(q: Tensor,
+                                            k: Tensor,
+                                            v: Tensor,
+                                            attn_mask: Optional[Tensor] = None,
+                                            dropout_p: float = 0.0,
+                                            **kwargs) -> Tuple[Tensor, Tensor]:
     """Rewrite for custom ops."""
     return ScaledDotProductAttentionTRT.apply(q, k, v, attn_mask)