[Spec Decode][Perf] Replicate DSpark Markov head across TP ranks (#49731)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Michael Goin committed
Jul 29, 2026 at 11:43 UTC
d6247d71731249c65c829b4cbfb782b647849e4e
4 files changed
+94
-23
tests/v1/sample/test_head_dtype.py
+52
-8
@@ -14,6 +14,7 @@ import torch
14
from vllm import LLM, SamplingParams
15
from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
from vllm.model_executor.layers.vocab_parallel_embedding import (
17
+ ParallelLMHead,
18
UnquantizedEmbeddingMethod,
19
)
20
@@ -28,6 +29,7 @@ class _FakeLmHead:
29
self.weight = weight
30
self.quant_method = object() if quantized else UnquantizedEmbeddingMethod()
31
self.shard_indices = shard_indices
32
+ self.tp_size = 1
33
34
35
def _build_processor(vocab_size: int) -> LogitsProcessor:
@@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
137
lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)
138
139
140
+def test_replicated_lm_head_skips_tp_communication_and_preserves_processing(
141
+ default_vllm_config,
142
+):
143
+ from unittest import mock
144
+
145
+ vocab_size, hidden_size = 12, 8
146
+ soft_cap, scale = 2.0, 0.5
147
+ lp = LogitsProcessor(
148
+ vocab_size,
149
+ soft_cap=soft_cap,
150
+ scale=scale,
151
+ )
152
+ lp.head_dtype = torch.float32
153
+
154
+ hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
155
+ weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
156
+ world_size_getter = (
157
+ "vllm.model_executor.layers.vocab_parallel_embedding."
158
+ "get_tensor_model_parallel_world_size"
159
+ )
160
+ with mock.patch(world_size_getter, return_value=2):
161
+ lm_head = ParallelLMHead(
162
+ vocab_size,
163
+ hidden_size,
164
+ params_dtype=torch.bfloat16,
165
+ disable_tp=True,
166
+ )
167
+ lm_head.weight_loader(lm_head.weight, weight)
168
+ assert lm_head.tp_size == 1
169
+
170
+ with mock.patch.object(lp, "_gather_logits") as gather_mock:
171
+ logits = lp(lm_head, hidden_states)
172
+
173
+ gather_mock.assert_not_called()
174
+
175
+ expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
176
+ expected = torch.tanh(expected / soft_cap) * soft_cap * scale
177
+ torch.testing.assert_close(logits, expected)
178
+
179
+ all_gather_path = (
180
+ "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather"
181
+ )
182
+ with mock.patch(all_gather_path) as all_gather:
183
+ top = lp.get_top_tokens(lm_head, hidden_states)
184
+
185
+ all_gather.assert_not_called()
186
+ assert torch.equal(top, expected.argmax(dim=-1))
187
+
188
+
189
def test_get_top_tokens_honors_head_dtype(default_vllm_config):
190
# The spec-decode local-argmax path (get_top_tokens) must run the lm_head
191
# in head_dtype too, not just _get_logits.
192
import types
142
- from unittest import mock
193
194
vocab_size, hidden_size = 64, 16
195
lp = _build_processor(vocab_size)
@@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):
204
),
205
)
206
157
- with mock.patch(
158
- "vllm.model_executor.layers.logits_processor."
159
- "get_tensor_model_parallel_world_size",
160
- return_value=1,
161
- ):
162
- top = lp.get_top_tokens(lm_head, hidden_states, None)
163
-
207
+ top = lp.get_top_tokens(lm_head, hidden_states, None)
208
expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax(
209
dim=-1
210
)
vllm/model_executor/layers/logits_processor.py
+3
-3
@@ -7,7 +7,6 @@ import torch.nn.functional as F
7
8
from vllm.config import get_current_vllm_config
9
from vllm.distributed import (
10
- get_tensor_model_parallel_world_size,
10
tensor_model_parallel_all_gather,
11
tensor_model_parallel_gather,
12
)
@@ -145,7 +144,8 @@ class LogitsProcessor(PluggableLayer):
144
logits = self._apply_head(lm_head, hidden_states, embedding_bias)
145
146
# Gather logits for TP
148
- logits = self._gather_logits(logits)
147
+ if lm_head.tp_size > 1:
148
+ logits = self._gather_logits(logits)
149
150
# Remove paddings in vocab (if any).
151
if logits is not None:
@@ -169,7 +169,7 @@ class LogitsProcessor(PluggableLayer):
169
"The local argmax reduction optimization is not supported for "
170
"non-positive logit scaling factors."
171
)
172
- tp_size = get_tensor_model_parallel_world_size()
172
+ tp_size = lm_head.tp_size
173
174
logits = self._apply_head(lm_head, hidden_states, embedding_bias)
175
if self.soft_cap is not None:
vllm/model_executor/layers/vocab_parallel_embedding.py
+24
-5
@@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer):
232
padding_size: padding size for the vocabulary.
233
quant_config: quant config for the layer
234
prefix: full name of the layer in the state dict
235
+ disable_tp: If true, tensor parallelism will be disabled for this layer.
236
""" # noqa: E501
237
238
# --8<-- [end:vocab_parallel_embedding]
@@ -245,12 +246,19 @@ class VocabParallelEmbedding(PluggableLayer):
246
padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
247
quant_config: QuantizationConfig | None = None,
248
prefix: str = "",
249
+ *,
250
+ disable_tp: bool = False,
251
):
252
super().__init__()
253
254
# Keep the input dimensions.
252
- tp_rank = get_tensor_model_parallel_rank()
253
- self.tp_size = get_tensor_model_parallel_world_size()
255
+ self.disable_tp = disable_tp
256
+ if disable_tp:
257
+ tp_rank, self.tp_size = 0, 1
258
+ else:
259
+ tp_rank = get_tensor_model_parallel_rank()
260
+ self.tp_size = get_tensor_model_parallel_world_size()
261
+ self.tp_rank = tp_rank
262
self.num_embeddings = num_embeddings
263
self.padding_size = padding_size
264
self.org_vocab_size = org_num_embeddings or num_embeddings
@@ -323,6 +331,13 @@ class VocabParallelEmbedding(PluggableLayer):
331
params_dtype=params_dtype,
332
weight_loader=self.weight_loader,
333
)
334
+ self.update_param_tp_status()
335
+
336
+ def update_param_tp_status(self):
337
+ for param in self.parameters():
338
+ if isinstance(param, BasevLLMParameter):
339
+ param.tp_rank = self.tp_rank
340
+ param.tp_size = self.tp_size
341
342
@classmethod
343
def _get_indices(
@@ -487,9 +502,9 @@ class VocabParallelEmbedding(PluggableLayer):
502
# Mask the output embedding.
503
if self.tp_size > 1:
504
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)
490
- # Reduce across all the model parallel GPUs.
491
- output = tensor_model_parallel_all_reduce(output_parallel)
492
- return output
505
+ # Reduce across all the model parallel GPUs.
506
+ return tensor_model_parallel_all_reduce(output_parallel)
507
+ return output_parallel
508
509
def extra_repr(self) -> str:
510
s = f"num_embeddings={self.num_embeddings_per_partition}"
@@ -516,6 +531,7 @@ class ParallelLMHead(VocabParallelEmbedding):
531
params_dtype: type of the parameters.
532
org_num_embeddings: original vocabulary size (without LoRA).
533
padding_size: padding size for the vocabulary.
534
+ disable_tp: If true, tensor parallelism will be disabled for this layer.
535
"""
536
537
# --8<-- [end:parallel_lm_head]
@@ -530,6 +546,8 @@ class ParallelLMHead(VocabParallelEmbedding):
546
padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
547
quant_config: QuantizationConfig | None = None,
548
prefix: str = "",
549
+ *,
550
+ disable_tp: bool = False,
551
):
552
super().__init__(
553
num_embeddings,
@@ -539,6 +557,7 @@ class ParallelLMHead(VocabParallelEmbedding):
557
padding_size,
558
quant_config,
559
prefix,
560
+ disable_tp=disable_tp,
561
)
562
self.quant_config = quant_config
563
if bias:
vllm/model_executor/models/qwen3_dspark.py
+15
-7
@@ -24,7 +24,6 @@ from vllm.logger import init_logger
24
from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
from vllm.model_executor.layers.vocab_parallel_embedding import (
26
ParallelLMHead,
27
- VocabParallelEmbedding,
27
)
28
29
from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model
@@ -40,6 +39,10 @@ class DSparkMarkovHead(nn.Module):
39
``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
40
(``draft_vocab_size``) added to the base draft logits. The two sizes
41
coincide for full-vocab drafts.
42
+
43
+ Both weights are replicated because the head runs sequentially for every
44
+ draft position. Sharding them would add an all-reduce and a full-vocab
45
+ gather to each position.
46
"""
47
48
def __init__(
@@ -50,19 +53,24 @@ class DSparkMarkovHead(nn.Module):
53
prefix: str,
54
) -> None:
55
super().__init__()
53
- # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
54
- self.markov_w1 = VocabParallelEmbedding(
55
- vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
56
- )
56
+ self.markov_w1 = nn.Embedding(vocab_size, markov_rank)
57
self.markov_w2 = ParallelLMHead(
58
- draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
58
+ draft_vocab_size,
59
+ markov_rank,
60
+ bias=False,
61
+ prefix=maybe_prefix(prefix, "markov_w2"),
62
+ disable_tp=True,
63
)
64
65
def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
66
"""r-dim Markov embedding of ``token_ids`` ([B] -> [B, r])."""
67
return self.markov_w1(token_ids)
68
65
- def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor:
69
+ def bias(
70
+ self,
71
+ markov_embed: torch.Tensor,
72
+ logits_processor: LogitsProcessor,
73
+ ) -> torch.Tensor:
74
"""Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
75
return logits_processor(self.markov_w2, markov_embed)
76