[Fix] avoid stream sync in _compute_mrope_positions (#14956)
This commit is contained in:
@@ -664,7 +664,7 @@ class ForwardBatch:
|
|||||||
self, model_runner: ModelRunner, batch: ModelWorkerBatch
|
self, model_runner: ModelRunner, batch: ModelWorkerBatch
|
||||||
):
|
):
|
||||||
# batch_size * [3 * seq_len]
|
# batch_size * [3 * seq_len]
|
||||||
batch_size = self.seq_lens.shape[0]
|
batch_size = self.seq_lens_cpu.shape[0]
|
||||||
mrope_positions_list = [[]] * batch_size
|
mrope_positions_list = [[]] * batch_size
|
||||||
for batch_idx in range(batch_size):
|
for batch_idx in range(batch_size):
|
||||||
mm_input = batch.multimodal_inputs[batch_idx]
|
mm_input = batch.multimodal_inputs[batch_idx]
|
||||||
@@ -673,13 +673,13 @@ class ForwardBatch:
|
|||||||
if mm_input is None:
|
if mm_input is None:
|
||||||
mrope_positions_list[batch_idx] = torch.full(
|
mrope_positions_list[batch_idx] = torch.full(
|
||||||
(3, 1),
|
(3, 1),
|
||||||
self.seq_lens[batch_idx] - 1,
|
self.seq_lens_cpu[batch_idx] - 1,
|
||||||
dtype=torch.int64,
|
dtype=torch.int64,
|
||||||
device=model_runner.device,
|
device=model_runner.device,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
mrope_positions = self._expand_mrope_from_input(
|
mrope_positions = self._expand_mrope_from_input(
|
||||||
mm_input, self.seq_lens[batch_idx], model_runner.device
|
mm_input, self.seq_lens_cpu[batch_idx], model_runner.device
|
||||||
)
|
)
|
||||||
mrope_positions_list[batch_idx] = mrope_positions
|
mrope_positions_list[batch_idx] = mrope_positions
|
||||||
elif self.forward_mode.is_extend():
|
elif self.forward_mode.is_extend():
|
||||||
|
|||||||
Reference in New Issue
Block a user