A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 05:00 UTC
[Scheduler][ASR] Fix CrossAttn blocks per-request for Variable length encoder inputs (#31058)
Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Ekagra Ranjan committed
Feb 16, 2026 at 06:08 UTC
cd81cdb399e5ead89ac10eb3f8eff1fa85b427a1
2 files changed
+305
-13
tests/v1/core/test_scheduler.py
+294
@@ -3676,6 +3676,300 @@ def test_abort_request_finished_recving():
3676
assert not scheduler.finished_recving_kv_req_ids
3677
3678
3679
+# ==============================================================================
3680
+# Variable-length encoder cross-attention block allocation tests
3681
+# ==============================================================================
3682
+
3683
+
3684
+def _create_encoder_decoder_scheduler(
3685
+ block_size: int = 16,
3686
+ num_blocks: int = 10000,
3687
+ max_num_batched_tokens: int = 8192,
3688
+ max_num_seqs: int = 16,
3689
+) -> Scheduler:
3690
+ """Create a scheduler configured for encoder-decoder cross-attention
3691
+ block allocation testing.
3692
+
3693
+ Constructs a scheduler with both FullAttentionSpec (self-attention) and
3694
+ CrossAttentionSpec (cross-attention) KV cache groups, then patches it
3695
+ to behave as an encoder-decoder model.
3696
+ """
3697
+ from vllm.v1.core.encoder_cache_manager import EncoderDecoderCacheManager
3698
+ from vllm.v1.kv_cache_interface import CrossAttentionSpec
3699
+
3700
+ model_config = ModelConfig(
3701
+ model="facebook/opt-125m",
3702
+ trust_remote_code=True,
3703
+ dtype="float16",
3704
+ seed=42,
3705
+ )
3706
+ scheduler_config = SchedulerConfig(
3707
+ max_num_seqs=max_num_seqs,
3708
+ max_num_batched_tokens=max_num_batched_tokens,
3709
+ max_model_len=max_num_batched_tokens,
3710
+ # is_encoder_decoder disables chunked prefill and prefix caching
3711
+ is_encoder_decoder=True,
3712
+ )
3713
+ cache_config = CacheConfig(
3714
+ block_size=block_size,
3715
+ gpu_memory_utilization=0.9,
3716
+ swap_space=0,
3717
+ cache_dtype="auto",
3718
+ enable_prefix_caching=False,
3719
+ )
3720
+ cache_config.num_gpu_blocks = num_blocks
3721
+
3722
+ vllm_config = VllmConfig(
3723
+ scheduler_config=scheduler_config,
3724
+ model_config=model_config,
3725
+ cache_config=cache_config,
3726
+ )
3727
+
3728
+ # KV cache config with both self-attention and cross-attention groups,
3729
+ # mirroring an encoder-decoder model like Whisper.
3730
+ kv_cache_config = KVCacheConfig(
3731
+ num_blocks=num_blocks,
3732
+ kv_cache_tensors=[],
3733
+ kv_cache_groups=[
3734
+ KVCacheGroupSpec(
3735
+ ["self_attn_layer"],
3736
+ FullAttentionSpec(
3737
+ block_size=block_size,
3738
+ num_kv_heads=1,
3739
+ head_size=1,
3740
+ dtype=torch.float32,
3741
+ ),
3742
+ ),
3743
+ KVCacheGroupSpec(
3744
+ ["cross_attn_layer"],
3745
+ CrossAttentionSpec(
3746
+ block_size=block_size,
3747
+ num_kv_heads=1,
3748
+ head_size=1,
3749
+ dtype=torch.float32,
3750
+ ),
3751
+ ),
3752
+ ],
3753
+ )
3754
+
3755
+ # Construct the scheduler. Since opt-125m is not truly encoder-decoder,
3756
+ # the __init__ won't set up encoder-decoder internals. We patch them
3757
+ # after construction.
3758
+ scheduler = Scheduler(
3759
+ vllm_config=vllm_config,
3760
+ kv_cache_config=kv_cache_config,
3761
+ block_size=block_size,
3762
+ structured_output_manager=StructuredOutputManager(vllm_config),
3763
+ )
3764
+
3765
+ # Patch to enable encoder-decoder behavior in the scheduling loop.
3766
+ scheduler.is_encoder_decoder = True
3767
+ scheduler.max_num_encoder_input_tokens = max_num_batched_tokens
3768
+ scheduler.encoder_cache_manager = EncoderDecoderCacheManager(
3769
+ cache_size=max_num_batched_tokens
3770
+ )
3771
+
3772
+ return scheduler
3773
+
3774
+
3775
+def _get_num_cross_attn_blocks(scheduler: Scheduler, request_id: str) -> int:
3776
+ """Get the number of cross-attention blocks allocated for a request."""
3777
+ from vllm.v1.core.single_type_kv_cache_manager import CrossAttentionManager
3778
+
3779
+ coordinator = scheduler.kv_cache_manager.coordinator
3780
+ for manager in coordinator.single_type_managers:
3781
+ if isinstance(manager, CrossAttentionManager):
3782
+ blocks = manager.req_to_blocks.get(request_id, [])
3783
+ return len(blocks)
3784
+ raise AssertionError("No CrossAttentionManager found in coordinator")
3785
+
3786
+
3787
+def test_variable_length_cross_attn_block_allocation():
3788
+ """Test that cross-attention blocks are allocated per-request based on
3789
+ actual encoder input length, not a fixed maximum.
3790
+
3791
+ Fixed max-encoder-length allocation would assign
3792
+ `ceil(max_encoder_tokens / block_size)` blocks to
3793
+ every request whereas with dynamic allocation, exactly
3794
+ `ceil(actual_encoder_tokens / block_size)` blocks are assigned
3795
+ to each request.
3796
+ """
3797
+ block_size = 16
3798
+ scheduler = _create_encoder_decoder_scheduler(block_size=block_size)
3799
+
3800
+ # Create requests with distinctly different encoder input lengths,
3801
+ # simulating variable-length audio inputs to a model like Whisper.
3802
+ encoder_lengths = [500, 1000, 200]
3803
+ num_prompt_tokens = 100 # Decoder prompt tokens
3804
+
3805
+ requests = []
3806
+ for i, enc_len in enumerate(encoder_lengths):
3807
+ req = create_requests(
3808
+ num_requests=1,
3809
+ num_tokens=num_prompt_tokens,
3810
+ mm_hashes_list=[[f"enc_hash_{i}"]],
3811
+ mm_positions=[[PlaceholderRange(offset=0, length=enc_len)]],
3812
+ req_ids=[f"req_{i}"],
3813
+ )[0]
3814
+ requests.append(req)
3815
+
3816
+ # Add and schedule all requests.
3817
+ for req in requests:
3818
+ scheduler.add_request(req)
3819
+
3820
+ output = scheduler.schedule()
3821
+
3822
+ # All requests should be scheduled.
3823
+ assert len(output.scheduled_new_reqs) == len(requests)
3824
+
3825
+ # Verify cross-attention blocks per request match the actual encoder length.
3826
+ from math import ceil
3827
+
3828
+ for req, enc_len in zip(requests, encoder_lengths):
3829
+ expected_blocks = ceil(enc_len / block_size)
3830
+ actual_blocks = _get_num_cross_attn_blocks(scheduler, req.request_id)
3831
+
3832
+ assert actual_blocks == expected_blocks, (
3833
+ f"Request {req.request_id} with {enc_len} encoder tokens: "
3834
+ f"expected {expected_blocks} cross-attn blocks, "
3835
+ f"got {actual_blocks}"
3836
+ )
3837
+
3838
+ # Verify that different encoder lengths produce different block counts,
3839
+ # confirming variable-length (not fixed-max) allocation.
3840
+ block_counts = [
3841
+ _get_num_cross_attn_blocks(scheduler, req.request_id) for req in requests
3842
+ ]
3843
+ assert len(set(block_counts)) > 1, (
3844
+ "All requests have the same number of cross-attn blocks, "
3845
+ "suggesting static max-based allocation instead of per-request"
3846
+ )
3847
+
3848
+
3849
+def test_cross_attn_blocks_not_over_allocated():
3850
+ """Test that cross-attention blocks are not over-allocated compared to
3851
+ what each request actually needs."""
3852
+ from math import ceil
3853
+
3854
+ block_size = 16
3855
+ max_encoder_tokens = 1500 # e.g., Whisper's max mel-spectrogram length
3856
+ scheduler = _create_encoder_decoder_scheduler(block_size=block_size)
3857
+
3858
+ # Request with a small encoder input (much less than the max).
3859
+ small_enc_len = 200
3860
+ request = create_requests(
3861
+ num_requests=1,
3862
+ num_tokens=100,
3863
+ mm_hashes_list=[["enc_small"]],
3864
+ mm_positions=[[PlaceholderRange(offset=0, length=small_enc_len)]],
3865
+ req_ids=["req_small"],
3866
+ )[0]
3867
+
3868
+ scheduler.add_request(request)
3869
+ output = scheduler.schedule()
3870
+
3871
+ assert len(output.scheduled_new_reqs) == 1
3872
+
3873
+ actual_blocks = _get_num_cross_attn_blocks(scheduler, request.request_id)
3874
+ expected_blocks = ceil(small_enc_len / block_size)
3875
+ max_blocks = ceil(max_encoder_tokens / block_size)
3876
+
3877
+ # Blocks should match the actual encoder length.
3878
+ assert actual_blocks == expected_blocks, (
3879
+ f"Expected {expected_blocks} blocks for {small_enc_len} encoder tokens, "
3880
+ f"got {actual_blocks}"
3881
+ )
3882
+
3883
+ # Blocks should be strictly less than what max-based allocation would give.
3884
+ assert actual_blocks < max_blocks, (
3885
+ f"Cross-attn blocks ({actual_blocks}) should be less than max "
3886
+ f"({max_blocks}), indicating no over-allocation"
3887
+ )
3888
+
3889
+
3890
+def test_cross_attn_blocks_not_under_allocated():
3891
+ """Test that cross-attention blocks are sufficient for each request's
3892
+ actual encoder input length. Every encoder token must have a slot.
3893
+
3894
+ Tests various edge cases including exact block boundaries, off-by-one,
3895
+ and the minimum/maximum encoder input sizes.
3896
+ """
3897
+ from math import ceil
3898
+
3899
+ block_size = 16
3900
+
3901
+ # Test various encoder lengths including edge cases around block boundaries.
3902
+ test_cases = [
3903
+ 1, # Minimum: single encoder token
3904
+ block_size - 1, # Just under one full block
3905
+ block_size, # Exactly one full block
3906
+ block_size + 1, # Just over one block (needs 2 blocks)
3907
+ block_size * 10, # Exact multiple of block size
3908
+ block_size * 10 + 1, # One over exact multiple
3909
+ 1500, # Whisper's typical max
3910
+ ]
3911
+
3912
+ for enc_len in test_cases:
3913
+ scheduler = _create_encoder_decoder_scheduler(block_size=block_size)
3914
+
3915
+ request = create_requests(
3916
+ num_requests=1,
3917
+ num_tokens=100,
3918
+ mm_hashes_list=[[f"enc_{enc_len}"]],
3919
+ mm_positions=[[PlaceholderRange(offset=0, length=enc_len)]],
3920
+ req_ids=[f"req_{enc_len}"],
3921
+ )[0]
3922
+
3923
+ scheduler.add_request(request)
3924
+ output = scheduler.schedule()
3925
+
3926
+ assert len(output.scheduled_new_reqs) == 1
3927
+
3928
+ actual_blocks = _get_num_cross_attn_blocks(scheduler, request.request_id)
3929
+ expected_blocks = ceil(enc_len / block_size)
3930
+
3931
+ # Number of blocks must be exactly ceil(enc_len / block_size).
3932
+ assert actual_blocks == expected_blocks, (
3933
+ f"Encoder length {enc_len}: expected {expected_blocks} blocks, "
3934
+ f"got {actual_blocks}"
3935
+ )
3936
+
3937
+ # Total available slots must be >= encoder tokens (no under-allocation).
3938
+ total_slots = actual_blocks * block_size
3939
+ assert total_slots >= enc_len, (
3940
+ f"Encoder length {enc_len}: total slots {total_slots} < "
3941
+ f"needed {enc_len} (under-allocation)"
3942
+ )
3943
+
3944
+
3945
+def test_cross_attn_zero_blocks_without_encoder_inputs():
3946
+ """Test that requests without encoder inputs get zero cross-attention
3947
+ blocks, even when the scheduler is configured for encoder-decoder."""
3948
+ block_size = 16
3949
+ scheduler = _create_encoder_decoder_scheduler(block_size=block_size)
3950
+
3951
+ # Create a text-only request (no mm_features).
3952
+ request = create_requests(
3953
+ num_requests=1,
3954
+ num_tokens=100,
3955
+ req_ids=["req_text_only"],
3956
+ )[0]
3957
+
3958
+ # Text-only request has no encoder inputs.
3959
+ assert not request.has_encoder_inputs
3960
+
3961
+ scheduler.add_request(request)
3962
+ output = scheduler.schedule()
3963
+
3964
+ assert len(output.scheduled_new_reqs) == 1
3965
+
3966
+ # No cross-attention blocks should be allocated.
3967
+ actual_blocks = _get_num_cross_attn_blocks(scheduler, request.request_id)
3968
+ assert actual_blocks == 0, (
3969
+ f"Text-only request should have 0 cross-attn blocks, got {actual_blocks}"
3970
+ )
3971
+
3972
+
3973
def test_eagle3_mm_encoder_cache_with_shift():
3974
"""Test EAGLE3 encoder scheduling accounts for shift_computed_tokens.
3975
vllm/v1/core/sched/scheduler.py
+11
-13
@@ -201,14 +201,6 @@ class Scheduler(SchedulerInterface):
201
if self.is_encoder_decoder
202
else EncoderCacheManager(cache_size=encoder_cache_size)
203
)
204
- # For encoder-decoder models, allocate the maximum number of tokens for Cross
205
- # Attn blocks, as for Whisper its input is always padded to the maximum length.
206
- # TODO (NickLucche): Generalize to models with variable-length encoder inputs.
207
- self._num_encoder_max_input_tokens = (
208
- mm_budget.mm_max_toks_per_item[mm_budget.get_modality_with_max_tokens()]
209
- if mm_budget and mm_budget.mm_max_toks_per_item
210
- else 0
211
- )
204
205
speculative_config = vllm_config.speculative_config
206
self.use_eagle = False
@@ -715,11 +707,17 @@ class Scheduler(SchedulerInterface):
707
0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens
708
)
709
718
- num_encoder_tokens = (
719
- self._num_encoder_max_input_tokens
720
- if self.is_encoder_decoder and request.has_encoder_inputs
721
- else 0
722
- )
710
+ # Determine if we need to allocate cross-attention blocks.
711
+ num_encoder_tokens = 0
712
+ if (
713
+ self.is_encoder_decoder
714
+ and request.has_encoder_inputs
715
+ and encoder_inputs_to_schedule
716
+ ):
717
+ num_encoder_tokens = sum(
718
+ request.get_num_encoder_embeds(i)
719
+ for i in encoder_inputs_to_schedule
720
+ )
721
722
new_blocks = self.kv_cache_manager.allocate_slots(
723
request,