A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 17:30 UTC
[Attention] Make seq_lens_cpu optional in CommonAttentionMetadata to enable true async spec-decode (#29624)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
Lucas Wilkinson committed
Dec 9, 2025 at 20:18 UTC
abe93bce5952ed8adf90d4b77af6ed3515958620
9 files changed
+200
-37
tests/v1/attention/utils.py
+2
-2
@@ -106,8 +106,8 @@ def create_common_attn_metadata(
106
query_start_loc=query_start_loc,
107
query_start_loc_cpu=query_start_loc_cpu,
108
seq_lens=seq_lens,
109
- seq_lens_cpu=seq_lens_cpu,
110
- num_computed_tokens_cpu=num_computed_tokens_cpu,
109
+ _seq_lens_cpu=seq_lens_cpu,
110
+ _num_computed_tokens_cpu=num_computed_tokens_cpu,
111
num_reqs=batch_spec.batch_size,
112
num_actual_tokens=num_tokens,
113
max_query_len=max_query_len,
tests/v1/e2e/test_async_spec_decode.py
new
+131
@@ -0,0 +1,131 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Test that verifies no implicit GPU-CPU synchronization occurs during
5
+speculative decoding generation under expected conditions.
6
+"""
7
+
8
+import multiprocessing
9
+import sys
10
+import traceback
11
+
12
+import pytest
13
+import torch
14
+
15
+
16
+@pytest.fixture
17
+def sync_tracker():
18
+ """
19
+ Fixture that patches CommonAttentionMetadata.seq_lens_cpu to detect
20
+ lazy init syncs. Prints stack traces immediately when syncs occur.
21
+ """
22
+ from vllm.v1.attention.backends.utils import CommonAttentionMetadata
23
+
24
+ # Shared counter for cross-process communication (inherited by fork)
25
+ sync_count = multiprocessing.Value("i", 0)
26
+
27
+ # Save original property
28
+ original_prop = CommonAttentionMetadata.seq_lens_cpu
29
+ original_fget = original_prop.fget
30
+
31
+ # Create tracking wrapper
32
+ def tracking_seq_lens_cpu(self):
33
+ if self._seq_lens_cpu is None:
34
+ # Increment counter
35
+ with sync_count.get_lock():
36
+ sync_count.value += 1
37
+ count = sync_count.value
38
+ # Print stack trace immediately (shows in subprocess output)
39
+ print(f"\n{'=' * 60}", file=sys.stderr)
40
+ print(f"SYNC #{count}: seq_lens_cpu lazy init triggered!", file=sys.stderr)
41
+ print(f"{'=' * 60}", file=sys.stderr)
42
+ traceback.print_stack(file=sys.stderr)
43
+ print(f"{'=' * 60}\n", file=sys.stderr)
44
+ sys.stderr.flush()
45
+ return original_fget(self)
46
+
47
+ # Apply patch
48
+ CommonAttentionMetadata.seq_lens_cpu = property(tracking_seq_lens_cpu)
49
+
50
+ class SyncTracker:
51
+ @property
52
+ def count(self) -> int:
53
+ return sync_count.value
54
+
55
+ def assert_no_sync(self, msg: str = ""):
56
+ count = sync_count.value
57
+ assert count == 0, (
58
+ f"Unexpected GPU-CPU sync: seq_lens_cpu lazy init triggered "
59
+ f"{count} times. See stack traces above. {msg}"
60
+ )
61
+
62
+ yield SyncTracker()
63
+
64
+ # Restore original property
65
+ CommonAttentionMetadata.seq_lens_cpu = original_prop
66
+ torch._dynamo.reset()
67
+
68
+
69
+# Test configurations: (model, spec_model, method, num_spec_tokens, backend_env)
70
+SPEC_DECODE_CONFIGS = [
71
+ pytest.param(
72
+ "meta-llama/Llama-3.2-1B-Instruct",
73
+ "nm-testing/Llama3_2_1B_speculator.eagle3",
74
+ "eagle3",
75
+ 2,
76
+ id="eagle3-llama",
77
+ ),
78
+ pytest.param(
79
+ "eagle618/deepseek-v3-random",
80
+ "eagle618/eagle-deepseek-v3-random",
81
+ "eagle",
82
+ 2,
83
+ id="eagle-mla-deepseek",
84
+ ),
85
+]
86
+
87
+
88
+@pytest.mark.parametrize(
89
+ "model,spec_model,method,num_spec_tokens",
90
+ SPEC_DECODE_CONFIGS,
91
+)
92
+def test_no_sync_with_spec_decode(
93
+ sync_tracker,
94
+ model: str,
95
+ spec_model: str,
96
+ method: str,
97
+ num_spec_tokens: int,
98
+):
99
+ """
100
+ Test that no implicit GPU-CPU sync occurs during speculative decoding
101
+ generation.
102
+ """
103
+ # Import vLLM AFTER sync_tracker fixture has applied the patch
104
+ from vllm import LLM, SamplingParams
105
+ from vllm.distributed import cleanup_dist_env_and_memory
106
+
107
+ llm = LLM(
108
+ model=model,
109
+ max_model_len=256,
110
+ speculative_config={
111
+ "method": method,
112
+ "num_speculative_tokens": num_spec_tokens,
113
+ "model": spec_model,
114
+ },
115
+ enforce_eager=True,
116
+ async_scheduling=True,
117
+ )
118
+
119
+ outputs = llm.generate(
120
+ ["Hello, my name is"],
121
+ SamplingParams(temperature=0, max_tokens=10),
122
+ )
123
+
124
+ assert len(outputs) == 1
125
+ assert len(outputs[0].outputs[0].text) > 0
126
+
127
+ del llm
128
+ torch.cuda.empty_cache()
129
+ cleanup_dist_env_and_memory()
130
+
131
+ sync_tracker.assert_no_sync()
tests/v1/spec_decode/test_tree_attention.py
+2
-2
@@ -88,8 +88,8 @@ def forward_attention(
88
query_start_loc=query_start_loc,
89
query_start_loc_cpu=query_start_loc.cpu(),
90
seq_lens=seq_lens,
91
- seq_lens_cpu=seq_lens.cpu(),
92
- num_computed_tokens_cpu=context_lens.cpu(),
91
+ _seq_lens_cpu=seq_lens.cpu(),
92
+ _num_computed_tokens_cpu=context_lens.cpu(),
93
num_reqs=batch_size,
94
num_actual_tokens=num_actual_tokens,
95
max_query_len=max_query_len,
vllm/attention/layers/cross_attention.py
+1
-1
@@ -103,7 +103,7 @@ def create_cross_attention_backend(
103
# needed here to know how many tokens to attend to from the cached
104
# cross-attention KV cache.
105
new_metadata.seq_lens = common_attn_metadata.encoder_seq_lens
106
- new_metadata.seq_lens_cpu = torch.from_numpy(
106
+ new_metadata._seq_lens_cpu = torch.from_numpy(
107
common_attn_metadata.encoder_seq_lens_cpu
108
)
109
vllm/v1/attention/backends/gdn_attn.py
+1
-1
@@ -370,6 +370,6 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
370
371
num_accepted_tokens = torch.diff(m.query_start_loc)
372
num_decode_draft_tokens_cpu = (num_accepted_tokens - 1).cpu()
373
- m.num_computed_tokens_cpu = m.seq_lens_cpu - num_accepted_tokens.cpu()
373
+ m._num_computed_tokens_cpu = m.seq_lens_cpu - num_accepted_tokens.cpu()
374
375
return self.build(0, m, num_accepted_tokens, num_decode_draft_tokens_cpu)
vllm/v1/attention/backends/utils.py
+49
-17
@@ -18,7 +18,7 @@ from typing import (
18
19
import numpy as np
20
import torch
21
-from typing_extensions import runtime_checkable
21
+from typing_extensions import deprecated, runtime_checkable
22
23
from vllm.config import VllmConfig, get_layers_from_vllm_config
24
from vllm.utils.math_utils import cdiv
@@ -66,11 +66,6 @@ class CommonAttentionMetadata:
66
"""(batch_size + 1,), the start location of each request in query Tensor"""
67
68
seq_lens: torch.Tensor
69
- seq_lens_cpu: torch.Tensor
70
- """(batch_size,), the length of each request including both computed tokens
71
- and newly scheduled tokens"""
72
-
73
- num_computed_tokens_cpu: torch.Tensor
69
"""(batch_size,), the number of computed tokens for each request"""
70
71
num_reqs: int
@@ -81,7 +76,7 @@ class CommonAttentionMetadata:
76
max_query_len: int
77
"""Longest query in batch"""
78
max_seq_len: int
84
- """Longest context length in batch"""
79
+ """Longest context length (may be an upper bound)"""
80
81
block_table_tensor: torch.Tensor
82
slot_mapping: torch.Tensor
@@ -100,6 +95,40 @@ class CommonAttentionMetadata:
95
dcp_local_seq_lens_cpu: torch.Tensor | None = None
96
"""Sequence lengths of the local rank in decode context parallelism world"""
97
98
+ # WARNING: Deprecated fields. Will be removed in a future release (v0.14.0)
99
+ _seq_lens_cpu: torch.Tensor | None = None
100
+ _num_computed_tokens_cpu: torch.Tensor | None = None
101
+
102
+ @property
103
+ @deprecated(
104
+ """
105
+ Prefer using device seq_lens directly to avoid implicit H<>D sync.
106
+ If a CPU copy is needed, use `seq_lens.cpu()` instead.
107
+ Will be removed in a future release (v0.14.0)
108
+ """
109
+ )
110
+ def seq_lens_cpu(self) -> torch.Tensor:
111
+ if self._seq_lens_cpu is None:
112
+ self._seq_lens_cpu = self.seq_lens.to("cpu")
113
+ return self._seq_lens_cpu
114
+
115
+ @property
116
+ @deprecated(
117
+ """
118
+ Prefer using device seq_lens directly to avoid implicit H<>D sync which breaks full
119
+ async scheduling. If a CPU copy is needed, it can be derived from
120
+ query_start_loc_cpu and seq_lens.
121
+ Will be removed in a future release (v0.14.0)
122
+ """
123
+ )
124
+ def num_computed_tokens_cpu(self) -> torch.Tensor:
125
+ if self._num_computed_tokens_cpu is None:
126
+ query_seq_lens = (
127
+ self.query_start_loc_cpu[1:] - self.query_start_loc_cpu[:-1]
128
+ )
129
+ self._num_computed_tokens_cpu = self.seq_lens_cpu - query_seq_lens
130
+ return self._num_computed_tokens_cpu
131
+
132
# TODO(lucas): remove once we have FULL-CG spec-decode support
133
def unpadded(
134
self, num_actual_tokens: int, num_actual_reqs: int
@@ -109,8 +138,12 @@ class CommonAttentionMetadata:
138
query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
139
query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
140
seq_lens=self.seq_lens[:num_actual_reqs],
112
- seq_lens_cpu=self.seq_lens_cpu[:num_actual_reqs],
113
- num_computed_tokens_cpu=self.num_computed_tokens_cpu[:num_actual_reqs],
141
+ _seq_lens_cpu=self._seq_lens_cpu[:num_actual_reqs]
142
+ if self._seq_lens_cpu is not None
143
+ else None,
144
+ _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
145
+ if self._num_computed_tokens_cpu is not None
146
+ else None,
147
num_reqs=num_actual_reqs,
148
num_actual_tokens=num_actual_tokens,
149
max_query_len=self.max_query_len,
@@ -224,14 +257,14 @@ def _make_metadata_with_slice(
257
query_start_loc=query_start_loc,
258
query_start_loc_cpu=query_start_loc_cpu,
259
seq_lens=seq_lens,
227
- seq_lens_cpu=seq_lens_cpu,
228
- num_computed_tokens_cpu=num_computed_tokens_cpu,
260
num_reqs=num_requests,
261
num_actual_tokens=num_actual_tokens,
262
max_query_len=max_query_len,
263
max_seq_len=max_seq_len,
264
block_table_tensor=block_table_tensor,
265
slot_mapping=slot_mapping,
266
+ _seq_lens_cpu=seq_lens_cpu,
267
+ _num_computed_tokens_cpu=num_computed_tokens_cpu,
268
)
269
270
@@ -689,9 +722,7 @@ def make_local_attention_virtual_batches(
722
return CommonAttentionMetadata(
723
query_start_loc_cpu=query_start_loc_cpu,
724
query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
692
- seq_lens_cpu=seq_lens_cpu,
725
seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),
694
- num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),
726
num_reqs=len(seq_lens_cpu),
727
num_actual_tokens=common_attn_metadata.num_actual_tokens,
728
max_query_len=seqlens_q_local.max(),
@@ -699,6 +730,8 @@ def make_local_attention_virtual_batches(
730
block_table_tensor=block_table_local,
731
slot_mapping=common_attn_metadata.slot_mapping,
732
causal=True,
733
+ _seq_lens_cpu=seq_lens_cpu,
734
+ _num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local),
735
)
736
737
@@ -719,7 +752,6 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
752
logits_indices = logits_indices_padded[:num_logits_indices]
753
num_reqs = common_attn_metadata.num_reqs
754
query_start_loc = common_attn_metadata.query_start_loc
722
- seq_lens = common_attn_metadata.seq_lens
755
# Example inputs
756
# num_reqs: 3
757
# generation_indices: [14, 18, 19, 27]
@@ -748,9 +780,7 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
780
common_attn_metadata = CommonAttentionMetadata(
781
query_start_loc=decode_query_start_loc,
782
query_start_loc_cpu=decode_query_start_loc.to("cpu", non_blocking=True),
751
- seq_lens=seq_lens,
752
- seq_lens_cpu=seq_lens.to("cpu", non_blocking=True),
753
- num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,
783
+ seq_lens=common_attn_metadata.seq_lens,
784
num_reqs=num_reqs,
785
num_actual_tokens=total_num_decode_tokens,
786
max_query_len=decode_max_query_len,
@@ -758,6 +788,8 @@ def make_kv_sharing_fast_prefill_common_attn_metadata(
788
block_table_tensor=common_attn_metadata.block_table_tensor,
789
slot_mapping=common_attn_metadata.slot_mapping,
790
causal=True,
791
+ _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,
792
+ _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
793
)
794
return common_attn_metadata
795
vllm/v1/spec_decode/eagle.py
+10
-10
@@ -440,16 +440,16 @@ class EagleProposer:
440
# of main model.
441
# Increment the sequence lengths.
442
common_attn_metadata.seq_lens += 1
443
- # This is an out-of-place operation to avoid modifying the original tensor.
444
- common_attn_metadata.seq_lens_cpu = common_attn_metadata.seq_lens_cpu + 1
443
# For the requests that exceed the max model length, we set the
444
# sequence length to 1 to minimize their overheads in attention.
447
-
445
common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
446
450
- common_attn_metadata.num_computed_tokens_cpu = (
451
- common_attn_metadata.seq_lens_cpu - 1
452
- )
447
+ # Also update the CPU-side shadow; NOTE: this is hacky and should be
448
+ # removed in when common_attn_metadata.seq_lens_cpu is deprecated.
449
+ if common_attn_metadata._seq_lens_cpu is not None:
450
+ common_attn_metadata._seq_lens_cpu += 1
451
+ if common_attn_metadata._num_computed_tokens_cpu is not None:
452
+ common_attn_metadata._num_computed_tokens_cpu += 1
453
454
# Compute the slot mapping.
455
if self.uses_mrope:
@@ -656,8 +656,8 @@ class EagleProposer:
656
query_start_loc=common_attn_metadata.query_start_loc,
657
seq_lens=common_attn_metadata.seq_lens,
658
query_start_loc_cpu=query_start_loc_cpu,
659
- seq_lens_cpu=common_attn_metadata.seq_lens_cpu,
660
- num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,
659
+ _seq_lens_cpu=common_attn_metadata._seq_lens_cpu,
660
+ _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
661
num_reqs=common_attn_metadata.num_reqs,
662
num_actual_tokens=total_num_tokens,
663
max_query_len=new_query_len_per_req.max().item(),
@@ -932,8 +932,8 @@ class EagleProposer:
932
query_start_loc=new_query_start_loc_cpu.to(device, non_blocking=True),
933
seq_lens=new_seq_lens_cpu.to(device, non_blocking=True),
934
query_start_loc_cpu=new_query_start_loc_cpu,
935
- seq_lens_cpu=new_seq_lens_cpu,
936
- num_computed_tokens_cpu=common_attn_metadata.num_computed_tokens_cpu,
935
+ _seq_lens_cpu=new_seq_lens_cpu,
936
+ _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu,
937
num_reqs=common_attn_metadata.num_reqs,
938
num_actual_tokens=total_num_tokens,
939
max_query_len=new_query_len_per_req.max().item(),
vllm/v1/worker/gpu/attn_utils.py
+2
-2
@@ -168,9 +168,9 @@ def build_attn_metadata(
168
query_start_loc=query_start_loc_gpu,
169
query_start_loc_cpu=query_start_loc_cpu,
170
seq_lens=seq_lens,
171
- seq_lens_cpu=seq_lens_cpu,
171
+ _seq_lens_cpu=seq_lens_cpu,
172
max_seq_len=max_seq_len,
173
- num_computed_tokens_cpu=num_computed_tokens_cpu,
173
+ _num_computed_tokens_cpu=num_computed_tokens_cpu,
174
num_reqs=num_reqs,
175
num_actual_tokens=num_tokens,
176
max_query_len=max_query_len,
vllm/v1/worker/gpu_model_runner.py
+2
-2
@@ -1626,8 +1626,8 @@ class GPUModelRunner(
1626
query_start_loc=query_start_loc,
1627
query_start_loc_cpu=query_start_loc_cpu,
1628
seq_lens=seq_lens,
1629
- seq_lens_cpu=seq_lens_cpu,
1630
- num_computed_tokens_cpu=num_computed_tokens_cpu,
1629
+ _seq_lens_cpu=seq_lens_cpu,
1630
+ _num_computed_tokens_cpu=num_computed_tokens_cpu,
1631
num_actual_tokens=num_tokens_padded,
1632
num_reqs=num_reqs_padded,
1633
max_query_len=max_query_len,