Update modeling_mmMamba_embedding.py

modeling_mmMamba_embedding.py

@@ -410,7 +410,7 @@ class MHA_LM(nn.Module):
         ):
             if self.rotary_emb_dim > 0:
                 q, kv = self.rotary_emb(
-                    q, kv, seqlen_offset=seqlen_offset, max_seqlen=rotary_max_seqlen
+                    q, kv, seqlen_offset=seqlen_offset[:bsz,...], max_seqlen=rotary_max_seqlen
                 )
             if inference_params is None:
                 k, v = kv.unbind(dim=-3)
@@ -538,7 +538,9 @@ class Mamba2_LM(nn.Module):
         conv_state, ssm_state = None, None
         if inference_params is not None:
             conv_state, ssm_state = self._get_states_from_cache(inference_params, batch)
-
+        conv_state = conv_state[:batch, ...]
+        ssm_state = ssm_state[:batch, ...]
+        
         if use_cache and inference_params.seqlen_offset==0:
             vkq, new_conv_states = causal_conv1d_fn(
                 vkq.transpose(1, 2),