Skip to content

Commit

Permalink
Bugfix, the dest index didn't run on the seq ids
Browse files Browse the repository at this point in the history
  • Loading branch information
mzusman committed Jul 24, 2024
1 parent bda9876 commit f1e792d
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions vllm/model_executor/models/jamba.py
Original file line number Diff line number Diff line change
Expand Up @@ -760,16 +760,18 @@ def _prepare_current_run_mamba_cache(
self, request_ids_to_seq_ids: Dict[str, list[int]],
batch_size: int, finished_requests_ids: List[str]):
running_indices = []
for dest_index, (request_id,
seqs_id) in enumerate(request_ids_to_seq_ids.items()):
dest_index = 0
for (request_id, seqs_id) in request_ids_to_seq_ids.items():
if request_id in finished_requests_ids:
# Do not allocate cache index for requests that run
# and finish right after
dest_index += 1
continue
for seq_id in seqs_id:
self._assign_seq_id_to_mamba_cache_in_specific_dest(
request_id, seq_id, dest_index)
running_indices.append(dest_index)
dest_index += 1

self._clean_up_first_bs_blocks(batch_size, running_indices)
conv_state = self.mamba_cache[0][:, :batch_size]
Expand Down

0 comments on commit f1e792d

Please sign in to comment.