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 11, 00:00 UTC
[Mamba] - Consolidate Mambas Attention Logic (#28133)
Asaf Joseph Gardin committed
Dec 23, 2025 at 22:57 UTC
34916ae37fb91a234413830a58bebca42015833c
5 files changed
+305
-448
vllm/model_executor/layers/mamba/short_conv.py
+1
-5
@@ -118,6 +118,7 @@ class ShortConv(MambaBase, CustomOp):
118
conv_state = self_kv_cache[0].transpose(-1, -2)
119
state_indices_tensor = attn_metadata.state_indices_tensor
120
has_initial_states_p = attn_metadata.has_initial_states_p
121
+ query_start_loc_p = attn_metadata.query_start_loc_p
122
123
BCx, _ = self.in_proj(hidden_states)
124
@@ -165,11 +166,6 @@ class ShortConv(MambaBase, CustomOp):
166
[num_decodes, num_prefills],
167
dim=0,
168
)
168
- query_start_loc_p = (
169
- attn_metadata.query_start_loc[-num_prefills - 1 :] - num_decodes
170
- if has_prefill
171
- else None
172
- )
169
170
conv_output_list = []
171
vllm/v1/attention/backends/mamba1_attn.py
+7
-138
@@ -3,17 +3,11 @@
3
4
from dataclasses import dataclass
5
6
-import torch
7
-
6
from vllm.attention.backends.abstract import AttentionBackend
9
-from vllm.attention.backends.utils import PAD_SLOT_ID
10
-from vllm.config import VllmConfig
11
-from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder
12
-from vllm.v1.attention.backends.utils import (
13
- CommonAttentionMetadata,
14
- split_decodes_and_prefills,
7
+from vllm.v1.attention.backends.mamba_attn import (
8
+ BaseMambaAttentionMetadata,
9
+ BaseMambaAttentionMetadataBuilder,
10
)
16
-from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
11
12
13
class Mamba1AttentionBackend(AttentionBackend):
@@ -23,137 +17,12 @@ class Mamba1AttentionBackend(AttentionBackend):
17
18
19
@dataclass
26
-class Mamba1AttentionMetadata:
27
- query_start_loc_p: torch.Tensor
28
- state_indices_tensor: torch.Tensor
29
- has_initial_states_p: torch.Tensor | None
30
- num_prefills: int
31
- num_prefill_tokens: int
32
- num_decodes: int
33
- num_decode_tokens: int
34
-
35
- block_idx_last_scheduled_token: torch.Tensor # shape: [batch,]
36
- block_idx_first_scheduled_token_p: torch.Tensor # shape: [batch,]
37
- block_idx_last_computed_token: torch.Tensor # shape: [batch,]
38
- num_computed_tokens_p: torch.Tensor # shape: [batch,]
20
+class Mamba1AttentionMetadata(BaseMambaAttentionMetadata):
21
+ pass
22
23
24
class Mamba1AttentionMetadataBuilder(
25
BaseMambaAttentionMetadataBuilder[Mamba1AttentionMetadata]
26
):
44
- def __init__(
45
- self,
46
- kv_cache_spec: AttentionSpec,
47
- layer_names: list[str],
48
- vllm_config: VllmConfig,
49
- device: torch.device,
50
- ):
51
- super().__init__(kv_cache_spec, layer_names, vllm_config, device)
52
- assert isinstance(kv_cache_spec, MambaSpec)
53
-
54
- def build(
55
- self,
56
- common_prefix_len: int,
57
- common_attn_metadata: CommonAttentionMetadata,
58
- fast_build: bool = False,
59
- ) -> Mamba1AttentionMetadata:
60
- num_reqs = common_attn_metadata.num_reqs
61
-
62
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
63
- split_decodes_and_prefills(
64
- common_attn_metadata, decode_threshold=self.reorder_batch_threshold
65
- )
66
- )
67
-
68
- has_initial_states_p = None
69
- query_start_loc_p = None
70
- num_computed_tokens, num_computed_tokens_p = None, None
71
- block_idx_first_scheduled_token = None
72
- block_idx_first_scheduled_token_p = None
73
-
74
- # TODO(@Josephasafg) Mamba1 and Mamba2 have a lot of code in common here.
75
- # We should consolidate this code
76
- if self.vllm_config.cache_config.enable_prefix_caching:
77
- # Return a tensor of shape (#requests, #max blocks)
78
- state_indices_tensor = common_attn_metadata.block_table_tensor
79
- mamba_block_size = self.kv_cache_spec.block_size
80
- num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
81
- self.device
82
- )
83
- (
84
- block_idx_last_computed_token,
85
- block_idx_first_scheduled_token,
86
- block_idx_last_scheduled_token,
87
- ) = self._compute_prefix_caching_block_indices(
88
- common_attn_metadata, mamba_block_size
89
- )
90
- else:
91
- # Always return just a single block per each request:
92
- state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
93
- block_idx_last_scheduled_token = None
94
- block_idx_last_computed_token = None
95
-
96
- if num_prefills > 0:
97
- query_start_loc_p = (
98
- common_attn_metadata.query_start_loc[-num_prefills - 1 :]
99
- - num_decode_tokens
100
- )
101
- has_initial_states_cpu = (
102
- common_attn_metadata.num_computed_tokens_cpu[
103
- num_reqs - num_prefills : num_reqs
104
- ]
105
- > 0
106
- )
107
- has_initial_states_p = has_initial_states_cpu.to(
108
- common_attn_metadata.query_start_loc.device
109
- )
110
-
111
- if self.vllm_config.cache_config.enable_prefix_caching:
112
- assert num_computed_tokens is not None
113
- num_computed_tokens_p = num_computed_tokens[
114
- num_reqs - num_prefills : num_reqs
115
- ]
116
- assert block_idx_first_scheduled_token is not None
117
- block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[
118
- num_reqs - num_prefills : num_reqs
119
- ]
120
-
121
- elif (
122
- num_decodes > 0
123
- and num_decodes <= self.decode_cudagraph_max_bs
124
- and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
125
- ):
126
- self.state_indices_tensor[:num_decodes].copy_(
127
- state_indices_tensor, non_blocking=True
128
- )
129
- state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
130
- state_indices_tensor[num_decodes:] = PAD_SLOT_ID
131
-
132
- if self.vllm_config.cache_config.enable_prefix_caching:
133
- self.block_idx_last_scheduled_token[:num_decodes].copy_(
134
- block_idx_last_scheduled_token, non_blocking=True
135
- )
136
- block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
137
- :num_decode_tokens
138
- ]
139
-
140
- self.block_idx_last_computed_token[:num_decodes].copy_(
141
- block_idx_last_computed_token, non_blocking=True
142
- )
143
- block_idx_last_computed_token = self.block_idx_last_computed_token[
144
- :num_decode_tokens
145
- ]
146
-
147
- return Mamba1AttentionMetadata(
148
- query_start_loc_p=query_start_loc_p,
149
- has_initial_states_p=has_initial_states_p,
150
- state_indices_tensor=state_indices_tensor,
151
- num_prefills=num_prefills,
152
- num_prefill_tokens=num_prefill_tokens,
153
- num_decodes=num_decodes,
154
- num_decode_tokens=num_decode_tokens,
155
- block_idx_last_scheduled_token=block_idx_last_scheduled_token,
156
- block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
157
- block_idx_last_computed_token=block_idx_last_computed_token,
158
- num_computed_tokens_p=num_computed_tokens_p,
159
- )
27
+ metadata_cls = Mamba1AttentionMetadata
28
+ supports_update_block_table: bool = False
vllm/v1/attention/backends/mamba2_attn.py
+100
-220
@@ -1,19 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import copy
3
import itertools
5
-from dataclasses import dataclass
4
+from dataclasses import dataclass, replace
5
6
import torch
7
8
from vllm.attention.backends.abstract import AttentionBackend
9
from vllm.config import VllmConfig
10
from vllm.utils.math_utils import cdiv
12
-from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder
11
+from vllm.v1.attention.backends.mamba_attn import (
12
+ BaseMambaAttentionMetadata,
13
+ BaseMambaAttentionMetadataBuilder,
14
+)
15
from vllm.v1.attention.backends.utils import (
16
CommonAttentionMetadata,
15
- compute_causal_conv1d_metadata,
16
- split_decodes_and_prefills,
17
)
18
from vllm.v1.kv_cache_interface import AttentionSpec
19
@@ -94,48 +94,26 @@ class Mamba2AttentionBackend(AttentionBackend):
94
95
96
@dataclass
97
-class Mamba2AttentionMetadata:
98
- num_prefills: int
99
- num_prefill_tokens: int
100
- num_decodes: int
101
- num_decode_tokens: int
102
- query_start_loc_p: torch.Tensor
103
- seq_lens: torch.Tensor
104
-
105
- prep_initial_states: bool
106
- chunk_size: int
107
-
108
- # The following tensors only contain prefill requests and will be None if
109
- # the batch has no prefill request.
110
- has_initial_states_p: torch.Tensor | None
111
- seq_idx_p: torch.Tensor | None
97
+class Mamba2AttentionMetadata(BaseMambaAttentionMetadata):
98
+ prep_initial_states: bool = False
99
+ chunk_size: int = 0
100
101
+ # Chunk-related metadata (only for prefill)
102
+ seq_idx_p: torch.Tensor | None = None
103
# cu_chunk_seqlen_p is a tensor of shape (nchunks+1,) that contains, for
104
# each chunk, its offests into the varlen sequence dimension. It is defined
105
# such that the i-th chunk contains tokens from cu_chunk_seqlen_p[i] to
106
# cu_chunk_seqlen_p[i+1].
117
- cu_chunk_seqlen_p: torch.Tensor | None
118
-
107
+ cu_chunk_seqlen_p: torch.Tensor | None = None
108
# last_chunk_indices_p is a tensor of shape (batch,) that contains the
109
# index of the last chunk for every sequence in the (prefill) batch.
121
- last_chunk_indices_p: torch.Tensor | None
122
-
123
- state_indices_tensor: torch.Tensor # shape: [batch,]
124
- block_idx_last_scheduled_token: torch.Tensor # shape: [batch,]
125
- block_idx_first_scheduled_token_p: torch.Tensor # shape: [batch,]
126
- block_idx_last_computed_token: torch.Tensor # shape: [batch,]
127
- num_computed_tokens_p: torch.Tensor # shape: [batch,]
128
-
129
- # The following attributes are for triton implementation of causal_conv1d
130
- nums_dict: dict | None = None
131
- batch_ptr: torch.Tensor | None = None
132
- token_chunk_offset_ptr: torch.Tensor | None = None
110
+ last_chunk_indices_p: torch.Tensor | None = None
111
112
113
class Mamba2AttentionMetadataBuilder(
114
BaseMambaAttentionMetadataBuilder[Mamba2AttentionMetadata]
115
):
138
- supports_update_block_table: bool = True
116
+ metadata_cls = Mamba2AttentionMetadata
117
118
def __init__(
119
self,
@@ -150,87 +128,93 @@ class Mamba2AttentionMetadataBuilder(
128
"chunk_size needs to be set in the model config for Mamba2 models"
129
)
130
131
+ def _compute_chunk_metadata(
132
+ self,
133
+ num_prefills: int,
134
+ num_computed_tokens_p_cpu: torch.Tensor,
135
+ query_start_loc_p_cpu: torch.Tensor,
136
+ ) -> tuple[list[int], list[int], list[int]]:
137
+ """
138
+ Compute chunk-specific metadata for Mamba2.
139
+
140
+ The code below carefully constructs the chunks such that:
141
+ 1. Chunks contain tokens from a *single* sequence only.
142
+ 2. For every sequence, we are guaranteed that we can
143
+ retrieve the mamba state *every* chunk_size tokens.
144
+ Constraint (1) dramatically simplifies the mamba2 kernels.
145
+ Constraint (2) dramatically simplifies the implementation
146
+ of prefix caching for mamba2 (wip). We need to take care
147
+ of the interaction with chunked prefill in order to
148
+ satisfy constraint (2).
149
+ """
150
+ # TODO (tdoublep): This code could probably be optimized.
151
+ cu_chunk_seqlen = []
152
+ seq_idx = []
153
+ last_chunk_indices = []
154
+ seqlen_pos = 0
155
+
156
+ for req_idx in range(num_prefills):
157
+ this_num_computed = num_computed_tokens_p_cpu[req_idx].item()
158
+ this_new_tokens = (
159
+ query_start_loc_p_cpu[req_idx + 1].item()
160
+ - query_start_loc_p_cpu[req_idx].item()
161
+ )
162
+
163
+ # if computed tokens are not chunk-aligned, use the first
164
+ # chunk to finish it off
165
+ if this_num_computed % self.chunk_size != 0:
166
+ seq_idx.append(req_idx)
167
+ cu_chunk_seqlen.append(seqlen_pos)
168
+ # how many tokens to finish the chunk?
169
+ chunk_len = (
170
+ cdiv(this_num_computed, self.chunk_size) * self.chunk_size
171
+ - this_num_computed
172
+ )
173
+ # we can only use at most this_new_tokens
174
+ chunk_len = min(chunk_len, this_new_tokens)
175
+ seqlen_pos += chunk_len
176
+ this_new_tokens -= chunk_len
177
+
178
+ n_chunks = cdiv(this_new_tokens, self.chunk_size)
179
+ for chunk in range(n_chunks):
180
+ seq_idx.append(req_idx)
181
+ cu_chunk_seqlen.append(seqlen_pos)
182
+ chunk_len = min(self.chunk_size, this_new_tokens)
183
+ seqlen_pos += chunk_len
184
+ this_new_tokens -= chunk_len
185
+
186
+ assert this_new_tokens == 0
187
+ last_chunk_indices.append(len(cu_chunk_seqlen) - 1)
188
+
189
+ cu_chunk_seqlen.append(seqlen_pos)
190
+
191
+ return cu_chunk_seqlen, seq_idx, last_chunk_indices
192
+
193
def build(
194
self,
195
common_prefix_len: int,
196
common_attn_metadata: CommonAttentionMetadata,
197
fast_build: bool = False,
198
) -> Mamba2AttentionMetadata:
159
- num_reqs = common_attn_metadata.num_reqs
160
- seq_lens = common_attn_metadata.seq_lens
199
+ common = self._compute_common_metadata(common_attn_metadata)
200
162
- query_start_loc_p = None
201
seq_idx_p = None
202
cu_chunk_seqlen_p = None
203
last_chunk_indices_p = None
166
-
167
- # Need flags to indicate if there are initial states
168
- has_initial_states_p = None
204
prep_initial_states = False
205
171
- # for causal_conv1d
172
- nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
173
-
174
- num_computed_tokens, num_computed_tokens_p = None, None
175
- block_idx_first_scheduled_token = None
176
- block_idx_first_scheduled_token_p = None
177
-
178
- if self.vllm_config.cache_config.enable_prefix_caching:
179
- # Return a tensor of shape (#requests, #max blocks)
180
- state_indices_tensor = common_attn_metadata.block_table_tensor
181
- # Additional cache-related varaiables:
182
- mamba_block_size = self.kv_cache_spec.block_size
183
- num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
184
- self.device
185
- )
186
- (
187
- block_idx_last_computed_token,
188
- block_idx_first_scheduled_token,
189
- block_idx_last_scheduled_token,
190
- ) = self._compute_prefix_caching_block_indices(
191
- common_attn_metadata, mamba_block_size
192
- )
193
- else:
194
- # Always return just a single block per each request:
195
- state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
196
- # Additional cache-related varaiables:
197
- block_idx_last_scheduled_token = None
198
- block_idx_last_computed_token = None
199
-
200
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
201
- split_decodes_and_prefills(
202
- common_attn_metadata, decode_threshold=self.reorder_batch_threshold
203
- )
204
- )
205
-
206
# Compute seq_idx for prefill only
207
- if num_prefills > 0:
208
- # [batch,]
209
- has_initial_states_cpu = (
210
- common_attn_metadata.num_computed_tokens_cpu[
211
- num_reqs - num_prefills : num_reqs
212
- ]
213
- > 0
214
- )
215
- prep_initial_states = torch.any(has_initial_states_cpu).item()
216
- has_initial_states_p = has_initial_states_cpu.to(
217
- common_attn_metadata.query_start_loc.device
207
+ if common.num_prefills > 0:
208
+ prep_initial_states = (
209
+ torch.any(common.has_initial_states_p).item()
210
+ if common.has_initial_states_p is not None
211
+ else False
212
)
213
220
- query_start_loc_p = (
221
- common_attn_metadata.query_start_loc[-num_prefills - 1 :]
222
- - num_decode_tokens
223
- )
214
+ num_reqs = common.num_reqs
215
+ num_prefills = common.num_prefills
216
+ num_decode_tokens = common.num_decode_tokens
217
225
- if self.vllm_config.cache_config.enable_prefix_caching:
226
- assert num_computed_tokens is not None
227
- num_computed_tokens_p = num_computed_tokens[
228
- num_reqs - num_prefills : num_reqs
229
- ]
230
- assert block_idx_first_scheduled_token is not None
231
- block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[
232
- num_reqs - num_prefills : num_reqs
233
- ]
218
num_computed_tokens_p_cpu = common_attn_metadata.num_computed_tokens_cpu[
219
num_reqs - num_prefills : num_reqs
220
]
@@ -239,137 +223,33 @@ class Mamba2AttentionMetadataBuilder(
223
- num_decode_tokens
224
)
225
242
- # The code below carefully constructs the chunks such that:
243
- # 1. Chunks contain tokens from a *single* sequence only.
244
- # 2. For every sequence, we are guaranteed that we can
245
- # retrieve the mamba state *every* chunk_size tokens.
246
- # Constraint (1) dramatically simplifies the mamba2 kernels.
247
- # Constraint (2) dramatically simplifies the implementation
248
- # of prefix caching for mamba2 (wip). We need to take care
249
- # of the interaction with chunked prefill in order to
250
- # satisfy constraint (2).
251
- # TODO (tdoublep): This code could probably be optimized.
252
- cu_chunk_seqlen = []
253
- seq_idx = []
254
- last_chunk_indices = []
255
- seqlen_pos = 0
256
- for req_idx in range(num_prefills):
257
- this_num_computed = num_computed_tokens_p_cpu[req_idx].item()
258
- this_new_tokens = (
259
- query_start_loc_p_cpu[req_idx + 1].item()
260
- - query_start_loc_p_cpu[req_idx].item()
261
- )
262
-
263
- # if computed tokens are not chunk-aligned, use the first
264
- # chunk to finish it off
265
- if this_num_computed % self.chunk_size != 0:
266
- seq_idx.append(req_idx)
267
- cu_chunk_seqlen.append(seqlen_pos)
268
- # how many tokens to finish the chunk?
269
- chunk_len = (
270
- cdiv(this_num_computed, self.chunk_size) * self.chunk_size
271
- - this_num_computed
272
- )
273
- # we can only use at most this_new_tokens
274
- chunk_len = min(chunk_len, this_new_tokens)
275
- seqlen_pos += chunk_len
276
- this_new_tokens -= chunk_len
277
-
278
- n_chunks = cdiv(this_new_tokens, self.chunk_size)
279
- for chunk in range(n_chunks):
280
- seq_idx.append(req_idx)
281
- cu_chunk_seqlen.append(seqlen_pos)
282
- chunk_len = min(self.chunk_size, this_new_tokens)
283
- seqlen_pos += chunk_len
284
- this_new_tokens -= chunk_len
285
-
286
- assert this_new_tokens == 0
287
- last_chunk_indices.append(len(cu_chunk_seqlen) - 1)
288
-
289
- cu_chunk_seqlen.append(seqlen_pos)
226
+ cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata(
227
+ num_prefills,
228
+ num_computed_tokens_p_cpu,
229
+ query_start_loc_p_cpu,
230
+ )
231
232
seq_idx_p = torch.as_tensor(
292
- seq_idx, device=query_start_loc_p.device, dtype=torch.int32
233
+ seq_idx,
234
+ device=common_attn_metadata.query_start_loc.device,
235
+ dtype=torch.int32,
236
)
237
cu_chunk_seqlen_p = torch.as_tensor(
295
- cu_chunk_seqlen, device=query_start_loc_p.device, dtype=torch.int32
238
+ cu_chunk_seqlen,
239
+ device=common_attn_metadata.query_start_loc.device,
240
+ dtype=torch.int32,
241
)
242
last_chunk_indices_p = torch.as_tensor(
298
- last_chunk_indices, device=query_start_loc_p.device, dtype=torch.int32
299
- )
300
-
301
- nums_dict, batch_ptr, token_chunk_offset_ptr = (
302
- compute_causal_conv1d_metadata(query_start_loc_p)
303
- )
304
-
305
- elif (
306
- num_decodes <= self.decode_cudagraph_max_bs
307
- and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
308
- ):
309
- self.state_indices_tensor[:num_decodes].copy_(
310
- state_indices_tensor, non_blocking=True
243
+ last_chunk_indices,
244
+ device=common_attn_metadata.query_start_loc.device,
245
+ dtype=torch.int32,
246
)
312
- state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
313
-
314
- if self.vllm_config.cache_config.enable_prefix_caching:
315
- self.block_idx_last_scheduled_token[:num_decodes].copy_(
316
- block_idx_last_scheduled_token, non_blocking=True
317
- )
318
- block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
319
- :num_decode_tokens
320
- ]
247
322
- self.block_idx_last_computed_token[:num_decodes].copy_(
323
- block_idx_last_computed_token, non_blocking=True
324
- )
325
- block_idx_last_computed_token = self.block_idx_last_computed_token[
326
- :num_decode_tokens
327
- ]
328
-
329
- attn_metadata = Mamba2AttentionMetadata(
330
- num_prefills=num_prefills,
331
- num_prefill_tokens=num_prefill_tokens,
332
- num_decodes=num_decodes,
333
- num_decode_tokens=num_decode_tokens,
334
- query_start_loc_p=query_start_loc_p,
335
- seq_lens=seq_lens,
248
+ return replace(
249
+ common,
250
prep_initial_states=prep_initial_states,
251
chunk_size=self.chunk_size,
338
- has_initial_states_p=has_initial_states_p,
252
seq_idx_p=seq_idx_p,
340
- state_indices_tensor=state_indices_tensor,
253
cu_chunk_seqlen_p=cu_chunk_seqlen_p,
254
last_chunk_indices_p=last_chunk_indices_p,
343
- nums_dict=nums_dict,
344
- batch_ptr=batch_ptr,
345
- token_chunk_offset_ptr=token_chunk_offset_ptr,
346
- block_idx_last_scheduled_token=block_idx_last_scheduled_token,
347
- block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
348
- block_idx_last_computed_token=block_idx_last_computed_token,
349
- num_computed_tokens_p=num_computed_tokens_p,
255
)
351
- return attn_metadata
352
-
353
- def update_block_table(
354
- self,
355
- metadata: Mamba2AttentionMetadata,
356
- blk_table: torch.Tensor,
357
- slot_mapping: torch.Tensor,
358
- ) -> Mamba2AttentionMetadata:
359
- new_metadata = copy.copy(metadata)
360
- prefix_caching = self.vllm_config.cache_config.enable_prefix_caching
361
- state_indices_t = blk_table if prefix_caching else blk_table[:, 0]
362
- num_reqs = blk_table.shape[0]
363
-
364
- # For CUDA graphs, copy to persistent buffer
365
- if (
366
- metadata.num_prefills == 0
367
- and num_reqs <= self.decode_cudagraph_max_bs
368
- and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
369
- ):
370
- persistent_state_indices_t = self.state_indices_tensor[:num_reqs]
371
- persistent_state_indices_t.copy_(state_indices_t, non_blocking=True)
372
- state_indices_t = persistent_state_indices_t
373
-
374
- new_metadata.state_indices_tensor = state_indices_t
375
- return new_metadata
vllm/v1/attention/backends/mamba_attn.py
+191
-1
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import abc
5
+import copy
6
+from dataclasses import dataclass
7
from typing import ClassVar, TypeVar
8
9
import torch
@@ -9,20 +11,52 @@ import torch
11
from vllm.config import VllmConfig
12
from vllm.utils.math_utils import cdiv
13
from vllm.v1.attention.backends.utils import (
14
+ PAD_SLOT_ID,
15
AttentionCGSupport,
16
AttentionMetadataBuilder,
17
CommonAttentionMetadata,
18
+ compute_causal_conv1d_metadata,
19
+ split_decodes_and_prefills,
20
)
21
from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
22
18
-M = TypeVar("M")
23
+M = TypeVar("M", bound="BaseMambaAttentionMetadata")
24
+
25
+
26
+@dataclass
27
+class BaseMambaAttentionMetadata:
28
+ num_prefills: int
29
+ num_prefill_tokens: int
30
+ num_decodes: int
31
+ num_decode_tokens: int
32
+ num_reqs: int
33
+
34
+ # The following tensors only contain prefill requests and will be None if
35
+ # the batch has no prefill request.
36
+ has_initial_states_p: torch.Tensor | None
37
+ query_start_loc_p: torch.Tensor | None
38
+ num_computed_tokens_p: torch.Tensor | None
39
+
40
+ state_indices_tensor: torch.Tensor
41
+
42
+ # The following tensors are only used for prefix caching and are None if disabled
43
+ block_idx_last_scheduled_token: torch.Tensor | None
44
+ block_idx_first_scheduled_token_p: torch.Tensor | None
45
+ block_idx_last_computed_token: torch.Tensor | None
46
+
47
+ # The following attributes are for triton implementation of causal_conv1d
48
+ nums_dict: dict | None = None
49
+ batch_ptr: torch.Tensor | None = None
50
+ token_chunk_offset_ptr: torch.Tensor | None = None
51
52
53
class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
54
+ metadata_cls: type[M]
55
reorder_batch_threshold: int = 1
56
_cudagraph_support: ClassVar[AttentionCGSupport] = (
57
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
58
)
59
+ supports_update_block_table: bool = True
60
61
def __init__(
62
self,
@@ -87,6 +121,18 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
121
122
return self.build(0, m)
123
124
+ def build(
125
+ self,
126
+ common_prefix_len: int,
127
+ common_attn_metadata: CommonAttentionMetadata,
128
+ fast_build: bool = False,
129
+ ) -> M:
130
+ """
131
+ Default build implementation for Mamba-like attention backends.
132
+ Subclasses (e.g., Mamba2) can override to add additional metadata.
133
+ """
134
+ return self._compute_common_metadata(common_attn_metadata)
135
+
136
def _compute_prefix_caching_block_indices(
137
self,
138
common_attn_metadata: CommonAttentionMetadata,
@@ -115,3 +161,147 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
161
block_idx_first_scheduled_token,
162
block_idx_last_scheduled_token,
163
)
164
+
165
+ def _compute_common_metadata(
166
+ self,
167
+ common_attn_metadata: CommonAttentionMetadata,
168
+ ) -> M:
169
+ """
170
+ Compute metadata common to both Mamba1 and Mamba2.
171
+ """
172
+ num_reqs = common_attn_metadata.num_reqs
173
+
174
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
175
+ split_decodes_and_prefills(
176
+ common_attn_metadata, decode_threshold=self.reorder_batch_threshold
177
+ )
178
+ )
179
+
180
+ # Need flags to indicate if there are initial states
181
+ has_initial_states_p = None
182
+ query_start_loc_p = None
183
+ num_computed_tokens = None
184
+ num_computed_tokens_p = None
185
+
186
+ # for prefix caching
187
+ block_idx_first_scheduled_token = None
188
+ block_idx_first_scheduled_token_p = None
189
+ block_idx_last_computed_token = None
190
+ block_idx_last_scheduled_token = None
191
+
192
+ # for causal_conv1d
193
+ nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
194
+
195
+ if self.vllm_config.cache_config.enable_prefix_caching:
196
+ # Return a tensor of shape (#requests, #max blocks)
197
+ state_indices_tensor = common_attn_metadata.block_table_tensor
198
+ # Additional cache-related varaiables:
199
+ mamba_block_size = self.kv_cache_spec.block_size
200
+ num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
201
+ self.device
202
+ )
203
+ (
204
+ block_idx_last_computed_token,
205
+ block_idx_first_scheduled_token,
206
+ block_idx_last_scheduled_token,
207
+ ) = self._compute_prefix_caching_block_indices(
208
+ common_attn_metadata, mamba_block_size
209
+ )
210
+ else:
211
+ # Always return just a single block per each request:
212
+ state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
213
+
214
+ if num_prefills > 0:
215
+ query_start_loc_p = (
216
+ common_attn_metadata.query_start_loc[-num_prefills - 1 :]
217
+ - num_decode_tokens
218
+ )
219
+ has_initial_states_cpu = (
220
+ common_attn_metadata.num_computed_tokens_cpu[
221
+ num_reqs - num_prefills : num_reqs
222
+ ]
223
+ > 0
224
+ )
225
+ has_initial_states_p = has_initial_states_cpu.to(
226
+ common_attn_metadata.query_start_loc.device
227
+ )
228
+
229
+ nums_dict, batch_ptr, token_chunk_offset_ptr = (
230
+ compute_causal_conv1d_metadata(query_start_loc_p)
231
+ )
232
+
233
+ if self.vllm_config.cache_config.enable_prefix_caching:
234
+ assert num_computed_tokens is not None
235
+ num_computed_tokens_p = num_computed_tokens[
236
+ num_reqs - num_prefills : num_reqs
237
+ ]
238
+ assert block_idx_first_scheduled_token is not None
239
+ block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[
240
+ num_reqs - num_prefills : num_reqs
241
+ ]
242
+ elif (
243
+ num_decodes <= self.decode_cudagraph_max_bs
244
+ and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
245
+ ):
246
+ self.state_indices_tensor[:num_decodes].copy_(
247
+ state_indices_tensor, non_blocking=True
248
+ )
249
+ state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
250
+ state_indices_tensor[num_decodes:] = PAD_SLOT_ID
251
+
252
+ if self.vllm_config.cache_config.enable_prefix_caching:
253
+ self.block_idx_last_scheduled_token[:num_decodes].copy_(
254
+ block_idx_last_scheduled_token, non_blocking=True
255
+ )
256
+ block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
257
+ :num_decode_tokens
258
+ ]
259
+
260
+ self.block_idx_last_computed_token[:num_decodes].copy_(
261
+ block_idx_last_computed_token, non_blocking=True
262
+ )
263
+ block_idx_last_computed_token = self.block_idx_last_computed_token[
264
+ :num_decode_tokens
265
+ ]
266
+
267
+ return self.metadata_cls(
268
+ num_prefills=num_prefills,
269
+ num_prefill_tokens=num_prefill_tokens,
270
+ num_decodes=num_decodes,
271
+ num_decode_tokens=num_decode_tokens,
272
+ query_start_loc_p=query_start_loc_p,
273
+ has_initial_states_p=has_initial_states_p,
274
+ state_indices_tensor=state_indices_tensor,
275
+ block_idx_last_scheduled_token=block_idx_last_scheduled_token,
276
+ block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
277
+ block_idx_last_computed_token=block_idx_last_computed_token,
278
+ num_computed_tokens_p=num_computed_tokens_p,
279
+ num_reqs=num_reqs,
280
+ nums_dict=nums_dict,
281
+ batch_ptr=batch_ptr,
282
+ token_chunk_offset_ptr=token_chunk_offset_ptr,
283
+ )
284
+
285
+ def update_block_table(
286
+ self,
287
+ metadata: M,
288
+ blk_table: torch.Tensor,
289
+ slot_mapping: torch.Tensor,
290
+ ) -> M:
291
+ new_metadata = copy.copy(metadata)
292
+ prefix_caching = self.vllm_config.cache_config.enable_prefix_caching
293
+ state_indices_t = blk_table if prefix_caching else blk_table[:, 0]
294
+ num_reqs = blk_table.shape[0]
295
+
296
+ # For CUDA graphs, copy to persistent buffer
297
+ if (
298
+ metadata.num_prefills == 0
299
+ and num_reqs <= self.decode_cudagraph_max_bs
300
+ and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
301
+ ):
302
+ persistent_state_indices_t = self.state_indices_tensor[:num_reqs]
303
+ persistent_state_indices_t.copy_(state_indices_t, non_blocking=True)
304
+ state_indices_t = persistent_state_indices_t
305
+
306
+ new_metadata.state_indices_tensor = state_indices_t
307
+ return new_metadata
vllm/v1/attention/backends/short_conv_attn.py
+6
-84
@@ -2,15 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
5
-import torch
6
-
5
from vllm.attention.backends.abstract import AttentionBackend
8
-from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadataBuilder
9
-from vllm.v1.attention.backends.utils import (
10
- PAD_SLOT_ID,
11
- CommonAttentionMetadata,
12
- compute_causal_conv1d_metadata,
13
- split_decodes_and_prefills,
6
+from vllm.v1.attention.backends.mamba_attn import (
7
+ BaseMambaAttentionMetadata,
8
+ BaseMambaAttentionMetadataBuilder,
9
)
10
11
@@ -21,84 +16,11 @@ class ShortConvAttentionBackend(AttentionBackend):
16
17
18
@dataclass
24
-class ShortConvAttentionMetadata:
25
- num_prefills: int
26
- num_prefill_tokens: int
27
- num_decodes: int
28
- num_decode_tokens: int
29
-
30
- query_start_loc: torch.Tensor
31
- state_indices_tensor: torch.Tensor
32
- has_initial_states_p: torch.Tensor | None
33
-
34
- # For causal_conv1d
35
- nums_dict: dict | None = None
36
- batch_ptr: torch.Tensor | None = None
37
- token_chunk_offset_ptr: torch.Tensor | None = None
19
+class ShortConvAttentionMetadata(BaseMambaAttentionMetadata):
20
+ pass
21
22
23
class ShortConvAttentionMetadataBuilder(
24
BaseMambaAttentionMetadataBuilder[ShortConvAttentionMetadata]
25
):
43
- def build(
44
- self,
45
- common_prefix_len: int,
46
- common_attn_metadata: CommonAttentionMetadata,
47
- fast_build: bool = False,
48
- ) -> ShortConvAttentionMetadata:
49
- num_reqs = common_attn_metadata.num_reqs
50
- query_start_loc = common_attn_metadata.query_start_loc
51
- state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
52
-
53
- # for causal_conv1d
54
- nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
55
-
56
- num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
57
- split_decodes_and_prefills(
58
- common_attn_metadata, decode_threshold=self.reorder_batch_threshold
59
- )
60
- )
61
-
62
- has_initial_states_p = None
63
- if num_prefills > 0:
64
- has_initial_states_cpu = (
65
- common_attn_metadata.num_computed_tokens_cpu[
66
- num_reqs - num_prefills : num_reqs
67
- ]
68
- > 0
69
- )
70
- has_initial_states_p = has_initial_states_cpu.to(query_start_loc.device)
71
-
72
- query_start_loc_p = (
73
- common_attn_metadata.query_start_loc[-num_prefills - 1 :]
74
- - num_decode_tokens
75
- )
76
-
77
- nums_dict, batch_ptr, token_chunk_offset_ptr = (
78
- compute_causal_conv1d_metadata(query_start_loc_p)
79
- )
80
-
81
- elif (
82
- num_decodes > 0
83
- and num_decodes <= self.decode_cudagraph_max_bs
84
- and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
85
- ):
86
- self.state_indices_tensor[:num_decodes].copy_(
87
- state_indices_tensor, non_blocking=True
88
- )
89
- state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
90
- state_indices_tensor[num_decodes:] = PAD_SLOT_ID
91
-
92
- attn_metadata = ShortConvAttentionMetadata(
93
- query_start_loc=query_start_loc,
94
- state_indices_tensor=state_indices_tensor,
95
- has_initial_states_p=has_initial_states_p,
96
- num_prefills=num_prefills,
97
- num_prefill_tokens=num_prefill_tokens,
98
- num_decodes=num_decodes,
99
- num_decode_tokens=num_decode_tokens,
100
- nums_dict=nums_dict,
101
- batch_ptr=batch_ptr,
102
- token_chunk_offset_ptr=token_chunk_offset_ptr,
103
- )
104
- return attn_metadata
26
+ metadata_cls = ShortConvAttentionMetadata