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 11, 07:30 UTC
[Fix] uniform decode batch check (#30747)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Dec 17, 2025 at 03:58 UTC
6e9dbcc50e35af75ec76bf033ee6402697c02609
2 files changed
+121
-8
tests/v1/worker/test_gpu_model_runner.py
+84
@@ -1110,3 +1110,87 @@ def test_hybrid_cache_integration(model_runner, dist_init):
1110
runner._update_states(scheduler_output)
1111
assert _is_req_scheduled(runner, req_id)
1112
assert _is_req_state_block_table_match(runner, req_id)
1113
+
1114
+
1115
+def test_is_uniform_decode() -> None:
1116
+ # Normal
1117
+ assert GPUModelRunner._is_uniform_decode(
1118
+ max_num_scheduled_tokens=1,
1119
+ uniform_decode_query_len=1,
1120
+ num_tokens=16,
1121
+ num_reqs=16,
1122
+ )
1123
+ assert not GPUModelRunner._is_uniform_decode(
1124
+ max_num_scheduled_tokens=2,
1125
+ uniform_decode_query_len=1,
1126
+ num_tokens=16,
1127
+ num_reqs=16,
1128
+ )
1129
+ assert not GPUModelRunner._is_uniform_decode(
1130
+ max_num_scheduled_tokens=1,
1131
+ uniform_decode_query_len=1,
1132
+ num_tokens=16,
1133
+ num_reqs=15,
1134
+ )
1135
+ # Spec decoding
1136
+ assert GPUModelRunner._is_uniform_decode(
1137
+ max_num_scheduled_tokens=5,
1138
+ uniform_decode_query_len=5,
1139
+ num_tokens=30,
1140
+ num_reqs=6,
1141
+ )
1142
+ assert not GPUModelRunner._is_uniform_decode(
1143
+ max_num_scheduled_tokens=5,
1144
+ uniform_decode_query_len=4,
1145
+ num_tokens=30,
1146
+ num_reqs=6,
1147
+ )
1148
+ assert not GPUModelRunner._is_uniform_decode(
1149
+ max_num_scheduled_tokens=5,
1150
+ uniform_decode_query_len=5,
1151
+ num_tokens=30,
1152
+ num_reqs=7,
1153
+ )
1154
+ # Force uniform decode
1155
+ assert GPUModelRunner._is_uniform_decode(
1156
+ max_num_scheduled_tokens=1,
1157
+ uniform_decode_query_len=1,
1158
+ num_tokens=16,
1159
+ num_reqs=16,
1160
+ force_uniform_decode=True,
1161
+ )
1162
+ assert GPUModelRunner._is_uniform_decode(
1163
+ max_num_scheduled_tokens=2,
1164
+ uniform_decode_query_len=1,
1165
+ num_tokens=16,
1166
+ num_reqs=16,
1167
+ force_uniform_decode=True,
1168
+ )
1169
+ assert GPUModelRunner._is_uniform_decode(
1170
+ max_num_scheduled_tokens=1,
1171
+ uniform_decode_query_len=1,
1172
+ num_tokens=16,
1173
+ num_reqs=15,
1174
+ force_uniform_decode=True,
1175
+ )
1176
+ assert not GPUModelRunner._is_uniform_decode(
1177
+ max_num_scheduled_tokens=1,
1178
+ uniform_decode_query_len=1,
1179
+ num_tokens=16,
1180
+ num_reqs=16,
1181
+ force_uniform_decode=False,
1182
+ )
1183
+ assert not GPUModelRunner._is_uniform_decode(
1184
+ max_num_scheduled_tokens=2,
1185
+ uniform_decode_query_len=1,
1186
+ num_tokens=16,
1187
+ num_reqs=16,
1188
+ force_uniform_decode=False,
1189
+ )
1190
+ assert not GPUModelRunner._is_uniform_decode(
1191
+ max_num_scheduled_tokens=1,
1192
+ uniform_decode_query_len=1,
1193
+ num_tokens=16,
1194
+ num_reqs=15,
1195
+ force_uniform_decode=False,
1196
+ )
vllm/v1/worker/gpu_model_runner.py
+37
-8
@@ -2777,6 +2777,27 @@ class GPUModelRunner(
2777
**model_kwargs,
2778
)
2779
2780
+ @staticmethod
2781
+ def _is_uniform_decode(
2782
+ max_num_scheduled_tokens: int,
2783
+ uniform_decode_query_len: int,
2784
+ num_tokens: int,
2785
+ num_reqs: int,
2786
+ force_uniform_decode: bool | None = None,
2787
+ ) -> bool:
2788
+ """
2789
+ Checks if it's a decode batch with same amount scheduled tokens
2790
+ across all requests.
2791
+ """
2792
+ return (
2793
+ (
2794
+ (max_num_scheduled_tokens == uniform_decode_query_len)
2795
+ and (num_tokens == max_num_scheduled_tokens * num_reqs)
2796
+ )
2797
+ if force_uniform_decode is None
2798
+ else force_uniform_decode
2799
+ )
2800
+
2801
def _determine_batch_execution_and_padding(
2802
self,
2803
num_tokens: int,
@@ -2798,14 +2819,12 @@ class GPUModelRunner(
2819
torch.Tensor | None,
2820
CUDAGraphStat | None,
2821
]:
2801
- num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens)
2802
- uniform_decode = (
2803
- (
2804
- (max_num_scheduled_tokens == self.uniform_decode_query_len)
2805
- and (num_tokens_padded == max_num_scheduled_tokens * num_reqs)
2806
- )
2807
- if force_uniform_decode is None
2808
- else force_uniform_decode
2822
+ uniform_decode = self._is_uniform_decode(
2823
+ max_num_scheduled_tokens=max_num_scheduled_tokens,
2824
+ uniform_decode_query_len=self.uniform_decode_query_len,
2825
+ num_tokens=num_tokens,
2826
+ num_reqs=num_reqs,
2827
+ force_uniform_decode=force_uniform_decode,
2828
)
2829
# Encoder-decoder models only support CG for decoder_step > 0 (no enc_output
2830
# is present). Also, chunked-prefill is disabled, so batch are uniform.
@@ -2819,6 +2838,7 @@ class GPUModelRunner(
2838
else force_has_lora
2839
)
2840
2841
+ num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens)
2842
dispatch_cudagraph = (
2843
lambda num_tokens, disable_full: self.cudagraph_dispatcher.dispatch(
2844
num_tokens=num_tokens,
@@ -2834,6 +2854,15 @@ class GPUModelRunner(
2854
num_tokens_padded, use_cascade_attn or has_encoder_output
2855
)
2856
num_tokens_padded = batch_descriptor.num_tokens
2857
+ if self.compilation_config.pass_config.enable_sp:
2858
+ assert (
2859
+ batch_descriptor.num_tokens
2860
+ % self.vllm_config.parallel_config.tensor_parallel_size
2861
+ == 0
2862
+ ), (
2863
+ "Sequence parallelism requires num_tokens to be "
2864
+ "a multiple of tensor parallel size"
2865
+ )
2866
2867
# Extra coordination when running data-parallel since we need to coordinate
2868
# across ranks