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 15, 18:30 UTC
[V1] Add flag to disable cascade attention (#15243)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 20, 2025 at 15:24 UTC
2b22290ce01b033cc692e7dce159d74a43f6f2c5
3 files changed
+23
-5
vllm/config.py
+2
@@ -246,6 +246,7 @@ class ModelConfig:
246
max_seq_len_to_capture: Optional[int] = None,
247
max_logprobs: int = 20,
248
disable_sliding_window: bool = False,
249
+ disable_cascade_attn: bool = False,
250
skip_tokenizer_init: bool = False,
251
served_model_name: Optional[Union[str, list[str]]] = None,
252
limit_mm_per_prompt: Optional[Mapping[str, int]] = None,
@@ -322,6 +323,7 @@ class ModelConfig:
323
self.max_seq_len_to_capture = max_seq_len_to_capture
324
self.max_logprobs = max_logprobs
325
self.disable_sliding_window = disable_sliding_window
326
+ self.disable_cascade_attn = disable_cascade_attn
327
self.skip_tokenizer_init = skip_tokenizer_init
328
self.enable_sleep_mode = enable_sleep_mode
329
vllm/engine/arg_utils.py
+12
@@ -120,6 +120,7 @@ class EngineArgs:
120
block_size: Optional[int] = None
121
enable_prefix_caching: Optional[bool] = None
122
disable_sliding_window: bool = False
123
+ disable_cascade_attn: bool = False
124
use_v2_block_manager: bool = True
125
swap_space: float = 4 # GiB
126
cpu_offload_gb: float = 0 # GiB
@@ -1096,6 +1097,16 @@ class EngineArgs:
1097
"using. This is used to parse the reasoning content into OpenAI "
1098
"API format. Required for ``--enable-reasoning``.")
1099
1100
+ parser.add_argument(
1101
+ "--disable-cascade-attn",
1102
+ action="store_true",
1103
+ default=False,
1104
+ help="Disable cascade attention for V1. While cascade attention "
1105
+ "does not change the mathematical correctness, disabling it "
1106
+ "could be useful for preventing potential numerical issues. "
1107
+ "Note that even if this is set to False, cascade attention will be "
1108
+ "only used when the heuristic tells that it's beneficial.")
1109
+
1110
return parser
1111
1112
@classmethod
@@ -1141,6 +1152,7 @@ class EngineArgs:
1152
max_seq_len_to_capture=self.max_seq_len_to_capture,
1153
max_logprobs=self.max_logprobs,
1154
disable_sliding_window=self.disable_sliding_window,
1155
+ disable_cascade_attn=self.disable_cascade_attn,
1156
skip_tokenizer_init=self.skip_tokenizer_init,
1157
served_model_name=self.served_model_name,
1158
limit_mm_per_prompt=self.limit_mm_per_prompt,
vllm/v1/worker/gpu_model_runner.py
+9
-5
@@ -127,6 +127,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
127
128
self.attn_metadata_builder = self.attn_backend.get_builder_cls()(
129
weakref.proxy(self))
130
+ self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
131
132
# Multi-modal data support
133
self.input_registry = INPUT_REGISTRY
@@ -565,11 +566,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
566
self.positions_cpu[:total_num_scheduled_tokens],
567
non_blocking=True)
568
568
- # Prepare for cascade attention if needed.
569
- common_prefix_len = self._compute_cascade_attn_prefix_len(
570
- num_scheduled_tokens,
571
- scheduler_output.num_common_prefix_blocks,
572
- )
569
+ # Prepare for cascade attention if enabled & beneficial.
570
+ common_prefix_len = 0
571
+ if self.cascade_attn_enabled:
572
+ common_prefix_len = self._compute_cascade_attn_prefix_len(
573
+ num_scheduled_tokens,
574
+ scheduler_output.num_common_prefix_blocks,
575
+ )
576
+
577
attn_metadata = self.attn_metadata_builder.build(
578
num_reqs=num_reqs,
579
num_actual_tokens=total_num_scheduled_tokens,