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 10, 05:15 UTC
[6/N][Attention] Move utils to more appropriate locations (#32215)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Jan 13, 2026 at 08:38 UTC
98f60e5acbcbbc9b6e7d6370c4e19128fc3b0b32
14 files changed
+168
-178
tests/v1/attention/test_attention_splitting.py
+4
-2
@@ -7,13 +7,15 @@ import torch
7
from tests.v1.attention.test_attention_backends import BATCH_SPECS
8
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
9
from vllm.v1.attention.backends.utils import (
10
+ split_decodes_and_prefills,
11
+)
12
+from vllm.v1.worker.ubatch_utils import (
13
UBatchSlice,
14
_make_metadata_with_slice,
15
+ maybe_create_ubatch_slices,
16
slice_query_start_locs,
17
split_attn_metadata,
14
- split_decodes_and_prefills,
18
)
16
-from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
19
20
21
@pytest.fixture
vllm/model_executor/layers/attention/chunked_local_attention.py
+1
-1
@@ -13,10 +13,10 @@ from vllm.v1.attention.backend import (
13
AttentionCGSupport,
14
AttentionMetadataBuilder,
15
CommonAttentionMetadata,
16
+ subclass_attention_backend,
17
)
18
from vllm.v1.attention.backends.utils import (
19
make_local_attention_virtual_batches,
19
- subclass_attention_backend,
20
)
21
from vllm.v1.attention.selector import get_attn_backend
22
from vllm.v1.kv_cache_interface import (
vllm/model_executor/layers/attention/cross_attention.py
-2
@@ -15,8 +15,6 @@ from vllm.v1.attention.backend import (
15
AttentionMetadata,
16
AttentionType,
17
CommonAttentionMetadata,
18
-)
19
-from vllm.v1.attention.backends.utils import (
18
subclass_attention_backend,
19
)
20
from vllm.v1.attention.selector import get_attn_backend
vllm/model_executor/layers/attention/encoder_only_attention.py
-2
@@ -13,8 +13,6 @@ from vllm.v1.attention.backend import (
13
AttentionMetadata,
14
AttentionType,
15
CommonAttentionMetadata,
16
-)
17
-from vllm.v1.attention.backends.utils import (
16
subclass_attention_backend,
17
)
18
from vllm.v1.attention.selector import get_attn_backend
vllm/model_executor/layers/attention/static_sink_attention.py
-2
@@ -16,8 +16,6 @@ from vllm.v1.attention.backend import (
16
AttentionMetadata,
17
AttentionType,
18
CommonAttentionMetadata,
19
-)
20
-from vllm.v1.attention.backends.utils import (
19
subclass_attention_backend,
20
)
21
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
vllm/model_executor/models/whisper_utils.py
+1
-3
@@ -17,11 +17,9 @@ from vllm.v1.attention.backend import (
17
AttentionMetadata,
18
AttentionType,
19
CommonAttentionMetadata,
20
-)
21
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
22
-from vllm.v1.attention.backends.utils import (
20
subclass_attention_backend_with_overrides,
21
)
22
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
23
from vllm.v1.attention.selector import get_attn_backend
24
from vllm.v1.kv_cache_interface import AttentionSpec
25
vllm/v1/attention/backend.py
+25
-1
@@ -4,7 +4,7 @@
4
from abc import ABC, abstractmethod
5
from dataclasses import dataclass
6
from enum import Enum
7
-from typing import TYPE_CHECKING, ClassVar, Generic, Protocol, TypeVar, get_args
7
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, Protocol, TypeVar, get_args
8
9
import numpy as np
10
import torch
@@ -734,3 +734,27 @@ class MLAAttentionImpl(AttentionImpl[T], Generic[T]):
734
735
def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
736
return kv_cache_dtype != "auto"
737
+
738
+
739
+def subclass_attention_backend(
740
+ name_prefix: str,
741
+ attention_backend_cls: type[AttentionBackend],
742
+ builder_cls: type[AttentionMetadataBuilder[M]],
743
+) -> type[AttentionBackend]:
744
+ """
745
+ Return a new subclass where `get_builder_cls` returns `builder_cls`.
746
+ """
747
+ name: str = name_prefix + attention_backend_cls.__name__ # type: ignore
748
+
749
+ return type(
750
+ name, (attention_backend_cls,), {"get_builder_cls": lambda: builder_cls}
751
+ )
752
+
753
+
754
+def subclass_attention_backend_with_overrides(
755
+ name_prefix: str,
756
+ attention_backend_cls: type[AttentionBackend],
757
+ overrides: dict[str, Any],
758
+) -> type[AttentionBackend]:
759
+ name: str = name_prefix + attention_backend_cls.__name__ # type: ignore
760
+ return type(name, (attention_backend_cls,), overrides)
vllm/v1/attention/backends/utils.py
+1
-159
@@ -8,7 +8,6 @@ from typing import (
8
Any,
9
Literal,
10
Protocol,
11
- TypeVar,
11
get_args,
12
)
13
@@ -33,10 +32,9 @@ from vllm.v1.attention.backend import (
32
AttentionBackend,
33
AttentionImpl,
34
AttentionMetadata,
36
- AttentionMetadataBuilder,
35
CommonAttentionMetadata,
36
+ subclass_attention_backend,
37
)
39
-from vllm.v1.worker.ubatch_utils import UBatchSlice
38
39
logger = init_logger(__name__)
40
KVCacheLayoutType = Literal["NHD", "HND"]
@@ -49,135 +47,6 @@ def is_valid_kv_cache_layout(value: str) -> bool:
47
return value in get_args(KVCacheLayoutType)
48
49
52
-def slice_query_start_locs(
53
- query_start_loc: torch.Tensor,
54
- request_slice: slice,
55
-) -> torch.Tensor:
56
- """
57
- Creates a new query_start_loc that corresponds to the requests in
58
- request_slice.
59
-
60
- Note: This function creates a new tensor to hold the new query_start_locs.
61
- This will break cudagraph compatibility.
62
- """
63
- return (
64
- query_start_loc[request_slice.start : request_slice.stop + 1]
65
- - query_start_loc[request_slice.start]
66
- )
67
-
68
-
69
-def _make_metadata_with_slice(
70
- ubatch_slice: UBatchSlice, attn_metadata: CommonAttentionMetadata
71
-) -> CommonAttentionMetadata:
72
- """
73
- This function creates a new CommonAttentionMetadata that corresponds to
74
- the requests included in ubatch_slice
75
- """
76
-
77
- assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty"
78
-
79
- request_slice = ubatch_slice.request_slice
80
- token_slice = ubatch_slice.token_slice
81
-
82
- start_locs = attn_metadata.query_start_loc_cpu
83
- first_req = request_slice.start
84
- first_tok = token_slice.start
85
- last_req = request_slice.stop - 1
86
- last_tok = token_slice.stop - 1
87
-
88
- assert start_locs[first_req] <= first_tok < start_locs[first_req + 1], (
89
- "Token slice start outside of first request"
90
- )
91
- # NOTE: last token can be outside of the last request if we have CG padding.
92
-
93
- # If the request is split across ubatches, we have to adjust the metadata.
94
- # splits_first_request: The first request in this slice is the continuation of
95
- # a request that started in a previous slice.
96
- # splits_last_request: The last request in this slice continues into the
97
- # next slice.
98
- splits_first_request = first_tok > start_locs[first_req]
99
- splits_last_request = last_tok < start_locs[last_req + 1] - 1
100
-
101
- query_start_loc_cpu = slice_query_start_locs(start_locs, request_slice)
102
- query_start_loc = slice_query_start_locs(
103
- attn_metadata.query_start_loc, request_slice
104
- )
105
-
106
- assert len(query_start_loc) >= 2, (
107
- f"query_start_loc must have at least 2 elements, got {len(query_start_loc)}"
108
- )
109
-
110
- if splits_first_request:
111
- tokens_skipped = first_tok - start_locs[first_req]
112
- query_start_loc[1:] -= tokens_skipped
113
- query_start_loc_cpu[1:] -= tokens_skipped
114
- seq_lens = attn_metadata.seq_lens[request_slice]
115
- seq_lens_cpu = attn_metadata.seq_lens_cpu[request_slice]
116
-
117
- if splits_last_request:
118
- # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate
119
- # the tokens skipped because query_start_loc_cpu might have been modified
120
- # if splits_first_request is True.
121
- tokens_skipped = start_locs[last_req + 1] - token_slice.stop
122
- query_start_loc[-1] -= tokens_skipped
123
- query_start_loc_cpu[-1] -= tokens_skipped
124
-
125
- # Make sure we don't modify the seq_lens tensors
126
- # (not cudagraph compatible)
127
- seq_lens = seq_lens.clone()
128
- seq_lens_cpu = seq_lens_cpu.clone()
129
- seq_lens[-1] -= tokens_skipped
130
- seq_lens_cpu[-1] -= tokens_skipped
131
-
132
- max_seq_len = int(seq_lens_cpu.max())
133
- num_computed_tokens_cpu = attn_metadata.num_computed_tokens_cpu[request_slice]
134
-
135
- num_requests = request_slice.stop - request_slice.start
136
- num_actual_tokens = token_slice.stop - token_slice.start
137
- max_query_len = int(
138
- torch.max(torch.abs(query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])).item()
139
- )
140
-
141
- # This is to account for the case where we are in a dummy
142
- # run and query_start_loc_cpu is full of 0s
143
- if max_query_len == 0:
144
- max_query_len = attn_metadata.max_query_len
145
-
146
- block_table_tensor = attn_metadata.block_table_tensor[request_slice]
147
- slot_mapping = attn_metadata.slot_mapping[token_slice]
148
-
149
- return CommonAttentionMetadata(
150
- query_start_loc=query_start_loc,
151
- query_start_loc_cpu=query_start_loc_cpu,
152
- seq_lens=seq_lens,
153
- num_reqs=num_requests,
154
- num_actual_tokens=num_actual_tokens,
155
- max_query_len=max_query_len,
156
- max_seq_len=max_seq_len,
157
- block_table_tensor=block_table_tensor,
158
- slot_mapping=slot_mapping,
159
- _seq_lens_cpu=seq_lens_cpu,
160
- _num_computed_tokens_cpu=num_computed_tokens_cpu,
161
- )
162
-
163
-
164
-def split_attn_metadata(
165
- ubatch_slices: list[UBatchSlice],
166
- common_attn_metadata: CommonAttentionMetadata,
167
-) -> list[CommonAttentionMetadata]:
168
- """
169
- Creates a new CommonAttentionMetadata instance that corresponds to the
170
- requests for each UBatchSlice in ubatch_slices.
171
-
172
- Note: This function does not modify common_attn_metadata
173
- """
174
- results = []
175
- for ubatch_slice in ubatch_slices:
176
- results.append(_make_metadata_with_slice(ubatch_slice, common_attn_metadata))
177
-
178
- return results
179
-
180
-
50
@functools.lru_cache
51
def get_kv_cache_layout():
52
# Format specified by the code.
@@ -548,33 +417,6 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
417
return common_attn_metadata
418
419
551
-M = TypeVar("M")
552
-
553
-
554
-def subclass_attention_backend(
555
- name_prefix: str,
556
- attention_backend_cls: type[AttentionBackend],
557
- builder_cls: type[AttentionMetadataBuilder[M]],
558
-) -> type[AttentionBackend]:
559
- """
560
- Return a new subclass where `get_builder_cls` returns `builder_cls`.
561
- """
562
- name: str = name_prefix + attention_backend_cls.__name__ # type: ignore
563
-
564
- return type(
565
- name, (attention_backend_cls,), {"get_builder_cls": lambda: builder_cls}
566
- )
567
-
568
-
569
-def subclass_attention_backend_with_overrides(
570
- name_prefix: str,
571
- attention_backend_cls: type[AttentionBackend],
572
- overrides: dict[str, Any],
573
-) -> type[AttentionBackend]:
574
- name: str = name_prefix + attention_backend_cls.__name__ # type: ignore
575
- return type(name, (attention_backend_cls,), overrides)
576
-
577
-
420
def split_decodes_prefills_and_extends(
421
common_attn_metadata: CommonAttentionMetadata,
422
decode_threshold: int = 1,
vllm/v1/worker/gpu/cudagraph_utils.py
+1
-1
@@ -12,7 +12,7 @@ from vllm.config import VllmConfig
12
from vllm.config.compilation import CUDAGraphMode
13
from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14
from vllm.forward_context import set_forward_context
15
-from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
15
+from vllm.v1.attention.backend import AttentionMetadataBuilder
16
from vllm.v1.core.sched.output import SchedulerOutput
17
from vllm.v1.kv_cache_interface import KVCacheConfig
18
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
vllm/v1/worker/gpu/spec_decode/eagle.py
+1
-1
@@ -11,7 +11,7 @@ from vllm.forward_context import set_forward_context
11
from vllm.logger import init_logger
12
from vllm.model_executor.model_loader import get_model
13
from vllm.triton_utils import tl, triton
14
-from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
14
+from vllm.v1.attention.backend import AttentionMetadataBuilder
15
from vllm.v1.kv_cache_interface import KVCacheConfig
16
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
17
from vllm.v1.worker.gpu.block_table import BlockTables
vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py
+1
-1
@@ -6,7 +6,7 @@ import torch
6
7
from vllm.config import VllmConfig
8
from vllm.config.compilation import CUDAGraphMode
9
-from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
9
+from vllm.v1.attention.backend import AttentionMetadataBuilder
10
from vllm.v1.kv_cache_interface import KVCacheConfig
11
from vllm.v1.worker.gpu.block_table import BlockTables
12
from vllm.v1.worker.gpu.cudagraph_utils import (
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -112,7 +112,6 @@ from vllm.v1.attention.backends.utils import (
112
create_fast_prefill_custom_backend,
113
get_dcp_local_seq_lens,
114
reorder_batch_to_split_decodes_and_prefills,
115
- split_attn_metadata,
115
)
116
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
117
from vllm.v1.kv_cache_interface import (
@@ -165,6 +164,7 @@ from vllm.v1.worker.ubatch_utils import (
164
UBatchSlices,
165
check_ubatch_thresholds,
166
maybe_create_ubatch_slices,
167
+ split_attn_metadata,
168
)
169
from vllm.v1.worker.utils import is_residual_scattered_for_sp
170
from vllm.v1.worker.workspace import lock_workspace
vllm/v1/worker/ubatch_utils.py
+131
@@ -4,8 +4,10 @@ from dataclasses import dataclass
4
from typing import TypeAlias
5
6
import numpy as np
7
+import torch
8
9
from vllm.config import ParallelConfig
10
+from vllm.v1.attention.backend import CommonAttentionMetadata
11
12
13
@dataclass
@@ -110,3 +112,132 @@ def maybe_create_ubatch_slices(
112
assert sum(s.num_tokens for s in ubatch_slices_padded) == num_tokens_padded
113
114
return ubatch_slices, ubatch_slices_padded
115
+
116
+
117
+def slice_query_start_locs(
118
+ query_start_loc: torch.Tensor,
119
+ request_slice: slice,
120
+) -> torch.Tensor:
121
+ """
122
+ Creates a new query_start_loc that corresponds to the requests in
123
+ request_slice.
124
+
125
+ Note: This function creates a new tensor to hold the new query_start_locs.
126
+ This will break cudagraph compatibility.
127
+ """
128
+ return (
129
+ query_start_loc[request_slice.start : request_slice.stop + 1]
130
+ - query_start_loc[request_slice.start]
131
+ )
132
+
133
+
134
+def _make_metadata_with_slice(
135
+ ubatch_slice: UBatchSlice, attn_metadata: CommonAttentionMetadata
136
+) -> CommonAttentionMetadata:
137
+ """
138
+ This function creates a new CommonAttentionMetadata that corresponds to
139
+ the requests included in ubatch_slice
140
+ """
141
+
142
+ assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty"
143
+
144
+ request_slice = ubatch_slice.request_slice
145
+ token_slice = ubatch_slice.token_slice
146
+
147
+ start_locs = attn_metadata.query_start_loc_cpu
148
+ first_req = request_slice.start
149
+ first_tok = token_slice.start
150
+ last_req = request_slice.stop - 1
151
+ last_tok = token_slice.stop - 1
152
+
153
+ assert start_locs[first_req] <= first_tok < start_locs[first_req + 1], (
154
+ "Token slice start outside of first request"
155
+ )
156
+ # NOTE: last token can be outside of the last request if we have CG padding.
157
+
158
+ # If the request is split across ubatches, we have to adjust the metadata.
159
+ # splits_first_request: The first request in this slice is the continuation of
160
+ # a request that started in a previous slice.
161
+ # splits_last_request: The last request in this slice continues into the
162
+ # next slice.
163
+ splits_first_request = first_tok > start_locs[first_req]
164
+ splits_last_request = last_tok < start_locs[last_req + 1] - 1
165
+
166
+ query_start_loc_cpu = slice_query_start_locs(start_locs, request_slice)
167
+ query_start_loc = slice_query_start_locs(
168
+ attn_metadata.query_start_loc, request_slice
169
+ )
170
+
171
+ assert len(query_start_loc) >= 2, (
172
+ f"query_start_loc must have at least 2 elements, got {len(query_start_loc)}"
173
+ )
174
+
175
+ if splits_first_request:
176
+ tokens_skipped = first_tok - start_locs[first_req]
177
+ query_start_loc[1:] -= tokens_skipped
178
+ query_start_loc_cpu[1:] -= tokens_skipped
179
+ seq_lens = attn_metadata.seq_lens[request_slice]
180
+ seq_lens_cpu = attn_metadata.seq_lens_cpu[request_slice]
181
+
182
+ if splits_last_request:
183
+ # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate
184
+ # the tokens skipped because query_start_loc_cpu might have been modified
185
+ # if splits_first_request is True.
186
+ tokens_skipped = start_locs[last_req + 1] - token_slice.stop
187
+ query_start_loc[-1] -= tokens_skipped
188
+ query_start_loc_cpu[-1] -= tokens_skipped
189
+
190
+ # Make sure we don't modify the seq_lens tensors
191
+ # (not cudagraph compatible)
192
+ seq_lens = seq_lens.clone()
193
+ seq_lens_cpu = seq_lens_cpu.clone()
194
+ seq_lens[-1] -= tokens_skipped
195
+ seq_lens_cpu[-1] -= tokens_skipped
196
+
197
+ max_seq_len = int(seq_lens_cpu.max())
198
+ num_computed_tokens_cpu = attn_metadata.num_computed_tokens_cpu[request_slice]
199
+
200
+ num_requests = request_slice.stop - request_slice.start
201
+ num_actual_tokens = token_slice.stop - token_slice.start
202
+ max_query_len = int(
203
+ torch.max(torch.abs(query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])).item()
204
+ )
205
+
206
+ # This is to account for the case where we are in a dummy
207
+ # run and query_start_loc_cpu is full of 0s
208
+ if max_query_len == 0:
209
+ max_query_len = attn_metadata.max_query_len
210
+
211
+ block_table_tensor = attn_metadata.block_table_tensor[request_slice]
212
+ slot_mapping = attn_metadata.slot_mapping[token_slice]
213
+
214
+ return CommonAttentionMetadata(
215
+ query_start_loc=query_start_loc,
216
+ query_start_loc_cpu=query_start_loc_cpu,
217
+ seq_lens=seq_lens,
218
+ num_reqs=num_requests,
219
+ num_actual_tokens=num_actual_tokens,
220
+ max_query_len=max_query_len,
221
+ max_seq_len=max_seq_len,
222
+ block_table_tensor=block_table_tensor,
223
+ slot_mapping=slot_mapping,
224
+ _seq_lens_cpu=seq_lens_cpu,
225
+ _num_computed_tokens_cpu=num_computed_tokens_cpu,
226
+ )
227
+
228
+
229
+def split_attn_metadata(
230
+ ubatch_slices: list[UBatchSlice],
231
+ common_attn_metadata: CommonAttentionMetadata,
232
+) -> list[CommonAttentionMetadata]:
233
+ """
234
+ Creates a new CommonAttentionMetadata instance that corresponds to the
235
+ requests for each UBatchSlice in ubatch_slices.
236
+
237
+ Note: This function does not modify common_attn_metadata
238
+ """
239
+ results = []
240
+ for ubatch_slice in ubatch_slices:
241
+ results.append(_make_metadata_with_slice(ubatch_slice, common_attn_metadata))
242
+
243
+ return results
vllm/v1/worker/utils.py
+1
-2
@@ -16,8 +16,7 @@ from vllm.multimodal.cache import processor_only_cache_from_config
16
from vllm.multimodal.registry import MultiModalRegistry
17
from vllm.platforms import current_platform
18
from vllm.utils.mem_utils import MemorySnapshot, format_gib
19
-from vllm.v1.attention.backend import AttentionBackend
20
-from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
19
+from vllm.v1.attention.backend import AttentionBackend, AttentionMetadataBuilder
20
from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
21
from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec
22