A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 22:45 UTC
[Spec Decode] Reduce TP communication for speculative decoding draft token generation (#34049)
Signed-off-by: qizixi <qizixi@meta.com> Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com>
qizixi committed
Feb 22, 2026 at 14:59 UTC
2bcf71b9c0305e2a3d645e7a5ced4460262b4a6c
4 files changed
+114
-6
vllm/config/speculative.py
+5
@@ -109,6 +109,11 @@ class SpeculativeConfig:
109
speculative input batches can contain sequences of different lengths,
110
which may only be supported by certain attention backends. This currently
111
only affects the EAGLE method of speculation."""
112
+ use_local_argmax_reduction: bool = False
113
+ """Use vocab-parallel local argmax instead of all-gathering full logits
114
+ for draft token generation. Reduces communication from O(vocab_size) to
115
+ O(2 * tp_size) per token. Only applies to greedy draft selection in
116
+ non-tree speculation."""
117
118
# Ngram proposer configuration
119
prompt_lookup_max: int | None = Field(default=None, ge=1)
vllm/model_executor/layers/logits_processor.py
+53
@@ -5,6 +5,7 @@
5
import torch
6
7
from vllm.distributed import (
8
+ get_tensor_model_parallel_world_size,
9
tensor_model_parallel_all_gather,
10
tensor_model_parallel_gather,
11
)
@@ -102,6 +103,58 @@ class LogitsProcessor(CustomOp):
103
logits = logits[..., : self.org_vocab_size]
104
return logits
105
106
+ def get_top_tokens(
107
+ self,
108
+ lm_head: VocabParallelEmbedding,
109
+ hidden_states: torch.Tensor,
110
+ embedding_bias: torch.Tensor | None = None,
111
+ ) -> torch.Tensor:
112
+ """Vocab-parallel argmax without all-gathering full logits.
113
+
114
+ Each TP rank computes local argmax, then only the (value, index) pairs
115
+ are gathered and reduced. Communication: O(batch * 2 * tp_size) vs
116
+ O(batch * vocab_size).
117
+ """
118
+ if self.scale <= 0.0 and self.scale != 1.0:
119
+ raise ValueError(
120
+ "The local argmax reduction optimization is not supported for "
121
+ "non-positive logit scaling factors."
122
+ )
123
+ tp_size = get_tensor_model_parallel_world_size()
124
+
125
+ logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)
126
+ if self.soft_cap is not None:
127
+ logits = torch.tanh(logits / self.soft_cap) * self.soft_cap
128
+ if self.scale != 1.0:
129
+ logits = logits * self.scale
130
+
131
+ # Mask out padding entries beyond org_vocab_size on this shard.
132
+ num_pad = lm_head.shard_indices.num_org_vocab_padding
133
+ if num_pad > 0:
134
+ logits[..., -num_pad:] = -float("inf")
135
+
136
+ local_max_vals, local_max_indices = logits.max(dim=-1)
137
+
138
+ # Convert shard-local indices to global vocab indices.
139
+ vocab_start = lm_head.shard_indices.org_vocab_start_index
140
+ global_indices = local_max_indices + vocab_start
141
+
142
+ if tp_size == 1:
143
+ return global_indices
144
+
145
+ # All-gather (value, index) pairs, then reduce to global argmax.
146
+ # Use float32 to avoid bf16 precision loss on large vocab indices.
147
+ local_pair = torch.stack(
148
+ [local_max_vals.float(), global_indices.float()], dim=-1
149
+ )
150
+ # [batch, 2] -> [batch, 2 * tp_size]
151
+ gathered = tensor_model_parallel_all_gather(local_pair, dim=-1)
152
+ # [batch, tp_size, 2] where [:, :, 0]=values, [:, :, 1]=indices
153
+ gathered = gathered.view(hidden_states.shape[0], tp_size, 2)
154
+ max_rank_idx = gathered[:, :, 0].argmax(dim=-1, keepdim=True)
155
+ top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx)
156
+ return top_tokens.squeeze(-1).to(torch.int64)
157
+
158
def extra_repr(self) -> str:
159
s = f"vocab_size={self.vocab_size}"
160
s += f", org_vocab_size={self.org_vocab_size}"
vllm/model_executor/models/llama4_eagle.py
+17
@@ -208,6 +208,23 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
208
) -> tuple[torch.Tensor, torch.Tensor]:
209
return self.model(input_ids, positions, hidden_states, inputs_embeds)
210
211
+ def get_top_tokens(
212
+ self,
213
+ hidden_states: torch.Tensor,
214
+ ) -> torch.Tensor:
215
+ """Vocab-parallel argmax without all-gathering full logits.
216
+
217
+ Falls back to full logits when draft_id_to_target_id remapping is
218
+ active, since the shared lm_head covers the full target vocab but
219
+ the draft model only predicts over a subset (draft_vocab_size).
220
+ """
221
+ if (
222
+ hasattr(self, "draft_id_to_target_id")
223
+ and self.draft_id_to_target_id is not None
224
+ ):
225
+ return self.compute_logits(hidden_states).argmax(dim=-1)
226
+ return self.logits_processor.get_top_tokens(self.lm_head, hidden_states)
227
+
228
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None:
229
def transform(inputs):
230
name, loaded_weight = inputs
vllm/v1/spec_decode/eagle.py
+39
-6
@@ -99,6 +99,9 @@ class SpecDecodeBaseProposer:
99
self.parallel_drafting_hidden_state_tensor: torch.Tensor | None = None
100
if self.parallel_drafting:
101
self._init_parallel_drafting_params()
102
+ self.use_local_argmax_reduction: bool = (
103
+ self.speculative_config.use_local_argmax_reduction
104
+ )
105
106
max_batch_size = vllm_config.scheduler_config.max_num_seqs
107
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -369,6 +372,12 @@ class SpecDecodeBaseProposer:
372
373
self.cudagraph_dispatcher.initialize_cudagraph_keys(eagle_cudagraph_mode)
374
375
+ def _greedy_sample(self, hidden_states: torch.Tensor) -> torch.Tensor:
376
+ """Greedy-sample draft tokens from hidden states."""
377
+ if self.use_local_argmax_reduction:
378
+ return self.model.get_top_tokens(hidden_states)
379
+ return self.model.compute_logits(hidden_states).argmax(dim=-1)
380
+
381
def propose(
382
self,
383
# [num_tokens]
@@ -491,11 +500,10 @@ class SpecDecodeBaseProposer:
500
last_hidden_states, hidden_states = ret_hidden_states
501
502
sample_hidden_states = last_hidden_states[token_indices_to_sample]
494
- logits = self.model.compute_logits(sample_hidden_states)
503
504
# Early exit if there is only one draft token to be generated.
505
if self.num_speculative_tokens == 1 or self.parallel_drafting:
498
- draft_token_ids = logits.argmax(dim=-1)
506
+ draft_token_ids = self._greedy_sample(sample_hidden_states)
507
return draft_token_ids.view(-1, self.num_speculative_tokens)
508
509
if self.uses_mrope:
@@ -513,7 +521,8 @@ class SpecDecodeBaseProposer:
521
hidden_states = hidden_states[token_indices_to_sample]
522
523
if isinstance(attn_metadata, TreeAttentionMetadata):
516
- # Draft using tree attention.
524
+ # Draft using tree attention - requires full logits for top-k
525
+ logits = self.model.compute_logits(sample_hidden_states)
526
draft_token_ids_list = self.propose_tree(
527
batch_size=batch_size,
528
logits=logits,
@@ -525,7 +534,7 @@ class SpecDecodeBaseProposer:
534
# [batch_size, num_tree_tokens]
535
return torch.cat(draft_token_ids_list, dim=1)
536
528
- draft_token_ids = logits.argmax(dim=-1)
537
+ draft_token_ids = self._greedy_sample(sample_hidden_states)
538
539
if self.allowed_attn_types is not None and not isinstance(
540
attn_metadata, self.allowed_attn_types
@@ -690,8 +699,7 @@ class SpecDecodeBaseProposer:
699
last_hidden_states, hidden_states = ret_hidden_states
700
701
hidden_states = hidden_states[:batch_size]
693
- logits = self.model.compute_logits(last_hidden_states[:batch_size])
694
- draft_token_ids = logits.argmax(dim=-1)
702
+ draft_token_ids = self._greedy_sample(last_hidden_states[:batch_size])
703
draft_token_ids_list.append(draft_token_ids)
704
705
# [batch_size, num_speculative_tokens]
@@ -1521,6 +1529,31 @@ class SpecDecodeBaseProposer:
1529
"Shared target model lm_head with MTP shared_head.head."
1530
)
1531
1532
+ if self.use_local_argmax_reduction:
1533
+ if not hasattr(self.model, "get_top_tokens"):
1534
+ raise ValueError(
1535
+ "use_local_argmax_reduction is enabled but draft model "
1536
+ f"{self.model.__class__.__name__} does not implement "
1537
+ "get_top_tokens()."
1538
+ )
1539
+ # Warn if draft model has vocab remapping, which forces fallback
1540
+ # to the full-logits path (negating the optimization).
1541
+ if (
1542
+ hasattr(self.model, "draft_id_to_target_id")
1543
+ and self.model.draft_id_to_target_id is not None
1544
+ ):
1545
+ logger.warning(
1546
+ "use_local_argmax_reduction is enabled but draft model "
1547
+ "uses draft_id_to_target_id vocab remapping. The "
1548
+ "optimization will be bypassed (falling back to full "
1549
+ "logits gather + argmax)."
1550
+ )
1551
+ else:
1552
+ logger.info(
1553
+ "Using local argmax reduction for draft token generation "
1554
+ "(communication: O(2*tp_size) vs O(vocab_size))."
1555
+ )
1556
+
1557
@torch.inference_mode()
1558
def dummy_run(
1559
self,