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 14, 05:30 UTC
[BugFix] AssertionError: Do not capture num_reqs > max_num_reqs for uniform batch (#25505)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Sep 23, 2025 at 20:00 UTC
dc464a3d3937e30267514e1fc5b988a35dd9dbdf
1 file changed
+16
-16
vllm/v1/worker/gpu_model_runner.py
+16
-16
@@ -2828,7 +2828,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2828
def _dummy_run(
2829
self,
2830
num_tokens: int,
2831
- cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
2831
+ cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
2832
force_attention: bool = False,
2833
uniform_decode: bool = False,
2834
allow_microbatching: bool = True,
@@ -2844,6 +2844,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2844
Args:
2845
num_tokens: Number of tokens to run the dummy forward pass.
2846
cudagraph_runtime_mode: used to control the behavior.
2847
+ - if not set will determine the cudagraph mode based on using
2848
+ the self.cudagraph_dispatcher.
2849
- CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
2850
- CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
2851
- CUDAGraphMode.FULL: Full cudagraph, attention metadata is
@@ -2857,7 +2859,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2859
(1 token) and prefill (multiple tokens) requests.
2860
remove_lora: If False, dummy LoRAs are not destroyed after the run
2861
"""
2860
- assert cudagraph_runtime_mode in {
2862
+ assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in {
2863
CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL
2864
}
2865
@@ -2899,10 +2901,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2901
elif uniform_decode:
2902
assert not create_mixed_batch
2903
num_reqs = cdiv(num_tokens, max_query_len)
2902
- assert num_reqs <= max_num_reqs, \
2903
- f"Do not capture num_reqs {num_reqs} > max_num_reqs " \
2904
- f"{max_num_reqs} for uniform batch. Num tokens: " \
2905
- f"{num_tokens}, max_query_len: {max_query_len}"
2904
num_scheduled_tokens_list = [max_query_len] * num_reqs
2905
if num_tokens % max_query_len != 0:
2906
num_scheduled_tokens_list[-1] = num_tokens % max_query_len
@@ -3043,18 +3041,20 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3041
3042
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
3043
num_tokens, None, False)
3046
- if cudagraph_runtime_mode == CUDAGraphMode.NONE:
3047
- batch_descriptor = None
3048
- else:
3049
- # filter out the valid batch descriptor
3050
- _cg_mode, batch_descriptor = \
3051
- self.cudagraph_dispatcher.dispatch(
3052
- BatchDescriptor(num_tokens=num_tokens,
3053
- uniform_decode=uniform_decode))
3054
- # sanity check
3055
- assert cudagraph_runtime_mode == _cg_mode, (
3044
+
3045
+ # filter out the valid batch descriptor
3046
+ _cg_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(
3047
+ BatchDescriptor(num_tokens=num_tokens,
3048
+ uniform_decode=uniform_decode))
3049
+ if cudagraph_runtime_mode is not None:
3050
+ # we allow forcing NONE when the dispatcher disagrees to support
3051
+ # warm ups for cudagraph capture
3052
+ assert cudagraph_runtime_mode == CUDAGraphMode.NONE or \
3053
+ cudagraph_runtime_mode == _cg_mode, (
3054
f"Cudagraph runtime mode mismatch at dummy_run. "
3055
f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}.")
3056
+ else:
3057
+ cudagraph_runtime_mode = _cg_mode
3058
3059
if ubatch_slices is not None:
3060
num_tokens = num_tokens // 2