[Bugfix][Spec-Decode] Populate draft seq_lens_cpu_upper_bound for spec-decode attention metadata (#44492)
Signed-off-by: Oxana Korzh <okorzh@amd.com> Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Oxana Korzh committed
Jul 20, 2026 at 14:50 UTC
4ec199b66a791070348f3baf847b3a873c48cdd9
4 files changed
+150
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
new
+128
@@ -0,0 +1,128 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Unit tests for the EAGLE speculator's draft attention metadata builder.
4
+
5
+These tests guard the regression where ``_build_draft_attn_metadata`` did
6
+not populate ``seq_lens_cpu_upper_bound`` on the per-step
7
+``CommonAttentionMetadata``. Several downstream attention backends and
8
+helpers (``split_decodes_prefills_and_extends``, the MLA indexer,
9
+flex-attention, cross-attention) assert this field is non-None, so
10
+omitting it caused crashes at the start of draft decode for certain
11
+backends (e.g. ``ROCM_AITER_FA`` with eagle/eagle3 spec decode):
12
+
13
+ AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
14
+"""
15
+
16
+from types import SimpleNamespace
17
+from unittest.mock import patch
18
+
19
+import torch
20
+
21
+from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator
22
+from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator
23
+
24
+
25
+def _make_fake_speculator(
26
+ *,
27
+ max_num_reqs: int = 8,
28
+ max_num_tokens: int = 16,
29
+ max_model_len: int = 1024,
30
+ draft_max_seq_len: int = 1024,
31
+) -> SimpleNamespace:
32
+ """Build a fake EagleSpeculator with just the attributes used by
33
+ ``_build_draft_attn_metadata``. We deliberately avoid constructing a
34
+ real ``EagleSpeculator`` because that requires a full ``VllmConfig``
35
+ and a draft model.
36
+ """
37
+ fake_input_buffers = SimpleNamespace(
38
+ query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32),
39
+ seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
40
+ )
41
+ fake_block_tables = SimpleNamespace(
42
+ input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)],
43
+ slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64),
44
+ )
45
+ return SimpleNamespace(
46
+ arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),
47
+ block_tables=fake_block_tables,
48
+ input_buffers=fake_input_buffers,
49
+ attn_groups=[],
50
+ kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
51
+ max_model_len=max_model_len,
52
+ draft_max_seq_len=draft_max_seq_len,
53
+ )
54
+
55
+
56
+def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step):
57
+ captured: dict[str, object] = {}
58
+
59
+ def fake_build_attn_metadata(**kwargs):
60
+ captured.update(kwargs)
61
+ return {}
62
+
63
+ with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata):
64
+ EagleSpeculator._build_draft_attn_metadata(
65
+ fake, # type: ignore[arg-type]
66
+ num_reqs=num_reqs,
67
+ num_reqs_padded=num_reqs_padded,
68
+ num_tokens_padded=num_tokens_padded,
69
+ seq_lens_cpu_upper_bound=base,
70
+ step=step,
71
+ )
72
+ return captured
73
+
74
+
75
+def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():
76
+ """The fix: every per-step ``CommonAttentionMetadata`` carries a non-None
77
+ ``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus
78
+ the current draft-step offset. Padded entries are zeroed (matching the
79
+ main model runner's convention)."""
80
+ fake = _make_fake_speculator()
81
+ base = torch.tensor([100, 200, 300, 0], dtype=torch.int32)
82
+
83
+ captured = _run_build(
84
+ fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2
85
+ )
86
+
87
+ bound = captured["seq_lens_cpu_upper_bound"]
88
+ assert isinstance(bound, torch.Tensor), (
89
+ "seq_lens_cpu_upper_bound must be a tensor, not None"
90
+ )
91
+ assert bound.shape == (4,), (
92
+ f"expected shape (num_reqs_padded=4,), got {bound.shape}"
93
+ )
94
+ assert bound.device.type == "cpu"
95
+ assert bound.dtype == torch.int32
96
+ # base[:num_reqs] + step, padded tail zeroed.
97
+ assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32))
98
+
99
+
100
+def test_build_draft_attn_metadata_handles_zero_unpadded_reqs():
101
+ """Edge case: when ``num_reqs == 0`` the upper-bound tensor must
102
+ still be a valid all-zero tensor of length ``num_reqs_padded``."""
103
+ fake = _make_fake_speculator()
104
+ base = torch.zeros(2, dtype=torch.int32)
105
+
106
+ captured = _run_build(
107
+ fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1
108
+ )
109
+
110
+ bound = captured["seq_lens_cpu_upper_bound"]
111
+ assert isinstance(bound, torch.Tensor)
112
+ assert bound.shape == (2,)
113
+ assert torch.equal(bound, torch.zeros(2, dtype=torch.int32))
114
+
115
+
116
+def test_build_draft_attn_metadata_clamps_to_max_model_len():
117
+ """The per-request upper bound (target bound + step) is clamped to the
118
+ model length so it never exceeds the allocated KV range."""
119
+ fake = _make_fake_speculator(max_model_len=1024)
120
+ base = torch.tensor([1023, 500], dtype=torch.int32)
121
+
122
+ captured = _run_build(
123
+ fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3
124
+ )
125
+
126
+ bound = captured["seq_lens_cpu_upper_bound"]
127
+ # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
128
+ assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+4
@@ -268,6 +268,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
268
dummy_run and skip_attn_for_dummy_run,
269
decode_batch_desc,
270
num_tokens_across_dp,
271
+ input_batch.seq_lens_cpu_upper_bound,
272
)
273
274
return self.draft_tokens[:num_reqs]
@@ -376,6 +377,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
377
skip_attn: bool,
378
batch_desc: BatchExecutionDescriptor,
379
num_tokens_across_dp: torch.Tensor | None,
380
+ seq_lens_cpu_upper_bound: torch.Tensor,
381
) -> None:
382
positions = self.input_buffers.positions[:num_reqs]
383
query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
@@ -400,6 +402,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
402
num_reqs=num_reqs,
403
num_reqs_padded=batch_desc.num_reqs or num_reqs,
404
num_tokens_padded=batch_desc.num_tokens,
405
+ seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
406
+ step=step,
407
)
408
409
# Update the current draft step.
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+6
@@ -269,6 +269,8 @@ class DFlashSpeculator(DraftModelSpeculator):
269
num_reqs: int,
270
num_reqs_padded: int,
271
num_tokens_padded: int,
272
+ seq_lens_cpu_upper_bound: torch.Tensor,
273
+ step: int,
274
num_query_per_req: int | None = None,
275
causal: bool | Mapping[int, bool] = False,
276
) -> dict[str, Any] | None:
@@ -279,6 +281,8 @@ class DFlashSpeculator(DraftModelSpeculator):
281
num_reqs,
282
num_reqs_padded,
283
num_tokens_padded,
284
+ seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
285
+ step=step,
286
num_query_per_req=self.num_query_per_req,
287
causal=causal,
288
)
@@ -427,6 +431,8 @@ class DFlashSpeculator(DraftModelSpeculator):
431
num_reqs=num_reqs,
432
num_reqs_padded=num_reqs_padded,
433
num_tokens_padded=num_tokens_padded,
434
+ seq_lens_cpu_upper_bound=input_batch.seq_lens_cpu_upper_bound,
435
+ step=self.num_query_per_req,
436
causal=self._group_causal,
437
)
438
draft_slot_mappings_by_layer = build_slot_mappings_by_layer(
vllm/v1/worker/gpu/spec_decode/speculator.py
+12
@@ -210,6 +210,8 @@ class DraftModelSpeculator(BaseSpeculator):
210
num_reqs: int,
211
num_reqs_padded: int,
212
num_tokens_padded: int,
213
+ seq_lens_cpu_upper_bound: torch.Tensor,
214
+ step: int,
215
num_query_per_req: int = 1,
216
causal: bool | Mapping[int, bool] = True,
217
) -> dict[str, Any] | None:
@@ -224,6 +226,15 @@ class DraftModelSpeculator(BaseSpeculator):
226
x[:num_reqs_padded] for x in self.block_tables.input_block_tables
227
]
228
slot_mappings = self.block_tables.slot_mappings[:, :num_tokens_padded]
229
+ draft_seq_lens_cpu_upper_bound = torch.zeros(
230
+ num_reqs_padded, dtype=torch.int32, device="cpu"
231
+ )
232
+ torch.add(
233
+ seq_lens_cpu_upper_bound[:num_reqs],
234
+ step,
235
+ out=draft_seq_lens_cpu_upper_bound[:num_reqs],
236
+ )
237
+ draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
238
attn_metadata = build_attn_metadata(
239
attn_groups=self.attn_groups,
240
num_reqs=num_reqs_padded,
@@ -239,6 +250,7 @@ class DraftModelSpeculator(BaseSpeculator):
250
slot_mappings=slot_mappings,
251
kv_cache_config=self.kv_cache_config,
252
causal=causal,
253
+ seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound,
254
)
255
return attn_metadata
256