| 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) |
| 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)) |