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 6, 17:00 UTC
[Bugfix] Fix GDN attention crash with mixed decode/spec-decode batches (#34871)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Mar 16, 2026 at 17:30 UTC
116ed130f4d323cb0fb088490b52197683e875a8
2 files changed
+201
tests/v1/attention/test_gdn_metadata_builder.py
new
+191
@@ -0,0 +1,191 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for GDNAttentionMetadataBuilder.build() — specifically the
4
+reclassification of non-spec decodes as prefills when spec decodes exist.
5
+Covers the fix for https://github.com/vllm-project/vllm/issues/34845.
6
+"""
7
+
8
+from dataclasses import dataclass
9
+
10
+import pytest
11
+import torch
12
+
13
+from tests.v1.attention.utils import (
14
+ BatchSpec,
15
+ create_common_attn_metadata,
16
+ create_vllm_config,
17
+)
18
+from vllm.config import SpeculativeConfig
19
+from vllm.v1.attention.backends.gdn_attn import (
20
+ GDNAttentionMetadata,
21
+ GDNAttentionMetadataBuilder,
22
+)
23
+from vllm.v1.kv_cache_interface import MambaSpec
24
+
25
+BLOCK_SIZE = 16
26
+DEVICE = torch.device("cpu")
27
+
28
+
29
+@dataclass
30
+class GDNBuildTestCase:
31
+ """Specification for a GDN metadata builder classification test."""
32
+
33
+ seq_lens: list[int]
34
+ query_lens: list[int]
35
+ num_decode_draft_tokens: list[int] | None # None = no spec config
36
+ num_speculative_tokens: int
37
+ expected_num_decodes: int
38
+ expected_num_prefills: int
39
+ expected_num_prefill_tokens: int
40
+ expected_num_spec_decodes: int
41
+
42
+
43
+GDN_BUILD_TEST_CASES = {
44
+ # The original #34845 crash: non-spec query_len=1 + spec decode
45
+ "mixed_decode_and_spec_decode": GDNBuildTestCase(
46
+ seq_lens=[65, 20],
47
+ query_lens=[1, 3],
48
+ num_decode_draft_tokens=[-1, 2],
49
+ num_speculative_tokens=2,
50
+ expected_num_decodes=0,
51
+ expected_num_prefills=1,
52
+ expected_num_prefill_tokens=1,
53
+ expected_num_spec_decodes=1,
54
+ ),
55
+ # All requests are spec decodes — no reclassification needed
56
+ "pure_spec_decode": GDNBuildTestCase(
57
+ seq_lens=[50, 30],
58
+ query_lens=[3, 3],
59
+ num_decode_draft_tokens=[2, 2],
60
+ num_speculative_tokens=2,
61
+ expected_num_decodes=0,
62
+ expected_num_prefills=0,
63
+ expected_num_prefill_tokens=0,
64
+ expected_num_spec_decodes=2,
65
+ ),
66
+ # No speculative config at all — standard decode path
67
+ "pure_regular_decode": GDNBuildTestCase(
68
+ seq_lens=[40, 30, 20],
69
+ query_lens=[1, 1, 1],
70
+ num_decode_draft_tokens=None,
71
+ num_speculative_tokens=0,
72
+ expected_num_decodes=3,
73
+ expected_num_prefills=0,
74
+ expected_num_prefill_tokens=0,
75
+ expected_num_spec_decodes=0,
76
+ ),
77
+ # Multi-token prefill alongside spec decode — no decode to reclassify
78
+ "spec_decode_with_real_prefill": GDNBuildTestCase(
79
+ seq_lens=[100, 20],
80
+ query_lens=[50, 3],
81
+ num_decode_draft_tokens=[-1, 2],
82
+ num_speculative_tokens=2,
83
+ expected_num_decodes=0,
84
+ expected_num_prefills=1,
85
+ expected_num_prefill_tokens=50,
86
+ expected_num_spec_decodes=1,
87
+ ),
88
+ # All three types in one batch — decode gets reclassified
89
+ "prefill_decode_and_spec_decode": GDNBuildTestCase(
90
+ seq_lens=[100, 65, 20],
91
+ query_lens=[50, 1, 3],
92
+ num_decode_draft_tokens=[-1, -1, 2],
93
+ num_speculative_tokens=2,
94
+ expected_num_decodes=0,
95
+ expected_num_prefills=2,
96
+ expected_num_prefill_tokens=51,
97
+ expected_num_spec_decodes=1,
98
+ ),
99
+ # Multiple non-spec query_len=1 requests all reclassified
100
+ "multiple_decodes_reclassified": GDNBuildTestCase(
101
+ seq_lens=[40, 50, 60, 20],
102
+ query_lens=[1, 1, 1, 3],
103
+ num_decode_draft_tokens=[-1, -1, -1, 2],
104
+ num_speculative_tokens=2,
105
+ expected_num_decodes=0,
106
+ expected_num_prefills=3,
107
+ expected_num_prefill_tokens=3,
108
+ expected_num_spec_decodes=1,
109
+ ),
110
+ # Zero-length padded sequence excluded from counts
111
+ "zero_length_padding_with_spec": GDNBuildTestCase(
112
+ seq_lens=[16, 65, 20],
113
+ query_lens=[0, 1, 3],
114
+ num_decode_draft_tokens=[-1, -1, 2],
115
+ num_speculative_tokens=2,
116
+ expected_num_decodes=0,
117
+ expected_num_prefills=1,
118
+ expected_num_prefill_tokens=1,
119
+ expected_num_spec_decodes=1,
120
+ ),
121
+}
122
+
123
+
124
+def _create_gdn_builder(
125
+ num_speculative_tokens: int = 0,
126
+) -> GDNAttentionMetadataBuilder:
127
+ """Create a GDNAttentionMetadataBuilder with minimal config."""
128
+ vllm_config = create_vllm_config(block_size=BLOCK_SIZE)
129
+ if num_speculative_tokens > 0:
130
+ vllm_config.speculative_config = SpeculativeConfig(
131
+ method="ngram",
132
+ num_speculative_tokens=num_speculative_tokens,
133
+ )
134
+ mamba_spec = MambaSpec(
135
+ block_size=BLOCK_SIZE,
136
+ shapes=((16, 64),),
137
+ dtypes=(torch.float16,),
138
+ )
139
+ return GDNAttentionMetadataBuilder(
140
+ kv_cache_spec=mamba_spec,
141
+ layer_names=["layer.0"],
142
+ vllm_config=vllm_config,
143
+ device=DEVICE,
144
+ )
145
+
146
+
147
+def _build(
148
+ builder: GDNAttentionMetadataBuilder,
149
+ batch_spec: BatchSpec,
150
+ num_decode_draft_tokens: list[int] | None = None,
151
+) -> GDNAttentionMetadata:
152
+ """Build GDN attention metadata, optionally with spec-decode kwargs."""
153
+ common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE)
154
+ kwargs: dict = {}
155
+ if num_decode_draft_tokens is not None:
156
+ kwargs["num_decode_draft_tokens_cpu"] = torch.tensor(
157
+ num_decode_draft_tokens, dtype=torch.int32
158
+ )
159
+ kwargs["num_accepted_tokens"] = torch.ones(
160
+ batch_spec.batch_size, dtype=torch.int32, device=DEVICE
161
+ )
162
+ return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs)
163
+
164
+
165
+@pytest.mark.parametrize(
166
+ "test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys()
167
+)
168
+def test_gdn_build_classification(test_case: GDNBuildTestCase):
169
+ """Test that GDN metadata builder classifies requests correctly."""
170
+ builder = _create_gdn_builder(test_case.num_speculative_tokens)
171
+ batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens)
172
+ meta = _build(builder, batch, test_case.num_decode_draft_tokens)
173
+
174
+ assert meta.num_decodes == test_case.expected_num_decodes
175
+ assert meta.num_prefills == test_case.expected_num_prefills
176
+ assert meta.num_prefill_tokens == test_case.expected_num_prefill_tokens
177
+ assert meta.num_spec_decodes == test_case.expected_num_spec_decodes
178
+
179
+
180
+def test_has_initial_state_after_reclassification():
181
+ """After reclassification, num_prefills > 0 so the prefill kernel path
182
+ should compute has_initial_state. For the reclassified request with
183
+ context_lens > 0, the corresponding entry must be True."""
184
+ builder = _create_gdn_builder(num_speculative_tokens=2)
185
+ batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3])
186
+ meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2])
187
+
188
+ assert meta.num_prefills > 0, "reclassification should produce prefills"
189
+ assert meta.has_initial_state is not None
190
+ # req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True
191
+ assert meta.has_initial_state[0].item() is True
vllm/v1/attention/backends/gdn_attn.py
+10
@@ -220,6 +220,16 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
220
query_lens_cpu.sum().item() - num_prefill_tokens - num_decode_tokens
221
)
222
223
+ # num_decodes and num_spec_decodes are mutually exclusive.
224
+ # Reclassify non-spec decodes as prefills when spec decodes
225
+ # exist — the prefill kernel handles 1-token sequences with
226
+ # initial state correctly, producing identical results.
227
+ if num_decodes > 0 and num_spec_decodes > 0:
228
+ num_prefills += num_decodes
229
+ num_prefill_tokens += num_decode_tokens
230
+ num_decodes = 0
231
+ num_decode_tokens = 0
232
+
233
if num_prefills == 0 and num_decodes == 0:
234
spec_token_size = min(
235
num_spec_decodes * (self.num_spec + 1),