[None][feat] AutoDeploy: Perf improvement for mamba layers (NVIDIA#8991)

Signed-off-by: Chenghao Zhang <[email protected]>
Signed-off-by: Suyog Gupta <[email protected]>
Co-authored-by: Suyog Gupta <[email protected]>

Author: nvchenghaoz

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/cuda_backend_causal_conv.py

@@ -204,7 +204,9 @@ def _cuda_cached_causal_conv1d(
 
         if y_dec.dim() == 3:
             y_dec = y_dec.squeeze(-1)
-        y_flat[total_prefill_tokens : total_prefill_tokens + num_decode].copy_(y_dec)
+        y_flat[total_prefill_tokens : total_prefill_tokens + num_decode].copy_(
+            y_dec.to(y_flat.dtype)
+        )
 
     # Custom op must not return an alias of any input; return a fresh tensor
     return y
@@ -296,9 +298,5 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         stride, padding, dilation, groups, padding_mode = extract_op_args(
             source_attn_node, "stride", "padding", "dilation", "groups", "padding_mode"
         )
-        # activation parameter may not exist in the source node (added by fusion later)
-        try:
-            activation = extract_op_args(source_attn_node, "activation")[0]
-        except (RuntimeError, IndexError):
-            activation = None
-        return [stride, padding, dilation, groups, padding_mode, activation]
+        # None is for activation parameter, which may not exist in the source node (added by fusion later)
+        return [stride, padding, dilation, groups, padding_mode, None]

tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/torch_backend_causal_conv.py

@@ -355,9 +355,5 @@ def get_constants(cls, source_attn_node: Node) -> List[Constant]:
         stride, padding, dilation, groups, padding_mode = extract_op_args(
             source_attn_node, "stride", "padding", "dilation", "groups", "padding_mode"
         )
-        # activation parameter may not exist in the source node (added by fusion later)
-        try:
-            activation = extract_op_args(source_attn_node, "activation")[0]
-        except (RuntimeError, IndexError):
-            activation = None
-        return [stride, padding, dilation, groups, padding_mode, activation]
+        # None is for activation parameter, which may not exist in the source node (added by fusion later)
+        return [stride, padding, dilation, groups, padding_mode, None]