fix unpacking error (#507)

key_value_memory_dict in MHA module is a tuple of kv_cache and conv1d

Author: chiennv2000

mamba_ssm/modules/mha.py

@@ -180,7 +180,8 @@ def _update_kvcache_attention(self, q, kv, inference_params):
             ).transpose(1, 2)
         else:
             batch = q.shape[0]
-            kv_cache = inference_params.key_value_memory_dict[self.layer_idx][:batch]
+            kv_cache, _ = inference_params.key_value_memory_dict[self.layer_idx]
+            kv_cache = kv_cache[:batch]
             cache_seqlens = (
                 inference_params.lengths_per_sample[:batch]
                 if inference_params.lengths_per_sample is not None