diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py
index 963c0cb1b1a4..99e0cbb63789 100644
--- a/tensorrt_llm/_torch/compilation/utils.py
+++ b/tensorrt_llm/_torch/compilation/utils.py
@@ -183,7 +183,9 @@ def inplace_info():
             2: "output_sf"
         },
         "mla_custom_op_inplace": {
-            1: "output"
+            1: "output",
+            2: "dsv4_output",
+            3: "dsv4_output_sf",
         },
         "mla_dsa_attn_inplace": {
             1: "output"