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 2, 06:00 UTC
Fix DFlash prefix cache corruption due to missing lookahead block (#42971)
Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>
Shreyas Kulkarni committed
Jun 2, 2026 at 08:06 UTC
0eeba5eec17e9aba07eb804fbde2177d519c0524
2 files changed
+171
-3
tests/v1/spec_decode/test_dflash_lookahead.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import torch
7
+
8
+from tests.v1.core.utils import create_requests
9
+from vllm.config import (
10
+ CacheConfig,
11
+ ModelConfig,
12
+ ParallelConfig,
13
+ SchedulerConfig,
14
+ SpeculativeConfig,
15
+ VllmConfig,
16
+)
17
+from vllm.v1.core.sched.scheduler import Scheduler
18
+from vllm.v1.kv_cache_interface import (
19
+ FullAttentionSpec,
20
+ KVCacheConfig,
21
+ KVCacheGroupSpec,
22
+)
23
+from vllm.v1.structured_output import StructuredOutputManager
24
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
25
+
26
+# Matches defaults from tests/v1/spec_decode/test_eagle.py
27
+DFLASH_TARGET_DIR = "Qwen/Qwen3-8B"
28
+DFLASH_DRAFT_DIR = "z-lab/Qwen3-8B-DFlash-b16"
29
+
30
+BLOCK_SIZE = 16
31
+NUM_BLOCKS = 8
32
+NUM_SPECULATIVE_TOKENS = 3
33
+
34
+
35
+def _dflash_speculative_config(num_speculative_tokens: int) -> SpeculativeConfig:
36
+ model_config = ModelConfig(
37
+ model=DFLASH_TARGET_DIR,
38
+ runner="generate",
39
+ max_model_len=100,
40
+ trust_remote_code=True,
41
+ )
42
+ return SpeculativeConfig(
43
+ target_model_config=model_config,
44
+ target_parallel_config=ParallelConfig(),
45
+ model=DFLASH_DRAFT_DIR,
46
+ method="dflash",
47
+ num_speculative_tokens=num_speculative_tokens,
48
+ )
49
+
50
+
51
+def _create_dflash_scheduler(num_speculative_tokens: int) -> Scheduler:
52
+ speculative_config = _dflash_speculative_config(num_speculative_tokens)
53
+ model_config = speculative_config.target_model_config
54
+ scheduler_config = SchedulerConfig(
55
+ max_num_seqs=16,
56
+ max_num_batched_tokens=8192,
57
+ max_model_len=model_config.max_model_len,
58
+ is_encoder_decoder=model_config.is_encoder_decoder,
59
+ )
60
+ cache_config = CacheConfig(
61
+ block_size=BLOCK_SIZE,
62
+ gpu_memory_utilization=0.9,
63
+ cache_dtype="auto",
64
+ enable_prefix_caching=False,
65
+ )
66
+ vllm_config = VllmConfig(
67
+ scheduler_config=scheduler_config,
68
+ model_config=model_config,
69
+ cache_config=cache_config,
70
+ parallel_config=ParallelConfig(),
71
+ speculative_config=speculative_config,
72
+ )
73
+ kv_cache_config = KVCacheConfig(
74
+ num_blocks=NUM_BLOCKS,
75
+ kv_cache_tensors=[],
76
+ kv_cache_groups=[
77
+ KVCacheGroupSpec(
78
+ ["layer"],
79
+ FullAttentionSpec(
80
+ block_size=BLOCK_SIZE,
81
+ num_kv_heads=1,
82
+ head_size=1,
83
+ dtype=torch.float32,
84
+ ),
85
+ )
86
+ ],
87
+ )
88
+ cache_config.num_gpu_blocks = NUM_BLOCKS
89
+ return Scheduler(
90
+ vllm_config=vllm_config,
91
+ kv_cache_config=kv_cache_config,
92
+ block_size=BLOCK_SIZE,
93
+ log_stats=True,
94
+ structured_output_manager=StructuredOutputManager(vllm_config),
95
+ )
96
+
97
+
98
+def test_dflash_prefill_reserves_lookahead_blocks():
99
+ scheduler = _create_dflash_scheduler(NUM_SPECULATIVE_TOKENS)
100
+
101
+ assert scheduler.num_lookahead_tokens == NUM_SPECULATIVE_TOKENS + 1
102
+
103
+ (request,) = create_requests(
104
+ num_requests=1,
105
+ num_tokens=BLOCK_SIZE,
106
+ block_size=BLOCK_SIZE,
107
+ )
108
+ scheduler.add_request(request)
109
+
110
+ output = scheduler.schedule()
111
+
112
+ assert output.num_scheduled_tokens[request.request_id] == BLOCK_SIZE
113
+ # prefill block + one lookahead block
114
+ assert len(output.scheduled_new_reqs[0].block_ids[0]) == 2
115
+
116
+
117
+def test_dflash_first_prefill_query_window_fits_allocated_blocks():
118
+ scheduler = _create_dflash_scheduler(NUM_SPECULATIVE_TOKENS)
119
+
120
+ (request,) = create_requests(
121
+ num_requests=1,
122
+ num_tokens=BLOCK_SIZE,
123
+ block_size=BLOCK_SIZE,
124
+ )
125
+ scheduler.add_request(request)
126
+
127
+ output = scheduler.schedule()
128
+ block_ids = output.scheduled_new_reqs[0].block_ids[0]
129
+ query_positions = range(BLOCK_SIZE, BLOCK_SIZE + scheduler.num_lookahead_tokens)
130
+
131
+ assert all(pos // BLOCK_SIZE < len(block_ids) for pos in query_positions)
132
+
133
+
134
+def test_dflash_drafter_window_reserves_bonus_token():
135
+ # DFlash's drafter window is num_spec + 1 (the extra slot is the bonus token),
136
+ # so max_seq_len + num_spec + 1 must stay within the draft model's max len.
137
+ input_fits_in_drafter = GPUModelRunner._input_fits_in_drafter
138
+ dflash_runner = SimpleNamespace(
139
+ num_spec_tokens=NUM_SPECULATIVE_TOKENS,
140
+ effective_drafter_max_model_len=100,
141
+ speculative_config=_dflash_speculative_config(NUM_SPECULATIVE_TOKENS),
142
+ )
143
+ # window = 4, so 96 fits (96 + 4 == 100) but 97 does not (97 + 4 == 101)
144
+ assert input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=96))
145
+ assert not input_fits_in_drafter(dflash_runner, SimpleNamespace(max_seq_len=97))
146
+ assert not input_fits_in_drafter(dflash_runner, None) # no metadata
147
+
148
+ # Other drafters don't reserve the bonus token, so 97 fits (97 + 3 == 100).
149
+ plain_runner = SimpleNamespace(
150
+ num_spec_tokens=NUM_SPECULATIVE_TOKENS,
151
+ effective_drafter_max_model_len=100,
152
+ speculative_config=SimpleNamespace(use_dflash=lambda: False),
153
+ )
154
+ assert input_fits_in_drafter(plain_runner, SimpleNamespace(max_seq_len=97))
vllm/v1/worker/gpu_model_runner.py
+17
-3
@@ -4362,6 +4362,21 @@ class GPUModelRunner(
4362
4363
return None
4364
4365
+ def _input_fits_in_drafter(
4366
+ self, common_attn_metadata: CommonAttentionMetadata | None
4367
+ ) -> bool:
4368
+ if common_attn_metadata is None:
4369
+ return False
4370
+ assert self.speculative_config is not None
4371
+ # DFlash queries one extra token (the bonus token) beyond num_spec_tokens
4372
+ num_drafter_query_tokens = self.num_spec_tokens + (
4373
+ 1 if self.speculative_config.use_dflash() else 0
4374
+ )
4375
+ return (
4376
+ common_attn_metadata.max_seq_len + num_drafter_query_tokens
4377
+ <= self.effective_drafter_max_model_len
4378
+ )
4379
+
4380
@torch.inference_mode
4381
def sample_tokens(
4382
self, grammar_output: "GrammarOutput | None"
@@ -4441,9 +4456,8 @@ class GPUModelRunner(
4456
propose_drafts_after_bookkeeping = False
4457
if spec_config is not None:
4458
# Decide whether to run the drafter or zero out draft tokens.
4444
- input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
4445
- spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
4446
- <= self.effective_drafter_max_model_len
4459
+ input_fits_in_drafter = self._input_fits_in_drafter(
4460
+ spec_decode_common_attn_metadata
4461
)
4462
use_gpu_toks = (
4463
spec_config.use_eagle()