fix unpacking error (#507)

key_value_memory_dict in MHA module is a tuple of kv_cache and conv1d
state-spaces · Aug 6, 2024 · a07ff1b · a07ff1b
1 parent 49ddf83
commit a07ff1b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mamba_ssm/modules/mha.py b/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