@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 986537f1c3
Fork 0

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, 13:30 UTC

[V1] V1 FlashInfer Attention (#16684)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Aurick Qiao <qiao@aurick.net>

Michael Goin committed Apr 21, 2025 at 18:38 UTC 986537f1c3c8288635b1f00944409489ab1cf21b
7 files changed +668 -13
tests/v1/e2e/test_cascade_attention.py
+9 -1
@@ -1,13 +1,21 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 +import pytest
4 +
5 from vllm import LLM, SamplingParams
6
7 +from ...utils import fork_new_process_for_each_test
8 +
9
6 -def test_cascade_attention(example_system_message, monkeypatch):
10 +@fork_new_process_for_each_test
11 +@pytest.mark.parametrize("attn_backend",
12 + ["FLASH_ATTN_VLLM_V1", "FLASHINFER_VLLM_V1"])
13 +def test_cascade_attention(example_system_message, monkeypatch, attn_backend):
14 prompt = "\n<User>: Implement fibonacci sequence in Python.\n<Claude>:"
15
16 with monkeypatch.context() as m:
17 m.setenv("VLLM_USE_V1", "1")
18 + m.setenv("VLLM_ATTENTION_BACKEND", attn_backend)
19
20 llm = LLM(model="Qwen/Qwen2-1.5B-Instruct")
21 sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
vllm/engine/arg_utils.py
+10 -3
@@ -1474,10 +1474,17 @@ class EngineArgs:
1474 recommend_to_remove=False)
1475 return False
1476
1477 - # No FlashInfer or XFormers so far.
1477 + # No XFormers so far.
1478 V1_BACKENDS = [
1479 - "FLASH_ATTN_VLLM_V1", "FLASH_ATTN", "PALLAS", "PALLAS_VLLM_V1",
1480 - "TRITON_ATTN_VLLM_V1", "TRITON_MLA", "FLASHMLA"
1479 + "FLASH_ATTN_VLLM_V1",
1480 + "FLASH_ATTN",
1481 + "PALLAS",
1482 + "PALLAS_VLLM_V1",
1483 + "TRITON_ATTN_VLLM_V1",
1484 + "TRITON_MLA",
1485 + "FLASHMLA",
1486 + "FLASHINFER",
1487 + "FLASHINFER_VLLM_V1",
1488 ]
1489 if (envs.is_set("VLLM_ATTENTION_BACKEND")
1490 and envs.VLLM_ATTENTION_BACKEND not in V1_BACKENDS):
vllm/platforms/cuda.py
+3
@@ -213,6 +213,9 @@ class CudaPlatformBase(Platform):
213 return ("vllm.attention.backends."
214 "flashmla.FlashMLABackend")
215 if use_v1:
216 + if selected_backend == _Backend.FLASHINFER:
217 + logger.info_once("Using FlashInfer backend on V1 engine.")
218 + return "vllm.v1.attention.backends.flashinfer.FlashInferBackend"
219 if selected_backend == _Backend.TRITON_ATTN_VLLM_V1:
220 logger.info_once("Using Triton backend on V1 engine.")
221 return ("vllm.v1.attention.backends."
vllm/v1/attention/backends/flash_attn.py
+3 -4
@@ -64,10 +64,6 @@ class FlashAttentionBackend(AttentionBackend):
64 raise ValueError("Block size must be a multiple of 16.")
65 return (2, num_blocks, block_size, num_kv_heads, head_size)
66
67 - @staticmethod
68 - def use_cascade_attention(*args, **kwargs) -> bool:
69 - return use_cascade_attention(*args, **kwargs)
70 -
67
68 @dataclass
69 class FlashAttentionMetadata:
@@ -402,6 +398,9 @@ class FlashAttentionMetadataBuilder:
398 )
399 return attn_metadata
400
401 + def use_cascade_attention(self, *args, **kwargs) -> bool:
402 + return use_cascade_attention(*args, **kwargs)
403 +
404
405 class FlashAttentionImpl(AttentionImpl):
406
vllm/v1/attention/backends/flashinfer.py new
+639
@@ -0,0 +1,639 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""Attention layer with FlashInfer."""
3 +from __future__ import annotations
4 +
5 +from dataclasses import dataclass
6 +from typing import TYPE_CHECKING, Any, Optional
7 +
8 +import torch
9 +from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
10 + BatchPrefillWithPagedKVCacheWrapper,
11 + MultiLevelCascadeAttentionWrapper)
12 +
13 +import vllm.envs as envs
14 +from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
15 + AttentionType)
16 +from vllm.attention.layer import Attention
17 +from vllm.config import VllmConfig, get_current_vllm_config
18 +from vllm.logger import init_logger
19 +from vllm.v1.attention.backends.flash_attn import use_cascade_attention
20 +
21 +if TYPE_CHECKING:
22 + from vllm.v1.core.sched.output import SchedulerOutput
23 + from vllm.v1.worker.gpu_input_batch import InputBatch
24 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
25 +
26 +FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
27 +
28 +logger = init_logger(__name__)
29 +
30 +
31 +class FlashInferBackend(AttentionBackend):
32 +
33 + accept_output_buffer: bool = True
34 +
35 + @staticmethod
36 + def get_supported_head_sizes() -> list[int]:
37 + return [64, 128, 256]
38 +
39 + @staticmethod
40 + def get_name() -> str:
41 + return "FLASHINFER_VLLM_V1"
42 +
43 + @staticmethod
44 + def get_impl_cls() -> type[FlashInferImpl]:
45 + return FlashInferImpl
46 +
47 + @staticmethod
48 + def get_metadata_cls() -> type[FlashInferMetadata]:
49 + return FlashInferMetadata
50 +
51 + @staticmethod
52 + def get_builder_cls() -> type[FlashInferMetadataBuilder]:
53 + return FlashInferMetadataBuilder
54 +
55 + @staticmethod
56 + def get_kv_cache_shape(
57 + num_blocks: int,
58 + block_size: int,
59 + num_kv_heads: int,
60 + head_size: int,
61 + ) -> tuple[int, ...]:
62 + return (num_blocks, 2, block_size, num_kv_heads, head_size)
63 +
64 +
65 +@dataclass
66 +class PerLayerParameters:
67 + """
68 + Currently, FlashInfer backend only support models in which all layers share
69 + the same values for the following hyperparameters.
70 + """
71 +
72 + window_left: int
73 + logits_soft_cap: Optional[float]
74 + sm_scale: float
75 +
76 +
77 +def get_per_layer_parameters(
78 + vllm_config: VllmConfig) -> dict[str, PerLayerParameters]:
79 + """
80 + Scan all attention layers and determine some hyperparameters
81 + to use during `plan`.
82 + """
83 +
84 + layers = vllm_config.compilation_config.static_forward_context
85 + per_layer_params: dict[str, PerLayerParameters] = {}
86 +
87 + for key, layer in layers.items():
88 + assert isinstance(layer, Attention)
89 +
90 + impl = layer.impl
91 + assert isinstance(impl, FlashInferImpl)
92 +
93 + # Infer hyperparameters from the attention layer
94 + window_size = impl.sliding_window
95 + window_left = window_size[0] if window_size is not None else -1
96 + logits_soft_cap = impl.logits_soft_cap
97 + sm_scale = impl.scale
98 +
99 + per_layer_params[key] = PerLayerParameters(window_left,
100 + logits_soft_cap, sm_scale)
101 +
102 + return per_layer_params
103 +
104 +
105 +def infer_global_hyperparameters(
106 + per_layer_params: dict[str, PerLayerParameters]) -> PerLayerParameters:
107 + """
108 + Currently, FlashInfer backend only support models in which all layers share
109 + the same values for the following hyperparameters:
110 + - `window_left`
111 + - `logits_soft_cap`
112 + - `sm_scale`
113 +
114 + So this function asserts that all layers share the same values for these
115 + hyperparameters and returns the global values.
116 + """
117 +
118 + assert len(per_layer_params) > 0, "No attention layers found in the model."
119 +
120 + param_sets = list(per_layer_params.values())
121 + global_params = param_sets[0]
122 + for params in param_sets:
123 + assert params == global_params, (
124 + "FlashInfer backend currently only supports models in which all "
125 + "layers share the same values for the following hyperparameters: "
126 + "`window_left`, `logits_soft_cap`, `sm_scale`.")
127 +
128 + return global_params
129 +
130 +
131 +@dataclass
132 +class FlashInferMetadata:
133 +
134 + num_actual_tokens: int # Number of tokens excluding padding.
135 +
136 + # (batch_size + 1,). The cumulative subquery lengths of the sequences in
137 + # the batch, used to index into subquery. E.g., if the subquery length
138 + # is [4, 6], it is [0, 4, 10].
139 + qo_indptr: torch.Tensor
140 + # An example for paged_kv_indices, paged_kv_indptr:
141 + # request 1, page indices [0, 5, 8]
142 + # request 2, page indices [1, 6, 7]
143 + # request 3, page indices [3, 4]
144 + # paged_kv_indices is a concatenation of page indices of all requests:
145 + # [0, 5, 8, 1, 6, 7, 3, 4]
146 + # paged_kv_indptr is used to index into paged_kv_indices:
147 + # [0, 3, 6, 8]
148 + # The indptr of the paged kv cache, shape: [batch_size + 1]
149 + paged_kv_indptr: torch.Tensor
150 + # The page indices of the paged kv cache
151 + paged_kv_indices: torch.Tensor
152 + # The number of entries in the last page of each request in
153 + # the paged kv cache, shape: [batch_size]
154 + paged_kv_last_page_len: torch.Tensor
155 + # The number of query/output heads
156 + num_qo_heads: int
157 + # The number of key/value heads
158 + num_kv_heads: int
159 + # The dimension of the attention heads
160 + head_dim: int
161 + # Block size of vllm
162 + page_size: int
163 + # The data type of the paged kv cache
164 + data_type: torch.dtype
165 + # The data type of the query
166 + q_data_type: torch.dtype
167 +
168 + slot_mapping: torch.Tensor
169 +
170 + # For handling prefill decode split
171 + num_decodes: int
172 + num_decode_tokens: int
173 + num_prefills: int
174 + num_prefill_tokens: int
175 +
176 + # For cascade attention.
177 + use_cascade: bool
178 + shared_qo_indptr: Optional[torch.Tensor] = None
179 + shared_kv_page_indptr: Optional[torch.Tensor] = None
180 + shared_kv_page_indices: Optional[torch.Tensor] = None
181 + shared_kv_last_page_len: Optional[torch.Tensor] = None
182 +
183 + prefill_wrapper: Optional[BatchPrefillWithPagedKVCacheWrapper] = None
184 + decode_wrapper: Optional[BatchDecodeWithPagedKVCacheWrapper] = None
185 + cascade_wrapper: Optional[MultiLevelCascadeAttentionWrapper] = None
186 +
187 + # For logging.
188 + num_input_tokens: int = 0 # Number of tokens including padding.
189 +
190 + @property
191 + def query_start_loc(self):
192 + # The GPUModelRunner expects to be able to access this property.
193 + return self.qo_indptr
194 +
195 + def __post_init__(self):
196 + # Refer to
197 + # https://github.com/flashinfer-ai/flashinfer/blob/3d55c71a62052c590c130897d3a3db49b14fcc34/include/flashinfer/utils.cuh#L157
198 + supported_head_sizes = FlashInferBackend.get_supported_head_sizes()
199 + if self.head_dim is not None and self.head_dim \
200 + not in supported_head_sizes:
201 + raise ValueError(
202 + f"Only {supported_head_sizes} are supported for head_dim,",
203 + f" received {self.head_dim}.")
204 +
205 +
206 +class FlashInferMetadataBuilder:
207 +
208 + def __init__(self, runner: GPUModelRunner):
209 + self.runner = runner
210 + self._workspace_buffer = None
211 + self._prefill_wrapper = None # Wrapper for prefill/append
212 + self._decode_wrapper = None # Wrapper for decode
213 + self._cascade_wrapper = None # Wrapper for cascade attention
214 +
215 + # Global hyperparameters shared by all attention layers
216 + self.global_hyperparameters: Optional[PerLayerParameters] = None
217 +
218 + self.vllm_config = get_current_vllm_config()
219 +
220 + def reorder_batch(self, input_batch: InputBatch,
221 + scheduler_output: SchedulerOutput) -> bool:
222 + # We now want to reorder the batch so that the "decode" requests are and
223 + # the front and the "prefill" requests are at the using the least amount
224 + # swaps possible. (NOTE for now we loosely use "decode" to mean requests
225 + # where attention is likely memory-bound and "prefill" to mean requests
226 + # where attention is likely compute-bound, TODO(lucas): figure out a
227 + # better naming here)
228 + decodes = []
229 + prefills = []
230 + num_decode_tokens = 0
231 + num_prefill_tokens = 0
232 +
233 + for i, req_id in enumerate(input_batch.req_ids):
234 + num_tokens = scheduler_output.num_scheduled_tokens[req_id]
235 + # for now treat 1 scheduled token as "decode" even if its not,
236 + # we should update this to something like < 8 in the future but
237 + # currently the decode run only supports num_tokens = 1
238 + if num_tokens == 1:
239 + decodes.append(i)
240 + num_decode_tokens += num_tokens
241 + else:
242 + prefills.append(i)
243 + num_prefill_tokens += num_tokens
244 +
245 + # We hope that this is fairly minimal since decodes
246 + # should be around for a number of iterations so hopefully they are
247 + # relatively stationary (and new request are generally appended to the
248 + # persistent batch so already should be at the back)
249 + # To achieve this we loop over the decodes in descending order and
250 + # the prefills in ascending order. We swap decodes from the "back"
251 + # i.e. past where the last decode should be in the reodorered with
252 + # prefills from the front of the batch.
253 + # `decodes` and `prefills` are already in ascending order just based on
254 + # the above loop
255 + num_decodes = len(decodes)
256 + num_prefills = len(prefills)
257 + modified_batch = False
258 +
259 + for i in range(1, min(num_decodes, num_prefills) + 1):
260 + # If the decode is at the "back" of the batch, i, we can swap it
261 + # with the prefill closest to the front of the batch
262 + decode_idx = decodes[num_decodes - i]
263 + if decode_idx < num_decodes:
264 + break
265 +
266 + input_batch.swap_states(prefills[i - 1], decode_idx)
267 + modified_batch = True
268 +
269 + # Save for next `build` call
270 + # TODO(lucas): this is a bit of a hack, we should probably have a
271 + # better way of doing this
272 + self._num_decodes = num_decodes
273 + self._num_prefills = num_prefills
274 + self._num_decode_tokens = num_decode_tokens
275 + self._num_prefill_tokens = num_prefill_tokens
276 +
277 + return modified_batch
278 +
279 + def _get_workspace_buffer(self):
280 + if self._workspace_buffer is None:
281 + self._workspace_buffer = torch.empty(
282 + FLASHINFER_WORKSPACE_BUFFER_SIZE,
283 + dtype=torch.uint8,
284 + device=self.runner.device)
285 + return self._workspace_buffer
286 +
287 + def _get_prefill_wrapper(self):
288 + if self._prefill_wrapper is None:
289 + self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
290 + self._get_workspace_buffer(), "NHD")
291 + return self._prefill_wrapper
292 +
293 + def _get_decode_wrapper(self):
294 + if self._decode_wrapper is None:
295 + num_qo_heads = (self.runner.model_config.get_num_attention_heads(
296 + self.runner.parallel_config))
297 + num_kv_heads = self.runner.model_config.get_num_kv_heads(
298 + self.runner.parallel_config)
299 + use_tensor_cores = envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or (
300 + num_qo_heads // num_kv_heads > 4)
301 + self._decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
302 + self._get_workspace_buffer(),
303 + "NHD",
304 + use_tensor_cores=use_tensor_cores)
305 + return self._decode_wrapper
306 +
307 + def _get_cascade_wrapper(self):
308 + if self._cascade_wrapper is None:
309 + self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
310 + 2, self._get_workspace_buffer(), "NHD")
311 + return self._cascade_wrapper
312 +
313 + def _plan(self, attn_metadata: FlashInferMetadata):
314 + if self.global_hyperparameters is None:
315 + self.global_hyperparameters = infer_global_hyperparameters(
316 + get_per_layer_parameters(self.vllm_config))
317 + if attn_metadata.use_cascade:
318 + attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
319 + attn_metadata.cascade_wrapper.plan(
320 + [attn_metadata.shared_qo_indptr, attn_metadata.qo_indptr],
321 + [
322 + attn_metadata.shared_kv_page_indptr,
323 + attn_metadata.paged_kv_indptr
324 + ],
325 + [
326 + attn_metadata.shared_kv_page_indices,
327 + attn_metadata.paged_kv_indices
328 + ],
329 + [
330 + attn_metadata.shared_kv_last_page_len,
331 + attn_metadata.paged_kv_last_page_len
332 + ],
333 + attn_metadata.num_qo_heads,
334 + attn_metadata.num_kv_heads,
335 + attn_metadata.head_dim,
336 + attn_metadata.page_size,
337 + causal=True,
338 + sm_scale=self.global_hyperparameters.sm_scale,
339 + window_left=self.global_hyperparameters.window_left,
340 + logits_soft_cap=self.global_hyperparameters.logits_soft_cap,
341 + q_data_type=attn_metadata.q_data_type,
342 + )
343 + else:
344 + # Regular attention (common case).
345 + # Decodes are at the front and prefills are at the back,
346 + # according to reorder_batch()
347 + if self._num_prefills > 0:
348 + # Decodes are first so prefills start after the last decode
349 + prefill_start = self._num_decodes
350 + attn_metadata.prefill_wrapper = self._get_prefill_wrapper()
351 + assert attn_metadata.qo_indptr[prefill_start:].shape[
352 + 0] == self._num_prefills + 1
353 + assert attn_metadata.paged_kv_indptr[prefill_start:].shape[
354 + 0] == self._num_prefills + 1
355 + assert attn_metadata.paged_kv_last_page_len[
356 + prefill_start:].shape[0] == self._num_prefills
357 + # Since prefill_wrapper.run() will be called with
358 + # query[num_decode_tokens:] we need to adjust the qo_indptr
359 + # to be relative to the start of the prefill queries.
360 + qo_indptr = attn_metadata.qo_indptr[
361 + prefill_start:] - attn_metadata.qo_indptr[prefill_start]
362 + attn_metadata.prefill_wrapper.plan(
363 + qo_indptr,
364 + attn_metadata.paged_kv_indptr[prefill_start:],
365 + attn_metadata.paged_kv_indices,
366 + attn_metadata.paged_kv_last_page_len[prefill_start:],
367 + attn_metadata.num_qo_heads,
368 + attn_metadata.num_kv_heads,
369 + attn_metadata.head_dim,
370 + attn_metadata.page_size,
371 + causal=True,
372 + sm_scale=self.global_hyperparameters.sm_scale,
373 + window_left=self.global_hyperparameters.window_left,
374 + logits_soft_cap=self.global_hyperparameters.
375 + logits_soft_cap,
376 + q_data_type=attn_metadata.q_data_type,
377 + kv_data_type=attn_metadata.data_type,
378 + )
379 +
380 + if self._num_decodes > 0:
381 + attn_metadata.decode_wrapper = self._get_decode_wrapper()
382 + attn_metadata.decode_wrapper.plan(
383 + attn_metadata.paged_kv_indptr[:self._num_decodes + 1],
384 + attn_metadata.paged_kv_indices,
385 + attn_metadata.paged_kv_last_page_len[:self._num_decodes],
386 + attn_metadata.num_qo_heads,
387 + attn_metadata.num_kv_heads,
388 + attn_metadata.head_dim,
389 + attn_metadata.page_size,
390 + # Disable flashinfer's pos encoding and use vllm's rope.
391 + pos_encoding_mode="NONE",
392 + sm_scale=self.global_hyperparameters.sm_scale,
393 + window_left=self.global_hyperparameters.window_left,
394 + logits_soft_cap=self.global_hyperparameters.
395 + logits_soft_cap,
396 + q_data_type=attn_metadata.q_data_type,
397 + kv_data_type=attn_metadata.data_type,
398 + )
399 +
400 + def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
401 + common_prefix_len: int):
402 + assert self._num_decodes + self._num_prefills == num_reqs
403 + assert (self._num_decode_tokens +
404 + self._num_prefill_tokens == num_actual_tokens)
405 + page_size = self.runner.block_size
406 + device = self.runner.device
407 + qo_indptr = self.runner.query_start_loc_cpu[:num_reqs + 1].to(
408 + self.runner.device, non_blocking=True)
409 + seq_lens = self.runner.seq_lens_cpu[:num_reqs].to(self.runner.device,
410 + non_blocking=True)
411 + block_table = (
412 + self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])
413 + slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
414 + self.runner.device, non_blocking=True).long()
415 +
416 + block_table_bounds = (seq_lens + page_size - 1) // page_size
417 +
418 + use_cascade = common_prefix_len > 0
419 + if use_cascade:
420 + # Grab the blocks of the shared prefix from the first request.
421 + assert common_prefix_len % page_size == 0
422 + num_common_kv_blocks = common_prefix_len // page_size
423 + shared_qo_indptr = torch.tensor([0, num_actual_tokens],
424 + dtype=torch.int32,
425 + device=device)
426 + shared_kv_page_indptr = torch.tensor([0, num_common_kv_blocks],
427 + dtype=torch.int32,
428 + device=device)
429 + shared_kv_page_indices = block_table[0, :num_common_kv_blocks]
430 + shared_kv_last_page_len = torch.tensor([page_size],
431 + dtype=torch.int32,
432 + device=device)
433 + # Remove the blocks of the shared prefix from all requests.
434 + block_table = block_table[:, num_common_kv_blocks:]
435 + block_table_bounds -= num_common_kv_blocks
436 + else:
437 + shared_qo_indptr = None
438 + shared_kv_page_indptr = None
439 + shared_kv_page_indices = None
440 + shared_kv_last_page_len = None
441 +
442 + mask = (torch.arange(block_table.size(1),
443 + dtype=block_table.dtype,
444 + device=block_table.device).unsqueeze(0)
445 + < block_table_bounds.unsqueeze(1))
446 + paged_kv_indices = block_table[mask]
447 +
448 + paged_kv_indptr = torch.cat([
449 + torch.zeros(1,
450 + dtype=block_table_bounds.dtype,
451 + device=block_table_bounds.device),
452 + block_table_bounds.cumsum(dim=0, dtype=torch.int32)
453 + ])
454 +
455 + paged_kv_last_page_len = seq_lens % page_size
456 + paged_kv_last_page_len = torch.where(paged_kv_last_page_len == 0,
457 + page_size, paged_kv_last_page_len)
458 +
459 + attn_metadata = FlashInferMetadata(
460 + num_actual_tokens=num_actual_tokens,
461 + qo_indptr=qo_indptr,
462 + paged_kv_indptr=paged_kv_indptr,
463 + paged_kv_indices=paged_kv_indices,
464 + paged_kv_last_page_len=paged_kv_last_page_len,
465 + num_qo_heads=self.runner.num_query_heads,
466 + num_kv_heads=self.runner.num_kv_heads,
467 + head_dim=self.runner.head_size,
468 + page_size=page_size,
469 + data_type=self.runner.kv_cache_dtype,
470 + q_data_type=self.runner.dtype,
471 + slot_mapping=slot_mapping,
472 + num_decodes=self._num_decodes,
473 + num_decode_tokens=self._num_decode_tokens,
474 + num_prefills=self._num_prefills,
475 + num_prefill_tokens=self._num_prefill_tokens,
476 + use_cascade=use_cascade,
477 + shared_qo_indptr=shared_qo_indptr,
478 + shared_kv_page_indptr=shared_kv_page_indptr,
479 + shared_kv_page_indices=shared_kv_page_indices,
480 + shared_kv_last_page_len=shared_kv_last_page_len,
481 + )
482 +
483 + self._plan(attn_metadata)
484 +
485 + return attn_metadata
486 +
487 + def use_cascade_attention(self, *args, **kwargs) -> bool:
488 + if self.runner.kv_cache_dtype != self.runner.model_config.dtype:
489 + # TODO: The cascade wrapper currently does not support setting
490 + # kv cache dtype to something different from query dtype.
491 + return False
492 + return use_cascade_attention(*args, **kwargs)
493 +
494 +
495 +class FlashInferImpl(AttentionImpl):
496 +
497 + def __init__(
498 + self,
499 + num_heads: int,
500 + head_size: int,
501 + scale: float,
502 + num_kv_heads: int,
503 + alibi_slopes: Optional[list[float]],
504 + sliding_window: Optional[int],
505 + kv_cache_dtype: str,
506 + blocksparse_params: Optional[dict[str, Any]] = None,
507 + logits_soft_cap: Optional[float] = None,
508 + attn_type: AttentionType = AttentionType.DECODER,
509 + ) -> None:
510 + self.num_heads = num_heads
511 + self.head_size = head_size
512 + self.scale = float(scale)
513 + self.num_kv_heads = num_kv_heads
514 + if alibi_slopes is not None:
515 + alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
516 + self.alibi_slopes = alibi_slopes
517 + if sliding_window is None:
518 + self.sliding_window = (-1, -1)
519 + else:
520 + self.sliding_window = (sliding_window - 1, 0)
521 + self.kv_cache_dtype = kv_cache_dtype
522 + self.logits_soft_cap = logits_soft_cap
523 +
524 + assert self.num_heads % self.num_kv_heads == 0
525 + self.num_queries_per_kv = self.num_heads // self.num_kv_heads
526 +
527 + if attn_type != AttentionType.DECODER:
528 + raise NotImplementedError("Encoder self-attention and "
529 + "encoder/decoder cross-attention "
530 + "are not implemented for "
531 + "FlashInferImpl")
532 +
533 + def forward(
534 + self,
535 + layer: torch.nn.Module,
536 + query: torch.Tensor,
537 + key: torch.Tensor,
538 + value: torch.Tensor,
539 + kv_cache: torch.Tensor,
540 + attn_metadata: FlashInferMetadata,
541 + output: Optional[torch.Tensor] = None,
542 + ) -> torch.Tensor:
543 + """Forward pass with FlashInfer.
544 +
545 + Args:
546 + query: shape = [num_tokens, num_heads, head_size]
547 + key: shape = [num_tokens, num_kv_heads, head_size]
548 + value: shape = [num_tokens, num_kv_heads, head_size]
549 + kv_cache = [num_blocks, 2, block_size, num_kv_heads, head_size]
550 + attn_metadata: Metadata for attention.
551 + Returns:
552 + shape = [num_tokens, num_heads * head_size]
553 + """
554 + assert output is not None, "Output tensor must be provided."
555 +
556 + if attn_metadata is None:
557 + # Profiling run.
558 + return output
559 +
560 + # IMPORTANT!
561 + # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
562 + # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
563 + # in this method. For example, `view` and `slice` (or `[:n]`) operations
564 + # are surprisingly slow even in the case they do not invoke any GPU ops.
565 + # Minimize the PyTorch ops in this method as much as possible.
566 + # Whenever making a change in this method, please benchmark the
567 + # performance to make sure it does not introduce any overhead.
568 +
569 + num_actual_tokens = attn_metadata.num_actual_tokens
570 + # Reshape the input keys and values and store them in the cache.
571 + # NOTE(woosuk): Here, key and value are padded while slot_mapping is
572 + # not padded. However, we don't need to do key[:num_actual_tokens] and
573 + # value[:num_actual_tokens] because the reshape_and_cache_flash op uses
574 + # the slot_mapping's shape to determine the number of actual tokens.
575 + torch.ops._C_cache_ops.reshape_and_cache_flash(
576 + key,
577 + value,
578 + kv_cache[:, 0],
579 + kv_cache[:, 1],
580 + attn_metadata.slot_mapping,
581 + self.kv_cache_dtype,
582 + layer._k_scale,
583 + layer._v_scale,
584 + )
585 +
586 + window_left = (self.sliding_window[0]
587 + if self.sliding_window is not None else -1)
588 +
589 + # Inputs and outputs may be padded for CUDA graphs
590 + query = query[:num_actual_tokens]
591 + output_padded = output
592 + output = output[:num_actual_tokens]
593 +
594 + if attn_metadata.use_cascade:
595 + # Cascade attention (rare case).
596 + assert attn_metadata.cascade_wrapper is not None
597 + output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache))
598 + return output
599 +
600 + num_decode_tokens = attn_metadata.num_decode_tokens
601 + num_prefill_tokens = attn_metadata.num_prefill_tokens
602 +
603 + # Regular attention (common case).
604 + # Decodes are at the front and prefills are at the back,
605 + # according to reorder_batch()
606 + if prefill_wrapper := attn_metadata.prefill_wrapper:
607 + prefill_query = query[num_decode_tokens:]
608 + assert prefill_query.shape[0] == num_prefill_tokens
609 + assert prefill_wrapper is not None
610 + assert prefill_wrapper._causal
611 + assert prefill_wrapper._window_left == window_left
612 + assert prefill_wrapper._logits_soft_cap == (self.logits_soft_cap
613 + or 0.0)
614 + assert prefill_wrapper._sm_scale == self.scale
615 + prefill_wrapper.run(
616 + prefill_query,
617 + kv_cache,
618 + k_scale=layer._k_scale_float,
619 + v_scale=layer._v_scale_float,
620 + out=output[num_decode_tokens:],
621 + )
622 +
623 + if decode_wrapper := attn_metadata.decode_wrapper:
624 + decode_query = query[:num_decode_tokens]
625 + assert decode_query.shape[0] == num_decode_tokens
626 + assert decode_wrapper is not None
627 + assert decode_wrapper._window_left == window_left
628 + assert decode_wrapper._logits_soft_cap == (self.logits_soft_cap
629 + or 0.0)
630 + assert decode_wrapper._sm_scale == self.scale
631 + decode_wrapper.run(
632 + decode_query,
633 + kv_cache,
634 + k_scale=layer._k_scale_float,
635 + v_scale=layer._v_scale_float,
636 + out=output[:num_decode_tokens],
637 + )
638 +
639 + return output_padded
vllm/v1/attention/backends/mla/common.py
+3 -4
@@ -251,10 +251,6 @@ class MLACommonBackend(AttentionBackend):
251 def get_supported_head_sizes() -> list[int]:
252 return [576]
253
254 - @staticmethod
255 - def use_cascade_attention(*args, **kwargs) -> bool:
256 - return False
257 -
254
255 @dataclass
256 class MLACommonPrefillMetadata:
@@ -574,6 +570,9 @@ class MLACommonMetadataBuilder(Generic[M]):
570 decode=decode_metadata,
571 )
572
573 + def use_cascade_attention(self, *args, **kwargs) -> bool:
574 + return False
575 +
576
577 class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
578 """
vllm/v1/worker/gpu_model_runner.py
+1 -1
@@ -696,7 +696,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
696 # common_prefix_len should be a multiple of the block size.
697 common_prefix_len = (common_prefix_len // self.block_size *
698 self.block_size)
699 - use_cascade = self.attn_backend.use_cascade_attention(
699 + use_cascade = self.attn_metadata_builder.use_cascade_attention(
700 common_prefix_len=common_prefix_len,
701 query_lens=num_scheduled_tokens,
702 num_query_heads=self.num_query_heads,