[V1][Scheduler] Apply Mamba alignment before encoder caps (#51603)
Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com> Co-authored-by: Akshat Anand <40275336+cipheraxat@users.noreply.github.com> Co-authored-by: Codex <codex@openai.com>
Jiangyun Zhu committed
Aug 10, 2026 at 16:37 UTC
243c63baf5239f961305e25160c1bf6d34096df4
2 files changed
+105
-20
tests/v1/core/test_scheduler.py
+83
@@ -3612,6 +3612,89 @@ def test_scheduler_no_ec_connector_by_default():
3612
assert scheduler.ec_connector is None
3613
3614
3615
+def test_mamba_align_encoder_cache_cap_makes_progress():
3616
+ """Two individually cacheable images must not deadlock Mamba alignment."""
3617
+ block_size = 768
3618
+ encoder_cache_size = 600
3619
+ scheduler = create_scheduler(
3620
+ max_num_batched_tokens=8192,
3621
+ block_size=block_size,
3622
+ enable_prefix_caching=True,
3623
+ )
3624
+ scheduler.need_mamba_block_aligned_split = True
3625
+ scheduler.max_num_encoder_input_tokens = encoder_cache_size
3626
+ scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=encoder_cache_size)
3627
+
3628
+ first_image_end = 510
3629
+ request = create_requests(
3630
+ num_requests=1,
3631
+ num_tokens=1010,
3632
+ mm_positions=[
3633
+ [
3634
+ PlaceholderRange(offset=0, length=500),
3635
+ PlaceholderRange(offset=first_image_end, length=500),
3636
+ ]
3637
+ ],
3638
+ max_tokens=1,
3639
+ block_size=block_size,
3640
+ req_ids=["req"],
3641
+ )[0]
3642
+ scheduler.add_request(request)
3643
+
3644
+ output = scheduler.schedule()
3645
+ assert output.num_scheduled_tokens[request.request_id] == first_image_end
3646
+ assert output.scheduled_encoder_inputs[request.request_id] == [0]
3647
+
3648
+ scheduler.update_from_output(
3649
+ output,
3650
+ ModelRunnerOutput(
3651
+ req_ids=[request.request_id],
3652
+ req_id_to_index={request.request_id: 0},
3653
+ sampled_token_ids=[[]],
3654
+ logprobs=None,
3655
+ prompt_logprobs_dict={},
3656
+ pooler_output=[],
3657
+ ),
3658
+ )
3659
+
3660
+ output = scheduler.schedule()
3661
+ next_block_boundary = block_size
3662
+ assert output.num_scheduled_tokens[request.request_id] == (
3663
+ next_block_boundary - first_image_end
3664
+ )
3665
+ assert output.scheduled_encoder_inputs[request.request_id] == [1]
3666
+
3667
+
3668
+def test_mamba_align_eagle_schedules_encoder_at_boundary():
3669
+ """EAGLE lookahead at an aligned MM boundary requires encoder cache."""
3670
+ block_size = 512
3671
+ scheduler = create_scheduler(
3672
+ max_num_batched_tokens=700,
3673
+ max_model_len=2048,
3674
+ block_size=block_size,
3675
+ enable_prefix_caching=True,
3676
+ )
3677
+ scheduler.need_mamba_block_aligned_split = True
3678
+ scheduler.use_eagle = True
3679
+ scheduler.max_num_encoder_input_tokens = 2048
3680
+ scheduler.encoder_cache_manager = EncoderCacheManager(cache_size=2048)
3681
+
3682
+ request = create_requests(
3683
+ num_requests=1,
3684
+ num_tokens=1200,
3685
+ mm_positions=[[PlaceholderRange(offset=block_size, length=100)]],
3686
+ max_tokens=1,
3687
+ block_size=block_size,
3688
+ req_ids=["req"],
3689
+ )[0]
3690
+ scheduler.add_request(request)
3691
+
3692
+ output = scheduler.schedule()
3693
+
3694
+ assert output.num_scheduled_tokens[request.request_id] == block_size
3695
+ assert output.scheduled_encoder_inputs[request.request_id] == [0]
3696
+
3697
+
3698
@pytest.mark.parametrize("use_kv_connector", [False, True])
3699
def test_ec_connector_text_only_request(use_kv_connector):
3700
"""Test text-only requests don't allocate encoder cache."""
vllm/v1/core/sched/scheduler.py
+22
-20
@@ -530,6 +530,12 @@ class Scheduler(SchedulerInterface):
530
- self.num_sampled_tokens_per_step,
531
)
532
533
+ # Apply Mamba alignment before encoder caps.
534
+ if self.need_mamba_block_aligned_split:
535
+ num_new_tokens = self._mamba_block_aligned_split(
536
+ request, num_new_tokens
537
+ )
538
+
539
# Schedule encoder inputs.
540
encoder_inputs_to_schedule = None
541
external_load_encoder_input: list[int] = []
@@ -548,11 +554,6 @@ class Scheduler(SchedulerInterface):
554
shift_computed_tokens=1 if self.use_eagle else 0,
555
)
556
551
- if self.need_mamba_block_aligned_split:
552
- num_new_tokens = self._mamba_block_aligned_split(
553
- request, num_new_tokens
554
- )
555
-
557
if num_new_tokens == 0:
558
# The request cannot be scheduled because one of the following
559
# reasons:
@@ -906,6 +907,17 @@ class Scheduler(SchedulerInterface):
907
num_new_tokens = min(num_new_tokens, token_budget)
908
assert num_new_tokens > 0
909
910
+ # Apply Mamba alignment before encoder caps.
911
+ if self.need_mamba_block_aligned_split:
912
+ num_new_tokens = self._mamba_block_aligned_split(
913
+ request,
914
+ num_new_tokens,
915
+ num_new_local_computed_tokens,
916
+ num_external_computed_tokens,
917
+ )
918
+ if num_new_tokens == 0:
919
+ break
920
+
921
# Schedule encoder inputs.
922
if request.has_encoder_inputs:
923
(
@@ -924,17 +936,6 @@ class Scheduler(SchedulerInterface):
936
# The request cannot be scheduled.
937
break
938
927
- # Skip block alignment when setting up async receive (no local work).
928
- if self.need_mamba_block_aligned_split and not load_kv_async:
929
- num_new_tokens = self._mamba_block_aligned_split(
930
- request,
931
- num_new_tokens,
932
- num_new_local_computed_tokens,
933
- num_external_computed_tokens,
934
- )
935
- if num_new_tokens == 0:
936
- break
937
-
939
# During async KV load, no forward pass is run yet.
940
# Allocate speculative lookahead slots later to avoid
941
# mismatching local and remote block counts.
@@ -1501,10 +1502,13 @@ class Scheduler(SchedulerInterface):
1502
mm_hashes_to_schedule = set()
1503
num_embeds_to_schedule = 0
1504
1505
+ encoder_window_end = (
1506
+ num_computed_tokens + num_new_tokens + shift_computed_tokens
1507
+ )
1508
lo, hi = get_mm_features_in_window(
1509
mm_features,
1510
start=num_computed_tokens,
1507
- end=num_computed_tokens + num_new_tokens + shift_computed_tokens,
1511
+ end=encoder_window_end,
1512
)
1513
# For encoder-decoder, all inputs sit at start_pos=0, so lo=0 always.
1514
if self.is_encoder_decoder:
@@ -1586,9 +1590,7 @@ class Scheduler(SchedulerInterface):
1590
# Calculate the number of embeddings to schedule in the current range
1591
# of scheduled encoder placeholder tokens.
1592
start_idx_rel = max(0, num_computed_tokens - start_pos)
1589
- end_idx_rel = min(
1590
- num_encoder_tokens, num_computed_tokens + num_new_tokens - start_pos
1591
- )
1593
+ end_idx_rel = min(num_encoder_tokens, encoder_window_end - start_pos)
1594
curr_embeds_start, curr_embeds_end = (
1595
mm_feature.mm_position.get_embeds_indices_in_range(
1596
start_idx_rel, end_idx_rel