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, 01:00 UTC
[Attention] Refactor attention metadata builder interface (#20466)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jul 17, 2025 at 00:44 UTC
76b494444fd864ffc53a623420668d1865c804b9
18 files changed
+1447
-778
tests/v1/attention/test_attention_backends.py
new
+466
@@ -0,0 +1,466 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for v1 attention backends without GPUModelRunner dependency."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from tests.v1.attention.utils import (BatchSpec, _Backend,
9
+ create_common_attn_metadata,
10
+ create_standard_kv_cache_spec,
11
+ create_vllm_config,
12
+ get_attention_backend)
13
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
14
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
15
+from vllm.v1.kv_cache_interface import FullAttentionSpec
16
+
17
+BACKENDS_TO_TEST = [
18
+ _Backend.FLASH_ATTN_VLLM_V1, _Backend.FLASHINFER_VLLM_V1,
19
+ _Backend.FLEX_ATTENTION, _Backend.TRITON_ATTN_VLLM_V1
20
+]
21
+
22
+# Remove flashinfer from the list if it's not available
23
+try:
24
+ import flashinfer # noqa: F401
25
+except ImportError:
26
+ BACKENDS_TO_TEST.remove(_Backend.FLASHINFER_VLLM_V1)
27
+
28
+
29
+def _convert_dtype_to_torch(dtype):
30
+ """Convert ModelDType to torch.dtype."""
31
+ if isinstance(dtype, str):
32
+ if dtype == "auto":
33
+ return torch.float16 # Default dtype for testing
34
+ elif dtype in STR_DTYPE_TO_TORCH_DTYPE:
35
+ return STR_DTYPE_TO_TORCH_DTYPE[dtype]
36
+ else:
37
+ raise ValueError(f"Unknown dtype: {dtype}")
38
+ elif isinstance(dtype, torch.dtype):
39
+ return dtype
40
+ else:
41
+ raise ValueError(f"Unknown dtype: {dtype}")
42
+
43
+
44
+# Define common batch configurations
45
+BATCH_SPECS = {
46
+ "small_decode":
47
+ BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
48
+ "small_prefill":
49
+ BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
50
+ "mixed_small":
51
+ BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
52
+ "medium_decode":
53
+ BatchSpec(seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
54
+ query_lens=[1, 1, 1, 1, 1, 1, 1, 1]),
55
+ "medium_prefill":
56
+ BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]),
57
+ "mixed_medium":
58
+ BatchSpec(seq_lens=[512, 1024, 2048, 512, 1024, 2048],
59
+ query_lens=[1, 1, 1, 7, 7, 7]),
60
+ "large_decode":
61
+ BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
62
+ "large_prefill":
63
+ BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
64
+ "single_decode":
65
+ BatchSpec(seq_lens=[1024], query_lens=[1]),
66
+ "single_prefill":
67
+ BatchSpec(seq_lens=[1024], query_lens=[64]),
68
+}
69
+
70
+
71
+def create_dummy_kv_cache(kv_cache_spec: FullAttentionSpec,
72
+ device: torch.device,
73
+ num_blocks: int = 100) -> torch.Tensor:
74
+ """Create a dummy KV cache tensor for testing."""
75
+ kv_cache = torch.randn(
76
+ 2, # K and V
77
+ num_blocks,
78
+ kv_cache_spec.block_size,
79
+ kv_cache_spec.num_kv_heads,
80
+ kv_cache_spec.head_size,
81
+ dtype=_convert_dtype_to_torch(kv_cache_spec.dtype),
82
+ device=device,
83
+ )
84
+ return kv_cache
85
+
86
+
87
+def create_and_prepopulate_kv_cache(
88
+ k_contexts: list[torch.Tensor],
89
+ v_contexts: list[torch.Tensor],
90
+ block_size: int,
91
+ num_kv_heads: int,
92
+ head_size: int,
93
+ dtype: torch.dtype,
94
+ device: torch.device,
95
+ num_blocks: int,
96
+ common_attn_metadata: CommonAttentionMetadata,
97
+ randomize_blocks: bool = True) -> torch.Tensor:
98
+ """Create and prepopulate a KV cache with context data.
99
+
100
+ Args:
101
+ k_contexts: List of key context tensors for each sequence
102
+ v_contexts: List of value context tensors for each sequence
103
+ seq_lens: List of sequence lengths
104
+ block_size: Size of each block
105
+ num_kv_heads: Number of KV heads
106
+ head_size: Size of each head
107
+ dtype: Data type for the cache
108
+ device: Device to create the cache on
109
+ num_blocks: Total number of blocks in the cache
110
+ block_table: Block table tensor to populate
111
+ randomize_blocks: Whether to randomly permute blocks
112
+ or use sequential order
113
+
114
+ Returns:
115
+ Tuple of (kv_cache, updated_block_table)
116
+ """
117
+ batch_size = len(k_contexts)
118
+ seq_lens = common_attn_metadata.seq_lens_cpu
119
+ query_lens = common_attn_metadata.query_start_loc_cpu[
120
+ 1:] - common_attn_metadata.query_start_loc_cpu[:-1]
121
+ context_lens = common_attn_metadata.num_computed_tokens_cpu
122
+ block_table = common_attn_metadata.block_table_tensor
123
+ slot_mapping = common_attn_metadata.slot_mapping
124
+
125
+ # Create KV cache
126
+ kv_cache = torch.empty(2,
127
+ num_blocks,
128
+ block_size,
129
+ num_kv_heads,
130
+ head_size,
131
+ dtype=dtype,
132
+ device=device)
133
+ kv_cache_flat = kv_cache.view(2, -1, num_kv_heads, head_size)
134
+
135
+ # Populate the cache with the context tokens
136
+ # Start from block_id=1 since block_id=0 is considered the null block
137
+ start_block_idx = 1
138
+ for i in range(batch_size):
139
+ k_context, v_context = k_contexts[i], v_contexts[i]
140
+ start = start_block_idx * block_size
141
+ end = start + k_context.shape[0]
142
+ kv_cache_flat[0, start:end, ...] = k_context
143
+ kv_cache_flat[1, start:end, ...] = v_context
144
+
145
+ # Stay block aligned and allocate enough blocks for the new tokens
146
+ start_block_idx += cdiv(int(seq_lens[i]), block_size)
147
+
148
+ blocks_end = start_block_idx
149
+
150
+ # Permute the context blocks (excluding block 0 which is null)
151
+ if randomize_blocks:
152
+ perm = torch.randperm(
153
+ blocks_end - 1) + 1 # Random permutation starting from block 1
154
+ else:
155
+ perm = torch.arange(
156
+ 1, blocks_end) # Sequential order starting from block 1
157
+
158
+ inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
159
+ inv_perm[1:] = torch.argsort(
160
+ perm) + 1 # Add 1 to account for starting from block 1
161
+ kv_cache[:, 1:blocks_end, ...] = kv_cache[:, perm, ...]
162
+
163
+ # Construct the right block table
164
+ # Start from block_id=1 since block_id=0 is considered the null block
165
+ start_block_idx = 1
166
+ for i in range(batch_size):
167
+ num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
168
+ start = start_block_idx
169
+ end = start + num_blocks_for_seq
170
+ block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
171
+ start_block_idx += num_blocks_for_seq
172
+
173
+ # Create a realistic slot mapping that corresponds to the block table
174
+ for i in range(batch_size):
175
+ token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
176
+ block_indices = token_offsets // block_size
177
+ token_inter_block_offsets = token_offsets % block_size
178
+ start = common_attn_metadata.query_start_loc_cpu[i]
179
+ end = common_attn_metadata.query_start_loc_cpu[i + 1]
180
+ slot_mapping[start:end] = block_table[
181
+ i,
182
+ block_indices] * block_size + token_inter_block_offsets.to(device)
183
+
184
+ return kv_cache
185
+
186
+
187
+class MockAttentionLayer:
188
+ """A mock attention layer for testing."""
189
+
190
+ def __init__(self, device: torch.device):
191
+ self._q_scale = torch.tensor(1.0, device=device)
192
+ self._k_scale = torch.tensor(1.0, device=device)
193
+ self._v_scale = torch.tensor(1.0, device=device)
194
+ # Add float versions for flashinfer
195
+ self._k_scale_float = 1.0
196
+ self._v_scale_float = 1.0
197
+
198
+
199
+def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
200
+ vllm_config, device: torch.device,
201
+ common_attn_metadata: CommonAttentionMetadata,
202
+ query: torch.Tensor, key: torch.Tensor,
203
+ value: torch.Tensor,
204
+ kv_cache: torch.Tensor) -> torch.Tensor:
205
+ """Run attention computation using the specified backend's AttentionImpl."""
206
+
207
+ builder_cls, impl_cls = get_attention_backend(backend)
208
+
209
+ # Mock flashinfer's get_per_layer_parameters if needed
210
+ if backend == _Backend.FLASHINFER_VLLM_V1:
211
+ import unittest.mock
212
+
213
+ from vllm.v1.attention.backends.flashinfer import PerLayerParameters
214
+
215
+ def mock_get_per_layer_parameters(vllm_config):
216
+ # Return mock parameters for a single layer
217
+ head_size = vllm_config.model_config.get_head_size()
218
+ return {
219
+ "mock_layer":
220
+ PerLayerParameters(
221
+ window_left=-1, # No sliding window
222
+ logits_soft_cap=0.0, # No soft cap
223
+ sm_scale=1.0 / (head_size**0.5) # Standard scale
224
+ )
225
+ }
226
+
227
+ with unittest.mock.patch(
228
+ 'vllm.v1.attention.backends.flashinfer.get_per_layer_parameters',
229
+ mock_get_per_layer_parameters):
230
+ builder = builder_cls(kv_cache_spec, vllm_config, device)
231
+ attn_metadata = builder.build(
232
+ common_prefix_len=0,
233
+ common_attn_metadata=common_attn_metadata,
234
+ )
235
+ else:
236
+ # Build metadata
237
+ builder = builder_cls(kv_cache_spec, vllm_config, device)
238
+ attn_metadata = builder.build(
239
+ common_prefix_len=0,
240
+ common_attn_metadata=common_attn_metadata,
241
+ )
242
+
243
+ # Instantiate implementation
244
+ num_heads = vllm_config.model_config.get_num_attention_heads(
245
+ vllm_config.parallel_config)
246
+ num_kv_heads = vllm_config.model_config.get_num_kv_heads(
247
+ vllm_config.parallel_config)
248
+ head_size = vllm_config.model_config.get_head_size()
249
+ scale = 1.0 / (head_size**0.5)
250
+ impl = impl_cls(
251
+ num_heads=num_heads,
252
+ head_size=head_size,
253
+ scale=scale,
254
+ num_kv_heads=num_kv_heads,
255
+ alibi_slopes=None,
256
+ sliding_window=None,
257
+ kv_cache_dtype="auto",
258
+ )
259
+
260
+ # Create mock layer and output buffer
261
+ mock_layer = MockAttentionLayer(device)
262
+ output = torch.empty_like(query)
263
+
264
+ # Run forward pass
265
+ # NOTE: The query, key, and value are already shaped correctly
266
+ # in the calling test function.
267
+ output = impl.forward(mock_layer,
268
+ query,
269
+ key,
270
+ value,
271
+ kv_cache,
272
+ attn_metadata,
273
+ output=output)
274
+
275
+ return output
276
+
277
+
278
+@pytest.mark.parametrize("batch_spec_name", [
279
+ "small_decode", "small_prefill", "mixed_small", "medium_decode",
280
+ "medium_prefill", "mixed_medium"
281
+])
282
+@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
283
+def test_backend_correctness(batch_spec_name: str, model: str):
284
+ """
285
+ Test that all backends produce similar outputs to a reference implementation
286
+ using torch.nn.functional.scaled_dot_product_attention.
287
+
288
+ This test works by:
289
+ 1. Generating a batch of sequences with specified context and query lengths.
290
+ 2. Computing a ground-truth attention output using torch.sdpa on
291
+ contiguous Q, K, and V tensors.
292
+ 3. Simulating vLLM's paged KV cache: It takes the context portion of the
293
+ K/V tensors and manually places them into a paged buffer according to
294
+ the test's (randomly generated) block table.
295
+ 4. Running each vLLM attention backend with the new queries and the
296
+ simulated paged KV cache.
297
+ 5. Comparing the vLLM backend's output to the ground-truth SDPA output.
298
+ """
299
+ batch_spec = BATCH_SPECS[batch_spec_name]
300
+ vllm_config = create_vllm_config(model_name=model)
301
+ device = torch.device("cuda:0")
302
+
303
+ kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
304
+
305
+ # 1. Setup
306
+ batch_size = batch_spec.batch_size
307
+ seq_lens = batch_spec.seq_lens
308
+ query_lens = batch_spec.query_lens
309
+ num_q_heads = vllm_config.model_config.get_num_attention_heads(
310
+ vllm_config.parallel_config)
311
+ num_kv_heads = vllm_config.model_config.get_num_kv_heads(
312
+ vllm_config.parallel_config)
313
+ head_size = vllm_config.model_config.get_head_size()
314
+ dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
315
+ block_size = vllm_config.cache_config.block_size
316
+ scale = 1.0 / (head_size**0.5)
317
+
318
+ # 2. Generate data and compute SDPA reference output
319
+ all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
320
+ all_sdpa_outputs = []
321
+ k_contexts, v_contexts = [], []
322
+
323
+ for i in range(batch_size):
324
+ s_len = seq_lens[i]
325
+ q_len = query_lens[i]
326
+ context_len = s_len - q_len
327
+
328
+ # Generate Q, K, V for the whole sequence to be used in SDPA
329
+ q = torch.randn(q_len,
330
+ num_q_heads,
331
+ head_size,
332
+ dtype=dtype,
333
+ device=device)
334
+ k_full = torch.randn(s_len,
335
+ num_kv_heads,
336
+ head_size,
337
+ dtype=dtype,
338
+ device=device)
339
+ v_full = torch.randn(s_len,
340
+ num_kv_heads,
341
+ head_size,
342
+ dtype=dtype,
343
+ device=device)
344
+
345
+ # SDPA expects (N, H, L, D), so unsqueeze batch and permute
346
+ q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
347
+ k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
348
+ v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
349
+
350
+ if num_q_heads != num_kv_heads:
351
+ assert num_q_heads % num_kv_heads == 0, (
352
+ f"num_q_heads ({num_q_heads}) must be divisible by "
353
+ f"num_kv_heads ({num_kv_heads})")
354
+ repeats = num_q_heads // num_kv_heads
355
+ k_sdpa_in = k_sdpa_in.repeat_interleave(repeats, dim=1)
356
+ v_sdpa_in = v_sdpa_in.repeat_interleave(repeats, dim=1)
357
+
358
+ # Create causal mask: query token i attends to positions 0 to
359
+ # (context_len + i)
360
+ kv_len = s_len
361
+ offset = context_len
362
+ attn_mask = torch.full((q_len, kv_len),
363
+ float('-inf'),
364
+ device=device,
365
+ dtype=dtype)
366
+ for i in range(q_len):
367
+ attn_mask[i, :offset + i + 1] = 0.0
368
+
369
+ sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
370
+ q_sdpa_in,
371
+ k_sdpa_in,
372
+ v_sdpa_in,
373
+ attn_mask=attn_mask,
374
+ scale=scale,
375
+ enable_gqa=True)
376
+ # Convert back to (L, H, D)
377
+ all_sdpa_outputs.append(sdpa_out_i.transpose(1, 2).squeeze(0))
378
+
379
+ # Inputs for vLLM backends are just the new tokens
380
+ all_q_vllm.append(q)
381
+ all_k_vllm.append(k_full[context_len:])
382
+ all_v_vllm.append(v_full[context_len:])
383
+
384
+ # Contextual K/V data used to populate the paged cache
385
+ k_contexts.append(k_full[:context_len])
386
+ v_contexts.append(v_full[:context_len])
387
+
388
+ query_vllm = torch.cat(all_q_vllm, dim=0)
389
+ key_vllm = torch.cat(all_k_vllm, dim=0)
390
+ value_vllm = torch.cat(all_v_vllm, dim=0)
391
+ sdpa_output = torch.cat(all_sdpa_outputs, dim=0)
392
+
393
+ common_attn_metadata = create_common_attn_metadata(
394
+ batch_spec, vllm_config.cache_config.block_size, device)
395
+
396
+ # 3. Simulate Paged KV Cache and a realistic slot_mapping
397
+ kv_cache = create_and_prepopulate_kv_cache(
398
+ k_contexts=k_contexts,
399
+ v_contexts=v_contexts,
400
+ block_size=block_size,
401
+ num_kv_heads=num_kv_heads,
402
+ head_size=head_size,
403
+ dtype=dtype,
404
+ device=device,
405
+ num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
406
+ common_attn_metadata=common_attn_metadata,
407
+ randomize_blocks=True)
408
+
409
+ # 4. Run vLLM backends and compare
410
+ # Note: flex_attention has known Triton kernel compatibility issues
411
+ # with test infrastructures
412
+ for backend_name in BACKENDS_TO_TEST:
413
+ # FlashAttentionm + FlexAttention:
414
+ # [2, num_blocks, block_size, num_kv_heads, head_size]
415
+ # FlashInfer:
416
+ # [num_blocks, 2, block_size, num_kv_heads, head_size]
417
+ # Select the appropriate KV cache format for each backend
418
+ kv_cache_for_backend = kv_cache
419
+ if backend_name == _Backend.FLASHINFER_VLLM_V1:
420
+ kv_cache_for_backend = kv_cache.transpose(0, 1)
421
+
422
+ backend_output = run_attention_backend(backend_name, kv_cache_spec,
423
+ vllm_config, device,
424
+ common_attn_metadata,
425
+ query_vllm, key_vllm,
426
+ value_vllm,
427
+ kv_cache_for_backend)
428
+
429
+ # Check shape and dtype consistency
430
+ assert backend_output.shape == sdpa_output.shape, (
431
+ f"[{backend_name}] shape {backend_output.shape} != "
432
+ f"SDPA shape {sdpa_output.shape}")
433
+ assert backend_output.dtype == sdpa_output.dtype, (
434
+ f"[{backend_name}] dtype {backend_output.dtype} != "
435
+ f"SDPA dtype {sdpa_output.dtype}")
436
+
437
+ assert torch.isfinite(backend_output).all(), (
438
+ f"[{backend_name}] produced non-finite values")
439
+
440
+ # Check numerical similarity
441
+ rtol = 1e-2
442
+ atol = 5e-3
443
+
444
+ if backend_name == _Backend.FLEX_ATTENTION:
445
+ atol = 5e-1 # TODO: figure out why flex_attention has such large
446
+ # numerical differences for medium_decode, medium_prefill,
447
+ # mixed_medium
448
+
449
+ max_diff = torch.max(torch.abs(backend_output - sdpa_output)).item()
450
+ max_rel_diff = torch.max(
451
+ torch.abs(backend_output - sdpa_output) /
452
+ torch.abs(sdpa_output)).item()
453
+ all_close = torch.allclose(backend_output,
454
+ sdpa_output,
455
+ rtol=rtol,
456
+ atol=atol)
457
+
458
+ if not all_close:
459
+ print(f"[{backend_name}] output differs from SDPA baseline. "
460
+ f"Max diff: {max_diff:.6f} (rel: {max_rel_diff:.6f})")
461
+ print(f"[{backend_name}] output: {backend_output}")
462
+ print(f"[{backend_name}] SDPA baseline: {sdpa_output}")
463
+
464
+ assert all_close, (
465
+ f"[{backend_name}] output differs from SDPA baseline. "
466
+ f"Max diff: {max_diff:.6f} (rel: {max_rel_diff:.6f})")
tests/v1/attention/utils.py
new
+229
@@ -0,0 +1,229 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Utility functions for attention-related v1 tests."""
4
+
5
+from dataclasses import dataclass
6
+from typing import Union
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.config import (CacheConfig, CompilationConfig, DeviceConfig,
12
+ LoadConfig, ModelConfig, ModelDType, ParallelConfig,
13
+ SchedulerConfig, VllmConfig)
14
+from vllm.platforms import _Backend
15
+from vllm.utils import resolve_obj_by_qualname
16
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
17
+from vllm.v1.kv_cache_interface import FullAttentionSpec
18
+
19
+
20
+@dataclass
21
+class BatchSpec:
22
+ """Specification for a batch configuration (workload shape only)."""
23
+ seq_lens: list[int]
24
+ query_lens: list[int]
25
+
26
+ name: str = "unnamed"
27
+
28
+ @property
29
+ def batch_size(self):
30
+ return len(self.seq_lens)
31
+
32
+ def __post_init__(self):
33
+ assert len(self.seq_lens) == len(self.query_lens)
34
+
35
+ def compute_num_tokens(self):
36
+ return sum(self.query_lens)
37
+
38
+
39
+def create_common_attn_metadata(
40
+ batch_spec: BatchSpec,
41
+ block_size: int,
42
+ device: torch.device,
43
+ max_block_idx: int = 1000) -> CommonAttentionMetadata:
44
+ """Create CommonAttentionMetadata from a BatchSpec and ModelParams."""
45
+ # Create query start locations
46
+ query_start_loc = torch.zeros(batch_spec.batch_size + 1,
47
+ dtype=torch.int32,
48
+ device=device)
49
+ query_start_loc[1:] = torch.tensor(batch_spec.query_lens,
50
+ dtype=torch.int32,
51
+ device=device).cumsum(0)
52
+ query_start_loc_cpu = query_start_loc.cpu()
53
+ num_tokens = batch_spec.compute_num_tokens()
54
+
55
+ # Create sequence lengths
56
+ seq_lens = torch.tensor(batch_spec.seq_lens,
57
+ dtype=torch.int32,
58
+ device=device)
59
+ seq_lens_cpu = seq_lens.cpu()
60
+
61
+ # Create computed tokens (context length for each sequence)
62
+ context_lens = [
63
+ batch_spec.seq_lens[i] - batch_spec.query_lens[i]
64
+ for i in range(batch_spec.batch_size)
65
+ ]
66
+ num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
67
+
68
+ # Create block table (random for testing)
69
+ max_blocks = max(batch_spec.seq_lens) // block_size + 1
70
+ block_table_tensor = torch.randint(0,
71
+ max_block_idx,
72
+ (batch_spec.batch_size, max_blocks),
73
+ dtype=torch.int32,
74
+ device=device)
75
+
76
+ # Create slot mapping
77
+ slot_mapping = torch.randint(0,
78
+ max_block_idx, (num_tokens, ),
79
+ dtype=torch.int64,
80
+ device=device)
81
+
82
+ # Calculate max query length
83
+ max_query_len = max(batch_spec.query_lens)
84
+
85
+ return CommonAttentionMetadata(
86
+ query_start_loc=query_start_loc,
87
+ query_start_loc_cpu=query_start_loc_cpu,
88
+ seq_lens=seq_lens,
89
+ seq_lens_cpu=seq_lens_cpu,
90
+ num_computed_tokens_cpu=num_computed_tokens_cpu,
91
+ num_reqs=batch_spec.batch_size,
92
+ num_actual_tokens=num_tokens,
93
+ max_query_len=max_query_len,
94
+ block_table_tensor=block_table_tensor,
95
+ slot_mapping=slot_mapping,
96
+ )
97
+
98
+
99
+def get_attention_backend(backend_name: _Backend):
100
+ """Set up attention backend classes for testing.
101
+
102
+ Args:
103
+ backend_name: Name of the backend ("flash_attn", "flashinfer", etc.)
104
+ vllm_config: VllmConfig instance
105
+
106
+ Returns:
107
+ Tuple of (backend_builder_class, backend_impl_class)
108
+ """
109
+ backend_map = {
110
+ _Backend.FLASH_ATTN_VLLM_V1:
111
+ "vllm.v1.attention.backends.flash_attn.FlashAttentionBackend",
112
+ _Backend.FLASHINFER_VLLM_V1:
113
+ "vllm.v1.attention.backends.flashinfer.FlashInferBackend",
114
+ _Backend.FLEX_ATTENTION:
115
+ "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend",
116
+ _Backend.TRITON_ATTN_VLLM_V1:
117
+ "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend",
118
+ }
119
+
120
+ if backend_name not in backend_map:
121
+ raise ValueError(f"Unknown backend: {backend_name}")
122
+
123
+ backend_class_name = backend_map[backend_name]
124
+
125
+ try:
126
+ backend_class = resolve_obj_by_qualname(backend_class_name)
127
+ return backend_class.get_builder_cls(), backend_class.get_impl_cls()
128
+ except ImportError as e:
129
+ pytest.skip(f"{backend_name} not available: {e}")
130
+
131
+
132
+def create_standard_kv_cache_spec(
133
+ vllm_config: VllmConfig) -> FullAttentionSpec:
134
+ """Create a FullAttentionSpec from ModelParams only."""
135
+ return FullAttentionSpec(
136
+ block_size=vllm_config.cache_config.block_size,
137
+ num_kv_heads=vllm_config.model_config.get_num_kv_heads(
138
+ vllm_config.parallel_config),
139
+ head_size=vllm_config.model_config.get_head_size(),
140
+ dtype=vllm_config.model_config.dtype,
141
+ use_mla=vllm_config.model_config.use_mla,
142
+ sliding_window=vllm_config.model_config.get_sliding_window(),
143
+ )
144
+
145
+
146
+def create_vllm_config(model_name: str = "meta-llama/Meta-Llama-3-8B",
147
+ tensor_parallel_size: int = 1,
148
+ max_model_len: int = 1024,
149
+ dtype: Union[ModelDType, torch.dtype] = "auto",
150
+ block_size: int = 16,
151
+ max_num_seqs: int = 256,
152
+ max_num_batched_tokens: int = 8192,
153
+ add_mock_model_methods: bool = True) -> VllmConfig:
154
+ """Create a VllmConfig for testing with reasonable defaults."""
155
+
156
+ model_config = ModelConfig(
157
+ model=model_name,
158
+ tokenizer=model_name,
159
+ trust_remote_code=False,
160
+ dtype=dtype,
161
+ seed=0,
162
+ max_model_len=max_model_len,
163
+ )
164
+
165
+ cache_config = CacheConfig(
166
+ block_size=block_size,
167
+ cache_dtype="auto",
168
+ swap_space=0,
169
+ )
170
+ # Set cache blocks for testing
171
+ # (these may be set during initialization normally)
172
+ cache_config.num_gpu_blocks = 1000
173
+ cache_config.num_cpu_blocks = 0
174
+
175
+ parallel_config = ParallelConfig(
176
+ tensor_parallel_size=tensor_parallel_size, )
177
+
178
+ scheduler_config = SchedulerConfig(
179
+ max_num_seqs=max_num_seqs,
180
+ max_num_batched_tokens=max_num_batched_tokens,
181
+ )
182
+
183
+ device_config = DeviceConfig()
184
+ load_config = LoadConfig()
185
+ compilation_config = CompilationConfig()
186
+
187
+ if add_mock_model_methods:
188
+ # Add mock methods to satisfy backends that need them
189
+ # This is a workaround because tests don't build full, real models,
190
+ # but some backends expect to query the model for layer-specific
191
+ # parameters
192
+ import types
193
+ model_config.get_num_layers = types.MethodType(lambda self: 1,
194
+ model_config)
195
+ model_config.get_sliding_window_for_layer = types.MethodType(
196
+ lambda self, i: None, model_config)
197
+ model_config.get_logits_soft_cap_for_layer = types.MethodType(
198
+ lambda self, i: 0.0, model_config)
199
+ model_config.get_sm_scale_for_layer = types.MethodType(
200
+ lambda self, i: 1.0 / model_config.get_head_size()**0.5,
201
+ model_config)
202
+
203
+ return VllmConfig(
204
+ model_config=model_config,
205
+ cache_config=cache_config,
206
+ parallel_config=parallel_config,
207
+ scheduler_config=scheduler_config,
208
+ device_config=device_config,
209
+ load_config=load_config,
210
+ compilation_config=compilation_config,
211
+ )
212
+
213
+
214
+def create_dummy_kv_cache(block_size: int,
215
+ num_kv_heads: int,
216
+ head_size: int,
217
+ dtype: torch.dtype,
218
+ device: torch.device,
219
+ num_blocks: int = 100) -> torch.Tensor:
220
+ """Create a dummy KV cache tensor for testing."""
221
+ kv_cache = torch.randn(
222
+ num_blocks,
223
+ 2, # K and V
224
+ block_size,
225
+ num_kv_heads,
226
+ head_size,
227
+ dtype=dtype,
228
+ device=device)
229
+ return kv_cache
tests/v1/spec_decode/test_eagle.py
+44
-24
@@ -6,6 +6,10 @@ from unittest import mock
6
import pytest
7
import torch
8
9
+from tests.v1.attention.utils import (BatchSpec, _Backend,
10
+ create_common_attn_metadata,
11
+ create_standard_kv_cache_spec,
12
+ get_attention_backend)
13
from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, ModelConfig,
14
ParallelConfig, SchedulerConfig, SpeculativeConfig,
15
VllmConfig)
@@ -64,13 +68,19 @@ def test_prepare_inputs():
68
"""
69
device = torch.device(current_platform.device_type)
70
67
- # a = 4, b = 7, c = 5
71
+ # q1 = 4, q2 = 7, q3 = 5
72
# n1 = 1, n2 = 3, n3 = 2
73
70
- # Cumulative lengths: [0, 4, 11, 16]
71
- cu_target_query_lens = torch.tensor([0, 4, 11, 16],
72
- dtype=torch.int32,
73
- device=device)
74
+ batch_spec = BatchSpec(
75
+ seq_lens=[4, 7, 5],
76
+ query_lens=[4, 7, 5],
77
+ )
78
+
79
+ common_attn_metadata = create_common_attn_metadata(
80
+ batch_spec,
81
+ block_size=16,
82
+ device=device,
83
+ )
84
85
# Rejected tokens per request: [1, 3, 2]
86
num_rejected_tokens = torch.tensor([1, 3, 2],
@@ -104,15 +114,13 @@ def test_prepare_inputs():
114
],
115
dtype=torch.int32,
116
device=device)
117
+ proposer = _create_proposer("eagle", 1)
118
108
- # n1 + n2 + n3 - a - b -c
109
- num_tokens = cu_target_query_lens[-1].item() - num_rejected_tokens.sum(
110
- ).item()
119
+ updated_metadata, token_indices = proposer.prepare_inputs(
120
+ common_attn_metadata, num_rejected_tokens.cpu())
121
112
- cu_num_tokens, token_indices = EagleProposer.prepare_inputs(
113
- cu_target_query_lens, num_rejected_tokens, num_tokens)
114
-
115
- assert torch.equal(cu_num_tokens, expected_cu_num_tokens)
122
+ assert torch.equal(updated_metadata.query_start_loc,
123
+ expected_cu_num_tokens)
124
assert token_indices.shape[0] == expected_cu_num_tokens[-1].item()
125
assert torch.equal(token_indices, expected_token_indices)
126
@@ -209,6 +217,7 @@ def test_propose(num_speculative_tokens):
217
seq_len_2 = 3
218
total_tokens = seq_len_1 + seq_len_2
219
vocab_size = 100
220
+ seq_lens = [seq_len_1, seq_len_2]
221
222
# Create proposer first so we can use its actual hidden_size
223
proposer = _create_proposer("eagle", num_speculative_tokens)
@@ -270,9 +279,16 @@ def test_propose(num_speculative_tokens):
279
proposer.attn_layer_names = ["layer.0"]
280
281
# Create input tensors
273
- cu_num_tokens = torch.tensor([0, seq_len_1, total_tokens],
274
- dtype=torch.int32,
275
- device=device)
282
+ batch_spec = BatchSpec(
283
+ seq_lens=seq_lens,
284
+ query_lens=seq_lens,
285
+ )
286
+
287
+ common_attn_metadata = create_common_attn_metadata(
288
+ batch_spec,
289
+ block_size=16,
290
+ device=device,
291
+ )
292
293
target_token_ids = torch.randint(0,
294
vocab_size, (total_tokens, ),
@@ -284,25 +300,29 @@ def test_propose(num_speculative_tokens):
300
target_hidden_states = torch.randn(total_tokens,
301
hidden_size,
302
device=device)
287
- target_slot_mapping = torch.randint(0,
288
- 100, (total_tokens, ),
289
- device=device)
303
next_token_ids = torch.randint(0,
304
vocab_size, (batch_size, ),
305
dtype=torch.int32,
306
device=device)
294
- block_table = torch.randint(0, 10, (batch_size, 10), device=device)
295
-
307
sampling_metadata = mock.MagicMock()
308
298
- # Call the method under test
309
+ attn_metadata_builder_cls, _ = get_attention_backend(
310
+ _Backend.FLASH_ATTN_VLLM_V1)
311
+ attn_metadata_builder = attn_metadata_builder_cls(
312
+ kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config),
313
+ vllm_config=proposer.vllm_config,
314
+ device=device,
315
+ )
316
+
317
+ # Mock runner for attention metadata building
318
+ proposer.runner = mock.MagicMock()
319
+ proposer.runner.attn_metadata_builders = [attn_metadata_builder]
320
+
321
result = proposer.propose(target_token_ids=target_token_ids,
322
target_positions=target_positions,
323
target_hidden_states=target_hidden_states,
302
- target_slot_mapping=target_slot_mapping,
324
next_token_ids=next_token_ids,
304
- cu_num_tokens=cu_num_tokens,
305
- block_table=block_table,
325
+ common_attn_metadata=common_attn_metadata,
326
sampling_metadata=sampling_metadata)
327
328
assert result.shape == (batch_size, num_speculative_tokens)
vllm/v1/attention/backends/cpu_attn.py
+34
-31
@@ -12,13 +12,12 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
12
AttentionMetadata, AttentionType,
13
is_quantized_kv_cache)
14
from vllm.attention.backends.utils import CommonAttentionState
15
+from vllm.config import VllmConfig
16
from vllm.logger import init_logger
17
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
18
CommonAttentionMetadata)
19
from vllm.v1.core.sched.output import SchedulerOutput
20
from vllm.v1.kv_cache_interface import AttentionSpec
20
-from vllm.v1.worker.block_table import BlockTable
21
-from vllm.v1.worker.cpu_model_runner import CPUModelRunner
21
from vllm.v1.worker.gpu_input_batch import InputBatch
22
23
try:
@@ -316,19 +315,21 @@ class TorchSDPAMetadata(AttentionMetadata):
315
316
class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
317
319
- def __init__(self, runner: CPUModelRunner, kv_cache_spec: AttentionSpec,
320
- block_table: BlockTable) -> None:
321
- self.runner = runner
322
- self.block_table = block_table
318
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
319
+ device: torch.device) -> None:
320
+ self.kv_cache_spec = kv_cache_spec
321
+ self.vllm_config = vllm_config
322
+ self.scheduler_config = vllm_config.scheduler_config
323
+
324
# For reorder
324
- self.reorder_prompt_req_index_list = np.empty(self.runner.max_num_reqs,
325
- dtype=np.int64)
326
- self.reorder_decode_req_index_list = np.empty(self.runner.max_num_reqs,
327
- dtype=np.int64)
325
+ self.reorder_prompt_req_index_list = np.empty(
326
+ vllm_config.scheduler_config.max_num_seqs, dtype=np.int64)
327
+ self.reorder_decode_req_index_list = np.empty(
328
+ vllm_config.scheduler_config.max_num_seqs, dtype=np.int64)
329
self.num_prompt_req: int = 0
330
331
self.seq_start_loc_cpu = torch.zeros(
331
- runner.max_num_reqs + 1,
332
+ vllm_config.scheduler_config.max_num_seqs + 1,
333
dtype=torch.int32,
334
device="cpu",
335
)
@@ -378,15 +379,15 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
379
380
return True
381
381
- def build(self, common_prefix_len: int,
382
- common_attn_metadata: CommonAttentionMetadata):
382
+ def build(self,
383
+ common_prefix_len: int,
384
+ common_attn_metadata: CommonAttentionMetadata,
385
+ fast_build: bool = False) -> TorchSDPAMetadata:
386
num_reqs = common_attn_metadata.num_reqs
384
- num_actual_tokens = common_attn_metadata.num_actual_tokens
387
max_query_len = common_attn_metadata.max_query_len
388
387
- runner = self.runner
388
- block_table = self.block_table
389
- seq_lens_np = runner.seq_lens_np[:num_reqs]
389
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu
390
+ seq_lens_np = seq_lens_cpu.numpy()
391
num_prompt_req = self.num_prompt_req
392
max_prefill_seq_len = seq_lens_np[:num_prompt_req].max().item(
393
) if num_prompt_req > 0 else 0
@@ -394,34 +395,36 @@ class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
395
) if num_prompt_req < num_reqs else 0
396
self.seq_start_loc_np[0] = 0
397
np.cumsum(seq_lens_np, out=self.seq_start_loc_np[1:num_reqs + 1])
397
- num_prefill_tokens = runner.query_start_loc_np[num_prompt_req].item()
398
- num_decode_tokens = runner.query_start_loc_np[num_reqs].item(
399
- ) - num_prefill_tokens
400
- slot_mapping = block_table.slot_mapping_cpu[:num_actual_tokens].long()
401
- block_table_tensor = block_table.get_device_tensor()
398
+
399
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
400
+ num_prefill_tokens = int(query_start_loc_cpu[num_prompt_req].item())
401
+ num_decode_tokens = int(query_start_loc_cpu[num_reqs].item() -
402
+ num_prefill_tokens)
403
+
404
+ slot_mapping = common_attn_metadata.slot_mapping.long()
405
+ block_table_tensor = common_attn_metadata.block_table_tensor
406
+
407
attn_metadata = TorchSDPAMetadata(
408
num_prefills=num_prompt_req,
409
num_prefill_tokens=num_prefill_tokens,
410
num_decode_tokens=num_decode_tokens,
411
slot_mapping=slot_mapping,
412
# to ensure inference when chunked_prefill is disabled
408
- seq_lens=runner.seq_lens_cpu[:num_reqs].tolist(),
409
- seq_lens_tensor=runner.
410
- seq_lens_cpu[num_prompt_req:num_reqs], # decode
413
+ seq_lens=seq_lens_cpu.tolist(),
414
+ seq_lens_tensor=seq_lens_cpu[num_prompt_req:num_reqs], # decode
415
max_decode_seq_len=max_decode_seq_len, # decode
416
block_tables=block_table_tensor[num_prompt_req:num_reqs], # decode
413
- chunked_prefill=self.runner.scheduler_config.
414
- chunked_prefill_enabled,
417
+ chunked_prefill=self.scheduler_config.chunked_prefill_enabled,
418
max_query_len=max_query_len,
419
max_kv_len=max_prefill_seq_len,
417
- prefill_query_start_loc=runner.
418
- query_start_loc_cpu[:num_prompt_req + 1], # prefill
420
+ prefill_query_start_loc=query_start_loc_cpu[:num_prompt_req +
421
+ 1], # prefill
422
kv_start_loc=self.seq_start_loc_cpu[:num_prompt_req +
423
1], # prefill
424
prefill_block_tables=block_table_tensor[:
425
num_prompt_req], # prefill
423
- query_start_loc=runner.query_start_loc_cpu[:num_reqs +
424
- 1], # for logits index
426
+ query_start_loc=query_start_loc_cpu[:num_reqs +
427
+ 1], # for logits index
428
multi_modal_placeholder_index_maps=None,
429
enable_kv_scales_calculation=False,
430
)
vllm/v1/attention/backends/flash_attn.py
+48
-53
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with FlashAttention."""
4
from dataclasses import dataclass
5
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
5
+from typing import Any, ClassVar, Optional
6
7
import numpy as np
8
import torch
@@ -29,10 +29,6 @@ from vllm.v1.attention.backends.utils import (
29
AttentionMetadataBuilder, CommonAttentionMetadata, get_kv_cache_layout,
30
make_local_attention_virtual_batches)
31
from vllm.v1.kv_cache_interface import AttentionSpec
32
-from vllm.v1.worker.block_table import BlockTable
33
-
34
-if TYPE_CHECKING:
35
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner
32
33
logger = init_logger(__name__)
34
@@ -162,29 +158,30 @@ class FlashAttentionMetadataBuilder(
158
AttentionMetadataBuilder[FlashAttentionMetadata]):
159
full_cudagraph_supported: ClassVar[bool] = get_flash_attn_version() == 3
160
165
- def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
166
- block_table: BlockTable):
167
- model_config = runner.model_config
168
- compilation_config = runner.vllm_config.compilation_config
169
-
170
- self.runner = runner
171
- self.num_heads_q = model_config.get_num_attention_heads(
172
- runner.parallel_config)
173
- self.num_heads_kv = model_config.get_num_kv_heads(
174
- runner.parallel_config)
175
- self.headdim = model_config.get_head_size()
161
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
162
+ device: torch.device):
163
+ self.vllm_config = vllm_config
164
+ self.model_config = vllm_config.model_config
165
+ self.parallel_config = vllm_config.parallel_config
166
+ self.cache_config = vllm_config.cache_config
167
+ self.compilation_config = vllm_config.compilation_config
168
+ self.device = device
169
+
170
+ self.num_heads_q = self.model_config.get_num_attention_heads(
171
+ self.parallel_config)
172
+ self.num_heads_kv = self.model_config.get_num_kv_heads(
173
+ self.parallel_config)
174
+ self.headdim = self.model_config.get_head_size()
175
self.block_size = kv_cache_spec.block_size
177
- self.kv_cache_spec = kv_cache_spec
178
- self.block_table = block_table
176
177
self.max_num_splits = 0 # No upper bound on the number of splits.
178
self.aot_schedule = (get_flash_attn_version() == 3)
182
- self.use_full_cuda_graph = compilation_config.full_cuda_graph
179
+ self.use_full_cuda_graph = self.compilation_config.full_cuda_graph
180
if self.use_full_cuda_graph:
181
if not self.aot_schedule:
182
raise ValueError(
183
"AoT scheduling is required for full cuda graph.")
187
- capture_sizes = compilation_config.cudagraph_capture_sizes
184
+ capture_sizes = self.compilation_config.cudagraph_capture_sizes
185
if not capture_sizes:
186
raise ValueError(
187
"cudagraph_capture_sizes should not be None when "
@@ -198,9 +195,9 @@ class FlashAttentionMetadataBuilder(
195
"full cuda graph.")
196
197
self.scheduler_metadata = torch.zeros(
201
- self.runner.max_num_reqs + 1,
198
+ vllm_config.scheduler_config.max_num_seqs + 1,
199
dtype=torch.int32,
203
- device=self.runner.device,
200
+ device=self.device,
201
)
202
# When using cuda graph, we need to set the upper bound of the
203
# number of splits so that large enough intermediate buffers are
@@ -211,28 +208,27 @@ class FlashAttentionMetadataBuilder(
208
# populated on first build() call.
209
self.aot_sliding_window: Optional[tuple[int, int]] = None
210
214
- def build(
215
- self, common_prefix_len: int,
216
- common_attn_metadata: CommonAttentionMetadata
217
- ) -> FlashAttentionMetadata:
211
+ def build(self,
212
+ common_prefix_len: int,
213
+ common_attn_metadata: CommonAttentionMetadata,
214
+ fast_build: bool = False) -> FlashAttentionMetadata:
215
+ """
216
+ fast_build disables AOT scheduling, used when there will be few
217
+ iterations i.e. spec-decode
218
+ """
219
num_reqs = common_attn_metadata.num_reqs
220
num_actual_tokens = common_attn_metadata.num_actual_tokens
221
max_query_len = common_attn_metadata.max_query_len
221
-
222
- max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
222
+ max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
223
query_start_loc = common_attn_metadata.query_start_loc
224
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
225
seq_lens = common_attn_metadata.seq_lens
225
- block_table = self.block_table
226
- block_table_tensor = block_table.get_device_tensor()[:num_reqs]
227
-
228
- block_table.slot_mapping[:num_actual_tokens].copy_(
229
- block_table.slot_mapping_cpu[:num_actual_tokens],
230
- non_blocking=True)
231
- # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
232
- # mode.
233
- block_table.slot_mapping[num_actual_tokens:].fill_(-1)
226
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu
227
+ block_table_tensor = common_attn_metadata.block_table_tensor
228
+ slot_mapping = common_attn_metadata.slot_mapping
229
235
- slot_mapping = block_table.slot_mapping[:num_actual_tokens]
230
+ # the overhead of the aot schedule is not worth it for spec-decode
231
+ aot_schedule = self.aot_schedule and not fast_build
232
233
if self.aot_sliding_window is None:
234
self.aot_sliding_window = (-1, -1)
@@ -240,19 +236,20 @@ class FlashAttentionMetadataBuilder(
236
# constant for all layers to. We have to populate this on the first
237
# build() call so the layers are constructed (cannot populate)
238
# in __init__.
243
- if self.aot_schedule:
239
+ if aot_schedule:
240
sliding_window_configs = _get_sliding_window_configs(
245
- self.runner.vllm_config)
241
+ self.vllm_config)
242
if len(sliding_window_configs) == 1:
243
sliding_window_config = sliding_window_configs.pop()
244
if sliding_window_config is not None:
245
self.aot_sliding_window = sliding_window_config
246
elif len(sliding_window_configs) > 1:
247
self.aot_schedule = False
248
+ aot_schedule = False
249
250
def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
251
max_seq_len, causal):
255
- if self.aot_schedule:
252
+ if aot_schedule:
253
return get_scheduler_metadata(
254
batch_size=batch_size,
255
max_seqlen_q=max_query_len,
@@ -271,19 +268,19 @@ class FlashAttentionMetadataBuilder(
268
269
# for local attention
270
local_attn_metadata = None
274
- if self.runner.attention_chunk_size is not None:
271
+ if self.model_config.attention_chunk_size is not None:
272
seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
273
virt_block_table_tensor = make_local_attention_virtual_batches(
277
- self.runner.attention_chunk_size,
278
- self.runner.query_start_loc_np[:num_reqs + 1],
279
- self.runner.seq_lens_np[:num_reqs],
274
+ self.model_config.attention_chunk_size,
275
+ query_start_loc_cpu.numpy(),
276
+ seq_lens_cpu.numpy(),
277
block_table_tensor,
278
self.block_size,
279
)
280
local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
284
- self.runner.device, non_blocking=True)
281
+ self.device, non_blocking=True)
282
local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
286
- self.runner.device, non_blocking=True)
283
+ self.device, non_blocking=True)
284
local_max_query_len = seqlens_q_local_np.max()
285
local_max_seq_len = virt_k_seqlens_np.max()
286
local_scheduler_metadata = schedule(
@@ -308,14 +305,12 @@ class FlashAttentionMetadataBuilder(
305
if use_cascade:
306
cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],
307
dtype=torch.int32,
311
- device=self.runner.device)
308
+ device=self.device)
309
prefix_kv_lens = torch.tensor([common_prefix_len],
310
dtype=torch.int32,
314
- device=self.runner.device)
315
- suffix_kv_lens = (self.runner.seq_lens_np[:num_reqs] -
316
- common_prefix_len)
317
- suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(
318
- self.runner.device)
311
+ device=self.device)
312
+ suffix_kv_lens = (seq_lens_cpu[:num_reqs] - common_prefix_len).to(
313
+ self.device, non_blocking=True)
314
prefix_scheduler_metadata = schedule(
315
batch_size=1,
316
cu_query_lens=cu_prefix_query_lens,
vllm/v1/attention/backends/flashinfer.py
+51
-106
@@ -15,22 +15,20 @@ from flashinfer.decode import trtllm_batch_decode_with_kv_cache
15
import vllm.envs as envs
16
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
17
AttentionType)
18
+from vllm.config import VllmConfig
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
from vllm.v1.attention.backends.flash_attn import use_cascade_attention
21
-from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
22
- CommonAttentionMetadata,
23
- PerLayerParameters,
24
- get_kv_cache_layout,
25
- get_per_layer_parameters,
26
- infer_global_hyperparameters)
22
+from vllm.v1.attention.backends.utils import (
23
+ AttentionMetadataBuilder, CommonAttentionMetadata, PerLayerParameters,
24
+ get_kv_cache_layout, get_per_layer_parameters,
25
+ infer_global_hyperparameters, reorder_batch_to_split_decodes_and_prefills,
26
+ split_decodes_and_prefills)
27
from vllm.v1.kv_cache_interface import AttentionSpec
28
-from vllm.v1.worker.block_table import BlockTable
28
29
if TYPE_CHECKING:
30
from vllm.v1.core.sched.output import SchedulerOutput
31
from vllm.v1.worker.gpu_input_batch import InputBatch
33
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner
32
33
FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
34
@@ -226,9 +224,9 @@ class FlashInferMetadata:
224
225
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
226
229
- def __init__(self, runner: GPUModelRunner, kv_cache_spec: AttentionSpec,
230
- block_table: BlockTable):
231
- self.runner = runner
227
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
228
+ device: torch.device):
229
+ self.device = device
230
self._workspace_buffer = None
231
self._prefill_wrapper = None # Wrapper for prefill/append
232
self._decode_wrapper = None # Wrapper for decode
@@ -237,75 +235,22 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
235
# Global hyperparameters shared by all attention layers
236
self.global_hyperparameters: Optional[PerLayerParameters] = None
237
240
- self.vllm_config = runner.vllm_config
238
+ self.vllm_config = vllm_config
239
+ self.cache_config = vllm_config.cache_config
240
self.kv_cache_spec = kv_cache_spec
242
- self.block_table = block_table
241
242
def reorder_batch(self, input_batch: InputBatch,
243
scheduler_output: SchedulerOutput) -> bool:
246
- # We now want to reorder the batch so that the "decode" requests are and
247
- # the front and the "prefill" requests are at the using the least amount
248
- # swaps possible. (NOTE for now we loosely use "decode" to mean requests
249
- # where attention is likely memory-bound and "prefill" to mean requests
250
- # where attention is likely compute-bound, TODO(lucas): figure out a
251
- # better naming here)
252
- decodes = []
253
- prefills = []
254
- num_decode_tokens = 0
255
- num_prefill_tokens = 0
256
-
257
- for i, req_id in enumerate(input_batch.req_ids):
258
- num_tokens = scheduler_output.num_scheduled_tokens[req_id]
259
- # for now treat 1 scheduled token as "decode" even if its not,
260
- # we should update this to something like < 8 in the future but
261
- # currently the decode run only supports num_tokens = 1
262
- if num_tokens == 1:
263
- decodes.append(i)
264
- num_decode_tokens += num_tokens
265
- else:
266
- prefills.append(i)
267
- num_prefill_tokens += num_tokens
268
-
269
- # We hope that this is fairly minimal since decodes
270
- # should be around for a number of iterations so hopefully they are
271
- # relatively stationary (and new request are generally appended to the
272
- # persistent batch so already should be at the back)
273
- # To achieve this we loop over the decodes in descending order and
274
- # the prefills in ascending order. We swap decodes from the "back"
275
- # i.e. past where the last decode should be in the reodorered with
276
- # prefills from the front of the batch.
277
- # `decodes` and `prefills` are already in ascending order just based on
278
- # the above loop
279
- num_decodes = len(decodes)
280
- num_prefills = len(prefills)
281
- modified_batch = False
282
-
283
- for i in range(1, min(num_decodes, num_prefills) + 1):
284
- # If the decode is at the "back" of the batch, i, we can swap it
285
- # with the prefill closest to the front of the batch
286
- decode_idx = decodes[num_decodes - i]
287
- if decode_idx < num_decodes:
288
- break
289
-
290
- input_batch.swap_states(prefills[i - 1], decode_idx)
291
- modified_batch = True
292
-
293
- # Save for next `build` call
294
- # TODO(lucas): this is a bit of a hack, we should probably have a
295
- # better way of doing this
296
- self._num_decodes = num_decodes
297
- self._num_prefills = num_prefills
298
- self._num_decode_tokens = num_decode_tokens
299
- self._num_prefill_tokens = num_prefill_tokens
300
-
301
- return modified_batch
244
+ return reorder_batch_to_split_decodes_and_prefills(input_batch,
245
+ scheduler_output,
246
+ decode_threshold=1)
247
248
def _get_workspace_buffer(self):
249
if self._workspace_buffer is None:
250
self._workspace_buffer = torch.empty(
251
FLASHINFER_WORKSPACE_BUFFER_SIZE,
252
dtype=torch.uint8,
308
- device=self.runner.device)
253
+ device=self.device)
254
return self._workspace_buffer
255
256
def _get_prefill_wrapper(self):
@@ -316,10 +261,11 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
261
262
def _get_decode_wrapper(self):
263
if self._decode_wrapper is None:
319
- num_qo_heads = (self.runner.model_config.get_num_attention_heads(
320
- self.runner.parallel_config))
321
- num_kv_heads = self.runner.model_config.get_num_kv_heads(
322
- self.runner.parallel_config)
264
+ num_qo_heads = (
265
+ self.vllm_config.model_config.get_num_attention_heads(
266
+ self.vllm_config.parallel_config))
267
+ num_kv_heads = self.vllm_config.model_config.get_num_kv_heads(
268
+ self.vllm_config.parallel_config)
269
use_tensor_cores = envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or (
270
num_qo_heads // num_kv_heads > 4)
271
self._decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
@@ -334,7 +280,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
280
2, self._get_workspace_buffer(), get_kv_cache_layout())
281
return self._cascade_wrapper
282
337
- def _plan(self, attn_metadata: FlashInferMetadata):
283
+ def _plan(self, num_prefills: int, num_decodes: int,
284
+ attn_metadata: FlashInferMetadata):
285
if self.global_hyperparameters is None:
286
self.global_hyperparameters = infer_global_hyperparameters(
287
get_per_layer_parameters(self.vllm_config, FlashInferImpl))
@@ -369,16 +316,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
316
# Regular attention (common case).
317
# Decodes are at the front and prefills are at the back,
318
# according to reorder_batch()
372
- if self._num_prefills > 0:
319
+ if num_prefills > 0:
320
# Decodes are first so prefills start after the last decode
374
- prefill_start = self._num_decodes
321
+ prefill_start = num_decodes
322
attn_metadata.prefill_wrapper = self._get_prefill_wrapper()
323
assert attn_metadata.qo_indptr[prefill_start:].shape[
377
- 0] == self._num_prefills + 1
324
+ 0] == num_prefills + 1
325
assert attn_metadata.paged_kv_indptr[prefill_start:].shape[
379
- 0] == self._num_prefills + 1
326
+ 0] == num_prefills + 1
327
assert attn_metadata.paged_kv_last_page_len[
381
- prefill_start:].shape[0] == self._num_prefills
328
+ prefill_start:].shape[0] == num_prefills
329
# Since prefill_wrapper.run() will be called with
330
# query[num_decode_tokens:] we need to adjust the qo_indptr
331
# to be relative to the start of the prefill queries.
@@ -402,17 +349,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
349
kv_data_type=attn_metadata.kv_data_type,
350
)
351
405
- if self._num_decodes > 0:
352
+ if num_decodes > 0:
353
attn_metadata.decode_wrapper = self._get_decode_wrapper()
354
if not FlashInferBackend.use_trtllm_decode_attention(
408
- self._num_decodes, attn_metadata.max_seq_len,
355
+ num_decodes, attn_metadata.max_seq_len,
356
attn_metadata.kv_data_type, attn_metadata.num_qo_heads,
357
attn_metadata.num_kv_heads, attn_metadata.head_dim):
358
attn_metadata.decode_wrapper.plan(
412
- attn_metadata.paged_kv_indptr[:self._num_decodes + 1],
359
+ attn_metadata.paged_kv_indptr[:num_decodes + 1],
360
attn_metadata.paged_kv_indices,
414
- attn_metadata.paged_kv_last_page_len[:self.
415
- _num_decodes],
361
+ attn_metadata.paged_kv_last_page_len[:num_decodes],
362
attn_metadata.num_qo_heads,
363
attn_metadata.num_kv_heads,
364
attn_metadata.head_dim,
@@ -427,22 +373,20 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
373
kv_data_type=attn_metadata.kv_data_type,
374
)
375
430
- def build(self, common_prefix_len: int,
431
- common_attn_metadata: CommonAttentionMetadata):
432
- num_reqs = common_attn_metadata.num_reqs
376
+ def build(self,
377
+ common_prefix_len: int,
378
+ common_attn_metadata: CommonAttentionMetadata,
379
+ fast_build: bool = False) -> FlashInferMetadata:
380
num_actual_tokens = common_attn_metadata.num_actual_tokens
381
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens =\
382
+ split_decodes_and_prefills(common_attn_metadata)
383
435
- assert self._num_decodes + self._num_prefills == num_reqs
436
- assert (self._num_decode_tokens +
437
- self._num_prefill_tokens == num_actual_tokens)
384
page_size = self.kv_cache_spec.block_size
439
- device = self.runner.device
385
+ device = self.device
386
qo_indptr = common_attn_metadata.query_start_loc
441
- max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
387
+ max_seq_len = common_attn_metadata.seq_lens_cpu.max()
388
seq_lens = common_attn_metadata.seq_lens
443
- block_table_tensor = self.block_table.get_device_tensor()[:num_reqs]
444
- slot_mapping = self.block_table.slot_mapping_cpu[:num_actual_tokens].to(
445
- self.runner.device, non_blocking=True).long()
389
+ block_table_tensor = common_attn_metadata.block_table_tensor
390
391
block_table_bounds = (seq_lens + page_size - 1) // page_size
392
@@ -487,7 +431,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
431
paged_kv_last_page_len = seq_lens % page_size
432
paged_kv_last_page_len = torch.where(paged_kv_last_page_len == 0,
433
page_size, paged_kv_last_page_len)
490
- cache_dtype = self.runner.cache_config.cache_dtype
434
+ cache_dtype = self.cache_config.cache_dtype
435
if cache_dtype.startswith("fp8"):
436
kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
437
cache_dtype)
@@ -499,17 +443,18 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
443
paged_kv_indptr=paged_kv_indptr,
444
paged_kv_indices=paged_kv_indices,
445
paged_kv_last_page_len=paged_kv_last_page_len,
502
- num_qo_heads=self.runner.num_query_heads,
446
+ num_qo_heads=self.vllm_config.model_config.get_num_attention_heads(
447
+ self.vllm_config.parallel_config),
448
num_kv_heads=self.kv_cache_spec.num_kv_heads,
449
head_dim=self.kv_cache_spec.head_size,
450
page_size=page_size,
451
kv_data_type=kv_cache_dtype,
507
- q_data_type=self.runner.dtype,
508
- slot_mapping=slot_mapping,
509
- num_decodes=self._num_decodes,
510
- num_decode_tokens=self._num_decode_tokens,
511
- num_prefills=self._num_prefills,
512
- num_prefill_tokens=self._num_prefill_tokens,
452
+ q_data_type=self.vllm_config.model_config.dtype,
453
+ slot_mapping=common_attn_metadata.slot_mapping,
454
+ num_decodes=num_decodes,
455
+ num_decode_tokens=num_decode_tokens,
456
+ num_prefills=num_prefills,
457
+ num_prefill_tokens=num_prefill_tokens,
458
use_cascade=use_cascade,
459
shared_qo_indptr=shared_qo_indptr,
460
shared_kv_page_indptr=shared_kv_page_indptr,
@@ -521,12 +466,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
466
workspace_buffer=self._workspace_buffer,
467
)
468
524
- self._plan(attn_metadata)
469
+ self._plan(num_prefills, num_decodes, attn_metadata)
470
471
return attn_metadata
472
473
def use_cascade_attention(self, *args, **kwargs) -> bool:
529
- if self.kv_cache_spec.dtype != self.runner.model_config.dtype:
474
+ if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
475
# TODO: The cascade wrapper currently does not support setting
476
# kv cache dtype to something different from query dtype.
477
return False
vllm/v1/attention/backends/flex_attention.py
+26
-33
@@ -3,7 +3,7 @@
3
"""Attention layer with FlashAttention."""
4
from collections import defaultdict
5
from dataclasses import dataclass
6
-from typing import TYPE_CHECKING, Any, Optional
6
+from typing import Any, Optional
7
8
import torch
9
from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
@@ -14,18 +14,15 @@ from torch.nn.attention.flex_attention import (BlockMask, _mask_mod_signature,
14
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
15
AttentionMetadata, AttentionType,
16
is_quantized_kv_cache)
17
+from vllm.config import VllmConfig
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
21
CommonAttentionMetadata)
22
from vllm.v1.kv_cache_interface import AttentionSpec
22
-from vllm.v1.worker.block_table import BlockTable
23
24
logger = init_logger(__name__)
25
26
-if TYPE_CHECKING:
27
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner
28
-
26
create_block_mask_compiled = torch.compile(create_block_mask,
27
fullgraph=True,
28
mode="reduce-overhead")
@@ -261,36 +258,34 @@ class FlexAttentionMetadata:
258
class FlexAttentionMetadataBuilder(
259
AttentionMetadataBuilder[FlexAttentionMetadata]):
260
264
- def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
265
- block_table: BlockTable):
266
- model_config = runner.model_config
267
-
268
- self.runner = runner
269
- self.num_heads_q = model_config.get_num_attention_heads(
270
- runner.parallel_config)
271
- self.num_heads_kv = model_config.get_num_kv_heads(
272
- runner.parallel_config)
273
- self.headdim = model_config.get_head_size()
261
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
262
+ device: torch.device):
263
+ self.model_config = vllm_config.model_config
264
+ self.parallel_config = vllm_config.parallel_config
265
+ self.cache_config = vllm_config.cache_config
266
+
267
+ self.num_heads_q = self.model_config.get_num_attention_heads(
268
+ vllm_config.parallel_config)
269
+ self.num_heads_kv = self.model_config.get_num_kv_heads(
270
+ vllm_config.parallel_config)
271
+ self.headdim = self.model_config.get_head_size()
272
self.block_size = kv_cache_spec.block_size
273
self.kv_cache_spec = kv_cache_spec
276
- self.block_table = block_table
274
+ self.device = device
275
278
- def build(self, common_prefix_len: int,
279
- common_attn_metadata: CommonAttentionMetadata):
276
+ def build(self,
277
+ common_prefix_len: int,
278
+ common_attn_metadata: CommonAttentionMetadata,
279
+ fast_build: bool = False) -> FlexAttentionMetadata:
280
num_reqs = common_attn_metadata.num_reqs
281
num_actual_tokens = common_attn_metadata.num_actual_tokens
282
max_query_len = common_attn_metadata.max_query_len
283
284
- max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
284
+ max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
285
query_start_loc = common_attn_metadata.query_start_loc
286
seq_lens = common_attn_metadata.seq_lens
287
-
288
- block_table = self.block_table
289
- block_table_tensor = block_table.get_device_tensor()[:num_reqs]
290
- block_table.slot_mapping[:num_actual_tokens].copy_(
291
- block_table.slot_mapping_cpu[:num_actual_tokens],
292
- non_blocking=True)
293
- slot_mapping = block_table.slot_mapping[:num_actual_tokens]
287
+ block_table_tensor = common_attn_metadata.block_table_tensor
288
+ slot_mapping = common_attn_metadata.slot_mapping
289
290
use_cascade = common_prefix_len > 0
291
cu_prefix_query_lens = None
@@ -300,17 +295,15 @@ class FlexAttentionMetadataBuilder(
295
raise NotImplementedError("Not yet my friend")
296
297
block_size = self.kv_cache_spec.block_size
303
- max_possible_seq_len = self.runner.model_config.max_model_len
304
- total_cache_tokens = (self.runner.cache_config.num_gpu_blocks *
305
- block_size)
298
+ max_possible_seq_len = self.model_config.max_model_len
299
+ total_cache_tokens = self.cache_config.num_gpu_blocks * block_size
300
301
inverse_block_table = physical_to_logical_mapping(
308
- block_table_tensor, self.runner.cache_config.num_gpu_blocks)
302
+ block_table_tensor, self.cache_config.num_gpu_blocks)
303
304
# Get the original offset tensor
311
- offset_tensor = torch.tensor(
312
- self.runner.input_batch.num_computed_tokens_cpu[:num_reqs]).to(
313
- self.runner.device, non_blocking=True)
305
+ offset_tensor = common_attn_metadata.num_computed_tokens_cpu.to(
306
+ self.device, non_blocking=True)
307
308
out = FlexAttentionMetadata(
309
num_actual_tokens=num_actual_tokens,
vllm/v1/attention/backends/mamba_attn.py
+38
-92
@@ -7,15 +7,15 @@ from typing import TYPE_CHECKING, Optional
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
-from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
11
- CommonAttentionMetadata)
12
-from vllm.v1.kv_cache_interface import MambaSpec
13
-from vllm.v1.worker.block_table import BlockTable
10
+from vllm.config import VllmConfig
11
+from vllm.v1.attention.backends.utils import (
12
+ AttentionMetadataBuilder, CommonAttentionMetadata,
13
+ reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)
14
+from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
15
16
if TYPE_CHECKING:
17
from vllm.v1.core.sched.output import SchedulerOutput
18
from vllm.v1.worker.gpu_input_batch import InputBatch
18
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner
19
20
21
def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
@@ -87,80 +87,24 @@ class Mamba2AttentionMetadata:
87
class Mamba2AttentionMetadataBuilder(
88
AttentionMetadataBuilder[Mamba2AttentionMetadata]):
89
90
- def __init__(self, runner: "GPUModelRunner", kv_cache_spec: MambaSpec,
91
- block_table: BlockTable):
92
- self.runner = runner
90
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
91
+ device: torch.device):
92
+ assert isinstance(kv_cache_spec, MambaSpec)
93
self.kv_cache_spec = kv_cache_spec
94
- self.block_table = block_table
95
- self.chunk_size = runner.vllm_config.model_config.get_mamba_chunk_size(
96
- )
94
+ self.chunk_size = vllm_config.model_config.get_mamba_chunk_size()
95
assert self.chunk_size is not None, (
96
"chunk_size needs to be set in the model config for Mamba2 models")
97
98
def reorder_batch(self, input_batch: "InputBatch",
99
scheduler_output: "SchedulerOutput") -> bool:
102
- # NOTE (Chen): Copied from MLACommonMetadataBuilder and
103
- # FlashInferMetadataBuilder. Should be refactored later to avoid code
104
- # duplication of these 3 functions.
105
- # We now want to reorder the batch so that the "decode" requests are and
106
- # the front and the "prefill" requests are at the using the least amount
107
- # swaps possible. (NOTE for now we loosely use "decode" to mean requests
108
- # where attention is likely memory-bound and "prefill" to mean requests
109
- # where attention is likely compute-bound, TODO(lucas): figure out a
110
- # better naming here)
111
- decodes = []
112
- prefills = []
113
- num_decode_tokens = 0
114
- num_prefill_tokens = 0
115
-
116
- for i, req_id in enumerate(input_batch.req_ids):
117
- num_tokens = scheduler_output.num_scheduled_tokens[req_id]
118
- # for now treat 1 scheduled token as "decode" even if its not,
119
- # we should update this to something like < 8 in the future but
120
- # currently the decode run only supports num_tokens = 1
121
- if num_tokens == 1:
122
- decodes.append(i)
123
- num_decode_tokens += num_tokens
124
- else:
125
- prefills.append(i)
126
- num_prefill_tokens += num_tokens
127
-
128
- # We hope that this is fairly minimal since decodes
129
- # should be around for a number of iterations so hopefully they are
130
- # relatively stationary (and new request are generally appended to the
131
- # persistent batch so already should be at the back)
132
- # To achieve this we loop over the decodes in descending order and
133
- # the prefills in ascending order. We swap decodes from the "back"
134
- # i.e. past where the last decode should be in the reodorered with
135
- # prefills from the front of the batch.
136
- # `decodes` and `prefills` are already in ascending order just based on
137
- # the above loop
138
- num_decodes = len(decodes)
139
- num_prefills = len(prefills)
140
- modified_batch = False
141
-
142
- for i in range(1, min(num_decodes, num_prefills) + 1):
143
- # If the decode is at the "back" of the batch, i, we can swap it
144
- # with the prefill closest to the front of the batch
145
- decode_idx = decodes[num_decodes - i]
146
- if decode_idx < num_decodes:
147
- break
148
-
149
- input_batch.swap_states(prefills[i - 1], decode_idx)
150
- modified_batch = True
151
-
152
- # Save for next `build` call
153
- # TODO(lucas): this is a bit of a hack, we should probably have a
154
- # better way of doing this
155
- self._num_decodes = num_decodes
156
- self._num_prefills = num_prefills
157
- self._num_decode_tokens = num_decode_tokens
158
- self._num_prefill_tokens = num_prefill_tokens
159
-
160
- return modified_batch
161
-
162
- def build(self, common_prefix_len: int,
163
- common_attn_metadata: CommonAttentionMetadata):
100
+ return reorder_batch_to_split_decodes_and_prefills(input_batch,
101
+ scheduler_output,
102
+ decode_threshold=1)
103
+
104
+ def build(self,
105
+ common_prefix_len: int,
106
+ common_attn_metadata: CommonAttentionMetadata,
107
+ fast_build: bool = False) -> Mamba2AttentionMetadata:
108
num_reqs = common_attn_metadata.num_reqs
109
query_start_loc = common_attn_metadata.query_start_loc
110
seq_lens = common_attn_metadata.seq_lens
@@ -172,29 +116,31 @@ class Mamba2AttentionMetadataBuilder(
116
has_initial_states = None
117
prep_initial_states = False
118
175
- state_indices_tensor = self.block_table.block_table[:num_reqs, 0]
119
+ state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
120
+
121
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
122
+ split_decodes_and_prefills(common_attn_metadata,
123
+ decode_threshold=1))
124
125
# Compute seq_idx, chunk_indices and chunk_offsets for prefill only
178
- if self._num_prefills > 0:
126
+ if num_prefills > 0:
127
#[batch,]
128
has_initial_states_cpu = (
181
- self.runner.input_batch.
182
- num_computed_tokens_cpu_tensor[num_reqs -
183
- self._num_prefills:num_reqs]
184
- > 0)
129
+ common_attn_metadata.
130
+ num_computed_tokens_cpu[num_reqs - num_prefills:num_reqs] > 0)
131
prep_initial_states = torch.any(has_initial_states_cpu).item()
132
has_initial_states = has_initial_states_cpu.to(
133
query_start_loc.device)
134
135
query_start_loc_p = common_attn_metadata.query_start_loc[
190
- -self._num_prefills - 1:] - self._num_decode_tokens
191
-
192
- seq_idx = torch.repeat_interleave(
193
- torch.arange(self._num_prefills,
194
- dtype=torch.int32,
195
- device=query_start_loc_p.device),
196
- query_start_loc_p.diff(),
197
- output_size=self._num_prefill_tokens)
136
+ -num_prefills - 1:] - num_decode_tokens
137
+
138
+ seq_idx = torch.repeat_interleave(torch.arange(
139
+ num_prefills,
140
+ dtype=torch.int32,
141
+ device=query_start_loc_p.device),
142
+ query_start_loc_p.diff(),
143
+ output_size=num_prefill_tokens)
144
seq_idx.unsqueeze_(0)
145
146
# We compute metadata for chunked prefill once at the top level
@@ -204,13 +150,13 @@ class Mamba2AttentionMetadataBuilder(
150
chunk_indices, chunk_offsets = (
151
_query_start_loc_to_chunk_indices_offsets(
152
query_start_loc_p, self.chunk_size,
207
- self._num_prefill_tokens))
153
+ num_prefill_tokens))
154
155
attn_metadata = Mamba2AttentionMetadata(
210
- num_prefills=self._num_prefills,
211
- num_prefill_tokens=self._num_prefill_tokens,
212
- num_decodes=self._num_decodes,
213
- num_decode_tokens=self._num_decode_tokens,
156
+ num_prefills=num_prefills,
157
+ num_prefill_tokens=num_prefill_tokens,
158
+ num_decodes=num_decodes,
159
+ num_decode_tokens=num_decode_tokens,
160
query_start_loc=query_start_loc,
161
seq_lens=seq_lens,
162
has_initial_states=has_initial_states,
vllm/v1/attention/backends/mla/common.py
+64
-119
@@ -202,18 +202,18 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionLayer,
202
from vllm.attention.backends.utils import get_mla_dims
203
from vllm.attention.ops.merge_attn_states import merge_attn_states
204
from vllm.attention.utils.fa_utils import get_flash_attn_version
205
+from vllm.config import VllmConfig
206
from vllm.logger import init_logger
207
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
208
LinearBase,
209
UnquantizedLinearMethod)
210
from vllm.platforms import current_platform
211
from vllm.utils import cdiv, round_down
211
-from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
212
- CommonAttentionMetadata,
213
- get_per_layer_parameters,
214
- infer_global_hyperparameters)
212
+from vllm.v1.attention.backends.utils import (
213
+ AttentionMetadataBuilder, CommonAttentionMetadata,
214
+ get_per_layer_parameters, infer_global_hyperparameters,
215
+ reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)
216
from vllm.v1.kv_cache_interface import AttentionSpec
216
-from vllm.v1.worker.block_table import BlockTable
217
218
try:
219
from vllm.vllm_flash_attn import flash_attn_varlen_func
@@ -235,7 +235,6 @@ except ImportError:
235
if TYPE_CHECKING:
236
from vllm.v1.core.sched.output import SchedulerOutput
237
from vllm.v1.worker.gpu_input_batch import InputBatch
238
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner
238
239
logger = init_logger(__name__)
240
@@ -406,22 +405,23 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
405
"""
406
407
def __init__(self,
409
- runner: "GPUModelRunner",
408
kv_cache_spec: AttentionSpec,
411
- block_table: BlockTable,
409
+ vllm_config: VllmConfig,
410
+ device: torch.device,
411
metadata_cls: Optional[type[M]] = None):
412
self.metadata_cls = metadata_cls \
413
if metadata_cls is not None else MLACommonMetadata
415
- self.runner = runner
416
- scheduler_config = runner.scheduler_config
417
- model_config = runner.model_config
418
- cache_config = runner.cache_config
414
+ self.kv_cache_spec = kv_cache_spec
415
+ self.device = device
416
+ scheduler_config = vllm_config.scheduler_config
417
+ self.model_config = vllm_config.model_config
418
+ cache_config = vllm_config.cache_config
419
+ parallel_config = vllm_config.parallel_config
420
self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
420
- self.num_heads = model_config.get_num_attention_heads(
421
- runner.parallel_config)
422
- self.mla_dims = get_mla_dims(model_config)
421
+ self.num_heads = self.model_config.get_num_attention_heads(
422
+ parallel_config)
423
+ self.mla_dims = get_mla_dims(self.model_config)
424
self.aot_schedule = current_platform.is_cuda()
424
- self.kv_cache_spec = kv_cache_spec
425
426
# Dont try to access the runner on AMD
427
if self.aot_schedule:
@@ -432,7 +432,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
432
# Max sure there is enough for 8 full length request or at least
433
# 4 pages of cache per request
434
max(
435
- 8 * model_config.max_model_len, 4 *
435
+ 8 * self.model_config.max_model_len, 4 *
436
scheduler_config.max_num_seqs * cache_config.block_size),
437
# For long-context models try not to over-allocate limiting
438
# kv-cache space, limiting it to 64k tokens,
@@ -447,13 +447,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
447
scheduler_config.max_num_seqs * cache_config.block_size
448
self.chunked_prefill_workspace = torch.empty(
449
(self.chunked_prefill_workspace_size,
450
- model_config.get_head_size()),
451
- dtype=model_config.dtype,
452
- device=runner.device,
450
+ self.model_config.get_head_size()),
451
+ dtype=self.model_config.dtype,
452
+ device=device,
453
)
454
455
- self.block_table = block_table
456
-
455
self._use_cudnn_prefill = use_cudnn_prefill()
456
self._use_fi_prefill = use_flashinfer_prefill()
457
self.prefill_metadata_cls = (
@@ -465,7 +463,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
463
self._workspace_buffer = torch.empty(
464
FLASHINFER_WORKSPACE_BUFFER_SIZE,
465
dtype=torch.uint8,
468
- device=runner.device)
466
+ device=device)
467
468
self._fi_prefill_main: Optional[
469
BatchPrefillWithRaggedKVCacheWrapper] = None
@@ -473,13 +471,13 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
471
BatchPrefillWithRaggedKVCacheWrapper] = []
472
473
self._global_hyperparameters = infer_global_hyperparameters(
476
- get_per_layer_parameters(runner.vllm_config, MLACommonImpl))
474
+ get_per_layer_parameters(vllm_config, MLACommonImpl))
475
476
if self._use_cudnn_prefill:
477
self.cudnn_workspace = torch.empty(
478
CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs,
479
dtype=torch.int8,
482
- device=runner.device,
480
+ device=device,
481
)
482
483
def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):
@@ -505,7 +503,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
503
assert num_chunks <= len(self._fi_prefill_chunks)
504
505
# In MLA, the non-latent num_qo_heads == num_kv_heads
508
- num_qo_heads = self.runner.num_query_heads
506
+ num_qo_heads = self.num_heads
507
num_kv_heads = num_qo_heads
508
509
# Sanity: Verify that num_kv_heads == 1 since it is latent space
@@ -531,7 +529,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
529
sm_scale=self._global_hyperparameters.sm_scale,
530
window_left=self._global_hyperparameters.window_left,
531
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
534
- q_data_type=self.runner.dtype,
532
+ q_data_type=self.model_config.dtype,
533
kv_data_type=self.kv_cache_spec.dtype,
534
)
535
@@ -552,7 +550,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
550
window_left=self._global_hyperparameters.window_left,
551
logits_soft_cap=self._global_hyperparameters.
552
logits_soft_cap,
555
- q_data_type=self.runner.dtype,
553
+ q_data_type=self.model_config.dtype,
554
kv_data_type=self.kv_cache_spec.dtype,
555
)
556
@@ -561,63 +559,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
559
560
def reorder_batch(self, input_batch: "InputBatch",
561
scheduler_output: "SchedulerOutput") -> bool:
564
- # We now want to reorder the batch so that the "decode" requests are and
565
- # the front and the "prefill" requests are at the using the least amount
566
- # swaps possible. (NOTE for now we loosely use "decode" to mean requests
567
- # where attention is likely memory-bound and "prefill" to mean requests
568
- # where attention is likely compute-bound, TODO(lucas): figure out a
569
- # better naming here)
570
- decodes = []
571
- prefills = []
572
- num_decode_tokens = 0
573
- num_prefill_tokens = 0
574
-
575
- for i, req_id in enumerate(input_batch.req_ids):
576
- num_tokens = scheduler_output.num_scheduled_tokens[req_id]
577
- # for now treat 1 scheduled token as "decode" even if its not,
578
- # we should update this to something like < 8 in the future but
579
- # currently the TritonMLA._forward_decode only supports
580
- # num_tokens = 1
581
- if num_tokens == 1:
582
- decodes.append(i)
583
- num_decode_tokens += num_tokens
584
- else:
585
- prefills.append(i)
586
- num_prefill_tokens += num_tokens
587
-
588
- # We hope that this is fairly minimal since decodes
589
- # should be around for a number of iterations so hopefully they are
590
- # relatively stationary (and new request are generally appended to the
591
- # persistent batch so already should be at the back)
592
- # To achieve this we loop over the decodes in descending order and
593
- # the prefills in ascending order. We swap decodes from the "back"
594
- # i.e. past where the last decode should be in the reodorered with
595
- # prefills from the front of the batch.
596
- # `decodes` and `prefills` are already in ascending order just based on
597
- # the above loop
598
- num_decodes = len(decodes)
599
- num_prefills = len(prefills)
600
- modified_batch = False
601
-
602
- for i in range(1, min(num_decodes, num_prefills) + 1):
603
- # If the decode is at the "back" of the batch, i, we can swap it
604
- # with the prefill closest to the front of the batch
605
- decode_idx = decodes[num_decodes - i]
606
- if decode_idx < num_decodes:
607
- break
608
-
609
- input_batch.swap_states(prefills[i - 1], decode_idx)
610
- modified_batch = True
611
-
612
- # Save for next `build` call
613
- # TODO(lucas): this is a bit of a hack, we should probably have a
614
- # better way of doing this
615
- self._num_decodes = num_decodes
616
- self._num_prefills = num_prefills
617
- self._num_decode_tokens = num_decode_tokens
618
- self._num_prefill_tokens = num_prefill_tokens
619
-
620
- return modified_batch
562
+ return reorder_batch_to_split_decodes_and_prefills(input_batch,
563
+ scheduler_output,
564
+ decode_threshold=1)
565
566
def _build_decode(self, block_table_tensor: torch.Tensor,
567
seq_lens: torch.Tensor):
@@ -639,49 +583,50 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
583
584
m.max_query_len = 1 # decode-only
585
642
- # Update state usually set in reorder_batch.
643
- self._num_decodes = m.num_reqs
644
- self._num_decode_tokens = m.num_actual_tokens
645
- self._num_prefills = 0
646
- self._num_prefill_tokens = 0
586
return self.build(0, m)
587
649
- def build(self, common_prefix_len: int,
650
- common_attn_metadata: CommonAttentionMetadata) -> M:
588
+ def build(self,
589
+ common_prefix_len: int,
590
+ common_attn_metadata: CommonAttentionMetadata,
591
+ fast_build: bool = False) -> M:
592
num_reqs = common_attn_metadata.num_reqs
652
- num_actual_tokens = common_attn_metadata.num_actual_tokens
593
+ num_tokens = common_attn_metadata.num_actual_tokens
594
max_query_len = common_attn_metadata.max_query_len
595
655
- assert self._num_decodes + self._num_prefills == num_reqs
656
-
596
# Note(simon): be careful about the CPU <> GPU memory movement in this
597
# function. We should avoid GPU -> CPU sync as much as possible because
598
# it blocks on all previous kernels.
660
- device = self.runner.device
661
- block_table = self.block_table
662
- block_table_tensor = block_table.get_device_tensor()[:num_reqs]
663
- block_table.slot_mapping[:num_actual_tokens].copy_(
664
- block_table.slot_mapping_cpu[:num_actual_tokens],
665
- non_blocking=True)
666
- block_table.slot_mapping[num_actual_tokens:].fill_(-1)
667
- slot_mapping = block_table.slot_mapping[:num_actual_tokens]
599
+ device = self.device
600
+ block_table_tensor = common_attn_metadata.block_table_tensor
601
+ slot_mapping = common_attn_metadata.slot_mapping
602
603
query_start_loc = common_attn_metadata.query_start_loc
604
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
605
seq_lens = common_attn_metadata.seq_lens
606
607
+ query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
608
+
609
+ num_computed_tokens_cpu = (common_attn_metadata.seq_lens_cpu -
610
+ query_seq_lens_cpu)
611
+
612
+ num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = \
613
+ split_decodes_and_prefills(common_attn_metadata)
614
+
615
+ assert num_decodes + num_prefills == num_reqs
616
+ assert num_decode_tokens + num_prefill_tokens == num_tokens
617
+
618
prefill_metadata = None
673
- if self._num_prefills > 0:
674
- reqs_start = self._num_decodes # prefill_start
619
+ if num_prefills > 0:
620
+ reqs_start = num_decodes # prefill_start
621
676
- context_lens_cpu = self.runner.input_batch.\
677
- num_computed_tokens_cpu_tensor[reqs_start:num_reqs]
622
+ context_lens_cpu = num_computed_tokens_cpu[reqs_start:num_reqs]
623
max_context_len_cpu = context_lens_cpu.max().item()
624
num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
625
prefill_query_start_loc = query_start_loc[
626
reqs_start:] - query_start_loc[reqs_start]
627
628
chunked_context_metadata = None
684
- if self.chunked_prefill_enabled and self._num_prefills > 0 \
629
+ if self.chunked_prefill_enabled and num_prefills > 0 \
630
and max_context_len_cpu > 0:
631
# NOTE: it is recommend you read the `Chunked Prefill` section
632
# in the comment at the top of the file before trying to
@@ -712,14 +657,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
657
# of `to_list`.
658
chunk_starts = \
659
torch.arange(num_chunks, dtype=torch.int32) \
715
- .unsqueeze(1).expand(-1, self._num_prefills) \
660
+ .unsqueeze(1).expand(-1, num_prefills) \
661
* max_context_chunk
662
chunk_ends = torch.min(context_lens_cpu.unsqueeze(0),
663
chunk_starts + max_context_chunk)
664
chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
665
666
cu_seq_lens_cpu = torch.zeros(num_chunks,
722
- self._num_prefills + 1,
667
+ num_prefills + 1,
668
dtype=torch.int32,
669
pin_memory=True)
670
torch.cumsum(chunk_seq_lens,
@@ -762,28 +707,28 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
707
prefill_metadata.cudnn_workspace = self.cudnn_workspace
708
709
decode_metadata = None
765
- if self._num_decodes > 0:
710
+ if num_decodes > 0:
711
decode_metadata = self._build_decode(
767
- block_table_tensor=block_table_tensor[:self._num_decodes, ...],
768
- seq_lens=seq_lens[:self._num_decodes],
712
+ block_table_tensor=block_table_tensor[:num_decodes, ...],
713
+ seq_lens=seq_lens[:num_decodes],
714
)
715
716
attn_metadata = self.metadata_cls(
717
num_reqs=common_attn_metadata.num_reqs,
718
max_query_len=common_attn_metadata.max_query_len,
774
- num_actual_tokens=num_actual_tokens,
719
+ num_actual_tokens=num_tokens,
720
query_start_loc=query_start_loc,
721
slot_mapping=slot_mapping,
777
- head_dim=self.runner.model_config.get_head_size(),
722
+ head_dim=self.model_config.get_head_size(),
723
# MLACommonMetadata Chunk prefill specific
779
- num_decodes=self._num_decodes,
780
- num_decode_tokens=self._num_decode_tokens,
781
- num_prefills=self._num_prefills,
724
+ num_decodes=num_decodes,
725
+ num_decode_tokens=num_decode_tokens,
726
+ num_prefills=num_prefills,
727
prefill=prefill_metadata,
728
decode=decode_metadata,
729
)
730
786
- if self._use_fi_prefill and self._num_prefills > 0:
731
+ if self._use_fi_prefill and num_prefills > 0:
732
assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata)
733
self._build_fi_prefill_wrappers(attn_metadata.prefill)
734
vllm/v1/attention/backends/mla/flashmla.py
+8
-7
@@ -11,6 +11,7 @@ from vllm.attention.backends.abstract import (AttentionType,
11
from vllm.attention.ops.flashmla import (flash_mla_with_kvcache,
12
get_mla_metadata,
13
is_flashmla_supported)
14
+from vllm.config import VllmConfig
15
from vllm.logger import init_logger
16
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
17
MLACommonDecodeMetadata,
@@ -18,7 +19,6 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
19
MLACommonMetadata,
20
MLACommonMetadataBuilder)
21
from vllm.v1.kv_cache_interface import AttentionSpec
21
-from vllm.v1.worker.block_table import BlockTable
22
23
logger = init_logger(__name__)
24
@@ -56,12 +56,13 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
56
class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
57
full_cudagraph_supported: ClassVar[bool] = True # Decode-only
58
59
- def __init__(self, runner, kv_cache_spec: AttentionSpec,
60
- block_table: BlockTable):
61
- super().__init__(runner, kv_cache_spec, block_table, FlashMLAMetadata)
59
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
60
+ device: torch.device):
61
+ super().__init__(kv_cache_spec, vllm_config, device, FlashMLAMetadata)
62
63
- self.num_q_heads = self.runner.model_config.get_num_attention_heads(
64
- self.runner.parallel_config)
63
+ self.compilation_config = vllm_config.compilation_config
64
+ self.num_q_heads = vllm_config.model_config.get_num_attention_heads(
65
+ vllm_config.parallel_config)
66
67
self.cg_buf_tile_scheduler_metadata = None
68
self.cg_buf_num_splits = None
@@ -75,7 +76,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
76
1, # MQA for the decode path
77
)
78
78
- if self.runner.full_cuda_graph:
79
+ if self.compilation_config.full_cuda_graph:
80
# First time around (CUDAGraph capture), allocate the static buffer
81
if self.cg_buf_tile_scheduler_metadata is None:
82
self.cg_buf_tile_scheduler_metadata = tile_scheduler_metadata
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+19
-16
@@ -8,6 +8,8 @@ import torch
8
9
import vllm.envs as envs
10
from vllm.attention.ops.rocm_aiter_mla import aiter_mla_decode_fwd
11
+from vllm.config import VllmConfig
12
+from vllm.utils import cdiv
13
# yapf conflicts with isort for this docstring
14
# yapf: disable
15
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
@@ -16,7 +18,6 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
18
MLACommonMetadata,
19
MLACommonMetadataBuilder)
20
from vllm.v1.kv_cache_interface import AttentionSpec
19
-from vllm.v1.worker.block_table import BlockTable
21
22
# yapf: enable
23
@@ -65,24 +66,26 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
66
class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
67
full_cudagraph_supported: ClassVar[bool] = True # decode only
68
68
- def __init__(self, runner, kv_cache_spec: AttentionSpec,
69
- block_table: BlockTable):
70
- super().__init__(runner, kv_cache_spec, block_table, AiterMLAMetadata)
69
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
70
+ device: torch.device):
71
+ super().__init__(kv_cache_spec, vllm_config, device, AiterMLAMetadata)
72
assert self.kv_cache_spec.block_size == 1, "AITER MLA" \
73
"only supports block size 1."
74
75
+ self.compilation_config = vllm_config.compilation_config
76
+ max_num_pages_per_req = cdiv(vllm_config.model_config.max_model_len,
77
+ self.kv_cache_spec.block_size)
78
+ max_num_reqs = vllm_config.scheduler_config.max_num_seqs
79
+ max_num_pages = max_num_reqs * max_num_pages_per_req
80
+
81
# Preparing persistent buffers
75
- if self.runner.full_cuda_graph:
76
- device = self.runner.device
77
- max_num_reqs = self.runner.max_num_reqs
82
+ if vllm_config.compilation_config.full_cuda_graph:
83
self.paged_kv_indptr = torch.zeros(max_num_reqs + 1,
84
dtype=torch.int32,
85
device=device)
81
- self.paged_kv_indices = torch.zeros(
82
- block_table.get_device_tensor().numel(
83
- ), # max num pages possible
84
- dtype=torch.int32,
85
- device=device)
86
+ self.paged_kv_indices = torch.zeros(max_num_pages,
87
+ dtype=torch.int32,
88
+ device=device)
89
self.paged_kv_last_page_len = torch.zeros(max_num_reqs,
90
dtype=torch.int32,
91
device=device)
@@ -96,7 +99,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
99
seq_lens: torch.Tensor) -> AiterMLADecodeMetadata:
100
page_size = self.kv_cache_spec.block_size
101
block_table_bounds = (seq_lens + page_size - 1) // page_size
99
- device = self.runner.device
102
+ device = self.device
103
+ num_reqs = seq_lens.size(0)
104
105
mask = (torch.arange(block_table_tensor.size(1),
106
dtype=block_table_tensor.dtype,
@@ -113,8 +117,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
117
block_table_bounds.cumsum(dim=0, dtype=torch.int32)
118
])
119
116
- if self.runner.full_cuda_graph:
117
- num_reqs = self._num_decodes
120
+ if self.compilation_config.full_cuda_graph:
121
122
num_actual_pages = paged_kv_indices.size(0)
123
@@ -137,7 +140,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
140
141
else:
142
qo_indptr = torch.arange(0,
140
- self._num_decodes + 1,
143
+ num_reqs + 1,
144
step=1,
145
dtype=torch.int32,
146
device=device)
vllm/v1/attention/backends/rocm_aiter_fa.py
+39
-50
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with AiterFlashAttention."""
4
from dataclasses import dataclass
5
-from typing import TYPE_CHECKING, Any, Optional
5
+from typing import Any, Optional
6
7
import torch
8
@@ -10,18 +10,13 @@ from vllm import _custom_ops as ops
10
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionMetadata, AttentionType,
12
is_quantized_kv_cache)
13
+from vllm.config import VllmConfig
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
from vllm.v1.attention.backends.flash_attn import (
17
make_local_attention_virtual_batches)
18
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
19
from vllm.v1.kv_cache_interface import AttentionSpec
19
-from vllm.v1.worker.block_table import BlockTable
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
20
21
if current_platform.is_rocm():
22
import aiter
@@ -172,54 +167,49 @@ logger = init_logger(__name__)
167
168
class AiterFlashAttentionMetadataBuilder:
169
175
- def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
176
- block_table: BlockTable):
177
- model_config = runner.model_config
178
-
179
- self.runner = runner
180
- self.num_heads_q = model_config.get_num_attention_heads(
181
- runner.parallel_config)
182
- self.num_heads_kv = model_config.get_num_kv_heads(
183
- runner.parallel_config)
184
- self.headdim = model_config.get_head_size()
170
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
171
+ device: torch.device):
172
+ self.vllm_config = vllm_config
173
+ self.model_config = vllm_config.model_config
174
+ self.parallel_config = vllm_config.parallel_config
175
+ self.cache_config = vllm_config.cache_config
176
+ self.device = device
177
+
178
+ self.num_heads_q = self.model_config.get_num_attention_heads(
179
+ self.parallel_config)
180
+ self.num_heads_kv = self.model_config.get_num_kv_heads(
181
+ self.parallel_config)
182
+ self.headdim = self.model_config.get_head_size()
183
self.block_size = kv_cache_spec.block_size
184
self.kv_cache_spec = kv_cache_spec
187
- self.block_table = block_table
185
186
# Sliding window size to be used with the AOT scheduler will be
187
# populated on first build() call.
188
self.aot_sliding_window: Optional[tuple[int, int]] = None
189
193
- def reorder_batch(self, input_batch: "InputBatch",
194
- scheduler_output: "SchedulerOutput") -> bool:
190
+ def reorder_batch(self, input_batch, scheduler_output) -> bool:
191
return False
192
197
- def build(self, common_prefix_len: int,
198
- common_attn_metadata: CommonAttentionMetadata):
193
+ def build(self,
194
+ common_prefix_len: int,
195
+ common_attn_metadata: CommonAttentionMetadata,
196
+ fast_build: bool = False) -> 'AiterFlashAttentionMetadata':
197
200
- num_reqs = common_attn_metadata.num_reqs
198
num_actual_tokens = common_attn_metadata.num_actual_tokens
199
max_query_len = common_attn_metadata.max_query_len
200
204
- max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
205
- total_tokens = int(self.runner.seq_lens_np[:num_reqs].sum())
201
+ max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
202
+ total_tokens = int(common_attn_metadata.seq_lens_cpu.sum())
203
query_start_loc = common_attn_metadata.query_start_loc
204
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
205
seq_lens = common_attn_metadata.seq_lens
208
- block_table = self.block_table
209
- block_table_tensor = block_table.get_device_tensor()[:num_reqs]
210
-
211
- block_table.slot_mapping[:num_actual_tokens].copy_(
212
- block_table.slot_mapping_cpu[:num_actual_tokens],
213
- non_blocking=True)
214
- # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
215
- # mode.
216
- block_table.slot_mapping[num_actual_tokens:].fill_(-1)
217
-
218
- slot_mapping = block_table.slot_mapping[:num_actual_tokens]
206
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu
207
+ block_table_tensor = common_attn_metadata.block_table_tensor
208
+ slot_mapping = common_attn_metadata.slot_mapping
209
210
cu_seq_lens = torch.zeros(seq_lens.shape[0] + 1,
211
dtype=torch.int32,
222
- device="cuda")
212
+ device=self.device)
213
torch.cumsum(seq_lens,
214
dim=0,
215
dtype=cu_seq_lens.dtype,
@@ -231,21 +221,21 @@ class AiterFlashAttentionMetadataBuilder:
221
222
# for local attention
223
local_attn_metadata = None
234
- if self.runner.attention_chunk_size is not None:
224
+ if self.model_config.attention_chunk_size is not None:
225
seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
226
virt_block_table_tensor = make_local_attention_virtual_batches(
237
- self.runner.attention_chunk_size,
238
- self.runner.query_start_loc_np[:num_reqs + 1],
239
- self.runner.seq_lens_np[:num_reqs],
227
+ self.model_config.attention_chunk_size,
228
+ query_start_loc_cpu.numpy(),
229
+ seq_lens_cpu.numpy(),
230
block_table_tensor,
231
self.block_size,
232
)
233
local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
244
- self.runner.device, non_blocking=True)
234
+ self.device, non_blocking=True)
235
local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
246
- self.runner.device, non_blocking=True)
247
- local_max_query_len = int(seqlens_q_local_np.max())
248
- local_max_seq_len = int(virt_k_seqlens_np.max())
236
+ self.device, non_blocking=True)
237
+ local_max_query_len = seqlens_q_local_np.max().item()
238
+ local_max_seq_len = virt_k_seqlens_np.max().item()
239
local_scheduler_metadata = schedule(
240
batch_size=local_query_start_loc.shape[0] - 1,
241
cu_query_lens=local_query_start_loc,
@@ -256,12 +246,11 @@ class AiterFlashAttentionMetadataBuilder:
246
247
local_cu_seq_lens = torch.zeros(virt_k_seqlens_np.shape[0] + 1,
248
dtype=torch.int32,
259
- device=self.runner.device)
249
+ device=self.device)
250
local_cu_seq_lens[1:] = torch.cumsum(
261
- torch.from_numpy(virt_k_seqlens_np).to(
262
- device=self.runner.device,
263
- dtype=torch.int32,
264
- non_blocking=True),
251
+ torch.from_numpy(virt_k_seqlens_np).to(device=self.device,
252
+ dtype=torch.int32,
253
+ non_blocking=True),
254
dim=0)
255
256
vllm/v1/attention/backends/triton_attn.py
+34
-39
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer with PagedAttention and Triton prefix prefill."""
4
from dataclasses import dataclass
5
-from typing import TYPE_CHECKING, Any, ClassVar, Optional
5
+from typing import Any, ClassVar, Optional
6
7
import torch
8
@@ -14,6 +14,7 @@ from vllm.attention.ops.chunked_prefill_paged_decode import (
14
chunked_prefill_paged_decode)
15
from vllm.attention.ops.paged_attn import PagedAttention
16
from vllm.attention.ops.triton_unified_attention import unified_attention
17
+from vllm.config import VllmConfig
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
@@ -21,10 +22,6 @@ from vllm.v1.attention.backends.utils import (
22
AttentionMetadataBuilder, CommonAttentionMetadata,
23
make_local_attention_virtual_batches)
24
from vllm.v1.kv_cache_interface import AttentionSpec
24
-from vllm.v1.worker.block_table import BlockTable
25
-
26
-if TYPE_CHECKING:
27
- from vllm.v1.worker.gpu_model_runner import GPUModelRunner
25
26
logger = init_logger(__name__)
27
@@ -75,12 +72,21 @@ class TritonAttentionMetadataBuilder(
72
AttentionMetadataBuilder[TritonAttentionMetadata]):
73
full_cudagraph_supported: ClassVar[bool] = True
74
78
- def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
79
- block_table: BlockTable):
80
- self.runner = runner
75
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
76
+ device: torch.device):
77
+ self.device = device
78
self.block_size = kv_cache_spec.block_size
79
self.kv_cache_spec = kv_cache_spec
83
- self.block_table = block_table
80
+
81
+ model_config = vllm_config.model_config
82
+ self.num_heads_q = model_config.get_num_attention_heads(
83
+ vllm_config.parallel_config)
84
+ self.num_heads_kv = model_config.get_num_kv_heads(
85
+ vllm_config.parallel_config)
86
+ self.headdim = model_config.get_head_size()
87
+
88
+ self.attention_chunk_size = getattr(vllm_config.scheduler_config,
89
+ 'attention_chunk_size', None)
90
91
def build_for_cudagraph_capture(
92
self, common_attn_metadata: CommonAttentionMetadata
@@ -92,46 +98,36 @@ class TritonAttentionMetadataBuilder(
98
attn_metadata.seq_lens.fill_(1)
99
return attn_metadata
100
95
- def build(
96
- self, common_prefix_len: int,
97
- common_attn_metadata: CommonAttentionMetadata
98
- ) -> TritonAttentionMetadata:
99
- num_reqs = common_attn_metadata.num_reqs
101
+ def build(self,
102
+ common_prefix_len: int,
103
+ common_attn_metadata: CommonAttentionMetadata,
104
+ fast_build: bool = False) -> TritonAttentionMetadata:
105
num_actual_tokens = common_attn_metadata.num_actual_tokens
106
max_query_len = common_attn_metadata.max_query_len
107
103
- max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
108
+ max_seq_len = int(common_attn_metadata.seq_lens_cpu.max())
109
query_start_loc = common_attn_metadata.query_start_loc
110
seq_lens = common_attn_metadata.seq_lens
106
- block_table = self.block_table
107
- block_table_tensor = block_table.get_device_tensor()[:num_reqs]
108
-
109
- block_table.slot_mapping[:num_actual_tokens].copy_(
110
- block_table.slot_mapping_cpu[:num_actual_tokens],
111
- non_blocking=True)
112
- # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
113
- # mode.
114
- block_table.slot_mapping[num_actual_tokens:].fill_(-1)
115
-
116
- slot_mapping = block_table.slot_mapping[:num_actual_tokens]
111
+ block_table_tensor = common_attn_metadata.block_table_tensor
112
+ slot_mapping = common_attn_metadata.slot_mapping
113
114
# for local attention
115
local_attn_metadata = None
120
- if self.runner.attention_chunk_size is not None:
116
+ if self.attention_chunk_size is not None:
117
seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
118
virt_block_table_tensor = make_local_attention_virtual_batches(
123
- self.runner.attention_chunk_size,
124
- self.runner.query_start_loc_np[:num_reqs + 1],
125
- self.runner.seq_lens_np[:num_reqs],
119
+ self.attention_chunk_size,
120
+ common_attn_metadata.query_start_loc_cpu.numpy(),
121
+ common_attn_metadata.seq_lens_cpu.numpy(),
122
block_table_tensor,
123
self.block_size,
124
)
125
local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
130
- self.runner.device, non_blocking=True)
126
+ self.device, non_blocking=True)
127
local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
132
- self.runner.device, non_blocking=True)
133
- local_max_query_len = seqlens_q_local_np.max()
134
- local_max_seq_len = virt_k_seqlens_np.max()
128
+ self.device, non_blocking=True)
129
+ local_max_query_len = seqlens_q_local_np.max().item()
130
+ local_max_seq_len = virt_k_seqlens_np.max().item()
131
132
local_attn_metadata = TritonAttentionMetadata \
133
.LocalAttentionMetadata(
@@ -148,14 +144,13 @@ class TritonAttentionMetadataBuilder(
144
if use_cascade:
145
cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],
146
dtype=torch.int32,
151
- device=self.runner.device)
147
+ device=self.device)
148
prefix_kv_lens = torch.tensor([common_prefix_len],
149
dtype=torch.int32,
154
- device=self.runner.device)
155
- suffix_kv_lens = (self.runner.seq_lens_np[:num_reqs] -
150
+ device=self.device)
151
+ suffix_kv_lens = (common_attn_metadata.seq_lens_cpu -
152
common_prefix_len)
157
- suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(
158
- self.runner.device)
153
+ suffix_kv_lens = suffix_kv_lens.to(self.device)
154
else:
155
cu_prefix_query_lens = None
156
prefix_kv_lens = None
vllm/v1/attention/backends/utils.py
+138
-2
@@ -22,6 +22,7 @@ import vllm.envs as envs
22
from vllm.distributed.kv_transfer.kv_connector.utils import (
23
get_kv_connector_cache_layout)
24
from vllm.logger import init_logger
25
+from vllm.v1.kv_cache_interface import AttentionSpec
26
27
logger = init_logger(__name__)
28
_KV_CACHE_LAYOUT_OVERRIDE = None
@@ -32,14 +33,22 @@ class CommonAttentionMetadata:
33
"""
34
Per-batch attention metadata, shared across layers and backends.
35
AttentionMetadataBuilder instances use it to construct per-layer metadata.
36
+
37
+ For many of the tensors we keep both GPU and CPU versions.
38
"""
39
40
query_start_loc: torch.Tensor
41
+ query_start_loc_cpu: torch.Tensor
42
"""(batch_size + 1,), the start location of each request in query Tensor"""
43
+
44
seq_lens: torch.Tensor
45
+ seq_lens_cpu: torch.Tensor
46
"""(batch_size,), the length of each request including both computed tokens
47
and newly scheduled tokens"""
48
49
+ num_computed_tokens_cpu: torch.Tensor
50
+ """(batch_size,), the number of computed tokens for each request"""
51
+
52
num_reqs: int
53
"""Number of requests"""
54
num_actual_tokens: int
@@ -47,6 +56,14 @@ class CommonAttentionMetadata:
56
max_query_len: int
57
"""Longest query in batch"""
58
59
+ block_table_tensor: torch.Tensor
60
+ slot_mapping: torch.Tensor
61
+
62
+ def __post_init__(self):
63
+ # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
64
+ # mode.
65
+ self.slot_mapping[self.num_actual_tokens:].fill_(-1)
66
+
67
68
M = TypeVar("M")
69
@@ -56,11 +73,25 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
73
full_cudagraph_supported: ClassVar[bool] = False
74
75
@abstractmethod
59
- def build(self, common_prefix_len: int,
60
- common_attn_metadata: CommonAttentionMetadata) -> M:
76
+ def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,
77
+ device: torch.device):
78
+ self.kv_cache_spec = kv_cache_spec
79
+
80
+ @abstractmethod
81
+ def build(self,
82
+ common_prefix_len: int,
83
+ common_attn_metadata: CommonAttentionMetadata,
84
+ fast_build: bool = False) -> M:
85
"""
86
Central method that builds attention metadata.
87
Some builders (MLA) require reorder_batch to be called prior to build.
88
+
89
+ Args:
90
+ common_prefix_len: The length of the common prefix of the batch.
91
+ common_attn_metadata: The common attention metadata.
92
+ fast_build: The meta-data will prioritize speed of building over
93
+ then speed at execution. Can be used for spec-decode where the
94
+ result of a build call may only be used for few layers/iters.
95
"""
96
raise NotImplementedError
97
@@ -351,3 +382,108 @@ def make_local_attention_virtual_batches(
382
383
return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, \
384
block_table_local
385
+
386
+
387
+def split_decodes_and_prefills(
388
+ common_attn_metadata: CommonAttentionMetadata,
389
+ decode_threshold: int = 1,
390
+) -> tuple[int, int, int, int]:
391
+ """
392
+ Assuming a reordered batch, finds the boundary between prefill and decode
393
+ requests.
394
+
395
+ Args:
396
+ common_attn_metadata: CommonAttentionMetadata object containing the
397
+ batch metadata.
398
+ decode_threshold: The maximum query length to be considered a decode.
399
+
400
+ Returns:
401
+ num_decodes: The number of decode requests.
402
+ num_prefills: The number of prefill requests.
403
+ num_decode_tokens: The number of tokens in the decode requests.
404
+ num_prefill_tokens: The number of tokens in the prefill requests.
405
+ """
406
+ max_query_len = common_attn_metadata.max_query_len
407
+ num_reqs = common_attn_metadata.num_reqs
408
+ num_tokens = common_attn_metadata.num_actual_tokens
409
+ query_start_loc = common_attn_metadata.query_start_loc_cpu
410
+
411
+ if max_query_len <= decode_threshold:
412
+ return num_reqs, 0, num_tokens, 0
413
+
414
+ query_lens = query_start_loc[1:] - query_start_loc[:-1]
415
+ is_prefill = query_lens > decode_threshold
416
+ if not torch.any(is_prefill):
417
+ return num_reqs, 0, num_tokens, 0
418
+
419
+ first_prefill = is_prefill.int().argmax(dim=-1).item()
420
+ assert torch.all(query_lens[first_prefill:] > decode_threshold)
421
+ assert torch.all(query_lens[:first_prefill] <= decode_threshold)
422
+ num_decodes = first_prefill
423
+ num_prefills = num_reqs - num_decodes
424
+ num_decode_tokens = query_start_loc[first_prefill].item()
425
+ num_prefill_tokens = num_tokens - num_decode_tokens
426
+ return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens)
427
+
428
+
429
+def reorder_batch_to_split_decodes_and_prefills(
430
+ input_batch: "InputBatch",
431
+ scheduler_output: "SchedulerOutput",
432
+ decode_threshold: int = 1,
433
+) -> bool:
434
+ """
435
+ Reorders the batch to split into prefill and decode requests; places all
436
+ requests with <= decode_threshold tokens at the front of the batch.
437
+
438
+ Returns:
439
+ True if the batch was modified, False otherwise.
440
+ """
441
+ # We now want to reorder the batch so that the "decode" requests are at
442
+ # the front and the "prefill" requests are at the back using the least
443
+ # amount of swaps possible. (NOTE for now we loosely use "decode" to mean
444
+ # requests where attention is likely memory-bound and "prefill" to mean
445
+ # requests where attention is likely compute-bound, TODO(lucas): figure out
446
+ # a better naming here)
447
+ decodes = []
448
+ prefills = []
449
+ num_decode_tokens = 0
450
+ num_prefill_tokens = 0
451
+
452
+ for i, req_id in enumerate(input_batch.req_ids):
453
+ num_tokens = scheduler_output.num_scheduled_tokens[req_id]
454
+ # for now treat 1 scheduled token as "decode" even if its not,
455
+ # we should update this to something like < 8 in the future but
456
+ # currently the TritonMLA._forward_decode only supports
457
+ # num_tokens = 1
458
+ if num_tokens <= decode_threshold:
459
+ decodes.append(i)
460
+ num_decode_tokens += num_tokens
461
+ else:
462
+ prefills.append(i)
463
+ num_prefill_tokens += num_tokens
464
+
465
+ # We hope that this is fairly minimal since decodes
466
+ # should be around for a number of iterations so hopefully they are
467
+ # relatively stationary (and new request are generally appended to the
468
+ # persistent batch so already should be at the back)
469
+ # To achieve this we loop over the decodes in descending order and
470
+ # the prefills in ascending order. We swap decodes from the "back"
471
+ # i.e. past where the last decode should be in the reodorered with
472
+ # prefills from the front of the batch.
473
+ # `decodes` and `prefills` are already in ascending order just based on
474
+ # the above loop
475
+ num_decodes = len(decodes)
476
+ num_prefills = len(prefills)
477
+ modified_batch = False
478
+
479
+ for i in range(1, min(num_decodes, num_prefills) + 1):
480
+ # If the decode is at the "back" of the batch, i, we can swap it
481
+ # with the prefill closest to the front of the batch
482
+ decode_idx = decodes[num_decodes - i]
483
+ if decode_idx < num_decodes:
484
+ break
485
+
486
+ input_batch.swap_states(prefills[i - 1], decode_idx)
487
+ modified_batch = True
488
+
489
+ return modified_batch
vllm/v1/spec_decode/eagle.py
+108
-90
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
import torch
5
import torch.nn as nn
6
@@ -12,11 +13,11 @@ from vllm.logger import init_logger
13
from vllm.model_executor.model_loader import get_model
14
from vllm.model_executor.models import supports_multimodal
15
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
16
+from vllm.utils import is_pin_memory_available
17
from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
18
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
19
from vllm.v1.kv_cache_interface import KVCacheConfig
20
from vllm.v1.sample.metadata import SamplingMetadata
19
-from vllm.v1.spec_decode.utils import prepare_eagle_input_kernel
21
22
logger = init_logger(__name__)
23
@@ -37,7 +38,6 @@ class EagleProposer:
38
self.method = self.speculative_config.method
39
40
self.runner = runner
40
-
41
self.dtype = vllm_config.model_config.dtype
42
self.max_model_len = vllm_config.model_config.max_model_len
43
self.block_size = vllm_config.cache_config.block_size
@@ -45,6 +45,7 @@ class EagleProposer:
45
self.speculative_config.num_speculative_tokens)
46
self.max_num_tokens = (
47
vllm_config.scheduler_config.max_num_batched_tokens)
48
+ self.token_arange_np = np.arange(self.max_num_tokens)
49
# We need to get the hidden size from the draft model config because
50
# the draft model's hidden size can be different from the target model's
51
# hidden size (e.g., Llama 3.3 70B).
@@ -83,19 +84,14 @@ class EagleProposer:
84
target_positions: torch.Tensor,
85
# [num_tokens, hidden_size]
86
target_hidden_states: torch.Tensor,
86
- # [num_tokens]
87
- target_slot_mapping: torch.Tensor,
87
# [batch_size]
88
next_token_ids: torch.Tensor,
90
- # [batch_size + 1] starting with 0
91
- cu_num_tokens: torch.Tensor,
92
- # [batch_size, max_num_blocks_per_req]
93
- block_table: torch.Tensor,
89
+ common_attn_metadata: CommonAttentionMetadata,
90
sampling_metadata: SamplingMetadata,
91
) -> torch.Tensor:
92
num_tokens = target_token_ids.shape[0]
93
batch_size = next_token_ids.shape[0]
98
- last_token_indices = cu_num_tokens[1:] - 1
94
+ last_token_indices = common_attn_metadata.query_start_loc[1:] - 1
95
96
if self.method == "eagle3":
97
assert isinstance(self.model, Eagle3LlamaForCausalLM)
@@ -110,50 +106,14 @@ class EagleProposer:
106
# E.g., [b1, b2, c1, c2, c3, c3] -> [a2, b2, b3, c2, c3, c4]
107
self.input_ids[last_token_indices] = next_token_ids
108
113
- # FA requires seq_len to have dtype int32.
114
- seq_lens = (target_positions[last_token_indices] + 1).int()
115
-
116
- if self.method in ["eagle", "eagle3"]:
117
- # FIXME(woosuk): The below two ops cause synchronization. Optimize.
118
- max_seq_len = seq_lens.max().item()
119
- max_num_tokens = (cu_num_tokens[1:] -
120
- cu_num_tokens[:-1]).max().item()
121
- attn_metadata = FlashAttentionMetadata(
122
- num_actual_tokens=num_tokens,
123
- max_query_len=max_num_tokens,
124
- query_start_loc=cu_num_tokens,
125
- max_seq_len=max_seq_len,
126
- seq_lens=seq_lens,
127
- block_table=block_table,
128
- slot_mapping=target_slot_mapping,
129
- # TODO(woosuk): Support cascade attention.
130
- use_cascade=False,
131
- common_prefix_len=0,
132
- cu_prefix_query_lens=None,
133
- prefix_kv_lens=None,
134
- suffix_kv_lens=None,
135
- )
136
- elif self.method == "deepseek_mtp":
137
- query_lens = cu_num_tokens[1:] - cu_num_tokens[:-1]
138
- max_query_len = query_lens.max().item()
139
-
140
- common_attn_metadata = CommonAttentionMetadata(
141
- query_start_loc=cu_num_tokens,
142
- seq_lens=seq_lens,
143
- num_reqs=batch_size,
144
- num_actual_tokens=num_tokens,
145
- max_query_len=max_query_len,
146
- )
147
-
148
- assert self.runner is not None
109
+ assert self.runner is not None
110
150
- # FIXME: need to consider multiple kv_cache_groups
151
- attn_metadata = self.runner.attn_metadata_builders[0].build(
152
- common_prefix_len=0,
153
- common_attn_metadata=common_attn_metadata,
154
- )
155
- else:
156
- raise ValueError(f"Unsupported method: {self.method}")
111
+ # FIXME: need to consider multiple kv_cache_groups
112
+ attn_metadata = self.runner.attn_metadata_builders[0].build(
113
+ common_prefix_len=0,
114
+ common_attn_metadata=common_attn_metadata,
115
+ fast_build=True,
116
+ )
117
118
# At this moment, we assume all eagle layers belong to the same KV
119
# cache group, thus using the same attention metadata.
@@ -194,6 +154,11 @@ class EagleProposer:
154
# one layer. Adapt this code to support multiple layers once
155
# there's a multi-layer MTP module.
156
157
+ # Currently FlashAttention is the only backend that supports
158
+ # multi-token eagle spec decode. This is because the code below
159
+ # makes assumptions about attn_metadata attributes available.
160
+ assert isinstance(attn_metadata, FlashAttentionMetadata)
161
+
162
# Generate the remaining draft tokens.
163
draft_token_ids_list = [draft_token_ids]
164
@@ -238,8 +203,8 @@ class EagleProposer:
203
204
# Compute the slot mapping.
205
block_numbers = clamped_positions // self.block_size
241
- block_ids = block_table.gather(dim=1,
242
- index=block_numbers.view(-1, 1))
206
+ block_ids = attn_metadata.block_table.gather(
207
+ dim=1, index=block_numbers.view(-1, 1))
208
block_ids = block_ids.view(-1)
209
attn_metadata.slot_mapping = (block_ids * self.block_size +
210
clamped_positions % self.block_size)
@@ -275,46 +240,99 @@ class EagleProposer:
240
draft_token_ids = torch.stack(draft_token_ids_list, dim=1)
241
return draft_token_ids
242
278
- @staticmethod
243
def prepare_inputs(
280
- # [batch_size + 1]
281
- cu_target_query_lens: torch.Tensor,
244
+ self,
245
+ common_attn_metadata: CommonAttentionMetadata,
246
# [batch_size]
283
- num_rejected_tokens: torch.Tensor,
284
- num_tokens: int,
285
- ) -> tuple[torch.Tensor, torch.Tensor]:
286
- # cu_target_query_lens: [0, a, a + b, a + b + c]
287
- # num_rejected_tokens: [n1, n2, n3]
288
- # num_tokens_per_req: [a - n1, b - n2, c - n3]
289
- # cu_num_tokens: [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3]
290
- # token_indices: [0, 1, ..., a - n1 - 1,
291
- # a, a + 1, ..., a + b - n2 - 1,
292
- # a + b, a + b + 1, ..., a + b + c - n3 - 1]
293
-
294
- # [0, a, a + b, a + b + c] -> [a, b, c]
295
- query_len_per_req = (cu_target_query_lens[1:] -
296
- cu_target_query_lens[:-1])
297
- # [a, b, c] -> [a - n1, b - n2, c - n3]
298
- num_tokens_per_req = query_len_per_req - num_rejected_tokens
299
-
300
- # [a - n1, b - n2, c - n3] ->
301
- # [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3]
302
- cu_num_tokens = torch.zeros_like(cu_target_query_lens)
303
- torch.cumsum(num_tokens_per_req, dim=0, out=cu_num_tokens[1:])
304
- token_indices = torch.empty(
305
- num_tokens,
247
+ num_rejected_tokens: torch.Tensor
248
+ ) -> tuple[CommonAttentionMetadata, torch.Tensor]:
249
+ """
250
+ This function is used to prepare the inputs for the spec decode.
251
+ It updates to the common_attn_metadata to account for the rejected
252
+ tokens (and newly sampled tokens). It also returns the token indices
253
+ of the tokens that should be fed to the speculator.
254
+ """
255
+ # E.g.
256
+ # common_attn_metadata.query_start_loc{_cpu}:
257
+ # [0, q1, q1 + q2, q1 + q2 + q3]
258
+ # common_attn_metadata.seq_lens{_cpu}: [s1, s2, s3]
259
+ # num_rejected_tokens: [n1, n2, n3]
260
+ # This function computes the intermediate values:
261
+ # num_tokens_per_req: [q1 - n1, q2 - n2, q3 - n3]
262
+ # And returns:
263
+ # common_attn_metadata.query_start_loc{_cpu}:
264
+ # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
265
+ # common_attn_metadata.seq_lens{_cpu}:
266
+ # [s1 - n1 + 1, s2 - n2 + 1, s3 - n3 + 1]
267
+ # token_indices: [0, 1, ..., q1 - n1 - 1,
268
+ # q1, q1 + 1, ..., q1 + q2 - n2 - 1,
269
+ # q1 + q2, q1 + q2 + 1, ..., q1 + q2 + q3 - n3 - 1]
270
+
271
+ device = common_attn_metadata.query_start_loc.device
272
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
273
+ new_seq_lens_cpu = common_attn_metadata.seq_lens_cpu \
274
+ - num_rejected_tokens
275
+
276
+ # [0, q1, q1 + q2, q1 + q2 + q3] -> [q1, q2, q3]
277
+ new_query_len_per_req = (query_start_loc_cpu[1:] -
278
+ query_start_loc_cpu[:-1])
279
+ # [q1, q2, q3] -> [q1 - n1, q2 - n2, q3 - n3]
280
+ new_num_tokens_per_req = new_query_len_per_req - num_rejected_tokens
281
+ new_num_tokens_per_req_np = new_num_tokens_per_req.numpy()
282
+
283
+ # [q1 - n1, q2 - n2, q3 - n3] ->
284
+ # [0, q1 - n1, q1 + q2 - n1 - n2, q1 + q2 + q3 - n1 - n2 - n3]
285
+ new_query_start_loc_cpu = torch.zeros(
286
+ query_start_loc_cpu.shape,
287
dtype=torch.int32,
307
- device=cu_target_query_lens.device,
308
- )
309
- batch_size = num_rejected_tokens.shape[0]
310
- BLOCK_SIZE = 1024
311
- prepare_eagle_input_kernel[(batch_size, )](
312
- token_indices,
313
- cu_target_query_lens,
314
- cu_num_tokens,
315
- BLOCK_SIZE=BLOCK_SIZE,
288
+ pin_memory=is_pin_memory_available())
289
+ new_query_start_loc_np = new_query_start_loc_cpu.numpy()
290
+ np.cumsum(new_num_tokens_per_req_np, out=new_query_start_loc_np[1:])
291
+
292
+ total_num_tokens = new_query_start_loc_np[-1]
293
+ # Example assuming num_tokens_per_req_np = [2, 4, 3]
294
+ # this implies that `new_query_start_locs` is:
295
+ # [0, 2, 6, 9] ->
296
+ # [0, 0, 2, 2, 2, 2, 6, 6, 6]
297
+ # _r1_ ____r2____ ___r3__
298
+ new_query_start_locs_expanded = np.repeat(new_query_start_loc_np[:-1],
299
+ new_num_tokens_per_req_np)
300
+ # [0, 1, 2, 3, 4, 5, 6, 7, 8] ->
301
+ # [0, 1, 0, 1, 2, 3, 0, 1, 2]
302
+ # _r1_ ____r2____ ___r3__
303
+ token_offests = self.token_arange_np[:total_num_tokens] \
304
+ - new_query_start_locs_expanded
305
+
306
+ # Expand starting positions to match token pattern
307
+ # [0, q1, q1 + q2] ->
308
+ # [0, 0, q1, q1, q1, q1, q1 + q2, q1 + q2, q1 + q2]
309
+ # _r1_ _____r2_______ ___________r3____________
310
+ old_query_start_locs_expanded = np.repeat(
311
+ query_start_loc_cpu[:-1].numpy(), new_num_tokens_per_req_np)
312
+ # Final token indices are:
313
+ # [0, 1, // req 1
314
+ # q1 + 0, q1 + 1, q1 + 2, q1 + 3, // req 2
315
+ # q1 + q2 + 0, q1 + q2 + 1, q1 + q2 + 2] // req 3
316
+ token_indices_np = token_offests + old_query_start_locs_expanded
317
+ token_indices = torch.from_numpy(token_indices_np).to(
318
+ device, non_blocking=True)
319
+
320
+ spec_common_attn_metadata = CommonAttentionMetadata(
321
+ query_start_loc=new_query_start_loc_cpu.to(device,
322
+ non_blocking=True),
323
+ seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),
324
+ query_start_loc_cpu=new_query_start_loc_cpu,
325
+ seq_lens_cpu=new_seq_lens_cpu,
326
+ num_computed_tokens_cpu=common_attn_metadata.
327
+ num_computed_tokens_cpu,
328
+ num_reqs=common_attn_metadata.num_reqs,
329
+ num_actual_tokens=total_num_tokens,
330
+ max_query_len=new_query_len_per_req.max().item(),
331
+ block_table_tensor=common_attn_metadata.block_table_tensor,
332
+ slot_mapping=common_attn_metadata.slot_mapping[token_indices],
333
)
317
- return cu_num_tokens, token_indices
334
+
335
+ return spec_common_attn_metadata, token_indices
336
337
def load_model(self, target_model: nn.Module) -> None:
338
draft_model_config = \
vllm/v1/spec_decode/utils.py
-27
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from vllm.sampling_params import SamplingParams
4
-from vllm.triton_utils import tl, triton
4
5
_SAMPLING_EPS = 1e-5
6
@@ -13,29 +12,3 @@ def is_spec_decode_unsupported(sampling_params: SamplingParams) -> bool:
12
or sampling_params.repetition_penalty != 1.0
13
or sampling_params.min_p > _SAMPLING_EPS
14
or sampling_params.logprobs is not None)
16
-
17
-
18
-@triton.jit
19
-def prepare_eagle_input_kernel(
20
- out_ptr,
21
- cu_query_lens_ptr,
22
- cu_num_tokens_ptr,
23
- BLOCK_SIZE: tl.constexpr,
24
-):
25
- pid = tl.program_id(0)
26
-
27
- # [start_pos, end_pos)
28
- start_pos = tl.load(cu_num_tokens_ptr + pid)
29
- end_pos = tl.load(cu_num_tokens_ptr + pid + 1)
30
- num_tokens = end_pos - start_pos
31
-
32
- index_start = tl.load(cu_query_lens_ptr + pid)
33
-
34
- num_blocks = tl.cdiv(num_tokens, BLOCK_SIZE)
35
- for i in tl.range(num_blocks):
36
- offset = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
37
- tl.store(
38
- out_ptr + start_pos + offset,
39
- index_start + offset,
40
- mask=offset < num_tokens,
41
- )
vllm/v1/worker/block_table.py
+37
-4
@@ -14,12 +14,14 @@ class BlockTable:
14
15
def __init__(
16
self,
17
+ block_size: int,
18
max_num_reqs: int,
19
max_num_blocks_per_req: int,
20
max_num_batched_tokens: int,
21
pin_memory: bool,
22
device: torch.device,
23
):
24
+ self.block_size = block_size
25
self.max_num_reqs = max_num_reqs
26
self.max_num_blocks_per_req = max_num_blocks_per_req
27
self.max_num_batched_tokens = max_num_batched_tokens
@@ -79,10 +81,31 @@ class BlockTable:
81
82
self.block_table_np[[src, tgt]] = self.block_table_np[[tgt, src]]
83
82
- def commit(self, num_reqs: int) -> None:
84
+ def compute_slot_mapping(self, req_indices: np.ndarray,
85
+ positions: np.ndarray) -> None:
86
+ # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
87
+ # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
88
+ # where K is the max_num_blocks_per_req and the block size is 2.
89
+ # NOTE(woosuk): We can't simply use `token_indices // block_size`
90
+ # here because M (max_model_len) is not necessarily divisible by
91
+ # block_size.
92
+ block_table_indices = (req_indices * self.max_num_blocks_per_req +
93
+ positions // self.block_size)
94
+ block_table_cpu = self.get_cpu_tensor()
95
+ block_numbers = block_table_cpu.flatten()[block_table_indices].numpy()
96
+ block_offsets = positions % self.block_size
97
+ np.add(block_numbers * self.block_size,
98
+ block_offsets,
99
+ out=self.slot_mapping_np[:req_indices.shape[0]])
100
+
101
+ def commit_block_table(self, num_reqs: int) -> None:
102
self.block_table[:num_reqs].copy_(self.block_table_cpu[:num_reqs],
103
non_blocking=True)
104
105
+ def commit_slot_mapping(self, num_tokens: int) -> None:
106
+ self.slot_mapping[:num_tokens].copy_(
107
+ self.slot_mapping_cpu[:num_tokens], non_blocking=True)
108
+
109
def clear(self) -> None:
110
self.block_table.fill_(0)
111
self.block_table_cpu.fill_(0)
@@ -107,7 +130,8 @@ class MultiGroupBlockTable:
130
max_num_batched_tokens: int, pin_memory: bool,
131
device: torch.device, block_sizes: list[int]) -> None:
132
self.block_tables = [
110
- BlockTable(max_num_reqs, cdiv(max_model_len, block_size),
133
+ BlockTable(block_size, max_num_reqs, cdiv(max_model_len,
134
+ block_size),
135
max_num_batched_tokens, pin_memory, device)
136
for block_size in block_sizes
137
]
@@ -129,9 +153,18 @@ class MultiGroupBlockTable:
153
for block_table in self.block_tables:
154
block_table.swap_row(src, tgt)
155
132
- def commit(self, num_reqs: int) -> None:
156
+ def compute_slot_mapping(self, req_indices: np.ndarray,
157
+ positions: np.ndarray) -> None:
158
+ for block_table in self.block_tables:
159
+ block_table.compute_slot_mapping(req_indices, positions)
160
+
161
+ def commit_block_table(self, num_reqs: int) -> None:
162
+ for block_table in self.block_tables:
163
+ block_table.commit_block_table(num_reqs)
164
+
165
+ def commit_slot_mapping(self, num_tokens: int) -> None:
166
for block_table in self.block_tables:
134
- block_table.commit(num_reqs)
167
+ block_table.commit_slot_mapping(num_tokens)
168
169
def clear(self) -> None:
170
for block_table in self.block_tables:
vllm/v1/worker/gpu_model_runner.py
+64
-85
@@ -3,7 +3,6 @@
3
4
import gc
5
import time
6
-import weakref
6
from contextlib import contextmanager
7
from typing import TYPE_CHECKING, Any, Optional, Union
8
@@ -42,8 +41,7 @@ from vllm.pooling_params import PoolingParams
41
from vllm.sampling_params import SamplingType
42
from vllm.sequence import IntermediateTensors
43
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
45
- GiB_bytes, LazyLoader, async_tensor_h2d,
46
- check_use_alibi, get_dtype_size,
44
+ GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,
45
is_pin_memory_available, round_up)
46
from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
47
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
@@ -62,7 +60,6 @@ from vllm.v1.spec_decode.eagle import EagleProposer
60
from vllm.v1.spec_decode.medusa import MedusaProposer
61
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
62
from vllm.v1.spec_decode.ngram_proposer import NgramProposer
65
-from vllm.v1.worker.block_table import BlockTable
63
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
64
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
65
@@ -577,8 +574,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
574
def _prepare_inputs(
575
self,
576
scheduler_output: "SchedulerOutput",
580
- ) -> tuple[dict[str, Any], bool, torch.Tensor,
581
- Optional[SpecDecodeMetadata], np.ndarray]:
577
+ ) -> tuple[dict[str,
578
+ Any], bool, torch.Tensor, Optional[SpecDecodeMetadata],
579
+ np.ndarray, Optional[CommonAttentionMetadata]]:
580
"""
581
:return: tuple[
582
attn_metadata: layer-to-attention_metadata mapping,
@@ -593,7 +591,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
591
592
# OPTIMIZATION: Start copying the block table first.
593
# This way, we can overlap the copy with the following CPU operations.
596
- self.input_batch.block_table.commit(num_reqs)
594
+ self.input_batch.block_table.commit_block_table(num_reqs)
595
596
# Get the number of scheduled tokens for each request.
597
req_ids = self.input_batch.req_ids
@@ -637,29 +635,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
635
torch.from_numpy(token_indices),
636
out=self.input_ids_cpu[:total_num_scheduled_tokens])
637
640
- # Calculate the slot mapping for each KV cache group.
641
- for kv_cache_group_id, kv_cache_group_spec in enumerate(
642
- self.kv_cache_config.kv_cache_groups):
643
- block_size = kv_cache_group_spec.kv_cache_spec.block_size
644
- block_table: BlockTable = self.input_batch.block_table[
645
- kv_cache_group_id]
646
- # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
647
- # -> [0, 0, K, K, K + 1, K + 1, K + 2, 2 * K, 2 * K, 2 * K + 1]
648
- # where K is the max_num_blocks_per_req and the block size is 2.
649
- # NOTE(woosuk): We can't simply use `token_indices // block_size`
650
- # here because M (max_model_len) is not necessarily divisible by
651
- # block_size.
652
- block_table_indices = (
653
- req_indices * block_table.max_num_blocks_per_req +
654
- positions_np // block_size)
655
- block_table_cpu = block_table.get_cpu_tensor()
656
- block_numbers = block_table_cpu.flatten(
657
- )[block_table_indices].numpy()
658
- block_offsets = positions_np % block_size
659
- np.add(
660
- block_numbers * block_size,
661
- block_offsets,
662
- out=block_table.slot_mapping_np[:total_num_scheduled_tokens])
638
+ self.input_batch.block_table.compute_slot_mapping(
639
+ req_indices, positions_np)
640
+ self.input_batch.block_table.commit_slot_mapping(
641
+ total_num_scheduled_tokens)
642
643
# Prepare the attention metadata.
644
self.query_start_loc_np[0] = 0
@@ -696,15 +675,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
675
self.query_start_loc_cpu[num_reqs].item())
676
677
query_start_loc = self.query_start_loc[:num_reqs + 1]
699
- seq_lens = self.seq_lens[:num_reqs]
700
-
701
- common_attn_metadata = CommonAttentionMetadata(
702
- query_start_loc=query_start_loc,
703
- seq_lens=seq_lens,
704
- num_reqs=num_reqs,
705
- num_actual_tokens=total_num_scheduled_tokens,
706
- max_query_len=max_num_scheduled_tokens,
707
- )
678
+
679
+ spec_decode_common_attn_metadata = None
680
681
attn_metadata: dict[str, Any] = {}
682
# Prepare the attention metadata for each KV cache group and make layers
@@ -712,6 +684,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
684
for kv_cache_group_id, kv_cache_group_spec in enumerate(
685
self.kv_cache_config.kv_cache_groups):
686
687
+ blk_table = self.input_batch.block_table[kv_cache_group_id]
688
+ blk_table_tensor = blk_table.get_device_tensor()[:num_reqs]
689
+ slot_mapping = blk_table.slot_mapping[:total_num_scheduled_tokens]
690
+ common_attn_metadata = CommonAttentionMetadata(
691
+ query_start_loc=self.query_start_loc[:num_reqs + 1],
692
+ query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
693
+ seq_lens=self.seq_lens[:num_reqs],
694
+ seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
695
+ num_computed_tokens_cpu=self.input_batch.
696
+ num_computed_tokens_cpu_tensor[:num_reqs],
697
+ num_reqs=num_reqs,
698
+ num_actual_tokens=total_num_scheduled_tokens,
699
+ max_query_len=max_num_scheduled_tokens,
700
+ block_table_tensor=blk_table_tensor,
701
+ slot_mapping=slot_mapping,
702
+ )
703
+
704
+ if self.speculative_config and \
705
+ spec_decode_common_attn_metadata is None:
706
+ spec_decode_common_attn_metadata = common_attn_metadata
707
+
708
# Prepare for cascade attention if enabled & beneficial.
709
common_prefix_len = 0
710
builder = self.attn_metadata_builders[kv_cache_group_id]
@@ -765,7 +758,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
758
self.set_active_loras(self.input_batch, num_scheduled_tokens)
759
760
return (attn_metadata, attention_cuda_graphs, logits_indices,
768
- spec_decode_metadata, num_scheduled_tokens)
761
+ spec_decode_metadata, num_scheduled_tokens,
762
+ spec_decode_common_attn_metadata)
763
764
def _compute_cascade_attn_prefix_len(
765
self,
@@ -1286,8 +1280,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1280
1281
# Prepare the decoder inputs.
1282
(attn_metadata, attention_cuda_graphs, logits_indices,
1289
- spec_decode_metadata,
1290
- num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
1283
+ spec_decode_metadata, num_scheduled_tokens_np,
1284
+ spec_decode_common_attn_metadata) = (
1285
+ self._prepare_inputs(scheduler_output))
1286
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1287
if (self.use_cuda_graph
1288
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
@@ -1528,6 +1523,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1523
# Speculative decoding is not enabled.
1524
spec_token_ids = None
1525
else:
1526
+ assert spec_decode_common_attn_metadata is not None
1527
spec_token_ids = self.propose_draft_token_ids(
1528
scheduler_output,
1529
valid_sampled_token_ids,
@@ -1536,7 +1532,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1532
sample_hidden_states,
1533
aux_hidden_states,
1534
spec_decode_metadata,
1539
- attn_metadata,
1535
+ spec_decode_common_attn_metadata,
1536
)
1537
1538
self.eplb_step()
@@ -1561,7 +1557,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1557
sample_hidden_states: torch.Tensor,
1558
aux_hidden_states: Optional[torch.Tensor],
1559
spec_decode_metadata: Optional[SpecDecodeMetadata],
1564
- attn_metadata: dict[str, Any],
1560
+ common_attn_metadata: CommonAttentionMetadata,
1561
) -> list[list[int]]:
1562
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1563
if self.speculative_config.method == "ngram":
@@ -1608,16 +1604,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1604
next_token_ids = torch.tensor(next_token_ids,
1605
dtype=torch.int32,
1606
device=self.device)
1611
- # At this moment, we assume all eagle layers belong to the same KV
1612
- # cache group, thus using the same attention metadata.
1613
- eagle_attn_metadata = attn_metadata[
1614
- self.drafter.attn_layer_names[0]]
1615
-
1616
- # NOTE: deepseek_mtp uses MLA which does not have `block_table`
1617
- if hasattr(eagle_attn_metadata, "block_table"):
1618
- block_table = eagle_attn_metadata.block_table
1619
- else:
1620
- block_table = None
1607
1608
if spec_decode_metadata is None:
1609
# input_ids can be None for multimodal models.
@@ -1630,8 +1616,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1616
dim=-1)
1617
else:
1618
target_hidden_states = hidden_states[:num_scheduled_tokens]
1633
- target_slot_mapping = eagle_attn_metadata.slot_mapping
1634
- cu_num_tokens = eagle_attn_metadata.query_start_loc
1619
else:
1620
# TODO(woosuk): Refactor this.
1621
num_draft_tokens = spec_decode_metadata.num_draft_tokens
@@ -1639,17 +1623,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1623
n + 1 - len(sampled_token_ids[i]) if n > 0 else 0
1624
for i, n in enumerate(num_draft_tokens)
1625
]
1642
- num_rejected_tokens_tensor = async_tensor_h2d(
1643
- num_rejected_tokens,
1644
- dtype=torch.int32,
1645
- target_device=self.device,
1646
- pin_memory=True)
1647
- num_tokens = num_scheduled_tokens - sum(num_rejected_tokens)
1648
- cu_num_tokens, token_indices = self.drafter.prepare_inputs(
1649
- eagle_attn_metadata.query_start_loc,
1650
- num_rejected_tokens_tensor,
1651
- num_tokens,
1652
- )
1626
+ num_rejected_tokens_cpu = torch.tensor(num_rejected_tokens,
1627
+ dtype=torch.int32)
1628
+ common_attn_metadata, token_indices =\
1629
+ self.drafter.prepare_inputs(
1630
+ common_attn_metadata, num_rejected_tokens_cpu)
1631
+
1632
target_token_ids = self.input_ids[token_indices]
1633
# TODO(woosuk): Support M-RoPE.
1634
target_positions = self.positions[token_indices]
@@ -1658,17 +1637,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1637
[h[token_indices] for h in aux_hidden_states], dim=-1)
1638
else:
1639
target_hidden_states = hidden_states[token_indices]
1661
- target_slot_mapping = eagle_attn_metadata.slot_mapping[
1662
- token_indices]
1640
draft_token_ids = self.drafter.propose(
1641
target_token_ids=target_token_ids,
1642
target_positions=target_positions,
1643
target_hidden_states=target_hidden_states,
1667
- target_slot_mapping=target_slot_mapping,
1644
next_token_ids=next_token_ids,
1669
- cu_num_tokens=cu_num_tokens,
1670
- block_table=block_table,
1645
sampling_metadata=sampling_metadata,
1646
+ common_attn_metadata=common_attn_metadata,
1647
)
1648
spec_token_ids = draft_token_ids.tolist()
1649
return spec_token_ids
@@ -1970,24 +1945,29 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1945
if capture_attn_cudagraph:
1946
attn_metadata = {}
1947
1973
- query_start_loc = self.query_start_loc[:num_reqs + 1]
1948
# Make sure max_model_len is used at the graph capture time.
1949
self.seq_lens_np[:num_reqs] = self.max_model_len
1950
self.seq_lens_np[num_reqs:] = 0
1951
self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
1952
non_blocking=True)
1979
- seq_lens = self.seq_lens[:num_reqs]
1980
-
1981
- common_attn_metadata = CommonAttentionMetadata(
1982
- query_start_loc=query_start_loc,
1983
- seq_lens=seq_lens,
1984
- num_reqs=num_reqs,
1985
- num_actual_tokens=num_tokens,
1986
- max_query_len=num_tokens,
1987
- )
1953
1954
for kv_cache_group_id, kv_cache_group_spec in enumerate(
1955
self.kv_cache_config.kv_cache_groups):
1956
+ common_attn_metadata = CommonAttentionMetadata(
1957
+ query_start_loc=self.query_start_loc[:num_reqs + 1],
1958
+ query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
1959
+ 1],
1960
+ seq_lens=self.seq_lens[:num_reqs],
1961
+ seq_lens_cpu=self.seq_lens_cpu[:num_reqs],
1962
+ num_computed_tokens_cpu=self.input_batch.
1963
+ num_computed_tokens_cpu_tensor[:num_reqs],
1964
+ num_reqs=num_reqs,
1965
+ num_actual_tokens=num_tokens,
1966
+ max_query_len=num_tokens,
1967
+ block_table_tensor=self.input_batch.block_table[
1968
+ kv_cache_group_id].get_device_tensor()[:num_reqs],
1969
+ slot_mapping=self.input_batch.
1970
+ block_table[kv_cache_group_id].slot_mapping[:num_reqs])
1971
1972
attn_metadata_i = self.attn_metadata_builders[
1973
kv_cache_group_id].build_for_cudagraph_capture(
@@ -2339,11 +2319,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2319
raise ValueError(
2320
f"Unknown KV cache spec type: {type(kv_cache_spec)}")
2321
2342
- block_table_i = self.input_batch.block_table[i]
2322
attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
2344
- weakref.proxy(self),
2323
kv_cache_spec,
2346
- block_table_i,
2324
+ self.vllm_config,
2325
+ self.device,
2326
)
2327
2328
if (self.full_cuda_graph