diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index bec834965..eaecdc54b 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -715,15 +715,21 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): else: # target_verify or draft_decode seq_positions = batch.spec_info.positions.view(batch_size, -1) - mrope_deltas = [ - ( - torch.tensor([0], dtype=torch.int64) - if mm_inputs[i] is None - else mm_inputs[i].mrope_position_delta.squeeze(0) + # Split text-only and mixed batches here because SpecV2 text-only batches can avoid an extra D2H. + if all(mm_input is None for mm_input in mm_inputs): + mrope_delta_tensor = torch.zeros( + (batch_size, 1), dtype=torch.int64, device=device ) - for i in range(batch_size) - ] - mrope_delta_tensor = torch.stack(mrope_deltas, dim=0).to(device=device) + else: + mrope_deltas = [ + ( + torch.zeros(1, dtype=torch.int64) + if mm_inputs[i] is None + else mm_inputs[i].mrope_position_delta.squeeze(0) + ) + for i in range(batch_size) + ] + mrope_delta_tensor = torch.stack(mrope_deltas, dim=0).to(device=device) next_input_positions = ( (seq_positions + mrope_delta_tensor).flatten().unsqueeze(0).repeat(3, 1) ) diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 1224fbd33..1f348f9b1 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -234,14 +234,12 @@ class EagleVerifyInputV2Mixin: # Set mamba_track_indices for mamba prefix-cache state tracking if get_global_server_args().enable_mamba_extra_buffer(): - batch.mamba_track_indices = torch.tensor( + batch.mamba_track_indices = torch.stack( [ req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx] for req in batch.reqs - ], - dtype=torch.int64, - device=device, - ) + ] + ).to(torch.int64) batch.mamba_track_mask = None batch.mamba_track_seqlens = None