[Fix] avoid stream sync in _compute_mrope_positions (#14956)

This commit is contained in:
narutolhy
2025-12-13 18:50:57 -08:00
committed by GitHub
parent a7a4b1755d
commit 69cfb17b9a
@@ -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():