diff --git a/tests/v1/spec_decode/test_vocab_mapping.py b/tests/v1/spec_decode/test_vocab_mapping.py index b8773c5a55..25c99ad049 100644 --- a/tests/v1/spec_decode/test_vocab_mapping.py +++ b/tests/v1/spec_decode/test_vocab_mapping.py @@ -2,9 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from transformers import AutoTokenizer -from vllm.v1.spec_decode.vocab_mapping import _detect_space_prefix +import vllm.utils.gpu_sync_debug as gsd +from vllm.utils.gpu_sync_debug import with_gpu_sync_check +from vllm.v1.spec_decode.vocab_mapping import VocabMapping, _detect_space_prefix @pytest.mark.parametrize( @@ -46,3 +49,79 @@ def test_detect_space_prefix_empty_encode(): result = _detect_space_prefix(EmptyTokenizer()) assert result == ("Ġ", "▁") + + +class _FakeTokenizer: + """Minimal tokenizer exposing only what ``VocabMapping`` needs.""" + + def __init__(self, vocab, unk_token_id=None, eos_token_id=None): + self._vocab = dict(vocab) + self.unk_token_id = unk_token_id + self.eos_token_id = eos_token_id + + def get_vocab(self): + return dict(self._vocab) + + def encode(self, text, **kwargs): + raise RuntimeError("no encode") + + def convert_ids_to_tokens(self, ids): + raise RuntimeError("no convert") + + +@pytest.fixture(params=["cpu", "cuda"]) +def vocab_mapping(request): + device = request.param + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + # Intersection: {a, b}. x/y/z are target-only, u/v are draft-only + target = _FakeTokenizer( + vocab={"a": 0, "b": 1, "x": 2, "y": 3, "z": 4}, unk_token_id=0 + ) + draft = _FakeTokenizer(vocab={"a": 0, "b": 1, "u": 2, "v": 3}, unk_token_id=2) + return VocabMapping(target, draft, 5, 4, device) + + +@pytest.mark.parametrize( + "target_ids,expected", + [ + ([0, 1, 2, 3, 4], [0, 1, 2, 2, 2]), # mixed + ([0, 1], [0, 1]), # all present + ([2, 3, 4], [2, 2, 2]), # all missing + ], +) +def test_map_target_to_draft_ids(monkeypatch, vocab_mapping, target_ids, expected): + device = vocab_mapping.device + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error" if device == "cuda" else None) + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + ids = torch.tensor(target_ids, dtype=torch.int32, device=device) + out = with_gpu_sync_check(lambda: vocab_mapping.map_target_to_draft_ids(ids))() + + assert out.tolist() == expected + assert out.dtype == torch.int32, f"expected torch.int32, got {out.dtype}" + assert out.device.type == device, f"expected {device} device, got {out.device}" + assert (out != -1).all() + + +@pytest.mark.parametrize( + "draft_ids,expected", + [ + ([0, 1, 2, 3], [0, 1, 0, 0]), # mixed + ([0, 1], [0, 1]), # all present + ([2, 3], [0, 0]), # all missing + ], +) +def test_map_draft_to_target_ids(monkeypatch, vocab_mapping, draft_ids, expected): + device = vocab_mapping.device + monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error" if device == "cuda" else None) + monkeypatch.setattr(gsd, "_sync_check_enabled", True) + + ids = torch.tensor(draft_ids, dtype=torch.int32, device=device) + out = with_gpu_sync_check(lambda: vocab_mapping.map_draft_to_target_ids(ids))() + + assert out.tolist() == expected + assert out.dtype == torch.int32, f"expected torch.int32, got {out.dtype}" + assert out.device.type == device, f"expected {device} device, got {out.device}" + assert (out != -1).all() diff --git a/vllm/v1/spec_decode/vocab_mapping.py b/vllm/v1/spec_decode/vocab_mapping.py index 9a6bbe7cfa..dad059fa77 100644 --- a/vllm/v1/spec_decode/vocab_mapping.py +++ b/vllm/v1/spec_decode/vocab_mapping.py @@ -137,16 +137,12 @@ class VocabMapping: def map_target_to_draft_ids(self, target_ids): draft_ids = self.target_to_draft_ids[target_ids] # new tensor; no clone needed - missing = draft_ids == -1 - if missing.any(): - draft_ids[missing] = self.draft_unk_token_id + draft_ids[draft_ids == -1] = self.draft_unk_token_id return draft_ids.to(target_ids.dtype) def map_draft_to_target_ids(self, draft_ids): target_ids = self.draft_to_target_ids[draft_ids] # new tensor; no clone needed - missing = target_ids == -1 - if missing.any(): - target_ids[missing] = self.target_unk_token_id + target_ids[target_ids == -1] = self.target_unk_token_id return target_ids.to(draft_ids.dtype) def constrain_draft_logits(self, logits):