148 lines
6.0 KiB
Diff
148 lines
6.0 KiB
Diff
From 108e2c9e5b62d1990796ee3f710d39f5063b3759 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
Date: Tue, 22 Sep 2026 09:11:49 +0800
|
|
Subject: [PATCH 061/103] gdn: fix stateless first-chunk classification (port
|
|
of upstream #51565)
|
|
|
|
Port of the upstream vLLM change #51565 (Apache-2.0). A Gated DeltaNet prefill whose first chunk carries no prior state is
|
|
classified correctly.
|
|
---
|
|
.../v1/attention/test_gdn_metadata_builder.py | 83 +++++++++++++++++++
|
|
vllm/v1/attention/backends/gdn_attn.py | 27 +++++-
|
|
2 files changed, 109 insertions(+), 1 deletion(-)
|
|
|
|
diff --git a/tests/v1/attention/test_gdn_metadata_builder.py b/tests/v1/attention/test_gdn_metadata_builder.py
|
|
index 221f933..b4d3acf 100644
|
|
--- a/tests/v1/attention/test_gdn_metadata_builder.py
|
|
+++ b/tests/v1/attention/test_gdn_metadata_builder.py
|
|
@@ -221,3 +221,86 @@ def test_full_cudagraph_spec_metadata_uses_request_count():
|
|
assert meta.spec_query_start_loc.shape == (batch.batch_size + 1,)
|
|
assert meta.num_accepted_tokens is not None
|
|
assert meta.num_accepted_tokens.shape == (batch.batch_size,)
|
|
+
|
|
+
|
|
+def _build_non_spec(
|
|
+ batch: BatchSpec,
|
|
+ is_prefilling: list[bool] | None,
|
|
+ full_cuda_graph: bool = False,
|
|
+):
|
|
+ common_attn_metadata = create_common_attn_metadata(
|
|
+ batch, BLOCK_SIZE, DEVICE
|
|
+ ).replace(
|
|
+ is_prefilling=None
|
|
+ if is_prefilling is None
|
|
+ else torch.tensor(is_prefilling, dtype=torch.bool)
|
|
+ )
|
|
+ builder = _create_gdn_builder(full_cuda_graph=full_cuda_graph)
|
|
+ return builder, common_attn_metadata, builder.build(0, common_attn_metadata)
|
|
+
|
|
+
|
|
+@pytest.mark.parametrize(
|
|
+ ("seq_len", "query_len", "is_prefilling", "num_prefills"),
|
|
+ [
|
|
+ pytest.param(1, 1, True, 1, id="first-chunk"),
|
|
+ pytest.param(65, 1, True, 0, id="resumed-chunk"),
|
|
+ pytest.param(0, 0, True, 0, id="padding"),
|
|
+ pytest.param(1, 1, None, 0, id="missing-prefill-flag"),
|
|
+ ],
|
|
+)
|
|
+def test_one_token_chunk_classification(
|
|
+ seq_len: int,
|
|
+ query_len: int,
|
|
+ is_prefilling: bool | None,
|
|
+ num_prefills: int,
|
|
+):
|
|
+ """Only a real first chunk with a prefill flag needs state initialization."""
|
|
+ _, _, meta = _build_non_spec(
|
|
+ BatchSpec(seq_lens=[100, seq_len], query_lens=[1, query_len]),
|
|
+ is_prefilling=None if is_prefilling is None else [False, is_prefilling],
|
|
+ )
|
|
+
|
|
+ assert meta.num_prefills == num_prefills
|
|
+ assert meta.num_decodes == 2 - num_prefills
|
|
+ assert meta.num_prefill_tokens == num_prefills
|
|
+ assert meta.num_decode_tokens == 1 + query_len - num_prefills
|
|
+ if num_prefills:
|
|
+ assert meta.has_initial_state is not None
|
|
+ assert meta.has_initial_state.tolist() == [True, False]
|
|
+ else:
|
|
+ assert meta.has_initial_state is None
|
|
+
|
|
+
|
|
+def test_one_token_first_chunk_excludes_padding():
|
|
+ """Neither padding requests nor padding tokens count as prefill work."""
|
|
+ common = create_common_attn_metadata(
|
|
+ BatchSpec(seq_lens=[100, 1, 0, 0], query_lens=[1, 1, 0, 0]),
|
|
+ BLOCK_SIZE,
|
|
+ DEVICE,
|
|
+ ).replace(
|
|
+ is_prefilling=torch.tensor([False, True, False, False], dtype=torch.bool),
|
|
+ num_actual_tokens=4,
|
|
+ )
|
|
+ meta = _create_gdn_builder().build(0, common)
|
|
+
|
|
+ assert meta.num_decodes == 1
|
|
+ assert meta.num_prefills == 1
|
|
+ assert meta.num_decode_tokens == 1
|
|
+ assert meta.num_prefill_tokens == 1
|
|
+
|
|
+
|
|
+def test_cudagraph_capture_batch_stays_decode_only():
|
|
+ """Capture rows have no history, but must still select decode kernels."""
|
|
+ batch = BatchSpec(seq_lens=[1] * 4, query_lens=[1] * 4)
|
|
+ builder, common_attn_metadata, _ = _build_non_spec(
|
|
+ batch, [False] * 4, full_cuda_graph=True
|
|
+ )
|
|
+ meta = builder.build_for_cudagraph_capture(common_attn_metadata)
|
|
+
|
|
+ assert meta.num_prefills == 0
|
|
+ assert meta.num_decodes == 4
|
|
+ assert meta.has_initial_state is None
|
|
+ staged = meta.non_spec_state_indices_tensor
|
|
+ assert staged is not None
|
|
+ assert staged.data_ptr() == builder.non_spec_state_indices_tensor.data_ptr()
|
|
+ torch.testing.assert_close(staged, common_attn_metadata.block_table_tensor[:, 0])
|
|
diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py
|
|
index 23daff5..dfc04a8 100644
|
|
--- a/vllm/v1/attention/backends/gdn_attn.py
|
|
+++ b/vllm/v1/attention/backends/gdn_attn.py
|
|
@@ -247,9 +247,34 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
|
|
)
|
|
|
|
if spec_sequence_masks is None:
|
|
+ # V2 already excludes prefills from full decode graphs via
|
|
+ # has_prefill. Classify first chunks as prefills to mask recycled
|
|
+ # state; resumed one-token chunks can still use the decode kernels.
|
|
+ assert m.seq_lens_cpu_upper_bound is not None
|
|
+ query_lens_cpu = query_start_loc_cpu.diff()
|
|
+ no_prior_state = (query_lens_cpu > 0) & (
|
|
+ m.seq_lens_cpu_upper_bound <= query_lens_cpu
|
|
+ )
|
|
+ # Capture batches also have seq_len == query_len, but are not
|
|
+ # prefills.
|
|
+ if m.is_prefilling is not None:
|
|
+ no_prior_state &= m.is_prefilling
|
|
+ else:
|
|
+ no_prior_state = torch.zeros_like(no_prior_state)
|
|
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
|
|
- split_decodes_and_prefills(m, decode_threshold=1)
|
|
+ split_decodes_and_prefills(
|
|
+ m.replace(is_prefilling=no_prior_state),
|
|
+ decode_threshold=1,
|
|
+ treat_short_extends_as_decodes=False,
|
|
+ )
|
|
)
|
|
+ # Exclude trailing padding from both prefill counts.
|
|
+ if num_prefills:
|
|
+ num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum())
|
|
+ num_prefill_tokens = (
|
|
+ int(query_start_loc_cpu[num_decodes + num_prefills])
|
|
+ - num_decode_tokens
|
|
+ )
|
|
num_spec_decode_tokens = 0
|
|
spec_token_indx = None
|
|
non_spec_token_indx = None
|
|
--
|
|
2.43.0
|
|
|