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 12, 16:00 UTC
Refactor: Move CUDA graph dispatch logic earlier (#27382)
Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>
Yizhou committed
Nov 23, 2025 at 05:10 UTC
df78aeef084cf35eecc6ba52640de8c390c99543
1 file changed
+25
-25
vllm/v1/worker/gpu_model_runner.py
+25
-25
@@ -3740,6 +3740,31 @@ class GPUModelRunner(
3740
dp_rank = self.parallel_config.data_parallel_rank
3741
num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])
3742
3743
+ # filter out the valid batch descriptor
3744
+ _cg_mode, batch_descriptor = (
3745
+ self.cudagraph_dispatcher.dispatch(
3746
+ BatchDescriptor(
3747
+ num_tokens=num_tokens_after_padding,
3748
+ uniform_decode=uniform_decode,
3749
+ has_lora=activate_lora and self.lora_config is not None,
3750
+ )
3751
+ )
3752
+ if not is_profile
3753
+ else (CUDAGraphMode.NONE, None)
3754
+ )
3755
+ if cudagraph_runtime_mode is not None:
3756
+ # we allow forcing NONE when the dispatcher disagrees to support
3757
+ # warm ups for cudagraph capture
3758
+ assert (
3759
+ cudagraph_runtime_mode == CUDAGraphMode.NONE
3760
+ or cudagraph_runtime_mode == _cg_mode
3761
+ ), (
3762
+ f"Cudagraph runtime mode mismatch at dummy_run. "
3763
+ f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
3764
+ )
3765
+ else:
3766
+ cudagraph_runtime_mode = _cg_mode
3767
+
3768
attn_metadata: PerLayerAttnMetadata | None = None
3769
3770
# If force_attention is True, we always capture attention. Otherwise,
@@ -3814,31 +3839,6 @@ class GPUModelRunner(
3839
num_tokens_after_padding, None, False
3840
)
3841
3817
- # filter out the valid batch descriptor
3818
- _cg_mode, batch_descriptor = (
3819
- self.cudagraph_dispatcher.dispatch(
3820
- BatchDescriptor(
3821
- num_tokens=num_tokens_after_padding,
3822
- uniform_decode=uniform_decode,
3823
- has_lora=activate_lora and self.lora_config is not None,
3824
- )
3825
- )
3826
- if not is_profile
3827
- else (CUDAGraphMode.NONE, None)
3828
- )
3829
- if cudagraph_runtime_mode is not None:
3830
- # we allow forcing NONE when the dispatcher disagrees to support
3831
- # warm ups for cudagraph capture
3832
- assert (
3833
- cudagraph_runtime_mode == CUDAGraphMode.NONE
3834
- or cudagraph_runtime_mode == _cg_mode
3835
- ), (
3836
- f"Cudagraph runtime mode mismatch at dummy_run. "
3837
- f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
3838
- )
3839
- else:
3840
- cudagraph_runtime_mode = _cg_mode
3841
-
3842
if ubatch_slices is not None:
3843
# Adjust values to reflect a single ubatch.
3844
# TODO(sage,lucas): this is cruft that should be addressed in