[Bugfix] Avoid racy accepted counts in async spec decode (#45100)
Signed-off-by: Weiwei Sun <68775773+sunnweiwei@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Weiwei Sun committed
Jun 22, 2026 at 01:53 UTC
cec2ec11760f9f3beabd4c90451936078bf91533
3 files changed
+46
-5
tests/v1/attention/test_gdn_metadata_builder.py
+33
-1
@@ -16,6 +16,7 @@ from tests.v1.attention.utils import (
16
create_vllm_config,
17
)
18
from vllm.config import SpeculativeConfig
19
+from vllm.config.compilation import CUDAGraphMode
20
from vllm.v1.attention.backends.gdn_attn import (
21
GDNAttentionMetadata,
22
GDNAttentionMetadataBuilder,
@@ -123,9 +124,15 @@ GDN_BUILD_TEST_CASES = {
124
125
def _create_gdn_builder(
126
num_speculative_tokens: int = 0,
127
+ full_cuda_graph: bool = False,
128
) -> GDNAttentionMetadataBuilder:
129
"""Create a GDNAttentionMetadataBuilder with minimal config."""
128
- vllm_config = create_vllm_config(block_size=BLOCK_SIZE)
130
+ vllm_config = create_vllm_config(
131
+ model_name="Qwen/Qwen3.5-0.8B",
132
+ block_size=BLOCK_SIZE,
133
+ )
134
+ if full_cuda_graph:
135
+ vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
136
if num_speculative_tokens > 0:
137
vllm_config.speculative_config = SpeculativeConfig(
138
method="ngram",
@@ -189,3 +196,28 @@ def test_has_initial_state_after_reclassification():
196
assert meta.has_initial_state is not None
197
# req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True
198
assert meta.has_initial_state[0].item() is True
199
+
200
+
201
+def test_full_cudagraph_spec_metadata_uses_request_count():
202
+ """FULL cudagraph token padding must not pad request-indexed metadata."""
203
+ num_speculative_tokens = 3
204
+ builder = _create_gdn_builder(
205
+ num_speculative_tokens=num_speculative_tokens,
206
+ full_cuda_graph=True,
207
+ )
208
+ batch = BatchSpec(seq_lens=[80, 96], query_lens=[4, 4])
209
+ meta = _build(builder, batch, num_decode_draft_tokens=[3, 3])
210
+
211
+ assert meta.num_spec_decodes == batch.batch_size
212
+ assert meta.num_spec_decode_tokens == batch.compute_num_tokens()
213
+ assert meta.spec_state_indices_tensor is not None
214
+ assert meta.spec_state_indices_tensor.shape == (
215
+ batch.batch_size,
216
+ num_speculative_tokens + 1,
217
+ )
218
+ assert meta.spec_sequence_masks is not None
219
+ assert meta.spec_sequence_masks.shape == (batch.batch_size,)
220
+ assert meta.spec_query_start_loc is not None
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,)
vllm/v1/attention/backends/gdn_attn.py
+4
-3
@@ -410,9 +410,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
410
f"num_decodes: {num_decodes}, num_spec_decodes: {num_spec_decodes}"
411
)
412
413
- # Prepare tensors for cudagraph
414
- # Note: m.num_actual_tokens is already padded by the model runner for CUDAGraph
415
- batch_size = m.num_actual_tokens
413
+ # Prepare per-request tensors for cudagraph. m.num_actual_tokens is
414
+ # token-padded for FULL graph replay, but the GDN state/query/accepted
415
+ # metadata below is indexed by request.
416
+ batch_size = m.num_reqs
417
418
if (
419
self.use_full_cuda_graph
vllm/v1/worker/gpu_model_runner.py
+9
-1
@@ -2035,7 +2035,13 @@ class GPUModelRunner(
2035
2036
# Sync num_accepted_tokens from CPU (set by
2037
# _update_states_after_model_execute for hybrid models).
2038
- if self.num_accepted_tokens_event is not None:
2038
+ # Skipped under async scheduling (non-align): the CPU copy races with
2039
+ # the in-flight D2H copy and with input-batch row moves.
2040
+ needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and not (
2041
+ self.use_async_scheduling and self.cache_config.mamba_cache_mode != "align"
2042
+ )
2043
+ if needs_cpu_accepted_counts:
2044
+ assert self.num_accepted_tokens_event is not None
2045
self.num_accepted_tokens_event.synchronize()
2046
# Async mode: condense() reordered indices, use prev_positions mapping
2047
if self.use_async_scheduling and prev_req_id_to_index:
@@ -2058,6 +2064,8 @@ class GPUModelRunner(
2064
self.num_accepted_tokens.np[num_reqs:].fill(1)
2065
self.num_accepted_tokens.copy_to_gpu()
2066
else:
2067
+ # Default to 1; update_num_computed_tokens_for_batch_change below
2068
+ # corrects rows that had drafts from valid_sampled_token_count.
2069
self.num_accepted_tokens.np.fill(1)
2070
self.num_accepted_tokens.gpu.fill_(1)
2071