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, 02:15 UTC
[Attention] MLA - Flashinfer Ragged Prefill (#20034)
Alexander Matveev committed
Jul 10, 2025 at 23:17 UTC
5b032352cc7285ed0b0d5c2fcbb9b7deab85d6c6
10 files changed
+422
-215
tests/v1/kv_connector/__init__.py
tests/v1/kv_connector/unit/test_multi_connector.py
+15
-72
@@ -3,16 +3,10 @@
3
import filecmp
4
import shutil
5
import tempfile
6
-from collections import defaultdict
6
from pathlib import Path
7
8
from vllm import LLM, SamplingParams
10
-from vllm.config import KVTransferConfig, VllmConfig
11
-from vllm.distributed.kv_transfer.kv_connector.factory import (
12
- KVConnectorFactory)
13
-from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa
14
- SharedStorageConnector)
15
-from vllm.v1.core.kv_cache_manager import KVCacheBlocks
9
+from vllm.config import KVTransferConfig
10
11
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
12
@@ -25,65 +19,6 @@ PROMPTS = [
19
SAMPLING_PARAMS = SamplingParams(temperature=0, max_tokens=20)
20
21
28
-class TestSharedStorageConnector(SharedStorageConnector):
29
-
30
- def __init__(self, config: VllmConfig, role):
31
- self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
32
- self._connector = SharedStorageConnector(config, role)
33
- self.call_record: dict[str, int] = defaultdict(int)
34
- # Use a unique temp file per connector
35
- self._event_file = tempfile.gettempdir(
36
- ) + f"/connector_{self.name}-{self.role.name}_events.log"
37
- # Start with an empty file
38
- with open(self._event_file, "w") as _:
39
- pass
40
-
41
- def __getattribute__(self, name):
42
- if name in ("_connector", "call_record", "name", "_event_file",
43
- "__class__", "__dict__", "__getattribute__",
44
- "__init__"): # avoid recursion
45
- return object.__getattribute__(self, name)
46
- if not hasattr(self._connector, name):
47
- return object.__getattribute__(self, name)
48
- attr = getattr(self._connector, name)
49
-
50
- # Intercept calls to the connector interface and write an event
51
- # for each one to a file, which can be read back in the main test proc.
52
- if callable(attr):
53
-
54
- def wrapper(*args, **kwargs):
55
- self.call_record[name] += 1
56
-
57
- # Include args that we're interested in
58
- to_log = [name]
59
- for arg in args:
60
- if isinstance(arg, int):
61
- to_log.append(str(arg))
62
- elif isinstance(arg, KVCacheBlocks):
63
- to_log.append(
64
- f"num_blocks={[len(b) for b in arg.blocks]}")
65
-
66
- # Log the event as a line to the file
67
- try:
68
- with open(self._event_file, "a") as f:
69
- f.write(' '.join(to_log) + "\n")
70
- except Exception as e:
71
- print(f"[ERROR] Could not log event {name} "
72
- f"for {self.name}: {e}")
73
- return attr(*args, **kwargs)
74
-
75
- return wrapper
76
- return attr
77
-
78
-
79
-# This relies on "fork" multiprocessing method being used.
80
-# It's the default but vLLM may fall back to spawn if for example CUDA
81
-# is already initialized.
82
-KVConnectorFactory.register_connector("TestSharedStorageConnector",
83
- TestSharedStorageConnector.__module__,
84
- TestSharedStorageConnector.__name__)
85
-
86
-
22
# Helper function to compare directories recursively
23
def _compare_directories(dir1: Path, dir2: Path) -> bool:
24
"""Compares two directories recursively for identical content."""
@@ -118,19 +53,27 @@ def test_multi_shared_storage_connector_consistency():
53
kv_role="kv_both",
54
kv_connector_extra_config={
55
"connectors": [{
121
- "kv_connector": "TestSharedStorageConnector",
122
- "kv_role": "kv_both",
56
+ "kv_connector":
57
+ "TestSharedStorageConnector",
58
+ "kv_role":
59
+ "kv_both",
60
"kv_connector_extra_config": {
61
"shared_storage_path": str(storage_1_path),
62
"name": "storage1",
126
- }
63
+ },
64
+ "kv_connector_module_path":
65
+ "tests.v1.kv_connector.unit.utils",
66
}, {
128
- "kv_connector": "TestSharedStorageConnector",
129
- "kv_role": "kv_both",
67
+ "kv_connector":
68
+ "TestSharedStorageConnector",
69
+ "kv_role":
70
+ "kv_both",
71
"kv_connector_extra_config": {
72
"shared_storage_path": str(storage_2_path),
73
"name": "storage2",
133
- }
74
+ },
75
+ "kv_connector_module_path":
76
+ "tests.v1.kv_connector.unit.utils",
77
}]
78
},
79
)
tests/v1/kv_connector/unit/utils.py
+62
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import tempfile
4
+from collections import defaultdict
5
from typing import Any, Optional
6
7
import torch
@@ -7,6 +9,11 @@ import torch
9
from vllm import SamplingParams
10
from vllm.config import (CacheConfig, DeviceConfig, KVTransferConfig,
11
ModelConfig, SchedulerConfig, VllmConfig)
12
+from vllm.distributed.kv_transfer.kv_connector.factory import (
13
+ KVConnectorFactory)
14
+from vllm.distributed.kv_transfer.kv_connector.v1.shared_storage_connector import ( # noqa
15
+ SharedStorageConnector)
16
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
17
from vllm.v1.core.sched.scheduler import Scheduler
18
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
19
KVCacheGroupSpec)
@@ -187,3 +194,58 @@ def create_model_runner_output(
194
finished_sending=finished_sending,
195
finished_recving=finished_recving,
196
)
197
+
198
+
199
+class TestSharedStorageConnector(SharedStorageConnector):
200
+
201
+ def __init__(self, config: VllmConfig, role):
202
+ self.name = config.kv_transfer_config.kv_connector_extra_config["name"]
203
+ self._connector = SharedStorageConnector(config, role)
204
+ self.call_record: dict[str, int] = defaultdict(int)
205
+ # Use a unique temp file per connector
206
+ self._event_file = tempfile.gettempdir(
207
+ ) + f"/connector_{self.name}-{self.role.name}_events.log"
208
+ # Start with an empty file
209
+ with open(self._event_file, "w") as _:
210
+ pass
211
+
212
+ def __getattribute__(self, name):
213
+ if name in ("_connector", "call_record", "name", "_event_file",
214
+ "__class__", "__dict__", "__getattribute__",
215
+ "__init__"): # avoid recursion
216
+ return object.__getattribute__(self, name)
217
+ if not hasattr(self._connector, name):
218
+ return object.__getattribute__(self, name)
219
+ attr = getattr(self._connector, name)
220
+
221
+ # Intercept calls to the connector interface and write an event
222
+ # for each one to a file, which can be read back in the main test proc.
223
+ if callable(attr):
224
+
225
+ def wrapper(*args, **kwargs):
226
+ self.call_record[name] += 1
227
+
228
+ # Include args that we're interested in
229
+ to_log = [name]
230
+ for arg in args:
231
+ if isinstance(arg, int):
232
+ to_log.append(str(arg))
233
+ elif isinstance(arg, KVCacheBlocks):
234
+ to_log.append(
235
+ f"num_blocks={[len(b) for b in arg.blocks]}")
236
+
237
+ # Log the event as a line to the file
238
+ try:
239
+ with open(self._event_file, "a") as f:
240
+ f.write(' '.join(to_log) + "\n")
241
+ except Exception as e:
242
+ print(f"[ERROR] Could not log event {name} "
243
+ f"for {self.name}: {e}")
244
+ return attr(*args, **kwargs)
245
+
246
+ return wrapper
247
+ return attr
248
+
249
+
250
+KVConnectorFactory.register_connector("TestSharedStorageConnector", __name__,
251
+ TestSharedStorageConnector.__name__)
vllm/attention/layer.py
+1
-1
@@ -10,6 +10,7 @@ import torch.nn.functional as F
10
import vllm.envs as envs
11
from vllm.attention import AttentionType
12
from vllm.attention.selector import backend_name_to_enum, get_attn_backend
13
+from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
14
from vllm.config import CacheConfig, get_current_vllm_config
15
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
16
has_kv_transfer_group,
@@ -21,7 +22,6 @@ from vllm.model_executor.layers.quantization.base_config import (
22
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
23
from vllm.platforms import _Backend, current_platform
24
from vllm.utils import direct_register_custom_op
24
-from vllm.v1.attention.backends.utils import validate_kv_sharing_target
25
26
27
class Attention(nn.Module):
vllm/attention/utils/kv_sharing_utils.py
new
+33
@@ -0,0 +1,33 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+def validate_kv_sharing_target(current_layer_name, target_layer_name,
4
+ static_forward_context):
5
+ error_msg = (f"Specified KV sharing target layer for {current_layer_name} "
6
+ f"is not valid: target layer {target_layer_name} ")
7
+
8
+ if current_layer_name == target_layer_name:
9
+ raise ValueError(error_msg +
10
+ "cannot be the same as the current layer.")
11
+
12
+ if target_layer_name not in static_forward_context:
13
+ from vllm.model_executor.models.utils import extract_layer_index
14
+
15
+ # If target layer name is not in the static fwd context, it means either
16
+ # a) the target layer does not come BEFORE the current layer, or
17
+ # b) the target layer is not an Attention layer that exists in the model
18
+ current_layer_idx = extract_layer_index(current_layer_name)
19
+ target_layer_idx = extract_layer_index(target_layer_name)
20
+ if current_layer_idx <= target_layer_idx:
21
+ raise ValueError(error_msg + "must come before the current layer.")
22
+ else:
23
+ raise ValueError(error_msg +
24
+ "is not a valid Attention layer in the model.")
25
+
26
+ # Currently KV sharing is only supported between layers of the same type
27
+ target_layer_attn_type = static_forward_context[
28
+ target_layer_name].attn_type
29
+ expected = static_forward_context[current_layer_name].attn_type
30
+ if target_layer_attn_type != expected:
31
+ raise ValueError(
32
+ error_msg +
33
+ f"must be the same type as the current layer ({expected}).")
vllm/logger.py
+14
@@ -53,6 +53,12 @@ DEFAULT_LOGGING_CONFIG = {
53
}
54
55
56
+@lru_cache
57
+def _print_debug_once(logger: Logger, msg: str, *args: Hashable) -> None:
58
+ # Set the stacklevel to 2 to print the original caller's line info
59
+ logger.debug(msg, *args, stacklevel=2)
60
+
61
+
62
@lru_cache
63
def _print_info_once(logger: Logger, msg: str, *args: Hashable) -> None:
64
# Set the stacklevel to 2 to print the original caller's line info
@@ -74,6 +80,13 @@ class _VllmLogger(Logger):
80
`intel_extension_for_pytorch.utils._logger`.
81
"""
82
83
+ def debug_once(self, msg: str, *args: Hashable) -> None:
84
+ """
85
+ As [`debug`][logging.Logger.debug], but subsequent calls with
86
+ the same message are silently dropped.
87
+ """
88
+ _print_debug_once(self, msg, *args)
89
+
90
def info_once(self, msg: str, *args: Hashable) -> None:
91
"""
92
As [`info`][logging.Logger.info], but subsequent calls with
@@ -132,6 +145,7 @@ def init_logger(name: str) -> _VllmLogger:
145
logger = logging.getLogger(name)
146
147
methods_to_patch = {
148
+ "debug_once": _print_debug_once,
149
"info_once": _print_info_once,
150
"warning_once": _print_warning_once,
151
}
vllm/v1/attention/backends/flashinfer.py
+5
-68
@@ -14,13 +14,14 @@ from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
14
import vllm.envs as envs
15
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
16
AttentionType)
17
-from vllm.attention.layer import Attention
18
-from vllm.config import VllmConfig, get_layers_from_vllm_config
17
from vllm.logger import init_logger
18
from vllm.v1.attention.backends.flash_attn import use_cascade_attention
19
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
20
CommonAttentionMetadata,
23
- get_kv_cache_layout)
21
+ PerLayerParameters,
22
+ get_kv_cache_layout,
23
+ get_per_layer_parameters,
24
+ infer_global_hyperparameters)
25
from vllm.v1.kv_cache_interface import AttentionSpec
26
from vllm.v1.worker.block_table import BlockTable
27
@@ -93,70 +94,6 @@ class FlashInferBackend(AttentionBackend):
94
return stride_order
95
96
96
-@dataclass
97
-class PerLayerParameters:
98
- """
99
- Currently, FlashInfer backend only support models in which all layers share
100
- the same values for the following hyperparameters.
101
- """
102
-
103
- window_left: int
104
- logits_soft_cap: Optional[float]
105
- sm_scale: float
106
-
107
-
108
-def get_per_layer_parameters(
109
- vllm_config: VllmConfig) -> dict[str, PerLayerParameters]:
110
- """
111
- Scan all attention layers and determine some hyperparameters
112
- to use during `plan`.
113
- """
114
-
115
- layers = get_layers_from_vllm_config(vllm_config, Attention)
116
- per_layer_params: dict[str, PerLayerParameters] = {}
117
-
118
- for key, layer in layers.items():
119
- impl = layer.impl
120
- assert isinstance(impl, FlashInferImpl)
121
-
122
- # Infer hyperparameters from the attention layer
123
- window_size = impl.sliding_window
124
- window_left = window_size[0] if window_size is not None else -1
125
- logits_soft_cap = impl.logits_soft_cap
126
- sm_scale = impl.scale
127
-
128
- per_layer_params[key] = PerLayerParameters(window_left,
129
- logits_soft_cap, sm_scale)
130
-
131
- return per_layer_params
132
-
133
-
134
-def infer_global_hyperparameters(
135
- per_layer_params: dict[str, PerLayerParameters]) -> PerLayerParameters:
136
- """
137
- Currently, FlashInfer backend only support models in which all layers share
138
- the same values for the following hyperparameters:
139
- - `window_left`
140
- - `logits_soft_cap`
141
- - `sm_scale`
142
-
143
- So this function asserts that all layers share the same values for these
144
- hyperparameters and returns the global values.
145
- """
146
-
147
- assert len(per_layer_params) > 0, "No attention layers found in the model."
148
-
149
- param_sets = list(per_layer_params.values())
150
- global_params = param_sets[0]
151
- for params in param_sets:
152
- assert params == global_params, (
153
- "FlashInfer backend currently only supports models in which all "
154
- "layers share the same values for the following hyperparameters: "
155
- "`window_left`, `logits_soft_cap`, `sm_scale`.")
156
-
157
- return global_params
158
-
159
-
97
@dataclass
98
class FlashInferMetadata:
99
@@ -336,7 +273,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
273
def _plan(self, attn_metadata: FlashInferMetadata):
274
if self.global_hyperparameters is None:
275
self.global_hyperparameters = infer_global_hyperparameters(
339
- get_per_layer_parameters(self.vllm_config))
276
+ get_per_layer_parameters(self.vllm_config, FlashInferImpl))
277
if attn_metadata.use_cascade:
278
attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
279
attn_metadata.cascade_wrapper.plan(
vllm/v1/attention/backends/mla/common.py
+222
-40
@@ -189,8 +189,8 @@ return curr_o @ W_O
189
190
import functools
191
from abc import abstractmethod
192
-from dataclasses import dataclass
193
-from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar
192
+from dataclasses import dataclass, field
193
+from typing import TYPE_CHECKING, Any, Generic, Optional, TypeVar, Union
194
195
import torch
196
@@ -208,7 +208,9 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
208
from vllm.platforms import current_platform
209
from vllm.utils import cdiv, round_down
210
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
211
- CommonAttentionMetadata)
211
+ CommonAttentionMetadata,
212
+ get_per_layer_parameters,
213
+ infer_global_hyperparameters)
214
from vllm.v1.kv_cache_interface import AttentionSpec
215
from vllm.v1.worker.block_table import BlockTable
216
@@ -221,6 +223,12 @@ except ImportError:
223
from flash_attn import flash_attn_varlen_func
224
is_vllm_fa = False
225
226
+try:
227
+ from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
228
+ flashinfer_available = True
229
+except ImportError:
230
+ flashinfer_available = False
231
+
232
if TYPE_CHECKING:
233
from vllm.v1.core.sched.output import SchedulerOutput
234
from vllm.v1.worker.gpu_input_batch import InputBatch
@@ -290,6 +298,13 @@ class MLACommonPrefillMetadata:
298
chunked_context: Optional[ChunkedContextMetadata] = None
299
300
301
+@dataclass
302
+class FlashInferPrefillMetadata(MLACommonPrefillMetadata):
303
+ prefill_main: Optional['BatchPrefillWithRaggedKVCacheWrapper'] = None
304
+ prefill_chunks: list['BatchPrefillWithRaggedKVCacheWrapper'] = field(
305
+ default_factory=list)
306
+
307
+
308
@dataclass
309
class MLACommonDecodeMetadata:
310
block_table: torch.Tensor
@@ -328,7 +343,8 @@ class MLACommonMetadata(Generic[D]):
343
head_dim: Optional[int] = None
344
345
decode: Optional[D] = None
331
- prefill: Optional[MLACommonPrefillMetadata] = None
346
+ prefill: Optional[Union[MLACommonPrefillMetadata,
347
+ FlashInferPrefillMetadata]] = None
348
349
def __post_init__(self):
350
if self.head_dim is not None:
@@ -338,6 +354,20 @@ class MLACommonMetadata(Generic[D]):
354
M = TypeVar("M", bound=MLACommonMetadata)
355
356
357
+def use_flashinfer_prefill() -> bool:
358
+ if flashinfer_available:
359
+ # For blackwell default to flashinfer prefill if its available since
360
+ # its faster than FA2.
361
+ return current_platform.has_device_capability(100)
362
+ return False
363
+
364
+
365
+# Currently 394MB, this can be tuned based on GEMM sizes used.
366
+# Choosen to be the same as sglang:
367
+# https://github.com/sgl-project/sglang/blob/766392c6bda2558b61ce6d1c1bfd8081a549e1f1/python/sglang/global_config.py#L37
368
+FLASHINFER_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024
369
+
370
+
371
class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
372
"""
373
NOTE: Please read the comment at the top of the file before trying to
@@ -392,6 +422,101 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
422
)
423
self.block_table = block_table
424
425
+ self._use_fi_prefill = use_flashinfer_prefill()
426
+ self.prefill_metadata_cls = FlashInferPrefillMetadata \
427
+ if self._use_fi_prefill else MLACommonPrefillMetadata
428
+
429
+ if self._use_fi_prefill:
430
+ self._workspace_buffer = torch.empty(
431
+ FLASHINFER_WORKSPACE_BUFFER_SIZE,
432
+ dtype=torch.uint8,
433
+ device=runner.device)
434
+
435
+ self._fi_prefill_main: Optional[
436
+ BatchPrefillWithRaggedKVCacheWrapper] = None
437
+ self._fi_prefill_chunks: list[
438
+ BatchPrefillWithRaggedKVCacheWrapper] = []
439
+
440
+ self._global_hyperparameters = infer_global_hyperparameters(
441
+ get_per_layer_parameters(runner.vllm_config, MLACommonImpl))
442
+
443
+ def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):
444
+ qo_indptr = prefill.query_start_loc
445
+
446
+ has_context = False
447
+ if prefill.chunked_context is not None:
448
+ chunked_context = prefill.chunked_context
449
+ has_context = True
450
+
451
+ if self._fi_prefill_main is None:
452
+ self._fi_prefill_main = BatchPrefillWithRaggedKVCacheWrapper(
453
+ self._workspace_buffer, "NHD", backend="cutlass")
454
+
455
+ if has_context:
456
+ num_chunks = chunked_context.cu_seq_lens.shape[0]
457
+ # Allocate more prefill chunk wrappers if needed
458
+ if len(self._fi_prefill_chunks) < num_chunks:
459
+ for _ in range(len(self._fi_prefill_chunks), num_chunks):
460
+ self._fi_prefill_chunks.append(
461
+ BatchPrefillWithRaggedKVCacheWrapper(
462
+ self._workspace_buffer, "NHD", backend="cutlass"))
463
+ assert num_chunks <= len(self._fi_prefill_chunks)
464
+
465
+ # In MLA, the non-latent num_qo_heads == num_kv_heads
466
+ num_qo_heads = self.runner.num_query_heads
467
+ num_kv_heads = num_qo_heads
468
+
469
+ # Sanity: Verify that num_kv_heads == 1 since it is latent space
470
+ assert self.kv_cache_spec.num_kv_heads == 1
471
+
472
+ # Get non-latent head_dim_qk and head_dim_vo
473
+ head_dim_qk = (self.mla_dims.qk_nope_head_dim +
474
+ self.mla_dims.qk_rope_head_dim)
475
+ head_dim_vo = self.mla_dims.v_head_dim
476
+
477
+ # For main run, qo_indptr == kv_indptr
478
+ kv_indptr = qo_indptr.clone()
479
+
480
+ # Prepare main prefill
481
+ self._fi_prefill_main.plan(
482
+ qo_indptr=qo_indptr,
483
+ kv_indptr=kv_indptr,
484
+ num_qo_heads=num_qo_heads,
485
+ num_kv_heads=num_kv_heads,
486
+ head_dim_qk=head_dim_qk,
487
+ head_dim_vo=head_dim_vo,
488
+ causal=True, # This is main run
489
+ sm_scale=self._global_hyperparameters.sm_scale,
490
+ window_left=self._global_hyperparameters.window_left,
491
+ logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
492
+ q_data_type=self.runner.dtype,
493
+ kv_data_type=self.kv_cache_spec.dtype,
494
+ )
495
+
496
+ # Prepare context prefills
497
+ if has_context:
498
+ for i in range(num_chunks):
499
+ kv_indptr_chunk = chunked_context.cu_seq_lens[i]
500
+
501
+ self._fi_prefill_chunks[i].plan(
502
+ qo_indptr=qo_indptr,
503
+ kv_indptr=kv_indptr_chunk,
504
+ num_qo_heads=num_qo_heads,
505
+ num_kv_heads=num_kv_heads,
506
+ head_dim_qk=head_dim_qk,
507
+ head_dim_vo=head_dim_vo,
508
+ causal=False, # This is context run
509
+ sm_scale=self._global_hyperparameters.sm_scale,
510
+ window_left=self._global_hyperparameters.window_left,
511
+ logits_soft_cap=self._global_hyperparameters.
512
+ logits_soft_cap,
513
+ q_data_type=self.runner.dtype,
514
+ kv_data_type=self.kv_cache_spec.dtype,
515
+ )
516
+
517
+ prefill.prefill_main = self._fi_prefill_main
518
+ prefill.prefill_chunks = self._fi_prefill_chunks
519
+
520
def reorder_batch(self, input_batch: "InputBatch",
521
scheduler_output: "SchedulerOutput") -> bool:
522
# We now want to reorder the batch so that the "decode" requests are and
@@ -572,7 +697,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
697
assert max(chunked_context_metadata.max_seq_lens) <= \
698
self.chunked_prefill_workspace_size
699
575
- prefill_metadata = MLACommonPrefillMetadata(
700
+ prefill_metadata = self.prefill_metadata_cls(
701
block_table=block_table_tensor[reqs_start:, ...],
702
query_start_loc=prefill_query_start_loc,
703
max_query_len=max_query_len,
@@ -586,7 +711,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
711
seq_lens=seq_lens[:self._num_decodes],
712
)
713
589
- return self.metadata_cls(
714
+ attn_metadata = self.metadata_cls(
715
num_actual_tokens=num_actual_tokens,
716
query_start_loc=query_start_loc,
717
slot_mapping=slot_mapping,
@@ -599,6 +724,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
724
decode=decode_metadata,
725
)
726
727
+ if self._use_fi_prefill and self._num_prefills > 0:
728
+ assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata)
729
+ self._build_fi_prefill_wrappers(attn_metadata.prefill)
730
+
731
+ return attn_metadata
732
+
733
def can_run_in_cudagraph(
734
self, common_attn_metadata: CommonAttentionMetadata) -> bool:
735
return common_attn_metadata.max_query_len == 1
@@ -649,23 +780,34 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
780
self.v_head_dim = v_head_dim
781
self.kv_b_proj = kv_b_proj
782
652
- # Handle the differences between the flash_attn_varlen from flash_attn
653
- # and the one from vllm_flash_attn. The former is used on RoCM and the
654
- # latter has an additional parameter to control FA2 vs FA3
655
- self.flash_attn_varlen_func = flash_attn_varlen_func
656
- self.vllm_flash_attn_version = get_flash_attn_version()
657
- if self.vllm_flash_attn_version is not None:
658
- self.flash_attn_varlen_func = \
659
- functools.partial(flash_attn_varlen_func,
660
- fa_version=self.vllm_flash_attn_version)
661
-
662
- # For MLA the v head dim is smaller than qk head dim so we pad out
663
- # v with 0s to match the qk head dim for attention backends that do
664
- # not support different headdims
665
- # We don't need to pad V if we are on a hopper system with FA3
666
- self._pad_v = self.vllm_flash_attn_version is None or not (
667
- self.vllm_flash_attn_version == 3
668
- and current_platform.get_device_capability()[0] == 9)
783
+ if use_flashinfer_prefill():
784
+ logger.debug_once("Using FlashInfer prefill for MLA")
785
+ self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi
786
+ self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi
787
+ self._pad_v = False
788
+ else: # Use FlashAttention
789
+ logger.debug_once("Using FlashAttention prefill for MLA")
790
+ self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa
791
+ self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa
792
+
793
+ # Handle the differences between the flash_attn_varlen from
794
+ # flash_attn and the one from vllm_flash_attn. The former is used on
795
+ # RoCM and the latter has an additional parameter to control
796
+ # FA2 vs FA3
797
+ self.flash_attn_varlen_func = flash_attn_varlen_func
798
+ self.vllm_flash_attn_version = get_flash_attn_version()
799
+ if self.vllm_flash_attn_version is not None:
800
+ self.flash_attn_varlen_func = \
801
+ functools.partial(flash_attn_varlen_func,
802
+ fa_version=self.vllm_flash_attn_version)
803
+
804
+ # For MLA the v head dim is smaller than qk head dim so we pad out
805
+ # v with 0s to match the qk head dim for attention backends that do
806
+ # not support different headdims
807
+ # We don't need to pad V if we are on a hopper system with FA3
808
+ self._pad_v = self.vllm_flash_attn_version is None or not (
809
+ self.vllm_flash_attn_version == 3
810
+ and current_platform.get_device_capability()[0] == 9)
811
812
def _flash_attn_varlen_diff_headdims(self,
813
q,
@@ -705,6 +847,58 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
847
return attn_out, lse
848
return attn_out
849
850
+ def _run_prefill_new_tokens_fa(self, prefill: MLACommonPrefillMetadata, q,
851
+ k, v, return_softmax_lse):
852
+ return self._flash_attn_varlen_diff_headdims(
853
+ q=q,
854
+ k=k,
855
+ v=v,
856
+ cu_seqlens_q=prefill.query_start_loc,
857
+ cu_seqlens_k=prefill.query_start_loc,
858
+ max_seqlen_q=prefill.max_query_len,
859
+ max_seqlen_k=prefill.max_query_len,
860
+ softmax_scale=self.scale,
861
+ causal=True,
862
+ return_softmax_lse=return_softmax_lse,
863
+ )
864
+
865
+ def _run_prefill_new_tokens_fi(self, prefill: MLACommonPrefillMetadata, q,
866
+ k, v, return_softmax_lse):
867
+ assert isinstance(prefill, FlashInferPrefillMetadata)
868
+ assert prefill.prefill_main is not None
869
+ return prefill.prefill_main.run(
870
+ q=q,
871
+ k=k,
872
+ v=v,
873
+ return_lse=return_softmax_lse,
874
+ )
875
+
876
+ def _run_prefill_context_chunk_fa(self, prefill: MLACommonPrefillMetadata,
877
+ chunk_idx: int, q, k, v):
878
+ assert prefill.chunked_context is not None
879
+ return self._flash_attn_varlen_diff_headdims(
880
+ q=q,
881
+ k=k,
882
+ v=v,
883
+ cu_seqlens_q=prefill.query_start_loc,
884
+ cu_seqlens_k=prefill.chunked_context.cu_seq_lens[chunk_idx],
885
+ max_seqlen_q=prefill.max_query_len,
886
+ max_seqlen_k=prefill.chunked_context.max_seq_lens[chunk_idx],
887
+ softmax_scale=self.scale,
888
+ causal=False, # Context is unmasked
889
+ return_softmax_lse=True,
890
+ )
891
+
892
+ def _run_prefill_context_chunk_fi(self, prefill: MLACommonPrefillMetadata,
893
+ chunk_idx: int, q, k, v):
894
+ assert isinstance(prefill, FlashInferPrefillMetadata)
895
+ return prefill.prefill_chunks[chunk_idx].run(
896
+ q=q,
897
+ k=k,
898
+ v=v,
899
+ return_lse=True,
900
+ )
901
+
902
def _v_up_proj(self, x):
903
# Convert from (B, N, L) to (N, B, L)
904
x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
@@ -803,18 +997,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
997
k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),
998
dim=-1)
999
806
- attn_output, attn_softmax_lse = \
807
- self._flash_attn_varlen_diff_headdims(
1000
+ attn_output, attn_softmax_lse = self._run_prefill_context_chunk(
1001
+ prefill=prefill_metadata,
1002
+ chunk_idx=i,
1003
q=q,
1004
k=k,
1005
v=v,
811
- cu_seqlens_q=prefill_metadata.query_start_loc,
812
- cu_seqlens_k=prefill_metadata.chunked_context.cu_seq_lens[i],
813
- max_seqlen_q=prefill_metadata.max_query_len,
814
- max_seqlen_k=prefill_metadata.chunked_context.max_seq_lens[i],
815
- softmax_scale=self.scale,
816
- causal=False, # Context is unmasked
817
- return_softmax_lse=True,
1006
)
1007
1008
if output is None:
@@ -854,16 +1042,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1042
1043
k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
1044
857
- output = self._flash_attn_varlen_diff_headdims(
1045
+ output = self._run_prefill_new_tokens(
1046
+ prefill=attn_metadata.prefill,
1047
q=q,
1048
k=k,
1049
v=v,
861
- cu_seqlens_q=attn_metadata.prefill.query_start_loc,
862
- cu_seqlens_k=attn_metadata.prefill.query_start_loc,
863
- max_seqlen_q=attn_metadata.prefill.max_query_len,
864
- max_seqlen_k=attn_metadata.prefill.max_query_len,
865
- softmax_scale=self.scale,
866
- causal=True,
1050
return_softmax_lse=has_context,
1051
)
1052
@@ -908,7 +1091,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
1091
output: Optional[torch.Tensor] = None,
1092
output_scale: Optional[torch.Tensor] = None,
1093
) -> torch.Tensor:
911
-
1094
assert output is not None, "Output tensor must be provided."
1095
1096
if output_scale is not None:
vllm/v1/attention/backends/mla/cutlass_mla.py
+1
@@ -91,6 +91,7 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
91
# Clone q_nope and q_pe to make sure strides computation is correct.
92
q_nope = q_nope.clone()
93
q_pe = q_pe.clone()
94
+
95
ops.cutlass_mla_decode(o, q_nope, q_pe, kv_c_and_k_pe_cache,
96
attn_metadata.decode.seq_lens,
97
attn_metadata.decode.block_table, self.scale)
vllm/v1/attention/backends/utils.py
+69
-34
@@ -4,14 +4,17 @@ import abc
4
import functools
5
from abc import abstractmethod
6
from dataclasses import dataclass
7
-from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar
7
+from typing import TYPE_CHECKING, ClassVar, Generic, Optional, TypeVar
8
9
import numpy as np
10
import torch
11
12
+from vllm.attention.layer import Attention
13
+from vllm.config import VllmConfig, get_layers_from_vllm_config
14
from vllm.utils import cdiv
15
16
if TYPE_CHECKING:
17
+ from vllm.attention.backends.abstract import AttentionImpl
18
from vllm.v1.core.sched.output import SchedulerOutput
19
from vllm.v1.worker.gpu_input_batch import InputBatch
20
@@ -98,39 +101,6 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
101
return False
102
103
101
-def validate_kv_sharing_target(current_layer_name, target_layer_name,
102
- static_forward_context):
103
- error_msg = (f"Specified KV sharing target layer for {current_layer_name} "
104
- f"is not valid: target layer {target_layer_name} ")
105
-
106
- if current_layer_name == target_layer_name:
107
- raise ValueError(error_msg +
108
- "cannot be the same as the current layer.")
109
-
110
- if target_layer_name not in static_forward_context:
111
- from vllm.model_executor.models.utils import extract_layer_index
112
-
113
- # If target layer name is not in the static fwd context, it means either
114
- # a) the target layer does not come BEFORE the current layer, or
115
- # b) the target layer is not an Attention layer that exists in the model
116
- current_layer_idx = extract_layer_index(current_layer_name)
117
- target_layer_idx = extract_layer_index(target_layer_name)
118
- if current_layer_idx <= target_layer_idx:
119
- raise ValueError(error_msg + "must come before the current layer.")
120
- else:
121
- raise ValueError(error_msg +
122
- "is not a valid Attention layer in the model.")
123
-
124
- # Currently KV sharing is only supported between layers of the same type
125
- target_layer_attn_type = static_forward_context[
126
- target_layer_name].attn_type
127
- expected = static_forward_context[current_layer_name].attn_type
128
- if target_layer_attn_type != expected:
129
- raise ValueError(
130
- error_msg +
131
- f"must be the same type as the current layer ({expected}).")
132
-
133
-
104
@functools.lru_cache
105
def get_kv_cache_layout():
106
# Override with format specified by the user.
@@ -144,6 +114,71 @@ def get_kv_cache_layout():
114
return cache_layout
115
116
117
+@dataclass
118
+class PerLayerParameters:
119
+ """
120
+ Currently, FlashInfer backend only support models in which all layers share
121
+ the same values for the following hyperparameters.
122
+ """
123
+
124
+ window_left: int
125
+ logits_soft_cap: Optional[float]
126
+ sm_scale: float
127
+
128
+
129
+def get_per_layer_parameters(
130
+ vllm_config: VllmConfig,
131
+ cls_: type['AttentionImpl']) -> dict[str, PerLayerParameters]:
132
+ """
133
+ Scan all attention layers and determine some hyperparameters
134
+ to use during `plan`.
135
+ """
136
+
137
+ layers = get_layers_from_vllm_config(vllm_config, Attention)
138
+ per_layer_params: dict[str, PerLayerParameters] = {}
139
+
140
+ for key, layer in layers.items():
141
+ impl = layer.impl
142
+ assert isinstance(impl, cls_)
143
+
144
+ # Infer hyperparameters from the attention layer
145
+ window_size = getattr(impl, "sliding_window", None)
146
+ window_left = window_size[0] if window_size is not None else -1
147
+ logits_soft_cap = getattr(impl, "logits_soft_cap", None)
148
+ sm_scale = impl.scale
149
+
150
+ per_layer_params[key] = PerLayerParameters(window_left,
151
+ logits_soft_cap, sm_scale)
152
+
153
+ return per_layer_params
154
+
155
+
156
+def infer_global_hyperparameters(
157
+ per_layer_params: dict[str, PerLayerParameters]) -> PerLayerParameters:
158
+ """
159
+ Currently, FlashInfer backend only support models in which all layers share
160
+ the same values for the following hyperparameters:
161
+ - `window_left`
162
+ - `logits_soft_cap`
163
+ - `sm_scale`
164
+
165
+ So this function asserts that all layers share the same values for these
166
+ hyperparameters and returns the global values.
167
+ """
168
+
169
+ assert len(per_layer_params) > 0, "No attention layers found in the model."
170
+
171
+ param_sets = list(per_layer_params.values())
172
+ global_params = param_sets[0]
173
+ for params in param_sets:
174
+ assert params == global_params, (
175
+ "FlashInfer backend currently only supports models in which all "
176
+ "layers share the same values for the following hyperparameters: "
177
+ "`window_left`, `logits_soft_cap`, `sm_scale`.")
178
+
179
+ return global_params
180
+
181
+
182
#
183
# Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into
184
# local attention blocks, where each block is passed to the attention kernel