| 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 pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.utils.torch_utils import set_random_seed |
| 11 | from vllm.v1.sample.logits_processor import LogitsProcessors |
| 12 | from vllm.v1.sample.metadata import SamplingMetadata |
| 13 | from vllm.v1.spec_decode.llm_base_proposer import ( |
| 14 | SpecDecodeBaseProposer, |
| 15 | compute_probs_and_sample_next_token, |
| 16 | ) |
| 17 | |
| 18 | DEVICE_TYPE = current_platform.device_type |
| 19 | |
| 20 | |
| 21 | def _seed_default_generator(seed: int) -> None: |
| 22 | set_random_seed(seed) |
| 23 | |
| 24 | |
| 25 | def _make_sampling_metadata(batch_size: int) -> SamplingMetadata: |
| 26 | return SamplingMetadata( |
| 27 | temperature=torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE), |
| 28 | all_greedy=False, |
| 29 | all_random=True, |
| 30 | top_p=None, |
| 31 | top_k=None, |
| 32 | generators={}, |
| 33 | max_num_logprobs=None, |
| 34 | no_penalties=True, |
| 35 | prompt_token_ids=None, |
| 36 | frequency_penalties=torch.empty(0, device=DEVICE_TYPE), |
| 37 | presence_penalties=torch.empty(0, device=DEVICE_TYPE), |
| 38 | repetition_penalties=torch.empty(0, device=DEVICE_TYPE), |
| 39 | output_token_ids=[[] for _ in range(batch_size)], |
| 40 | spec_token_ids=[[] for _ in range(batch_size)], |
| 41 | allowed_token_ids_mask=None, |
| 42 | bad_words_token_ids={}, |
| 43 | logitsprocs=LogitsProcessors(), |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race(): |
| 48 | batch_size = 4 |
| 49 | vocab_size = 32 |
| 50 | generator = torch.Generator(device=DEVICE_TYPE).manual_seed(11) |
| 51 | logits = torch.randn( |
| 52 | batch_size, |
| 53 | vocab_size, |
| 54 | dtype=torch.float32, |
| 55 | device=DEVICE_TYPE, |
| 56 | generator=generator, |
| 57 | ) |
| 58 | metadata = _make_sampling_metadata(batch_size) |
| 59 | |
| 60 | _seed_default_generator(12345) |
| 61 | probs = logits.softmax(dim=-1, dtype=torch.float32) |
| 62 | q = torch.empty(probs.shape, dtype=torch.float64, device=probs.device) |
| 63 | q.exponential_() |
| 64 | expected_ids = q.reciprocal_().mul_(probs).argmax(dim=-1).view(-1) |
| 65 | |
| 66 | _seed_default_generator(12345) |
| 67 | actual_ids, actual_probs = compute_probs_and_sample_next_token( |
| 68 | logits.clone(), |
| 69 | metadata, |
| 70 | use_fp64_gumbel=True, |
| 71 | ) |
| 72 | |
| 73 | assert torch.equal(actual_ids, expected_ids) |
| 74 | assert torch.allclose(actual_probs, probs) |
| 75 | |
| 76 | |
| 77 | @pytest.mark.parametrize( |
| 78 | ("architecture", "expected"), |
| 79 | [ |
| 80 | ("DeepSeekMTPModel", True), |
| 81 | ("KimiK3MTPModel", True), |
| 82 | ("MiniMaxM3ForCausalLM", False), |
| 83 | ], |
| 84 | ) |
| 85 | def test_mtp_model_returns_tuple(architecture: str, expected: bool): |
| 86 | proposer = object.__new__(SpecDecodeBaseProposer) |
| 87 | proposer.method = "mtp" |
| 88 | proposer.draft_model_config = SimpleNamespace( |
| 89 | hf_config=SimpleNamespace(architectures=[architecture]) |
| 90 | ) |
| 91 | |
| 92 | assert proposer.model_returns_tuple() is expected |