[Bugfix][GDN] Fix stateless first-chunk classification (#51565)
Signed-off-by: taking-lying-flat <1615405@qq.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
cherry77-cloud committed
Sep 22, 2026 at 09:11 UTC
79468c20ef23227d4e051f807e10fd54fb24e24f
2 files changed
+109
-1
tests/v1/attention/test_gdn_metadata_builder.py
+83
@@ -221,3 +221,86 @@ def test_full_cudagraph_spec_metadata_uses_request_count():
221
assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,)
222
assert meta.num_accepted_tokens is not None
223
assert meta.num_accepted_tokens.shape == (batch.batch_size,)
224
+
225
+
226
+def _build_non_spec(
227
+ batch: BatchSpec,
228
+ is_prefilling: list[bool] | None,
229
+ full_cuda_graph: bool = False,
230
+):
231
+ common_attn_metadata = create_common_attn_metadata(
232
+ batch, BLOCK_SIZE, DEVICE
233
+ ).replace(
234
+ is_prefilling=None
235
+ if is_prefilling is None
236
+ else torch.tensor(is_prefilling, dtype=torch.bool)
237
+ )
238
+ builder = _create_gdn_builder(full_cuda_graph=full_cuda_graph)
239
+ return builder, common_attn_metadata, builder.build(0, common_attn_metadata)
240
+
241
+
242
+@pytest.mark.parametrize(
243
+ ("seq_len", "query_len", "is_prefilling", "num_prefills"),
244
+ [
245
+ pytest.param(1, 1, True, 1, id="first-chunk"),
246
+ pytest.param(65, 1, True, 0, id="resumed-chunk"),
247
+ pytest.param(0, 0, True, 0, id="padding"),
248
+ pytest.param(1, 1, None, 0, id="missing-prefill-flag"),
249
+ ],
250
+)
251
+def test_one_token_chunk_classification(
252
+ seq_len: int,
253
+ query_len: int,
254
+ is_prefilling: bool | None,
255
+ num_prefills: int,
256
+):
257
+ """Only a real first chunk with a prefill flag needs state initialization."""
258
+ _, _, meta = _build_non_spec(
259
+ BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]),
260
+ is_prefilling=None if is_prefilling is None else [False, is_prefilling],
261
+ )
262
+
263
+ assert meta.num_prefills == num_prefills
264
+ assert meta.num_decodes == 2 - num_prefills
265
+ assert meta.num_prefill_tokens == num_prefills
266
+ assert meta.num_decode_tokens == 1 + query_len - num_prefills
267
+ if num_prefills:
268
+ assert meta.has_initial_state is not None
269
+ assert meta.has_initial_state.tolist() == [True, False]
270
+ else:
271
+ assert meta.has_initial_state is None
272
+
273
+
274
+def test_one_token_first_chunk_excludes_padding():
275
+ """Neither padding requests nor padding tokens count as prefill work."""
276
+ common = create_common_attn_metadata(
277
+ BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]),
278
+ BLOCK_SIZE,
279
+ DEVICE,
280
+ ).replace(
281
+ is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool),
282
+ num_actual_tokens=4,
283
+ )
284
+ meta = _create_gdn_builder().build(0, common)
285
+
286
+ assert meta.num_decodes == 1
287
+ assert meta.num_prefills == 1
288
+ assert meta.num_decode_tokens == 1
289
+ assert meta.num_prefill_tokens == 1
290
+
291
+
292
+def test_cudagraph_capture_batch_stays_decode_only():
293
+ """Capture rows have no history, but must still select decode kernels."""
294
+ batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4)
295
+ builder, common_attn_metadata, _ = _build_non_spec(
296
+ batch, [False] * 4, full_cuda_graph=True
297
+ )
298
+ meta = builder.build_for_cudagraph_capture(common_attn_metadata)
299
+
300
+ assert meta.num_prefills == 0
301
+ assert meta.num_decodes == 4
302
+ assert meta.has_initial_state is None
303
+ staged = meta.non_spec_state_indices_tensor
304
+ assert staged is not None
305
+ assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()
306
+ torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
vllm/v1/attention/backends/gdn_attn.py
+26
-1
@@ -247,9 +247,34 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
247
)
248
249
if spec_sequence_masks is None:
250
+ # V2 already excludes prefills from full decode graphs via
251
+ # has_prefill. Classify first chunks as prefills to mask recycled
252
+ # state; resumed one-token chunks can still use the decode kernels.
253
+ assert m.seq_lens_cpu_upper_bound is not None
254
+ query_lens_cpu = query_start_loc_cpu.diff()
255
+ no_prior_state = (query_lens_cpu > 0) & (
256
+ m.seq_lens_cpu_upper_bound <= query_lens_cpu
257
+ )
258
+ # Capture batches also have seq_len == query_len, but are not
259
+ # prefills.
260
+ if m.is_prefilling is not None:
261
+ no_prior_state &= m.is_prefilling
262
+ else:
263
+ no_prior_state = torch.zeros_like(no_prior_state)
264
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
251
- split_decodes_and_prefills(m, decode_threshold=1)
265
+ split_decodes_and_prefills(
266
+ m.replace(is_prefilling=no_prior_state),
267
+ decode_threshold=1,
268
+ treat_short_extends_as_decodes=False,
269
+ )
270
)
271
+ # Exclude trailing padding from both prefill counts.
272
+ if num_prefills:
273
+ num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum())
274
+ num_prefill_tokens = (
275
+ int(query_start_loc_cpu[num_decodes + num_prefills])
276
+ - num_decode_tokens
277
+ )
278
num_spec_decode_tokens = 0
279
spec_token_indx = None
280
non_spec_token_indx = None