diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 14da7ad9e..9d5a96b85 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -701,12 +701,20 @@ class TboForwardBatchPreparer: "spec_algorithm", "capture_hidden_mode", "padded_static_len", - "mrope_positions", # only used by qwen2-vl, thus not care "split_index", # for split prefill "orig_seq_lens", # only used by qwen-1m, thus not care "return_pooled_hidden_states", ]: output_dict[key] = getattr(batch, key) + + mrope_positions = getattr(batch, "mrope_positions") + if mrope_positions is not None: + output_dict["mrope_positions"] = mrope_positions[ + :, start_token_index:end_token_index + ] + else: + output_dict["mrope_positions"] = None + if not batch.forward_mode.is_target_verify(): assert ( _compute_extend_num_tokens(batch.input_ids, batch.forward_mode) diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index 494d82d80..2ae120686 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -179,6 +179,9 @@ class TboAttnBackend(AttentionBackend): assert ans == child.get_cuda_graph_seq_len_fill_value() return ans + def forward(self, *args, **kwargs): + return self.primary.forward(*args, **kwargs) + def forward_extend(self, *args, **kwargs): return self.primary.forward_extend(*args, **kwargs)