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][CI] Move resolving cudagraph_mode before initializing attn_metadata_builder (#27427)
Signed-off-by: fhl2000 <63384265+fhl2000@users.noreply.github.com>
fhl2000 committed
Oct 24, 2025 at 11:31 UTC
85fee74b337522f7e0807fc100b9e00682ff45e1
3 files changed
+34
-19
docs/design/cuda_graphs.md
+1
-1
@@ -167,7 +167,7 @@ class AttentionCGSupport(enum.Enum):
167
"""NO CUDA Graphs support"""
168
```
169
170
-Suppose we have hybrid attention backends (e.g., in mamba mixer models). In that case, we seek the minimum capability of all backends to determine the final capability of the model, and we might resolve the incompatible CUDA Graphs mode by downgrading the mode to the best fit one. For example, downgrading `FULL` mode to `FULL_AND_PIECEWISE` mode if the minimum capability is `UNIFORM_BATCH`, or `PIECEWISE` mode if the minimum capability is `NEVER` for -O3 compilation mode. For the complete fallback policy, please see the code of [initialize_cudagraph_capture][vllm.v1.worker.gpu_model_runner.GPUModelRunner.initialize_cudagraph_capture].
170
+Suppose we have hybrid attention backends (e.g., in mamba mixer models). In that case, we seek the minimum capability of all backends to determine the final capability of the model, and we might resolve the incompatible CUDA Graphs mode by downgrading the mode to the best fit one. For example, downgrading `FULL` mode to `FULL_AND_PIECEWISE` mode if the minimum capability is `UNIFORM_BATCH`, or `PIECEWISE` mode if the minimum capability is `NEVER` for -O3 compilation mode. For the complete fallback policy, please see the code for [this][vllm.v1.worker.gpu_model_runner.GPUModelRunner._check_and_update_cudagraph_mode].
171
172
The following table lists backends that support full CUDA Graphs at the time of writing.
173
tests/compile/test_fusions_e2e.py
+3
@@ -132,6 +132,9 @@ def test_attn_quant(
132
mode = CUDAGraphMode.FULL_AND_PIECEWISE
133
splitting_ops: list[str] | None = None
134
else:
135
+ # FIXME: Llama-4-Scout-17B-16E-Instruct-FP8 + FlashInfer + Blackwell end at
136
+ # CUDAGraphMode.NONE here because it derives an attention backend that
137
+ # does not support full cudagraphs
138
mode = CUDAGraphMode.FULL_DECODE_ONLY
139
splitting_ops = []
140
vllm/v1/worker/gpu_model_runner.py
+30
-18
@@ -3751,8 +3751,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3751
"ensure `cudagraph_mode` was not manually set to `NONE`"
3752
)
3753
return 0
3754
- else:
3755
- self.initialize_cudagraph_capture()
3754
3755
compilation_counter.num_gpu_runner_capture_triggers += 1
3756
@@ -3926,7 +3924,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3924
3925
def get_attn_backends_for_group(
3926
kv_cache_group_spec: KVCacheGroupSpec,
3929
- ) -> dict[AttentionGroupKey, list[str]]:
3927
+ ) -> tuple[dict[AttentionGroupKey, list[str]], set[type[AttentionBackend]]]:
3928
layers = get_layers_from_vllm_config(
3929
self.vllm_config, AttentionLayerBase, kv_cache_group_spec.layer_names
3930
)
@@ -3955,7 +3953,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3953
attn_backend, layer_kv_cache_spec
3954
)
3955
attn_backend_layers[key].append(layer_name)
3958
- return {attn_backends[k]: v for k, v in attn_backend_layers.items()}
3956
+ return (
3957
+ {attn_backends[k]: v for k, v in attn_backend_layers.items()},
3958
+ set(group_key.attn_backend for group_key in attn_backends.values()),
3959
+ )
3960
3961
def create_attn_groups(
3962
attn_backends_map: dict[AttentionGroupKey, list[str]],
@@ -3976,14 +3977,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3977
attn_groups.append(attn_group)
3978
return attn_groups
3979
3980
+ attention_backend_maps = []
3981
+ attention_backend_set: set[type[AttentionBackend]] = set()
3982
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
3983
attn_backends = get_attn_backends_for_group(kv_cache_group_spec)
3981
- self.attn_groups.append(create_attn_groups(attn_backends))
3984
+ attention_backend_maps.append(attn_backends[0])
3985
+ attention_backend_set.update(attn_backends[1])
3986
+
3987
+ # Resolve cudagraph_mode before actually initialize metadata_builders
3988
+ self._check_and_update_cudagraph_mode(attention_backend_set)
3989
+
3990
+ for attn_backends_map in attention_backend_maps:
3991
+ self.attn_groups.append(create_attn_groups(attn_backends_map))
3992
3993
# Calculate reorder batch threshold (if needed)
3994
self.calculate_reorder_batch_threshold()
3995
3986
- def initialize_cudagraph_capture(self) -> None:
3996
+ def _check_and_update_cudagraph_mode(
3997
+ self, attention_backends: set[type[AttentionBackend]]
3998
+ ) -> None:
3999
"""
4000
Resolve the cudagraph_mode when there are multiple attention
4001
backends with potential conflicting CUDA graph support.
@@ -3991,13 +4003,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4003
cudagraph_mode.
4004
"""
4005
min_cg_support = AttentionCGSupport.ALWAYS
3994
- min_cg_builder_name = None
4006
+ min_cg_backend_name = None
4007
3996
- for attn_group in self._attn_group_iterator():
3997
- builder = attn_group.get_metadata_builder()
3998
- if builder.cudagraph_support.value < min_cg_support.value:
3999
- min_cg_support = builder.cudagraph_support
4000
- min_cg_builder_name = builder.__class__.__name__
4008
+ for attn_backend in attention_backends:
4009
+ builder_cls = attn_backend.get_builder_cls()
4010
+ if builder_cls.cudagraph_support.value < min_cg_support.value:
4011
+ min_cg_support = builder_cls.cudagraph_support
4012
+ min_cg_backend_name = attn_backend.__name__
4013
# Flexible resolve the cudagraph mode
4014
cudagraph_mode = self.compilation_config.cudagraph_mode
4015
# check cudagraph for mixed batch is supported
@@ -4007,7 +4019,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4019
):
4020
msg = (
4021
f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
4010
- f"with {min_cg_builder_name} backend (support: "
4022
+ f"with {min_cg_backend_name} backend (support: "
4023
f"{min_cg_support})"
4024
)
4025
if min_cg_support == AttentionCGSupport.NEVER:
@@ -4038,7 +4050,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4050
):
4051
msg = (
4052
f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
4041
- f"with {min_cg_builder_name} backend (support: "
4053
+ f"with {min_cg_backend_name} backend (support: "
4054
f"{min_cg_support})"
4055
)
4056
if self.compilation_config.mode == CompilationMode.VLLM_COMPILE and (
@@ -4072,7 +4084,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4084
msg = (
4085
f"CUDAGraphMode.{cudagraph_mode.name} is not supported"
4086
f" with spec-decode for attention backend "
4075
- f"{min_cg_builder_name} (support: {min_cg_support})"
4087
+ f"{min_cg_backend_name} (support: {min_cg_support})"
4088
)
4089
if self.compilation_config.splitting_ops_contain_attention():
4090
msg += "; setting cudagraph_mode=PIECEWISE"
@@ -4094,14 +4106,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
4106
):
4107
raise ValueError(
4108
f"CUDAGraphMode.{cudagraph_mode.name} is not "
4097
- f"supported with {min_cg_builder_name} backend ("
4109
+ f"supported with {min_cg_backend_name} backend ("
4110
f"support:{min_cg_support}) "
4111
"; please try cudagraph_mode=PIECEWISE, "
4112
"and make sure compilation mode is VLLM_COMPILE"
4113
)
4114
4103
- # Trigger cudagraph dispatching keys initialization here (after
4104
- # initializing attn backends).
4115
+ # Trigger cudagraph dispatching keys initialization after
4116
+ # resolved cudagraph mode.
4117
self.cudagraph_dispatcher.initialize_cudagraph_keys(
4118
self.compilation_config.cudagraph_mode, self.uniform_decode_query_len
4119
)