Fixed model_nested_mamba.py

modeling_nested_mamba.py

@@ -16,6 +16,29 @@ try:
 except Exception:
     InferenceParams = None
 
+
+def _new_mamba_inference_params(max_batch_size: int, max_seqlen: int):
+    if InferenceParams is None:
+        return None
+    try:
+        return InferenceParams(
+            max_seqlen=max_seqlen,
+            max_batch_size=max_batch_size,
+        )
+    except TypeError:
+        try:
+            return InferenceParams(
+                max_batch_size=max_batch_size,
+                max_seqlen=max_seqlen,
+            )
+        except TypeError:
+            return InferenceParams(max_seqlen, max_batch_size)
+
+
+def _set_inference_seqlen_offset(inference_params, offset: int) -> None:
+    if inference_params is not None and hasattr(inference_params, "seqlen_offset"):
+        inference_params.seqlen_offset = int(offset)
+
 MAMBA_IMPORT_ERROR: Optional[Exception] = None
 try:
     from mamba_ssm import Mamba, Mamba2
@@ -2052,4 +2075,3 @@ class ForwardBackwardRepairModel(nn.Module):
 
     def forward(self, batch, inference_params=None):
         return self.forward_model(**batch, inference_params=inference_params)
-