[Perf] Remove CPU-GPU sync in heterogeneous vocabulary speculative decoding (#57396)
Michael Lapshin committed
Sep 22, 2026 at 18:34 UTC
d421a2e8fc0fd17c46c131d5c0884af5cf075fca
2 files changed
+82
-7
tests/v1/spec_decode/test_vocab_mapping.py
+80
-1
@@ -2,9 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import pytest
5
+import torch
6
from transformers import AutoTokenizer
7
7
-from vllm.v1.spec_decode.vocab_mapping import _detect_space_prefix
8
+import vllm.utils.gpu_sync_debug as gsd
9
+from vllm.utils.gpu_sync_debug import with_gpu_sync_check
10
+from vllm.v1.spec_decode.vocab_mapping import VocabMapping, _detect_space_prefix
11
12
13
@pytest.mark.parametrize(
@@ -46,3 +49,79 @@ def test_detect_space_prefix_empty_encode():
49
50
result = _detect_space_prefix(EmptyTokenizer())
51
assert result == ("Ġ", "▁")
52
+
53
+
54
+class _FakeTokenizer:
55
+ """Minimal tokenizer exposing only what ``VocabMapping`` needs."""
56
+
57
+ def __init__(self, vocab, unk_token_id=None, eos_token_id=None):
58
+ self._vocab = dict(vocab)
59
+ self.unk_token_id = unk_token_id
60
+ self.eos_token_id = eos_token_id
61
+
62
+ def get_vocab(self):
63
+ return dict(self._vocab)
64
+
65
+ def encode(self, text, **kwargs):
66
+ raise RuntimeError("no encode")
67
+
68
+ def convert_ids_to_tokens(self, ids):
69
+ raise RuntimeError("no convert")
70
+
71
+
72
+@pytest.fixture(params=["cpu", "cuda"])
73
+def vocab_mapping(request):
74
+ device = request.param
75
+ if device == "cuda" and not torch.cuda.is_available():
76
+ pytest.skip("requires CUDA")
77
+
78
+ # Intersection: {a, b}. x/y/z are target-only, u/v are draft-only
79
+ target = _FakeTokenizer(
80
+ vocab={"a": 0, "b": 1, "x": 2, "y": 3, "z": 4}, unk_token_id=0
81
+ )
82
+ draft = _FakeTokenizer(vocab={"a": 0, "b": 1, "u": 2, "v": 3}, unk_token_id=2)
83
+ return VocabMapping(target, draft, 5, 4, device)
84
+
85
+
86
+@pytest.mark.parametrize(
87
+ "target_ids,expected",
88
+ [
89
+ ([0, 1, 2, 3, 4], [0, 1, 2, 2, 2]), # mixed
90
+ ([0, 1], [0, 1]), # all present
91
+ ([2, 3, 4], [2, 2, 2]), # all missing
92
+ ],
93
+)
94
+def test_map_target_to_draft_ids(monkeypatch, vocab_mapping, target_ids, expected):
95
+ device = vocab_mapping.device
96
+ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error" if device == "cuda" else None)
97
+ monkeypatch.setattr(gsd, "_sync_check_enabled", True)
98
+
99
+ ids = torch.tensor(target_ids, dtype=torch.int32, device=device)
100
+ out = with_gpu_sync_check(lambda: vocab_mapping.map_target_to_draft_ids(ids))()
101
+
102
+ assert out.tolist() == expected
103
+ assert out.dtype == torch.int32, f"expected torch.int32, got {out.dtype}"
104
+ assert out.device.type == device, f"expected {device} device, got {out.device}"
105
+ assert (out != -1).all()
106
+
107
+
108
+@pytest.mark.parametrize(
109
+ "draft_ids,expected",
110
+ [
111
+ ([0, 1, 2, 3], [0, 1, 0, 0]), # mixed
112
+ ([0, 1], [0, 1]), # all present
113
+ ([2, 3], [0, 0]), # all missing
114
+ ],
115
+)
116
+def test_map_draft_to_target_ids(monkeypatch, vocab_mapping, draft_ids, expected):
117
+ device = vocab_mapping.device
118
+ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error" if device == "cuda" else None)
119
+ monkeypatch.setattr(gsd, "_sync_check_enabled", True)
120
+
121
+ ids = torch.tensor(draft_ids, dtype=torch.int32, device=device)
122
+ out = with_gpu_sync_check(lambda: vocab_mapping.map_draft_to_target_ids(ids))()
123
+
124
+ assert out.tolist() == expected
125
+ assert out.dtype == torch.int32, f"expected torch.int32, got {out.dtype}"
126
+ assert out.device.type == device, f"expected {device} device, got {out.device}"
127
+ assert (out != -1).all()
vllm/v1/spec_decode/vocab_mapping.py
+2
-6
@@ -137,16 +137,12 @@ class VocabMapping:
137
138
def map_target_to_draft_ids(self, target_ids):
139
draft_ids = self.target_to_draft_ids[target_ids] # new tensor; no clone needed
140
- missing = draft_ids == -1
141
- if missing.any():
142
- draft_ids[missing] = self.draft_unk_token_id
140
+ draft_ids[draft_ids == -1] = self.draft_unk_token_id
141
return draft_ids.to(target_ids.dtype)
142
143
def map_draft_to_target_ids(self, draft_ids):
144
target_ids = self.draft_to_target_ids[draft_ids] # new tensor; no clone needed
147
- missing = target_ids == -1
148
- if missing.any():
149
- target_ids[missing] = self.target_unk_token_id
145
+ target_ids[target_ids == -1] = self.target_unk_token_id
146
return target_ids.to(draft_ids.dtype)
147
148
def constrain_draft_logits(self, logits):