[Session + MM] Fix text positions in session continuations (#39144)

Co-authored-by: Byron Hsu <byron+per@periodiclabs.ai>
Co-authored-by: Manik Singhal <3400497+Manikvsin@users.noreply.github.com>
This commit is contained in:
Byron Hsu
2026-09-12 10:31:55 -07:00
committed by GitHub
co-authored by Byron Hsu Manik Singhal
parent ae1acf822d
commit 25a5641cf2
@@ -1330,7 +1330,21 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
:,
extend_prefix_len : extend_prefix_len + extend_seq_len,
]
if mrope_positions.numel() == 0:
if (
batch.reqs[batch_idx].session is not None
and mrope_positions.shape[1] < extend_seq_len
):
# Session history includes generated and appended text that
# is not covered by the saved prompt positions.
tail_len = extend_seq_len - mrope_positions.shape[1]
tail_start = extend_prefix_len + mrope_positions.shape[1]
text_positions = self._expand_mrope_from_input(
mm_input, tail_start + 1
) + torch.arange(tail_len)
mrope_positions = torch.cat(
[mrope_positions, text_positions], dim=1
)
elif mrope_positions.numel() == 0:
mrope_positions = self._expand_mrope_from_input(
mm_input, seq_lens_cpu[batch_idx]
)