[Feature] Universal speculative decoding for heterogeneous vocabularies (TLI) (#38174)
Signed-off-by: wan-danfeng <wandanfeng0802@gmail.com> Signed-off-by: Wonderful <wandanfeng0802@gmail.com> Co-authored-by: Wan_DF <wonderful199082@126.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Wonderful committed
Jul 2, 2026 at 16:34 UTC
3af878955935dd356182f6dd2ea9660acb1757be
8 files changed
+364
-5
docs/features/speculative_decoding/README.md
+29
@@ -86,6 +86,7 @@ only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
86
| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
87
| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
88
| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
89
+ | `use_heterogeneous_vocab` | `boolean` | `false` | Allow draft and target models with different vocabularies. Builds a token-level intersection at initialisation and constrains draft logits to shared tokens only. Only compatible with `method=draft_model`. Probabilistic draft sampling (`draft_sample_method='probabilistic'`) is not yet supported when this option is enabled. |
90
91
!!! note
92
Gemma 4 assistant checkpoints are handled as Gemma 4 MTP speculators, not
@@ -142,6 +143,33 @@ vllm serve <target-model> \
143
}'
144
```
145
146
+#### Cross-Vocabulary Draft Models (TLI)
147
+
148
+ By default, vLLM requires the draft and target models to share the same
149
+ vocabulary. Setting `use_heterogeneous_vocab: true` enables the
150
+ **Token-Level Intersection (TLI)** algorithm, which allows draft models
151
+ from a different model family with a different tokenizer.
152
+
153
+ At initialisation, vLLM builds a mapping between the two vocabularies by
154
+ normalising token strings and computing their intersection. Draft logits are
155
+ constrained to the shared tokens before sampling, and the sampled token IDs
156
+ are translated to the target vocabulary before rejection sampling.
157
+
158
+ ```python
159
+ from vllm import LLM, SamplingParams
160
+
161
+ llm = LLM(
162
+ model="Qwen/Qwen3-8B",
163
+ speculative_config={
164
+ "method": "draft_model",
165
+ "model": "HuggingFaceTB/SmolLM2-135M-Instruct",
166
+ "num_speculative_tokens": 3,
167
+ "use_heterogeneous_vocab": True,
168
+ },
169
+ gpu_memory_utilization=0.5,
170
+ )
171
+```
172
+
173
### Notes
174
175
- `--speculative-config` expects a JSON object on the CLI. In YAML config
@@ -153,6 +181,7 @@ vllm serve <target-model> \
181
- Internal fields such as `target_model_config`, `draft_model_config`,
182
`target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
183
are populated by vLLM and are not intended to be set by users.
184
+- `use_heterogeneous_vocab` currently supports greedy draft sampling only. Probabilistic acceptance (temperature > 0 draft sampling) is not yet supported and will be added in a future release.
185
186
## Lossless guarantees of Speculative Decoding
187
docs/features/speculative_decoding/draft_model.md
+28
@@ -76,6 +76,34 @@ The code used to request as completions as a client remains unchanged:
76
print(completion)
77
```
78
79
+## Draft Model Method with heterogeneous vocabs
80
+
81
+ By default, vLLM requires the draft and target models to share the same vocabulary. Setting `use_heterogeneous_vocab: true` enables the **Token-Level Intersection (TLI)** algorithm, which allows draft models from a different model family with a different tokenizer.
82
+
83
+ Currently,`use_heterogeneous_vocab` currently requires `draft_sample_method='greedy'` (the default). Probabilistic draft sampling is not yet supported and will be added in a
84
+ future release.
85
+
86
+ ```python
87
+ from vllm import LLM, SamplingParams
88
+
89
+ llm = LLM(
90
+ model="Qwen/Qwen3-8B",
91
+ speculative_config={
92
+ "method": "draft_model",
93
+ "model": "HuggingFaceTB/SmolLM2-135M-Instruct",
94
+ "num_speculative_tokens": 3,
95
+ "use_heterogeneous_vocab": True,
96
+ },
97
+ gpu_memory_utilization=0.5,
98
+ )
99
+outputs = llm.generate(prompts，sampling_params)
100
+
101
+for output in outputs:
102
+ prompt = output.prompt
103
+ generated_text = output.outputs[0].text
104
+ print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
105
+```
106
+
107
!!! warning
108
Note: Please use `--speculative-config` to set all configurations related
109
to speculative decoding. The previous method of specifying the model
examples/features/speculative_decoding/spec_decode_offline.py
+2
@@ -72,6 +72,7 @@ def parse_args():
72
parser.add_argument("--max-num-seqs", type=int, default=None)
73
parser.add_argument("--parallel-drafting", action="store_true")
74
parser.add_argument("--allowed-local-media-path", type=str, default="")
75
+ parser.add_argument("--use-heterogeneous-vocab", action="store_true")
76
return parser.parse_args()
77
78
@@ -135,6 +136,7 @@ def main(args):
136
"enforce_eager": args.enforce_eager,
137
"max_model_len": args.max_model_len,
138
"parallel_drafting": args.parallel_drafting,
139
+ "use_heterogeneous_vocab": args.use_heterogeneous_vocab,
140
}
141
elif args.method == "mtp":
142
speculative_config = {
tests/v1/spec_decode/test_vocab_mapping.py
new
+48
@@ -0,0 +1,48 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+from transformers import AutoTokenizer
6
+
7
+from vllm.v1.spec_decode.vocab_mapping import _detect_space_prefix
8
+
9
+
10
+@pytest.mark.parametrize(
11
+ "model_name,expected_prefix",
12
+ [
13
+ # BPE tokenizer (GPT-2 family) uses Ġ (U+0120)
14
+ ("HuggingFaceTB/SmolLM2-135M-Instruct", ("Ġ",)),
15
+ # SentencePiece tokenizer (LLaMA family) uses ▁ (U+2581)
16
+ ("TinyLlama/TinyLlama-1.1B-Chat-v1.0", ("▁",)),
17
+ # BPE tokenizer (Qwen family) uses Ġ (U+0120)
18
+ ("Qwen/Qwen2.5-0.5B-Instruct", ("Ġ",)),
19
+ ],
20
+)
21
+def test_detect_space_prefix_real_tokenizers(model_name, expected_prefix):
22
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
23
+ result = _detect_space_prefix(tokenizer)
24
+ assert result == expected_prefix, (
25
+ f"{model_name}: expected {expected_prefix!r}, got {result!r}"
26
+ )
27
+
28
+
29
+def test_detect_space_prefix_fallback_on_failure():
30
+ """When tokenizer lacks encode(), fall back to both known prefixes."""
31
+
32
+ class BrokenTokenizer:
33
+ def encode(self, text, **kwargs):
34
+ raise RuntimeError("broken")
35
+
36
+ result = _detect_space_prefix(BrokenTokenizer())
37
+ assert result == ("Ġ", "▁")
38
+
39
+
40
+def test_detect_space_prefix_empty_encode():
41
+ """When encode returns empty list, fall back."""
42
+
43
+ class EmptyTokenizer:
44
+ def encode(self, text, **kwargs):
45
+ return []
46
+
47
+ result = _detect_space_prefix(EmptyTokenizer())
48
+ assert result == ("Ġ", "▁")
vllm/config/speculative.py
+25
-2
@@ -139,6 +139,12 @@ class SpeculativeConfig:
139
O(2 * tp_size) per token. Only applies to greedy draft selection in
140
non-tree speculation."""
141
142
+ use_heterogeneous_vocab: bool = False
143
+ """Allow draft and target models to use different vocabularies.
144
+ When enabled, builds a token-level intersection at init and constrains
145
+ draft logits to shared tokens only (TLI algorithm). Requires
146
+ method='draft_model'."""
147
+
148
# Ngram proposer configuration
149
prompt_lookup_max: int | None = Field(default=None, ge=1)
150
"""Maximum size of ngram token window when using Ngram proposer, required
@@ -734,7 +740,11 @@ class SpeculativeConfig:
740
self.draft_model_config = ModelConfig(
741
model=self.model,
742
runner="draft",
737
- tokenizer=self.target_model_config.tokenizer,
743
+ tokenizer=(
744
+ self.model
745
+ if self.use_heterogeneous_vocab
746
+ else self.target_model_config.tokenizer
747
+ ),
748
tokenizer_mode=self.target_model_config.tokenizer_mode,
749
trust_remote_code=self.target_model_config.trust_remote_code,
750
allowed_local_media_path=self.target_model_config.allowed_local_media_path,
@@ -1103,7 +1113,20 @@ class SpeculativeConfig:
1113
self.draft_parallel_config
1114
)
1115
1106
- self.verify_equal_vocab_size_if_draft_model()
1116
+ if self.use_heterogeneous_vocab and not self.uses_draft_model():
1117
+ raise ValueError(
1118
+ "use_heterogeneous_vocab only works with method='draft_model'"
1119
+ )
1120
+
1121
+ if self.use_heterogeneous_vocab and self.draft_sample_method != "greedy":
1122
+ raise ValueError(
1123
+ "use_heterogeneous_vocab currently only supports greedy draft "
1124
+ "sampling. Set draft_sample_method='greedy' (the default) or "
1125
+ "omit it."
1126
+ )
1127
+
1128
+ if not self.use_heterogeneous_vocab:
1129
+ self.verify_equal_vocab_size_if_draft_model()
1130
return self
1131
1132
def verify_equal_vocab_size_if_draft_model(self):
vllm/v1/spec_decode/draft_model.py
+28
-1
@@ -9,7 +9,9 @@ from vllm.config import VllmConfig
9
from vllm.config.utils import replace
10
from vllm.logger import init_logger
11
from vllm.model_executor.model_loader import get_model
12
+from vllm.tokenizers.registry import get_tokenizer
13
from vllm.v1.spec_decode.llm_base_proposer import SpecDecodeBaseProposer
14
+from vllm.v1.spec_decode.vocab_mapping import VocabMapping
15
16
logger = init_logger(__name__)
17
@@ -27,9 +29,34 @@ class DraftModelProposer(SpecDecodeBaseProposer):
29
pass_hidden_states_to_model=False,
30
runner=runner,
31
)
30
- self._raise_if_vocab_size_mismatch()
32
self._raise_if_draft_tp_mismatch()
33
34
+ self.use_heterogeneous_vocab = self.speculative_config.use_heterogeneous_vocab
35
+
36
+ spec = self.speculative_config
37
+ if self.use_heterogeneous_vocab:
38
+ # Heterogeneous vocabularies: build a VocabMapping to translate
39
+ # token IDs between the two tokenizers and constrain draft logits
40
+ # to the intersection so rejection sampling stays lossless.
41
+ target_tokenizer = get_tokenizer(
42
+ spec.target_model_config.tokenizer,
43
+ trust_remote_code=spec.target_model_config.trust_remote_code,
44
+ )
45
+ draft_tokenizer = get_tokenizer(
46
+ spec.draft_model_config.model,
47
+ trust_remote_code=spec.draft_model_config.trust_remote_code,
48
+ )
49
+ self.vocab_mapping: VocabMapping | None = VocabMapping(
50
+ target_tokenizer=target_tokenizer,
51
+ draft_tokenizer=draft_tokenizer,
52
+ target_vocab_size=spec.target_model_config.get_vocab_size(),
53
+ draft_vocab_size=spec.draft_model_config.get_vocab_size(),
54
+ device=device,
55
+ )
56
+ else:
57
+ self._raise_if_vocab_size_mismatch()
58
+ self.vocab_mapping = None
59
+
60
def _raise_if_vocab_size_mismatch(self):
61
self.speculative_config.verify_equal_vocab_size_if_draft_model()
62
vllm/v1/spec_decode/llm_base_proposer.py
+50
-2
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import dataclasses
4
from importlib.util import find_spec
5
-from typing import Any, cast
5
+from typing import TYPE_CHECKING, Any, cast
6
7
import numpy as np
8
import torch
@@ -15,6 +15,10 @@ from vllm.config import (
15
get_layers_from_vllm_config,
16
replace,
17
)
18
+
19
+if TYPE_CHECKING:
20
+ from vllm.v1.spec_decode.vocab_mapping import VocabMapping
21
+
22
from vllm.distributed.eplb.eplb_state import EplbState
23
from vllm.distributed.parallel_state import get_pp_group
24
from vllm.forward_context import set_forward_context
@@ -125,6 +129,11 @@ class SpecDecodeBaseProposer:
129
)
130
self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
131
132
+ self.use_heterogeneous_vocab: bool = (
133
+ self.speculative_config.use_heterogeneous_vocab
134
+ )
135
+ self.vocab_mapping: VocabMapping | None = None
136
+
137
self.max_batch_size = vllm_config.scheduler_config.max_num_seqs
138
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
139
self.token_arange_np = np.arange(self.max_num_tokens, dtype=np.int32)
@@ -419,6 +428,12 @@ class SpecDecodeBaseProposer:
428
"""Greedy-sample draft tokens from hidden states."""
429
if self.use_local_argmax_reduction:
430
return self.model.get_top_tokens(hidden_states)
431
+ if self.use_heterogeneous_vocab:
432
+ logits = self.model.compute_logits(hidden_states)
433
+ assert self.vocab_mapping is not None
434
+ logits = self.vocab_mapping.constrain_draft_logits(logits)
435
+ draft_token_ids = logits.argmax(dim=-1)
436
+ return self.vocab_mapping.map_draft_to_target_ids(draft_token_ids)
437
return self.model.compute_logits(hidden_states).argmax(dim=-1)
438
439
def _sample_from_logits(
@@ -457,7 +472,28 @@ class SpecDecodeBaseProposer:
472
if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy:
473
return self._greedy_sample(hidden_states), None
474
logits = self.model.compute_logits(hidden_states)
460
- return self._sample_from_logits(logits, sampling_metadata)
475
+ if self.use_heterogeneous_vocab:
476
+ assert self.vocab_mapping is not None
477
+ logits = self.vocab_mapping.constrain_draft_logits(logits)
478
+ draft_token_ids, draft_probs = self._sample_from_logits(
479
+ logits, sampling_metadata
480
+ )
481
+ if self.use_heterogeneous_vocab:
482
+ assert self.vocab_mapping is not None
483
+ draft_token_ids = self.vocab_mapping.map_draft_to_target_ids(
484
+ draft_token_ids
485
+ )
486
+ # Config validation ensures draft_sample_method == "greedy" when
487
+ # use_heterogeneous_vocab is True, so this branch should never be
488
+ # reached. Kept as a safety fallback until probabilistic rejection
489
+ # sampling with heterogeneous vocabularies is implemented.
490
+ # TODO: remap draft_probs to target-vocab space for lossless
491
+ # probabilistic rejection sampling with heterogeneous vocabularies.
492
+ assert draft_probs is None, (
493
+ "probabilistic draft sampling is not supported with "
494
+ "use_heterogeneous_vocab"
495
+ )
496
+ return draft_token_ids, draft_probs
497
498
def take_last_draft_probs(self) -> torch.Tensor | None:
499
return self._last_draft_probs
@@ -647,6 +683,11 @@ class SpecDecodeBaseProposer:
683
# tensor.argmax() returns int64 by default.
684
input_ids = draft_token_ids_list[-1].int()
685
686
+ if self.use_heterogeneous_vocab:
687
+ # Map target token IDs to draft vocab space (TLI algorithm)
688
+ assert self.vocab_mapping is not None
689
+ input_ids = self.vocab_mapping.map_target_to_draft_ids(input_ids)
690
+
691
if not self.constant_draft_positions:
692
positions = self._update_positions_dependent_metadata(
693
positions,
@@ -785,6 +826,13 @@ class SpecDecodeBaseProposer:
826
cad: CommonAttentionMetadata,
827
num_rejected_tokens_gpu: torch.Tensor | None,
828
) -> tuple[int, torch.Tensor, CommonAttentionMetadata]:
829
+ # Map target token IDs to draft vocab space (TLI algorithm)
830
+ if self.use_heterogeneous_vocab:
831
+ assert self.vocab_mapping is not None
832
+ target_token_ids = self.vocab_mapping.map_target_to_draft_ids(
833
+ target_token_ids
834
+ )
835
+ next_token_ids = self.vocab_mapping.map_target_to_draft_ids(next_token_ids)
836
if not self.needs_extra_input_slots:
837
# Default EAGLE pathway: no reshaping of input tensors needed.
838
# Simply rotate the input ids and leave the positions unchanged,
vllm/v1/spec_decode/vocab_mapping.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.logger import init_logger
6
+
7
+logger = init_logger(__name__)
8
+
9
+
10
+def _detect_space_prefix(tokenizer) -> tuple[str, ...]:
11
+ """Detect the space-prefix character(s) by tokenizing a literal space.
12
+
13
+ Different tokenizer families mark word-initial spaces differently:
14
+ BPE uses 'Ġ' (U+0120), SentencePiece uses '▁' (U+2581). Probing at
15
+ runtime avoids hardcoding assumptions and correctly handles mixed-family
16
+ pairs (e.g. BPE draft + SentencePiece target).
17
+ """
18
+ try:
19
+ space_ids = tokenizer.encode(" a", add_special_tokens=False)
20
+ if space_ids:
21
+ tok_str = tokenizer.convert_ids_to_tokens(space_ids[0])
22
+ if (
23
+ isinstance(tok_str, str)
24
+ and len(tok_str) > 1
25
+ and tok_str.endswith("a")
26
+ and tok_str[0] not in (" ", " ")
27
+ ):
28
+ return (tok_str[:-1],)
29
+ except Exception:
30
+ pass
31
+ # Fallback: cover both BPE (Ġ U+0120) and SentencePiece (▁ U+2581)
32
+ return ("\u0120", "\u2581")
33
+
34
+
35
+def _normalize_token(token: str, space_prefixes: tuple[str, ...]) -> str:
36
+ for prefix in space_prefixes:
37
+ if token.startswith(prefix):
38
+ return " " + token[len(prefix) :]
39
+ return token
40
+
41
+
42
+def _get_unk_token_id(tokenizer, role: str) -> int:
43
+ """Return a safe fallback token ID for out-of-intersection tokens.
44
+
45
+ Preferred: unk_token_id → eos_token_id → ValueError.
46
+ Checking with ``is not None`` is required because token ID 0 is a valid
47
+ (and common) unk ID on many tokenizers; using ``or 0`` would silently
48
+ mishandle those cases.
49
+ """
50
+ unk = getattr(tokenizer, "unk_token_id", None)
51
+ if unk is not None:
52
+ return unk
53
+ eos = getattr(tokenizer, "eos_token_id", None)
54
+ if eos is not None:
55
+ logger.warning(
56
+ "VocabMapping: %s has no unk_token_id; "
57
+ "falling back to eos_token_id=%d for out-of-intersection tokens",
58
+ role,
59
+ eos,
60
+ )
61
+ return eos
62
+ raise ValueError(
63
+ f"VocabMapping: {role} has neither unk_token_id nor eos_token_id; "
64
+ "cannot safely map out-of-intersection tokens"
65
+ )
66
+
67
+
68
+class VocabMapping:
69
+ def __init__(
70
+ self,
71
+ target_tokenizer,
72
+ draft_tokenizer,
73
+ target_vocab_size,
74
+ draft_vocab_size,
75
+ device,
76
+ ):
77
+ self.target_vocab_size = target_vocab_size
78
+ self.draft_vocab_size = draft_vocab_size
79
+ self.device = device
80
+ self.target_unk_token_id = _get_unk_token_id(
81
+ target_tokenizer, "target tokenizer"
82
+ )
83
+ self.draft_unk_token_id = _get_unk_token_id(draft_tokenizer, "draft tokenizer")
84
+
85
+ target_prefixes = _detect_space_prefix(target_tokenizer)
86
+ draft_prefixes = _detect_space_prefix(draft_tokenizer)
87
+
88
+ target_vocab = target_tokenizer.get_vocab()
89
+ draft_vocab = draft_tokenizer.get_vocab()
90
+
91
+ target_normalized = {}
92
+ for token, tid in target_vocab.items():
93
+ norm = _normalize_token(token, target_prefixes)
94
+ if norm not in target_normalized:
95
+ target_normalized[norm] = tid
96
+
97
+ draft_normalized = {}
98
+ for token, tid in draft_vocab.items():
99
+ norm = _normalize_token(token, draft_prefixes)
100
+ if norm not in draft_normalized:
101
+ draft_normalized[norm] = tid
102
+
103
+ common_tokens = set(target_normalized.keys()) & set(draft_normalized.keys())
104
+
105
+ draft_to_target = torch.full((draft_vocab_size,), -1, dtype=torch.long)
106
+ target_to_draft = torch.full((target_vocab_size,), -1, dtype=torch.long)
107
+ intersection_mask_draft = torch.zeros(draft_vocab_size, dtype=torch.bool)
108
+
109
+ for norm_token in common_tokens:
110
+ t_id = target_normalized[norm_token]
111
+ d_id = draft_normalized[norm_token]
112
+ if t_id < target_vocab_size and d_id < draft_vocab_size:
113
+ draft_to_target[d_id] = t_id
114
+ target_to_draft[t_id] = d_id
115
+ intersection_mask_draft[d_id] = True
116
+
117
+ self.draft_to_target_ids = draft_to_target.to(device)
118
+ self.target_to_draft_ids = target_to_draft.to(device)
119
+ self.intersection_mask_draft = intersection_mask_draft.to(device)
120
+ self.intersection_size = int(intersection_mask_draft.sum().item())
121
+
122
+ logger.info(
123
+ "VocabMapping initialized: target_vocab=%d, draft_vocab=%d, "
124
+ "intersection=%d (%.1f%% of draft, %.1f%% of target)",
125
+ target_vocab_size,
126
+ draft_vocab_size,
127
+ self.intersection_size,
128
+ 100.0 * self.intersection_size / max(draft_vocab_size, 1),
129
+ 100.0 * self.intersection_size / max(target_vocab_size, 1),
130
+ )
131
+
132
+ if self.intersection_size < 100:
133
+ logger.warning(
134
+ "Very small vocabulary intersection (%d tokens).",
135
+ self.intersection_size,
136
+ )
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
143
+ return draft_ids.to(target_ids.dtype)
144
+
145
+ def map_draft_to_target_ids(self, draft_ids):
146
+ 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
150
+ return target_ids.to(draft_ids.dtype)
151
+
152
+ def constrain_draft_logits(self, logits):
153
+ # masked_fill returns a new tensor; no clone needed
154
+ return logits.masked_fill(~self.intersection_mask_draft, float("-inf"))