From d6247d71731249c65c829b4cbfb782b647849e4e Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 29 Jul 2026 11:43:51 -0400 Subject: [PATCH] [Spec Decode][Perf] Replicate DSpark Markov head across TP ranks (#49731) Signed-off-by: mgoin Co-authored-by: OpenAI Codex --- tests/v1/sample/test_head_dtype.py | 60 ++++++++++++++++--- .../model_executor/layers/logits_processor.py | 6 +- .../layers/vocab_parallel_embedding.py | 29 +++++++-- vllm/model_executor/models/qwen3_dspark.py | 22 ++++--- 4 files changed, 94 insertions(+), 23 deletions(-) diff --git a/tests/v1/sample/test_head_dtype.py b/tests/v1/sample/test_head_dtype.py index 7531cdfc45..ce07919577 100644 --- a/tests/v1/sample/test_head_dtype.py +++ b/tests/v1/sample/test_head_dtype.py @@ -14,6 +14,7 @@ import torch from vllm import LLM, SamplingParams from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, UnquantizedEmbeddingMethod, ) @@ -28,6 +29,7 @@ class _FakeLmHead: self.weight = weight self.quant_method = object() if quantized else UnquantizedEmbeddingMethod() self.shard_indices = shard_indices + self.tp_size = 1 def _build_processor(vocab_size: int) -> LogitsProcessor: @@ -135,11 +137,59 @@ def test_fp32_head_rejects_quantized_lm_head(default_vllm_config): lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None) +def test_replicated_lm_head_skips_tp_communication_and_preserves_processing( + default_vllm_config, +): + from unittest import mock + + vocab_size, hidden_size = 12, 8 + soft_cap, scale = 2.0, 0.5 + lp = LogitsProcessor( + vocab_size, + soft_cap=soft_cap, + scale=scale, + ) + lp.head_dtype = torch.float32 + + hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16) + weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16) + world_size_getter = ( + "vllm.model_executor.layers.vocab_parallel_embedding." + "get_tensor_model_parallel_world_size" + ) + with mock.patch(world_size_getter, return_value=2): + lm_head = ParallelLMHead( + vocab_size, + hidden_size, + params_dtype=torch.bfloat16, + disable_tp=True, + ) + lm_head.weight_loader(lm_head.weight, weight) + assert lm_head.tp_size == 1 + + with mock.patch.object(lp, "_gather_logits") as gather_mock: + logits = lp(lm_head, hidden_states) + + gather_mock.assert_not_called() + + expected = torch.nn.functional.linear(hidden_states.float(), weight.float()) + expected = torch.tanh(expected / soft_cap) * soft_cap * scale + torch.testing.assert_close(logits, expected) + + all_gather_path = ( + "vllm.model_executor.layers.logits_processor.tensor_model_parallel_all_gather" + ) + with mock.patch(all_gather_path) as all_gather: + top = lp.get_top_tokens(lm_head, hidden_states) + + all_gather.assert_not_called() + assert torch.equal(top, expected.argmax(dim=-1)) + + def test_get_top_tokens_honors_head_dtype(default_vllm_config): # The spec-decode local-argmax path (get_top_tokens) must run the lm_head # in head_dtype too, not just _get_logits. import types - from unittest import mock vocab_size, hidden_size = 64, 16 lp = _build_processor(vocab_size) @@ -154,13 +204,7 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config): ), ) - with mock.patch( - "vllm.model_executor.layers.logits_processor." - "get_tensor_model_parallel_world_size", - return_value=1, - ): - top = lp.get_top_tokens(lm_head, hidden_states, None) - + top = lp.get_top_tokens(lm_head, hidden_states, None) expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax( dim=-1 ) diff --git a/vllm/model_executor/layers/logits_processor.py b/vllm/model_executor/layers/logits_processor.py index 496a1dd153..31298a66be 100644 --- a/vllm/model_executor/layers/logits_processor.py +++ b/vllm/model_executor/layers/logits_processor.py @@ -7,7 +7,6 @@ import torch.nn.functional as F from vllm.config import get_current_vllm_config from vllm.distributed import ( - get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, tensor_model_parallel_gather, ) @@ -145,7 +144,8 @@ class LogitsProcessor(PluggableLayer): logits = self._apply_head(lm_head, hidden_states, embedding_bias) # Gather logits for TP - logits = self._gather_logits(logits) + if lm_head.tp_size > 1: + logits = self._gather_logits(logits) # Remove paddings in vocab (if any). if logits is not None: @@ -169,7 +169,7 @@ class LogitsProcessor(PluggableLayer): "The local argmax reduction optimization is not supported for " "non-positive logit scaling factors." ) - tp_size = get_tensor_model_parallel_world_size() + tp_size = lm_head.tp_size logits = self._apply_head(lm_head, hidden_states, embedding_bias) if self.soft_cap is not None: diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 8d9a7ccbac..6eed2b5249 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -232,6 +232,7 @@ class VocabParallelEmbedding(PluggableLayer): padding_size: padding size for the vocabulary. quant_config: quant config for the layer prefix: full name of the layer in the state dict + disable_tp: If true, tensor parallelism will be disabled for this layer. """ # noqa: E501 # --8<-- [end:vocab_parallel_embedding] @@ -245,12 +246,19 @@ class VocabParallelEmbedding(PluggableLayer): padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + disable_tp: bool = False, ): super().__init__() # Keep the input dimensions. - tp_rank = get_tensor_model_parallel_rank() - self.tp_size = get_tensor_model_parallel_world_size() + self.disable_tp = disable_tp + if disable_tp: + tp_rank, self.tp_size = 0, 1 + else: + tp_rank = get_tensor_model_parallel_rank() + self.tp_size = get_tensor_model_parallel_world_size() + self.tp_rank = tp_rank self.num_embeddings = num_embeddings self.padding_size = padding_size self.org_vocab_size = org_num_embeddings or num_embeddings @@ -323,6 +331,13 @@ class VocabParallelEmbedding(PluggableLayer): params_dtype=params_dtype, weight_loader=self.weight_loader, ) + self.update_param_tp_status() + + def update_param_tp_status(self): + for param in self.parameters(): + if isinstance(param, BasevLLMParameter): + param.tp_rank = self.tp_rank + param.tp_size = self.tp_size @classmethod def _get_indices( @@ -487,9 +502,9 @@ class VocabParallelEmbedding(PluggableLayer): # Mask the output embedding. if self.tp_size > 1: output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) - # Reduce across all the model parallel GPUs. - output = tensor_model_parallel_all_reduce(output_parallel) - return output + # Reduce across all the model parallel GPUs. + return tensor_model_parallel_all_reduce(output_parallel) + return output_parallel def extra_repr(self) -> str: s = f"num_embeddings={self.num_embeddings_per_partition}" @@ -516,6 +531,7 @@ class ParallelLMHead(VocabParallelEmbedding): params_dtype: type of the parameters. org_num_embeddings: original vocabulary size (without LoRA). padding_size: padding size for the vocabulary. + disable_tp: If true, tensor parallelism will be disabled for this layer. """ # --8<-- [end:parallel_lm_head] @@ -530,6 +546,8 @@ class ParallelLMHead(VocabParallelEmbedding): padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, quant_config: QuantizationConfig | None = None, prefix: str = "", + *, + disable_tp: bool = False, ): super().__init__( num_embeddings, @@ -539,6 +557,7 @@ class ParallelLMHead(VocabParallelEmbedding): padding_size, quant_config, prefix, + disable_tp=disable_tp, ) self.quant_config = quant_config if bias: diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 219819759a..aa2a9395de 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -24,7 +24,6 @@ from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, - VocabParallelEmbedding, ) from .qwen3_dflash import DFlashQwen3ForCausalLM, DFlashQwen3Model @@ -40,6 +39,10 @@ class DSparkMarkovHead(nn.Module): ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias (``draft_vocab_size``) added to the base draft logits. The two sizes coincide for full-vocab drafts. + + Both weights are replicated because the head runs sequentially for every + draft position. Sharding them would add an all-reduce and a full-vocab + gather to each position. """ def __init__( @@ -50,19 +53,24 @@ class DSparkMarkovHead(nn.Module): prefix: str, ) -> None: super().__init__() - # TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard - self.markov_w1 = VocabParallelEmbedding( - vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1") - ) + self.markov_w1 = nn.Embedding(vocab_size, markov_rank) self.markov_w2 = ParallelLMHead( - draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2") + draft_vocab_size, + markov_rank, + bias=False, + prefix=maybe_prefix(prefix, "markov_w2"), + disable_tp=True, ) def embed(self, token_ids: torch.Tensor) -> torch.Tensor: """r-dim Markov embedding of ``token_ids`` ([B] -> [B, r]).""" return self.markov_w1(token_ids) - def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor: + def bias( + self, + markov_embed: torch.Tensor, + logits_processor: LogitsProcessor, + ) -> torch.Tensor: """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V]).""" return logits_processor(self.markov_w2, markov_embed)