| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Literal |
| 5 | from unittest import mock |
| 6 | |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.utils import get_attn_backend_list_based_on_platform |
| 12 | from tests.v1.attention.utils import ( |
| 13 | BatchSpec, |
| 14 | create_common_attn_metadata, |
| 15 | create_standard_kv_cache_spec, |
| 16 | try_get_attention_backend, |
| 17 | ) |
| 18 | from vllm.config import ( |
| 19 | AttentionConfig, |
| 20 | CacheConfig, |
| 21 | DeviceConfig, |
| 22 | ModelConfig, |
| 23 | ParallelConfig, |
| 24 | SchedulerConfig, |
| 25 | SpeculativeConfig, |
| 26 | VllmConfig, |
| 27 | ) |
| 28 | from vllm.config.load import LoadConfig |
| 29 | from vllm.model_executor.models.llama import LlamaForCausalLM |
| 30 | from vllm.platforms import current_platform |
| 31 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 32 | from vllm.v1.spec_decode.dflash import DFlashProposer |
| 33 | from vllm.v1.spec_decode.draft_model import DraftModelProposer |
| 34 | from vllm.v1.spec_decode.eagle import EagleProposer |
| 35 | from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer |
| 36 | from vllm.v1.spec_decode.metadata import SpecDecodeMetadata |
| 37 | from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch |
| 38 | |
| 39 | model_dir = "meta-llama/Llama-3.1-8B-Instruct" |
| 40 | eagle_dir = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B" |
| 41 | eagle3_dir = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" |
| 42 | ar_draft_model_dir = "amd/PARD-Llama-3.2-1B" # Compatible with parallel and AR drafting |
| 43 | dflash_target_dir = "Qwen/Qwen3-8B" |
| 44 | dflash_dir = "z-lab/Qwen3-8B-DFlash-b16" |
| 45 | |
| 46 | BLOCK_SIZE = 16 |
| 47 | DEVICE_TYPE = current_platform.device_type |
| 48 | |
| 49 | |
| 50 | def _create_proposer( |
| 51 | method: Literal["eagle", "eagle3", "draft_model", "dflash"], |
| 52 | num_speculative_tokens: int, |
| 53 | attention_backend: str | None = None, |
| 54 | parallel_drafting: bool = False, |
| 55 | rejection_sample_method: Literal["standard", "synthetic", "block"] = "standard", |
| 56 | draft_sample_method: Literal["greedy", "probabilistic"] = "greedy", |
| 57 | ) -> SpecDecodeBaseProposer: |
| 58 | # Method-dependent setup |
| 59 | if method == "eagle": |
| 60 | target_model_dir = model_dir |
| 61 | draft_model_dir = eagle_dir |
| 62 | elif method == "eagle3": |
| 63 | target_model_dir = model_dir |
| 64 | draft_model_dir = eagle3_dir |
| 65 | elif method == "draft_model": |
| 66 | target_model_dir = model_dir |
| 67 | draft_model_dir = ar_draft_model_dir |
| 68 | elif method == "dflash": |
| 69 | target_model_dir = dflash_target_dir |
| 70 | draft_model_dir = dflash_dir |
| 71 | else: |
| 72 | raise ValueError(f"Unknown method: {method}") |
| 73 | |
| 74 | model_config = ModelConfig( |
| 75 | model=target_model_dir, |
| 76 | runner="generate", |
| 77 | max_model_len=100, |
| 78 | trust_remote_code=(method == "dflash"), |
| 79 | ) |
| 80 | |
| 81 | speculative_config = SpeculativeConfig( |
| 82 | target_model_config=model_config, |
| 83 | target_parallel_config=ParallelConfig(), |
| 84 | model=draft_model_dir, |
| 85 | method=method, |
| 86 | num_speculative_tokens=num_speculative_tokens, |
| 87 | parallel_drafting=parallel_drafting, |
| 88 | rejection_sample_method=rejection_sample_method, |
| 89 | draft_sample_method=draft_sample_method, |
| 90 | ) |
| 91 | if parallel_drafting: |
| 92 | # Overwrite pard_token to avoid crash during init |
| 93 | speculative_config.draft_model_config.hf_config.pard_token = 0 |
| 94 | |
| 95 | device = DEVICE_TYPE |
| 96 | vllm_config = VllmConfig( |
| 97 | model_config=model_config, |
| 98 | cache_config=CacheConfig(block_size=16), |
| 99 | speculative_config=speculative_config, |
| 100 | device_config=DeviceConfig(device=device), |
| 101 | parallel_config=ParallelConfig(), |
| 102 | load_config=LoadConfig(), |
| 103 | scheduler_config=SchedulerConfig( |
| 104 | max_model_len=model_config.max_model_len, |
| 105 | is_encoder_decoder=model_config.is_encoder_decoder, |
| 106 | ), |
| 107 | attention_config=AttentionConfig( |
| 108 | backend=AttentionBackendEnum[attention_backend] |
| 109 | if attention_backend |
| 110 | else None |
| 111 | ), |
| 112 | ) |
| 113 | |
| 114 | proposer: SpecDecodeBaseProposer |
| 115 | if method == "dflash": |
| 116 | proposer = DFlashProposer(vllm_config=vllm_config, device=device) |
| 117 | elif "eagle" in method: |
| 118 | proposer = EagleProposer(vllm_config=vllm_config, device=device) |
| 119 | else: |
| 120 | proposer = DraftModelProposer(vllm_config=vllm_config, device=device) |
| 121 | proposer.block_size = BLOCK_SIZE |
| 122 | return proposer |
| 123 | |
| 124 | |
| 125 | def test_prepare_next_token_ids(): |
| 126 | """Test for prepare_next_token_ids_cpu and prepare_next_token_ids_padded. |
| 127 | Each will produce a device tensor of next_token_ids, taking as input |
| 128 | either the GPU tensor of sampled_token_ids with -1 for rejected tokens, |
| 129 | or the CPU python list[list[int]] with the rejected tokens removed. |
| 130 | """ |
| 131 | device = torch.device(DEVICE_TYPE) |
| 132 | |
| 133 | num_requests = 4 |
| 134 | num_speculative_tokens = 4 |
| 135 | req_ids = [f"req_{i + 1}" for i in range(num_requests)] |
| 136 | mock_input_batch = mock.MagicMock(spec=InputBatch) |
| 137 | mock_input_batch.req_ids = req_ids |
| 138 | mock_input_batch.num_reqs = num_requests |
| 139 | mock_input_batch.vocab_size = 100 |
| 140 | mock_input_batch.num_tokens_no_spec = np.array( |
| 141 | [num_speculative_tokens + 1] * num_requests |
| 142 | ) |
| 143 | |
| 144 | mock_num_scheduled_tokens = {req_id: 0 for req_id in req_ids} |
| 145 | mock_requests: dict[str, CachedRequestState] = {} |
| 146 | for req_id in req_ids: |
| 147 | mock_request = mock.MagicMock(spec=CachedRequestState) |
| 148 | # Each request will have a backup next token id of 10, 20, 30, 40 |
| 149 | mock_request.get_token_id.return_value = int(req_id.split("_")[1]) * 10 |
| 150 | mock_request.num_computed_tokens = 0 |
| 151 | mock_requests[req_id] = mock_request |
| 152 | |
| 153 | # explicitly discard the last request |
| 154 | discarded_req_mask = torch.tensor( |
| 155 | [False, False, False, True], dtype=torch.bool, device=device |
| 156 | ) |
| 157 | sampled_token_ids = [ |
| 158 | [0, 1, -1, -1, -1], # 1 accepted, 3 rejected, "1" sampled |
| 159 | [0, 1, 2, 3, 4], # all accepted, "4" sampled |
| 160 | [-1, -1, -1, -1, -1], # sampling skipped, use backup token "30" |
| 161 | [0, 1, 2, -1, -1], # explicitly discarded, sampling should be ignored |
| 162 | ] |
| 163 | sampled_token_ids_tensor = torch.tensor( |
| 164 | sampled_token_ids, dtype=torch.int32, device=device |
| 165 | ) |
| 166 | sampled_token_ids_cpu = [[i for i in seq if i != -1] for seq in sampled_token_ids] |
| 167 | for i in range(len(sampled_token_ids_cpu)): |
| 168 | if discarded_req_mask[i]: |
| 169 | sampled_token_ids_cpu[i] = [] |
| 170 | |
| 171 | expected_next_token_ids_cpu = [1, 4, 30, 40] |
| 172 | expected_next_token_ids_tensor = torch.tensor( |
| 173 | expected_next_token_ids_cpu, dtype=torch.int32, device=device |
| 174 | ) |
| 175 | |
| 176 | proposer = _create_proposer("eagle", num_speculative_tokens) |
| 177 | |
| 178 | next_token_ids_from_cpu = proposer.prepare_next_token_ids_cpu( |
| 179 | sampled_token_ids_cpu, |
| 180 | mock_requests, |
| 181 | mock_input_batch, |
| 182 | mock_num_scheduled_tokens, |
| 183 | ) |
| 184 | |
| 185 | assert torch.equal(next_token_ids_from_cpu, expected_next_token_ids_tensor) |
| 186 | |
| 187 | expected_valid_sampled_tokens_count = torch.tensor( |
| 188 | [2, 5, 0, 0], dtype=torch.int32, device=device |
| 189 | ) |
| 190 | |
| 191 | next_token_ids_from_padded, valid_sampled_tokens_count = ( |
| 192 | proposer.prepare_next_token_ids_padded( |
| 193 | sampled_token_ids_tensor, |
| 194 | mock_requests, |
| 195 | mock_input_batch, |
| 196 | discarded_req_mask, |
| 197 | ) |
| 198 | ) |
| 199 | |
| 200 | assert torch.equal(next_token_ids_from_padded, expected_next_token_ids_tensor) |
| 201 | assert torch.equal(valid_sampled_tokens_count, expected_valid_sampled_tokens_count) |
| 202 | |
| 203 | |
| 204 | def test_prepare_inputs(): |
| 205 | """cu_target_query_lens: [0, a, a + b, a + b + c] |
| 206 | num_rejected_tokens: [n1, n2, n3] |
| 207 | num_tokens_per_req: [a - n1, b - n2, c - n3] |
| 208 | cu_num_tokens: [0, a - n1, a + b - n1 - n2, a + b + c - n1 - n2 - n3] |
| 209 | token_indices: [0, 1, ..., a - n1 - 1, |
| 210 | a, a + 1, ..., a + b - n2 - 1, |
| 211 | a + b, a + b + 1, ..., a + b + c - n3 - 1] |
| 212 | """ |
| 213 | device = torch.device(DEVICE_TYPE) |
| 214 | |
| 215 | # q1 = 4, q2 = 7, q3 = 5 |
| 216 | # n1 = 1, n2 = 3, n3 = 2 |
| 217 | |
| 218 | batch_spec = BatchSpec( |
| 219 | seq_lens=[4, 7, 5], |
| 220 | query_lens=[4, 7, 5], |
| 221 | ) |
| 222 | |
| 223 | common_attn_metadata = create_common_attn_metadata( |
| 224 | batch_spec, |
| 225 | block_size=BLOCK_SIZE, |
| 226 | device=device, |
| 227 | ) |
| 228 | |
| 229 | # If there are `k` sampled tokens, then `k-1` tokens are draft tokens |
| 230 | # from the previous iteration, and the last token is the bonus token sampled |
| 231 | # from the base model. |
| 232 | num_draft_tokens = [3, 6, 4] # one less than query_lens |
| 233 | # num rejected tokens is [1, 3, 2] |
| 234 | ACCEPT_TOKEN = 0 |
| 235 | BONUS_TOKEN = 1 |
| 236 | REJECT_TOKEN = -1 |
| 237 | sampled_token_ids = [ |
| 238 | [ACCEPT_TOKEN, ACCEPT_TOKEN, REJECT_TOKEN, BONUS_TOKEN], |
| 239 | [ |
| 240 | ACCEPT_TOKEN, |
| 241 | ACCEPT_TOKEN, |
| 242 | ACCEPT_TOKEN, |
| 243 | REJECT_TOKEN, |
| 244 | REJECT_TOKEN, |
| 245 | REJECT_TOKEN, |
| 246 | BONUS_TOKEN, |
| 247 | ], |
| 248 | [ACCEPT_TOKEN, ACCEPT_TOKEN, REJECT_TOKEN, REJECT_TOKEN, BONUS_TOKEN], |
| 249 | ] |
| 250 | sampled_token_ids = [ |
| 251 | [i for i in seq if i != REJECT_TOKEN] for seq in sampled_token_ids |
| 252 | ] |
| 253 | |
| 254 | # Expected calculations: |
| 255 | # query_len_per_req = [4, 7, 5] |
| 256 | # num_tokens_per_req = [3, 4, 3] (after subtracting rejected tokens) |
| 257 | # Expected cumulative counts: [0, 3, 7, 10] |
| 258 | expected_cu_num_tokens = torch.tensor( |
| 259 | [0, 3, 7, 10], dtype=torch.int32, device=device |
| 260 | ) |
| 261 | |
| 262 | # Expected token indices (mapped from original positions): |
| 263 | # First request: indices 0, 1, 2 (keeping first 3 from positions 0-3) |
| 264 | # Second request: indices 4, 5, 6, 7 (keeping first 4 from positions 4-10) |
| 265 | # Third request: indices 11, 12, 13 (keeping first 3 from positions 11-15) |
| 266 | expected_token_indices = torch.tensor( |
| 267 | [ |
| 268 | 0, |
| 269 | 1, |
| 270 | 2, # First request: 3 tokens (4-1) |
| 271 | 4, |
| 272 | 5, |
| 273 | 6, |
| 274 | 7, # Second request: 4 tokens (7-3) |
| 275 | 11, |
| 276 | 12, |
| 277 | 13, # Third request: 3 tokens (5-2) |
| 278 | ], |
| 279 | dtype=torch.int32, |
| 280 | device=device, |
| 281 | ) |
| 282 | proposer = _create_proposer("eagle", 1) |
| 283 | |
| 284 | updated_metadata, token_indices = proposer.prepare_inputs( |
| 285 | common_attn_metadata, sampled_token_ids, num_draft_tokens |
| 286 | ) |
| 287 | |
| 288 | assert torch.equal(updated_metadata.query_start_loc, expected_cu_num_tokens) |
| 289 | assert token_indices.shape[0] == expected_cu_num_tokens[-1].item() |
| 290 | assert torch.equal(token_indices, expected_token_indices) |
| 291 | |
| 292 | |
| 293 | def test_prepare_inputs_padded(): |
| 294 | """Input scenario is 3 requests with num_speculative_tokens == 2 and: |
| 295 | - Request 1: query_len = 3, rejected = 1 |
| 296 | - Request 2: query_len = 3, rejected = 0 |
| 297 | - Request 3: query_len = 3, rejected = 2 |
| 298 | |
| 299 | Expected outputs: |
| 300 | token_indices_to_sample: [1, 5, 6] |
| 301 | Reason: After accounting for rejections, these are the valid token positions |
| 302 | from the original indices to sample from. |
| 303 | """ |
| 304 | device = torch.device(DEVICE_TYPE) |
| 305 | |
| 306 | expected_token_indices_to_sample = torch.tensor( |
| 307 | [1, 5, 6], dtype=torch.int32, device=device |
| 308 | ) |
| 309 | |
| 310 | num_speculative_tokens = 2 |
| 311 | batch_spec = BatchSpec( |
| 312 | seq_lens=[3, 3, 3], |
| 313 | query_lens=[3, 3, 3], |
| 314 | ) |
| 315 | |
| 316 | common_attn_metadata = create_common_attn_metadata( |
| 317 | batch_spec, |
| 318 | block_size=BLOCK_SIZE, |
| 319 | device=device, |
| 320 | ) |
| 321 | |
| 322 | # Needed for cu_num_draft_tokens, which is expected to be [3, 6, 9] |
| 323 | expected_query_start_loc = torch.tensor( |
| 324 | [0, 3, 6, 9], dtype=torch.int32, device=device |
| 325 | ) |
| 326 | spec_decode_metadata = SpecDecodeMetadata.make_dummy( |
| 327 | draft_token_ids=[[0] * num_speculative_tokens] * 3, |
| 328 | device=device, |
| 329 | ) |
| 330 | |
| 331 | # num_rejected_tokens = [1, 0, 2] |
| 332 | # num_draft_tokens = [2, 2, 2] |
| 333 | # valid_sampled_tokens_count = num_draft_tokens + 1 - num_rejected_tokens |
| 334 | valid_sampled_tokens_count = torch.tensor( |
| 335 | [2, 3, 1], dtype=torch.int32, device=device |
| 336 | ) |
| 337 | |
| 338 | proposer = _create_proposer("eagle", num_speculative_tokens) |
| 339 | |
| 340 | output_metadata, token_indices_to_sample, num_rejected_tokens_gpu = ( |
| 341 | proposer.prepare_inputs_padded( |
| 342 | common_attn_metadata, spec_decode_metadata, valid_sampled_tokens_count |
| 343 | ) |
| 344 | ) |
| 345 | |
| 346 | # Verify num_rejected_tokens_gpu is calculated correctly |
| 347 | expected_num_rejected = torch.tensor([1, 0, 2], dtype=torch.int32, device=device) |
| 348 | assert torch.equal(num_rejected_tokens_gpu, expected_num_rejected) |
| 349 | |
| 350 | assert output_metadata.max_query_len == 3 |
| 351 | assert torch.equal(output_metadata.query_start_loc, expected_query_start_loc) |
| 352 | assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) |
| 353 | |
| 354 | |
| 355 | def test_set_inputs_first_pass_default_eagle(): |
| 356 | """Test for set_inputs_first_pass without extra input slots (default EAGLE). |
| 357 | |
| 358 | This tests the path where needs_extra_input_slots=False, which is the |
| 359 | default EAGLE pathway. In this case: |
| 360 | - Input IDs are rotated (shifted by one) |
| 361 | - The next_token_ids are inserted at the last position of each request |
| 362 | - Positions are copied as-is |
| 363 | - Hidden states are copied as-is |
| 364 | - The CommonAttentionMetadata is returned unchanged |
| 365 | |
| 366 | Setup: |
| 367 | - 3 requests with query_lens [3, 2, 4] |
| 368 | - Tokens: [a1, a2, a3, b1, b2, c1, c2, c3, c4] |
| 369 | - After rotation: [a2, a3, -, b2, -, c2, c3, c4, -] |
| 370 | - After inserting next_tokens [100, 200, 300]: |
| 371 | [a2, a3, 100, b2, 200, c2, c3, c4, 300] |
| 372 | """ |
| 373 | device = torch.device(DEVICE_TYPE) |
| 374 | |
| 375 | num_speculative_tokens = 3 |
| 376 | proposer = _create_proposer("eagle", num_speculative_tokens) |
| 377 | |
| 378 | # Setup batch with 3 requests |
| 379 | batch_spec = BatchSpec( |
| 380 | seq_lens=[10, 8, 12], # Arbitrary context lengths |
| 381 | query_lens=[3, 2, 4], |
| 382 | ) |
| 383 | |
| 384 | common_attn_metadata = create_common_attn_metadata( |
| 385 | batch_spec, |
| 386 | block_size=BLOCK_SIZE, |
| 387 | device=device, |
| 388 | ) |
| 389 | |
| 390 | # Input tensors |
| 391 | # Request 0: tokens [10, 11, 12] at positions [7, 8, 9] |
| 392 | # Request 1: tokens [20, 21] at positions [6, 7] |
| 393 | # Request 2: tokens [30, 31, 32, 33] at positions [8, 9, 10, 11] |
| 394 | target_token_ids = torch.tensor( |
| 395 | [10, 11, 12, 20, 21, 30, 31, 32, 33], dtype=torch.int32, device=device |
| 396 | ) |
| 397 | target_positions = torch.tensor( |
| 398 | [7, 8, 9, 6, 7, 8, 9, 10, 11], dtype=torch.int64, device=device |
| 399 | ) |
| 400 | target_hidden_states = torch.randn( |
| 401 | 9, proposer.hidden_size, dtype=proposer.dtype, device=device |
| 402 | ) |
| 403 | next_token_ids = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) |
| 404 | |
| 405 | num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( |
| 406 | target_token_ids=target_token_ids, |
| 407 | next_token_ids=next_token_ids, |
| 408 | target_positions=target_positions, |
| 409 | target_hidden_states=target_hidden_states, |
| 410 | token_indices_to_sample=None, |
| 411 | cad=common_attn_metadata, |
| 412 | num_rejected_tokens_gpu=None, |
| 413 | ) |
| 414 | |
| 415 | assert num_tokens == 9 # Total tokens unchanged |
| 416 | |
| 417 | expected_token_indices_to_sample = torch.tensor( |
| 418 | [2, 4, 8], dtype=torch.int32, device=device |
| 419 | ) |
| 420 | assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) |
| 421 | |
| 422 | assert output_cad is common_attn_metadata |
| 423 | |
| 424 | # Verify input_ids are rotated and next_tokens inserted |
| 425 | # Original: [10, 11, 12, 20, 21, 30, 31, 32, 33] |
| 426 | # After shift by 1: [11, 12, 12, 21, 21, 31, 32, 33, 33] |
| 427 | # After inserting at last indices [2, 4, 8]: [11, 12, 100, 21, 200, 31, 32, 33, 300] |
| 428 | expected_input_ids = torch.tensor( |
| 429 | [11, 12, 100, 21, 200, 31, 32, 33, 300], dtype=torch.int32, device=device |
| 430 | ) |
| 431 | assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) |
| 432 | |
| 433 | # Verify positions are copied as-is |
| 434 | assert torch.equal(proposer.positions[:num_tokens], target_positions) |
| 435 | |
| 436 | # Verify hidden states are copied as-is |
| 437 | assert torch.equal(proposer.hidden_states[:num_tokens], target_hidden_states) |
| 438 | |
| 439 | |
| 440 | def test_set_inputs_first_pass_draft_model(): |
| 441 | """Test for set_inputs_first_pass with a draft model (extra input slots, |
| 442 | no shift). |
| 443 | |
| 444 | This tests the path where needs_extra_input_slots=True and |
| 445 | shift_input_ids=False (draft model case). In this case: |
| 446 | - Input IDs are NOT shifted |
| 447 | - Each request gets extra_slots_per_request (1) new slots |
| 448 | - The kernel handles copying tokens and inserting bonus/padding tokens |
| 449 | - A new CommonAttentionMetadata is returned with updated query_start_loc |
| 450 | |
| 451 | Setup: |
| 452 | - 2 requests |
| 453 | - Request 0: tokens [10, 11, 12] at positions [0, 1, 2] |
| 454 | - Only tokens [10, 11] are "valid" (query_end_loc=1), |
| 455 | token 12 is a rejected token from previous speculation |
| 456 | - Request 1: tokens [20, 21] at positions [0, 1], both valid. |
| 457 | - Note: this is less than num_speculative_tokens (2) to ensure |
| 458 | we handle variable lengths correctly. |
| 459 | - next_token_ids: [100, 200] (bonus tokens) |
| 460 | |
| 461 | With extra_slots_per_request=1 and shift=False: |
| 462 | Expected output layout: |
| 463 | Request 0 (indices 0-3): |
| 464 | - idx 0: token 10, pos 0 |
| 465 | - idx 1: token 11, pos 1 |
| 466 | - idx 2: token 100, pos 2 (bonus token) |
| 467 | - idx 3: padding_token_id, is_rejected=True |
| 468 | Request 1 (indices 4-6): |
| 469 | - idx 4: token 20, pos 0 |
| 470 | - idx 5: token 21, pos 1 |
| 471 | - idx 6: token 200, pos 2 (bonus token) |
| 472 | """ |
| 473 | device = torch.device(DEVICE_TYPE) |
| 474 | |
| 475 | num_speculative_tokens = 2 |
| 476 | block_size = BLOCK_SIZE |
| 477 | |
| 478 | # Create a proposer configured as a draft model (pass_hidden_states=False) |
| 479 | # We need to mock this since _create_proposer defaults to EAGLE |
| 480 | proposer = _create_proposer("draft_model", num_speculative_tokens) |
| 481 | |
| 482 | proposer.parallel_drafting_token_id = 0 |
| 483 | proposer.is_rejected_token_mask = torch.zeros( |
| 484 | proposer.max_num_tokens, dtype=torch.bool, device=device |
| 485 | ) |
| 486 | proposer.is_masked_token_mask = torch.zeros( |
| 487 | proposer.max_num_tokens, dtype=torch.bool, device=device |
| 488 | ) |
| 489 | |
| 490 | # Mock draft_attn_groups to avoid needing the full model setup |
| 491 | mock_kv_cache_spec = mock.MagicMock() |
| 492 | mock_kv_cache_spec.block_size = block_size |
| 493 | mock_attn_group = mock.MagicMock() |
| 494 | mock_attn_group.kv_cache_spec = mock_kv_cache_spec |
| 495 | proposer.draft_attn_groups = [mock_attn_group] |
| 496 | |
| 497 | # Request 0: query_len=3 (but 1 rejected), Request 1: query_len=2 |
| 498 | batch_spec = BatchSpec( |
| 499 | seq_lens=[3, 2], |
| 500 | query_lens=[3, 2], |
| 501 | ) |
| 502 | |
| 503 | common_attn_metadata = create_common_attn_metadata( |
| 504 | batch_spec, |
| 505 | block_size=block_size, |
| 506 | device=device, |
| 507 | arange_block_indices=True, # Use predictable block indices |
| 508 | ) |
| 509 | |
| 510 | # Input tensors |
| 511 | target_token_ids = torch.tensor( |
| 512 | [10, 11, 12, 20, 21], dtype=torch.int32, device=device |
| 513 | ) |
| 514 | target_positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device) |
| 515 | target_hidden_states = torch.randn( |
| 516 | 5, proposer.hidden_size, dtype=proposer.dtype, device=device |
| 517 | ) |
| 518 | next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device) |
| 519 | |
| 520 | num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device) |
| 521 | |
| 522 | num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( |
| 523 | target_token_ids=target_token_ids, |
| 524 | next_token_ids=next_token_ids, |
| 525 | target_positions=target_positions, |
| 526 | target_hidden_states=target_hidden_states, |
| 527 | token_indices_to_sample=None, |
| 528 | cad=common_attn_metadata, |
| 529 | num_rejected_tokens_gpu=num_rejected_tokens_gpu, |
| 530 | ) |
| 531 | |
| 532 | assert proposer.net_num_new_slots_per_request == 1 |
| 533 | assert proposer.needs_extra_input_slots |
| 534 | |
| 535 | # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size |
| 536 | assert num_tokens == 7 |
| 537 | |
| 538 | # Request 0: [10, 11, 100, padding_token (0)] |
| 539 | # Request 1: [20, 21, 200] |
| 540 | # Combined: [10, 11, 100, 0, 20, 21, 200] |
| 541 | expected_input_ids = torch.tensor( |
| 542 | [10, 11, 100, 0, 20, 21, 200], dtype=torch.int32, device=device |
| 543 | ) |
| 544 | assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) |
| 545 | |
| 546 | # Verify positions |
| 547 | # Request 0: [0, 1, 2, 0 (don't care)] |
| 548 | # Request 1: [0, 1, 2] |
| 549 | # Combined: [0, 1, 2, 0, 0, 1, 2] |
| 550 | expected_positions = torch.tensor( |
| 551 | [0, 1, 2, 0, 0, 1, 2], dtype=torch.int64, device=device |
| 552 | ) |
| 553 | assert torch.equal( |
| 554 | proposer.positions[:num_tokens], |
| 555 | expected_positions, |
| 556 | ) |
| 557 | |
| 558 | # Verify rejection mask |
| 559 | expected_is_rejected = torch.zeros(7, dtype=torch.bool, device=device) |
| 560 | expected_is_rejected[3] = True # padding token at index 3 |
| 561 | assert torch.equal( |
| 562 | proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected |
| 563 | ) |
| 564 | |
| 565 | # Verify masked token mask (should all be False for non-parallel drafting) |
| 566 | expected_is_masked = torch.zeros(7, dtype=torch.bool, device=device) |
| 567 | assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked) |
| 568 | |
| 569 | # Verify token_indices_to_sample (bonus tokens at indices 2 and 6) |
| 570 | expected_token_indices_to_sample = torch.tensor( |
| 571 | [2, 6], dtype=torch.int32, device=device |
| 572 | ) |
| 573 | assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) |
| 574 | |
| 575 | # Verify the new CAD has updated query_start_loc |
| 576 | # Original: [0, 3, 5] -> New: [0, 4, 7] (each request gains 1 slot) |
| 577 | expected_query_start_loc = torch.tensor([0, 4, 7], dtype=torch.int32, device=device) |
| 578 | assert torch.equal(output_cad.query_start_loc, expected_query_start_loc) |
| 579 | |
| 580 | |
| 581 | def test_set_inputs_first_pass_parallel_drafting(): |
| 582 | """Test for set_inputs_first_pass with parallel drafting (extra input slots, |
| 583 | with shift). |
| 584 | |
| 585 | This tests the path where needs_extra_input_slots=True and |
| 586 | shift_input_ids=True (parallel drafting case). In this case: |
| 587 | - Input IDs ARE shifted (like default EAGLE) |
| 588 | - Each request gets extra_slots_per_request (3) new slots |
| 589 | - Parallel drafting tokens are inserted and marked as masked |
| 590 | - Hidden states are mapped correctly |
| 591 | |
| 592 | Setup: |
| 593 | - 2 requests with query_lens [4, 4] (1 bonus + 3 spec tokens each) |
| 594 | - Request 0: tokens [10, 11, 12, 13] at positions [5, 6, 7, 8] |
| 595 | - Only tokens [10, 11, 12] are "valid", token 13 is rejected |
| 596 | - Request 1: tokens [20, 21, 22, 23] at positions [10, 11, 12, 13], all valid. |
| 597 | - next_token_ids: [100, 200] (bonus tokens) |
| 598 | |
| 599 | With shift_input_ids=True, extra_slots_per_request=3: |
| 600 | Expected output layout: |
| 601 | Request 0 (6 output slots = 4 - 1 + 3): |
| 602 | - idx 0-2: shifted tokens [11, 12, 100] |
| 603 | - idx 3-4: parallel_drafting_tokens, is_masked=True |
| 604 | - idx 5: padding_token, is_rejected=True |
| 605 | Request 1 (6 output slots = 4 - 1 + 3): |
| 606 | - idx 6-8: shifted tokens [21, 22, 23] |
| 607 | - idx 9: bonus token 200 |
| 608 | - idx 10-11: parallel_drafting_tokens, is_masked=True |
| 609 | """ |
| 610 | device = torch.device(DEVICE_TYPE) |
| 611 | |
| 612 | num_speculative_tokens = 3 |
| 613 | block_size = BLOCK_SIZE |
| 614 | |
| 615 | proposer = _create_proposer("eagle", num_speculative_tokens, parallel_drafting=True) |
| 616 | |
| 617 | # Override to simulate parallel drafting behavior |
| 618 | proposer.parallel_drafting_token_id = -2 |
| 619 | proposer.parallel_drafting_hidden_state_tensor = torch.zeros( |
| 620 | proposer.hidden_size, dtype=proposer.dtype, device=device |
| 621 | ) |
| 622 | proposer.is_rejected_token_mask = torch.zeros( |
| 623 | proposer.max_num_tokens, dtype=torch.bool, device=device |
| 624 | ) |
| 625 | proposer.is_masked_token_mask = torch.zeros( |
| 626 | proposer.max_num_tokens, dtype=torch.bool, device=device |
| 627 | ) |
| 628 | |
| 629 | # Mock draft_attn_groups |
| 630 | mock_kv_cache_spec = mock.MagicMock() |
| 631 | mock_kv_cache_spec.block_size = block_size |
| 632 | mock_attn_group = mock.MagicMock() |
| 633 | mock_attn_group.kv_cache_spec = mock_kv_cache_spec |
| 634 | proposer.draft_attn_groups = [mock_attn_group] |
| 635 | |
| 636 | # Request 0: query_len=4 (1 rejected), Request 1: query_len=4 (all valid) |
| 637 | batch_spec = BatchSpec( |
| 638 | seq_lens=[9, 14], |
| 639 | query_lens=[4, 4], |
| 640 | ) |
| 641 | |
| 642 | common_attn_metadata = create_common_attn_metadata( |
| 643 | batch_spec, |
| 644 | block_size=block_size, |
| 645 | device=device, |
| 646 | arange_block_indices=True, |
| 647 | ) |
| 648 | |
| 649 | # Input tensors |
| 650 | target_token_ids = torch.tensor( |
| 651 | [10, 11, 12, 13, 20, 21, 22, 23], dtype=torch.int32, device=device |
| 652 | ) |
| 653 | target_positions = torch.tensor( |
| 654 | [5, 6, 7, 8, 10, 11, 12, 13], dtype=torch.int64, device=device |
| 655 | ) |
| 656 | target_hidden_states = torch.arange( |
| 657 | 8 * proposer.hidden_size, dtype=proposer.dtype, device=device |
| 658 | ).view(8, proposer.hidden_size) |
| 659 | next_token_ids = torch.tensor([100, 200], dtype=torch.int32, device=device) |
| 660 | |
| 661 | num_rejected_tokens_gpu = torch.tensor([1, 0], dtype=torch.int32, device=device) |
| 662 | |
| 663 | num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( |
| 664 | target_token_ids=target_token_ids, |
| 665 | next_token_ids=next_token_ids, |
| 666 | target_positions=target_positions, |
| 667 | target_hidden_states=target_hidden_states, |
| 668 | token_indices_to_sample=None, |
| 669 | cad=common_attn_metadata, |
| 670 | num_rejected_tokens_gpu=num_rejected_tokens_gpu, |
| 671 | ) |
| 672 | |
| 673 | # total_output_tokens = total_input_tokens + net_num_new_slots * batch_size |
| 674 | # = 8 + 2 * 2 = 12 |
| 675 | assert num_tokens == 12 |
| 676 | |
| 677 | # Request 0: [11, 12, 100, -2, -2, 0(padding)] |
| 678 | # Request 1: [21, 22, 23, 200, -2, -2] |
| 679 | expected_input_ids = torch.tensor( |
| 680 | [11, 12, 100, -2, -2, 0, 21, 22, 23, 200, -2, -2], |
| 681 | dtype=torch.int32, |
| 682 | device=device, |
| 683 | ) |
| 684 | assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) |
| 685 | |
| 686 | # Verify positions |
| 687 | # Request 0: [5, 6, 7, 8, 9, 0 (don't care)] |
| 688 | # Request 1: [10, 11, 12, 13, 14, 15] |
| 689 | expected_positions = torch.tensor( |
| 690 | [5, 6, 7, 8, 9, 0, 10, 11, 12, 13, 14, 15], dtype=torch.int64, device=device |
| 691 | ) |
| 692 | assert torch.equal( |
| 693 | proposer.positions[:num_tokens], |
| 694 | expected_positions, |
| 695 | ) |
| 696 | |
| 697 | # Verify rejection mask |
| 698 | expected_is_rejected = torch.zeros(12, dtype=torch.bool, device=device) |
| 699 | expected_is_rejected[5] = True |
| 700 | assert torch.equal( |
| 701 | proposer.is_rejected_token_mask[:num_tokens], expected_is_rejected |
| 702 | ) |
| 703 | |
| 704 | # Verify masked token mask (parallel drafting slots should be masked) |
| 705 | expected_is_masked = torch.zeros(12, dtype=torch.bool, device=device) |
| 706 | expected_is_masked[3] = True |
| 707 | expected_is_masked[4] = True |
| 708 | expected_is_masked[10] = True |
| 709 | expected_is_masked[11] = True |
| 710 | assert torch.equal(proposer.is_masked_token_mask[:num_tokens], expected_is_masked) |
| 711 | |
| 712 | # Verify token_indices_to_sample (bonus + parallel drafting tokens) |
| 713 | # Request 0: bonus at 2, parallel at 3, 4 |
| 714 | # Request 1: bonus at 9, parallel at 10, 11 |
| 715 | expected_token_indices_to_sample = torch.tensor( |
| 716 | [2, 3, 4, 9, 10, 11], dtype=torch.int32, device=device |
| 717 | ) |
| 718 | assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) |
| 719 | |
| 720 | # Verify the new CAD has updated query_start_loc |
| 721 | # Original query_lens: [4, 4] -> Output: [6, 6] |
| 722 | expected_query_start_loc = torch.tensor( |
| 723 | [0, 6, 12], dtype=torch.int32, device=device |
| 724 | ) |
| 725 | assert torch.equal(output_cad.query_start_loc, expected_query_start_loc) |
| 726 | |
| 727 | # Verify masked positions have the parallel drafting hidden state (zeros) |
| 728 | parallel_drafting_hs = proposer.parallel_drafting_hidden_state_tensor |
| 729 | for i in range(num_tokens): |
| 730 | if expected_is_masked[i]: |
| 731 | assert torch.equal(proposer.hidden_states[i], parallel_drafting_hs), ( |
| 732 | f"Masked position {i} should have parallel drafting hidden state" |
| 733 | ) |
| 734 | |
| 735 | |
| 736 | @pytest.mark.parametrize("method", ["eagle", "eagle3"]) |
| 737 | @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) |
| 738 | @pytest.mark.parametrize("pp_size", [1, 2]) |
| 739 | @pytest.mark.parametrize("use_distinct_embed_tokens", [True, False]) |
| 740 | @pytest.mark.parametrize("use_distinct_lm_head", [True, False]) |
| 741 | @mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_pp_group") |
| 742 | @mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_layers_from_vllm_config") |
| 743 | @mock.patch("vllm.v1.spec_decode.llm_base_proposer.get_model") |
| 744 | def test_load_model( |
| 745 | mock_get_model, |
| 746 | mock_get_layers, |
| 747 | mock_get_pp_group, |
| 748 | method, |
| 749 | attn_backend, |
| 750 | pp_size, |
| 751 | use_distinct_embed_tokens, |
| 752 | use_distinct_lm_head, |
| 753 | monkeypatch, |
| 754 | ): |
| 755 | if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): |
| 756 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 757 | |
| 758 | # Setup draft model mock |
| 759 | mock_model = mock.MagicMock() |
| 760 | mock_model.model = mock.MagicMock() |
| 761 | mock_model.has_own_embed_tokens = use_distinct_embed_tokens |
| 762 | if use_distinct_embed_tokens: |
| 763 | mock_model.model.embed_tokens = mock.MagicMock() |
| 764 | mock_model.has_own_lm_head = use_distinct_lm_head |
| 765 | if use_distinct_lm_head: |
| 766 | mock_model.lm_head = mock.MagicMock() |
| 767 | |
| 768 | mock_get_model.return_value = mock_model |
| 769 | |
| 770 | # Setup mocks for attention layers |
| 771 | target_attn_layers = { |
| 772 | "target_attn_1": mock.MagicMock(), |
| 773 | "target_attn_2": mock.MagicMock(), |
| 774 | } |
| 775 | target_indx_layers: dict[str, mock.MagicMock] = {} |
| 776 | # Draft model has one extra attention layer compared to target model |
| 777 | all_attn_layers = {**target_attn_layers, "draft_extra_attn": mock.MagicMock()} |
| 778 | |
| 779 | all_indx_layers: dict[str, mock.MagicMock] = {} |
| 780 | |
| 781 | # Make mock_get_layers return different values for each call |
| 782 | mock_get_layers.side_effect = [ |
| 783 | target_attn_layers, |
| 784 | target_indx_layers, |
| 785 | all_attn_layers, |
| 786 | all_indx_layers, |
| 787 | ] |
| 788 | |
| 789 | # Setup mock for pp group to return the appropriate value for world size |
| 790 | mock_pp_group = mock.MagicMock() |
| 791 | mock_pp_group.world_size = pp_size |
| 792 | mock_get_pp_group.return_value = mock_pp_group |
| 793 | |
| 794 | # Set up the target model mock with a custom class so that |
| 795 | # isinstance() checks match the expected type. |
| 796 | class _TargetModelStub(LlamaForCausalLM): |
| 797 | model: mock.MagicMock |
| 798 | lm_head: mock.MagicMock |
| 799 | |
| 800 | target_model = mock.create_autospec(_TargetModelStub, instance=True) |
| 801 | target_model.model = mock.MagicMock() |
| 802 | target_model.lm_head = mock.MagicMock() |
| 803 | target_model.model.embed_tokens = mock.MagicMock() |
| 804 | |
| 805 | from vllm.model_executor.models import SupportsMultiModal |
| 806 | |
| 807 | assert not isinstance(target_model, SupportsMultiModal) |
| 808 | |
| 809 | # Create proposer using the helper function |
| 810 | proposer = _create_proposer( |
| 811 | method, num_speculative_tokens=8, attention_backend=attn_backend |
| 812 | ) |
| 813 | |
| 814 | # Call the method under test |
| 815 | proposer.load_model(target_model) |
| 816 | |
| 817 | # Verify common interactions |
| 818 | mock_get_model.assert_called_once() |
| 819 | |
| 820 | # Verify that the lm head is set correctly |
| 821 | if use_distinct_lm_head: |
| 822 | assert proposer.model.lm_head is not target_model.lm_head |
| 823 | else: |
| 824 | assert proposer.model.lm_head is target_model.lm_head |
| 825 | |
| 826 | # Verify that the embed tokens are set correctly |
| 827 | # If pp_size is > 1, the embed tokens should be distinct |
| 828 | if pp_size > 1 or use_distinct_embed_tokens: |
| 829 | assert proposer.model.model.embed_tokens is not target_model.model.embed_tokens |
| 830 | else: |
| 831 | assert proposer.model.model.embed_tokens is target_model.model.embed_tokens |
| 832 | |
| 833 | |
| 834 | @pytest.mark.parametrize("method", ["eagle", "eagle3"]) |
| 835 | @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) |
| 836 | @pytest.mark.parametrize("num_speculative_tokens", [1, 3, 8]) |
| 837 | def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): |
| 838 | if attn_backend == "TRITON_ATTN" and not current_platform.is_rocm(): |
| 839 | pytest.skip( |
| 840 | "TRITON_ATTN does not support " |
| 841 | "multi-token eagle spec decode on current platform" |
| 842 | ) |
| 843 | |
| 844 | if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): |
| 845 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 846 | |
| 847 | # Use GPU device |
| 848 | device = torch.device(DEVICE_TYPE) |
| 849 | |
| 850 | # Setup test parameters |
| 851 | batch_size = 2 |
| 852 | seq_len_1 = 5 |
| 853 | seq_len_2 = 3 |
| 854 | total_tokens = seq_len_1 + seq_len_2 |
| 855 | vocab_size = 100 |
| 856 | seq_lens = [seq_len_1, seq_len_2] |
| 857 | |
| 858 | # Create proposer first so we can use its actual hidden_size |
| 859 | proposer = _create_proposer( |
| 860 | "eagle", num_speculative_tokens, attention_backend=attn_backend |
| 861 | ) |
| 862 | # Get the hidden_size from the proposer to ensure consistency |
| 863 | hidden_size = proposer.hidden_size |
| 864 | |
| 865 | # Helper to create deterministic logits that will produce specific tokens |
| 866 | def create_deterministic_logits(token_ids): |
| 867 | logits = torch.full((batch_size, vocab_size), -100.0, device=device) |
| 868 | for i, token_id in enumerate(token_ids): |
| 869 | logits[i, token_id] = 100.0 |
| 870 | return logits |
| 871 | |
| 872 | # We mock a model that returns deterministic logits |
| 873 | # Sequence 1: 42, 43, 44, ... |
| 874 | # Sequence 2: 60, 61, 62, ... |
| 875 | base_token_ids = [42, 60] |
| 876 | |
| 877 | # Skip loading the model and replace it with a mock directly |
| 878 | # Create the mock model with deterministic outputs |
| 879 | model_mock = mock.MagicMock() |
| 880 | |
| 881 | # Setup for model forward calls |
| 882 | forward_returns = [] |
| 883 | for i in range(num_speculative_tokens): |
| 884 | if i == 0: |
| 885 | # First call uses all tokens |
| 886 | h_logits = torch.zeros(total_tokens, hidden_size, device=device) |
| 887 | h_states = torch.zeros(total_tokens, hidden_size, device=device) |
| 888 | else: |
| 889 | # Subsequent calls use batch_size tokens |
| 890 | h_logits = torch.zeros(batch_size, hidden_size, device=device) |
| 891 | h_states = torch.zeros(batch_size, hidden_size, device=device) |
| 892 | forward_returns.append((h_logits, h_states)) |
| 893 | |
| 894 | # For single token case, we only need the first item; |
| 895 | # for multi-token, we need the sequence |
| 896 | if num_speculative_tokens == 1: |
| 897 | model_mock.return_value = forward_returns[0] |
| 898 | else: |
| 899 | model_mock.side_effect = forward_returns |
| 900 | |
| 901 | # Setup for compute_logits calls |
| 902 | logits_returns = [] |
| 903 | for i in range(num_speculative_tokens): |
| 904 | # For each call, increment the base token IDs |
| 905 | current_tokens = [base_id + i for base_id in base_token_ids] |
| 906 | logits_returns.append(create_deterministic_logits(current_tokens)) |
| 907 | |
| 908 | if num_speculative_tokens == 1: |
| 909 | model_mock.compute_logits.return_value = logits_returns[0] |
| 910 | else: |
| 911 | model_mock.compute_logits.side_effect = logits_returns |
| 912 | |
| 913 | # Assign the mock to the proposer |
| 914 | proposer.model = model_mock |
| 915 | |
| 916 | # Assign draft attn_layer_names since load_model is not invoked |
| 917 | proposer._draft_attn_layer_names = {"layer.0"} |
| 918 | |
| 919 | # Create input tensors |
| 920 | batch_spec = BatchSpec( |
| 921 | seq_lens=seq_lens, |
| 922 | query_lens=seq_lens, |
| 923 | ) |
| 924 | |
| 925 | common_attn_metadata = create_common_attn_metadata( |
| 926 | batch_spec, |
| 927 | block_size=BLOCK_SIZE, |
| 928 | device=device, |
| 929 | ) |
| 930 | |
| 931 | target_token_ids = torch.randint(0, vocab_size, (total_tokens,), device=device) |
| 932 | target_positions = torch.cat( |
| 933 | [torch.arange(seq_len_1, device=device), torch.arange(seq_len_2, device=device)] |
| 934 | ) |
| 935 | target_hidden_states = torch.randn(total_tokens, hidden_size, device=device) |
| 936 | next_token_ids = torch.randint( |
| 937 | 0, vocab_size, (batch_size,), dtype=torch.int32, device=device |
| 938 | ) |
| 939 | sampling_metadata = mock.MagicMock() |
| 940 | |
| 941 | if attn_backend == "FLASH_ATTN": |
| 942 | attn_metadata_builder_cls, _ = try_get_attention_backend( |
| 943 | AttentionBackendEnum.FLASH_ATTN |
| 944 | ) |
| 945 | elif attn_backend == "TRITON_ATTN": |
| 946 | attn_metadata_builder_cls, _ = try_get_attention_backend( |
| 947 | AttentionBackendEnum.TRITON_ATTN |
| 948 | ) |
| 949 | elif attn_backend == "ROCM_AITER_FA": |
| 950 | attn_metadata_builder_cls, _ = try_get_attention_backend( |
| 951 | AttentionBackendEnum.ROCM_AITER_FA |
| 952 | ) |
| 953 | else: |
| 954 | raise ValueError(f"Unsupported attention backend: {attn_backend}") |
| 955 | |
| 956 | attn_metadata_builder = attn_metadata_builder_cls( |
| 957 | kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config), |
| 958 | layer_names=sorted(proposer._draft_attn_layer_names), |
| 959 | vllm_config=proposer.vllm_config, |
| 960 | device=device, |
| 961 | ) |
| 962 | |
| 963 | # Mock draft_attn_groups for attention metadata building. |
| 964 | mock_attn_group = mock.MagicMock() |
| 965 | mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder |
| 966 | mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) |
| 967 | mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec |
| 968 | proposer.draft_attn_groups = [mock_attn_group] |
| 969 | |
| 970 | result = proposer.propose( |
| 971 | num_speculative_tokens=num_speculative_tokens, |
| 972 | target_token_ids=target_token_ids, |
| 973 | target_positions=target_positions, |
| 974 | target_hidden_states=target_hidden_states, |
| 975 | next_token_ids=next_token_ids, |
| 976 | token_indices_to_sample=None, |
| 977 | common_attn_metadata=common_attn_metadata, |
| 978 | sampling_metadata=sampling_metadata, |
| 979 | ) |
| 980 | |
| 981 | assert result.shape == (batch_size, num_speculative_tokens) |
| 982 | |
| 983 | # Create expected tokens based on our token pattern |
| 984 | if num_speculative_tokens == 1: |
| 985 | # Example for num_speculative_tokens=1: |
| 986 | # [[42], [60]] |
| 987 | expected_tokens = torch.tensor( |
| 988 | [[base_token_ids[0]], [base_token_ids[1]]], device=device |
| 989 | ) |
| 990 | else: |
| 991 | # Example for num_speculative_tokens=3: |
| 992 | # [[42, 43, 44], [60, 61, 62]] |
| 993 | expected_tokens = torch.zeros( |
| 994 | (batch_size, num_speculative_tokens), dtype=torch.int64, device=device |
| 995 | ) |
| 996 | for i in range(batch_size): |
| 997 | for j in range(num_speculative_tokens): |
| 998 | expected_tokens[i, j] = base_token_ids[i] + j |
| 999 | |
| 1000 | # Verify all tokens match our expectations |
| 1001 | assert torch.equal(result, expected_tokens) |
| 1002 | |
| 1003 | |
| 1004 | @pytest.mark.parametrize( |
| 1005 | "attn_backend", |
| 1006 | ["ROCM_ATTN", "TRITON_ATTN"] if current_platform.is_rocm() else ["FLASH_ATTN"], |
| 1007 | ) |
| 1008 | def test_propose_stores_probabilistic_draft_probs(attn_backend, monkeypatch): |
| 1009 | device = torch.device(DEVICE_TYPE) |
| 1010 | batch_size = 2 |
| 1011 | seq_lens = [5, 3] |
| 1012 | total_tokens = sum(seq_lens) |
| 1013 | num_speculative_tokens = 3 |
| 1014 | vocab_size = 8 |
| 1015 | |
| 1016 | proposer = _create_proposer( |
| 1017 | "draft_model", |
| 1018 | num_speculative_tokens, |
| 1019 | rejection_sample_method="standard", |
| 1020 | draft_sample_method="probabilistic", |
| 1021 | ) |
| 1022 | hidden_size = proposer.hidden_size |
| 1023 | expanded_total_tokens = total_tokens + batch_size |
| 1024 | |
| 1025 | model_mock = mock.MagicMock() |
| 1026 | forward_returns = [] |
| 1027 | logits_returns = [] |
| 1028 | for step in range(num_speculative_tokens): |
| 1029 | token_count = expanded_total_tokens if step == 0 else batch_size |
| 1030 | forward_returns.append(torch.zeros(token_count, hidden_size, device=device)) |
| 1031 | logits = torch.full((batch_size, vocab_size), -10.0, device=device) |
| 1032 | logits[0, step + 1] = 5.0 |
| 1033 | logits[1, step + 3] = 4.0 |
| 1034 | logits_returns.append(logits) |
| 1035 | |
| 1036 | model_mock.side_effect = forward_returns |
| 1037 | model_mock.compute_logits.side_effect = logits_returns |
| 1038 | proposer.model = model_mock |
| 1039 | proposer._draft_attn_layer_names = {"layer.0"} |
| 1040 | |
| 1041 | def fake_compute_probs(logits, sampling_metadata, use_fp64_gumbel): |
| 1042 | assert use_fp64_gumbel == proposer.use_fp64_gumbel |
| 1043 | probs = torch.softmax(logits, dim=-1) |
| 1044 | return probs.argmax(dim=-1), probs |
| 1045 | |
| 1046 | monkeypatch.setattr( |
| 1047 | "vllm.v1.spec_decode.llm_base_proposer.compute_probs_and_sample_next_token", |
| 1048 | fake_compute_probs, |
| 1049 | ) |
| 1050 | |
| 1051 | batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=seq_lens) |
| 1052 | common_attn_metadata = create_common_attn_metadata( |
| 1053 | batch_spec, |
| 1054 | block_size=BLOCK_SIZE, |
| 1055 | device=device, |
| 1056 | ) |
| 1057 | |
| 1058 | attn_metadata_builder_cls, _ = try_get_attention_backend( |
| 1059 | AttentionBackendEnum[attn_backend] |
| 1060 | ) |
| 1061 | attn_metadata_builder = attn_metadata_builder_cls( |
| 1062 | kv_cache_spec=create_standard_kv_cache_spec(proposer.vllm_config), |
| 1063 | layer_names=sorted(proposer._draft_attn_layer_names), |
| 1064 | vllm_config=proposer.vllm_config, |
| 1065 | device=device, |
| 1066 | ) |
| 1067 | mock_attn_group = mock.MagicMock() |
| 1068 | mock_attn_group.get_metadata_builder.return_value = attn_metadata_builder |
| 1069 | mock_attn_group.layer_names = list(proposer._draft_attn_layer_names) |
| 1070 | mock_attn_group.kv_cache_spec = attn_metadata_builder.kv_cache_spec |
| 1071 | proposer.draft_attn_groups = [mock_attn_group] |
| 1072 | |
| 1073 | sampling_metadata = mock.MagicMock() |
| 1074 | sampling_metadata.all_greedy = False |
| 1075 | sampling_metadata.temperature = torch.ones(batch_size, device=device) |
| 1076 | |
| 1077 | result = proposer.propose( |
| 1078 | num_speculative_tokens=num_speculative_tokens, |
| 1079 | target_token_ids=torch.randint(0, vocab_size, (total_tokens,), device=device), |
| 1080 | target_positions=torch.cat( |
| 1081 | [ |
| 1082 | torch.arange(seq_lens[0], device=device), |
| 1083 | torch.arange(seq_lens[1], device=device), |
| 1084 | ] |
| 1085 | ), |
| 1086 | target_hidden_states=torch.randn(total_tokens, hidden_size, device=device), |
| 1087 | next_token_ids=torch.randint( |
| 1088 | 0, vocab_size, (batch_size,), dtype=torch.int32, device=device |
| 1089 | ), |
| 1090 | token_indices_to_sample=None, |
| 1091 | common_attn_metadata=common_attn_metadata, |
| 1092 | sampling_metadata=sampling_metadata, |
| 1093 | ) |
| 1094 | |
| 1095 | assert result.shape == (batch_size, num_speculative_tokens) |
| 1096 | |
| 1097 | draft_probs = proposer.take_last_draft_probs() |
| 1098 | assert draft_probs is not None |
| 1099 | assert draft_probs.shape == (batch_size, num_speculative_tokens, vocab_size) |
| 1100 | for step, expected_logits in enumerate(logits_returns): |
| 1101 | assert torch.allclose( |
| 1102 | draft_probs[:, step, :], |
| 1103 | torch.softmax(expected_logits, dim=-1), |
| 1104 | ) |
| 1105 | |
| 1106 | |
| 1107 | def test_set_inputs_first_pass_dflash(): |
| 1108 | """Test for DFlash set_inputs_first_pass. |
| 1109 | |
| 1110 | DFlash uses cross-attention: context tokens become K/V and only |
| 1111 | query tokens (bonus + mask) are Q. This tests the DFlash-specific |
| 1112 | input preparation where: |
| 1113 | - Context hidden states are stored by reference (no copy) |
| 1114 | - Query input_ids are [next_token, mask, mask, ...] per request |
| 1115 | - Context and query positions are written to separate buffers |
| 1116 | - token_indices_to_sample points to mask token positions only |
| 1117 | - A new CommonAttentionMetadata is returned with causal=False |
| 1118 | |
| 1119 | Setup: |
| 1120 | - 3 requests with query_lens [3, 2, 4] |
| 1121 | - num_speculative_tokens = 3 |
| 1122 | - num_query_per_req = 4 (1 bonus + 3 mask tokens) |
| 1123 | - next_token_ids: [100, 200, 300] |
| 1124 | |
| 1125 | Expected output layout (query tokens only, 12 total): |
| 1126 | Request 0 (indices 0-3): [100, mask, mask, mask] |
| 1127 | Request 1 (indices 4-7): [200, mask, mask, mask] |
| 1128 | Request 2 (indices 8-11): [300, mask, mask, mask] |
| 1129 | |
| 1130 | Expected positions layout (separate buffers): |
| 1131 | Context (_context_positions_buffer, 9 tokens): copied from target_positions |
| 1132 | Query (positions, 12 tokens): |
| 1133 | Request 0: last_pos=9, query=[10, 11, 12, 13] |
| 1134 | Request 1: last_pos=7, query=[8, 9, 10, 11] |
| 1135 | Request 2: last_pos=11, query=[12, 13, 14, 15] |
| 1136 | """ |
| 1137 | device = torch.device(current_platform.device_type) |
| 1138 | |
| 1139 | num_speculative_tokens = 3 |
| 1140 | proposer = _create_proposer("dflash", num_speculative_tokens) |
| 1141 | assert isinstance(proposer, DFlashProposer) |
| 1142 | mask_token_id = proposer.parallel_drafting_token_id |
| 1143 | |
| 1144 | # Setup batch with 3 requests |
| 1145 | batch_spec = BatchSpec( |
| 1146 | seq_lens=[10, 8, 12], |
| 1147 | query_lens=[3, 2, 4], |
| 1148 | ) |
| 1149 | |
| 1150 | common_attn_metadata = create_common_attn_metadata( |
| 1151 | batch_spec, |
| 1152 | block_size=BLOCK_SIZE, |
| 1153 | device=device, |
| 1154 | arange_block_indices=True, |
| 1155 | ) |
| 1156 | |
| 1157 | # Input tensors |
| 1158 | # Request 0: tokens [10, 11, 12] at positions [7, 8, 9] |
| 1159 | # Request 1: tokens [20, 21] at positions [6, 7] |
| 1160 | # Request 2: tokens [30, 31, 32, 33] at positions [8, 9, 10, 11] |
| 1161 | target_token_ids = torch.tensor( |
| 1162 | [10, 11, 12, 20, 21, 30, 31, 32, 33], dtype=torch.int32, device=device |
| 1163 | ) |
| 1164 | target_positions = torch.tensor( |
| 1165 | [7, 8, 9, 6, 7, 8, 9, 10, 11], dtype=torch.int64, device=device |
| 1166 | ) |
| 1167 | target_hidden_states = torch.randn( |
| 1168 | 9, proposer.hidden_size, dtype=proposer.dtype, device=device |
| 1169 | ) |
| 1170 | next_token_ids = torch.tensor([100, 200, 300], dtype=torch.int32, device=device) |
| 1171 | |
| 1172 | num_tokens, token_indices_to_sample, output_cad = proposer.set_inputs_first_pass( |
| 1173 | target_token_ids=target_token_ids, |
| 1174 | next_token_ids=next_token_ids, |
| 1175 | target_positions=target_positions, |
| 1176 | target_hidden_states=target_hidden_states, |
| 1177 | token_indices_to_sample=None, |
| 1178 | cad=common_attn_metadata, |
| 1179 | num_rejected_tokens_gpu=None, |
| 1180 | ) |
| 1181 | |
| 1182 | num_query_per_req = 1 + num_speculative_tokens # 4 |
| 1183 | num_context = 9 |
| 1184 | |
| 1185 | # num_tokens is the query-only count |
| 1186 | assert num_tokens == 3 * num_query_per_req # 12 |
| 1187 | |
| 1188 | # Verify input_ids (query tokens only) |
| 1189 | # Each request: [next_token, mask, mask, mask] |
| 1190 | M = mask_token_id |
| 1191 | expected_input_ids = torch.tensor( |
| 1192 | [100, M, M, M, 200, M, M, M, 300, M, M, M], |
| 1193 | dtype=torch.int32, |
| 1194 | device=device, |
| 1195 | ) |
| 1196 | assert torch.equal(proposer.input_ids[:num_tokens], expected_input_ids) |
| 1197 | |
| 1198 | # Verify context positions (separate buffer): copied from target_positions |
| 1199 | assert torch.equal( |
| 1200 | proposer._context_positions_buffer[:num_context], target_positions |
| 1201 | ) |
| 1202 | |
| 1203 | # Verify query positions (separate buffer, starts at index 0): |
| 1204 | # req0: last_pos=9, query=[10, 11, 12, 13] |
| 1205 | # req1: last_pos=7, query=[8, 9, 10, 11] |
| 1206 | # req2: last_pos=11, query=[12, 13, 14, 15] |
| 1207 | expected_query_positions = torch.tensor( |
| 1208 | [10, 11, 12, 13, 8, 9, 10, 11, 12, 13, 14, 15], |
| 1209 | dtype=torch.int64, |
| 1210 | device=device, |
| 1211 | ) |
| 1212 | assert torch.equal( |
| 1213 | proposer.positions[:num_tokens], |
| 1214 | expected_query_positions, |
| 1215 | ) |
| 1216 | |
| 1217 | # Verify token_indices_to_sample (mask tokens only, skip bonus at offset 0) |
| 1218 | # req0: query indices 0-3, mask at 1,2,3 |
| 1219 | # req1: query indices 4-7, mask at 5,6,7 |
| 1220 | # req2: query indices 8-11, mask at 9,10,11 |
| 1221 | expected_token_indices_to_sample = torch.tensor( |
| 1222 | [1, 2, 3, 5, 6, 7, 9, 10, 11], dtype=torch.int32, device=device |
| 1223 | ) |
| 1224 | assert torch.equal(token_indices_to_sample, expected_token_indices_to_sample) |
| 1225 | |
| 1226 | # Verify the new CAD has DFlash-specific properties |
| 1227 | assert output_cad.causal is False # DFlash requires non-causal attention |
| 1228 | assert output_cad.num_actual_tokens == num_tokens # query-only count |
| 1229 | assert output_cad.max_query_len == num_query_per_req |
| 1230 | |
| 1231 | expected_query_start_loc = torch.tensor( |
| 1232 | [0, 4, 8, 12], dtype=torch.int32, device=device |
| 1233 | ) |
| 1234 | assert torch.equal(output_cad.query_start_loc, expected_query_start_loc) |
| 1235 | |
| 1236 | # Verify hidden states (stored by reference, not copied) |
| 1237 | assert proposer._dflash_hidden_states is target_hidden_states |