[Model Runner V2][Spec Decode] Reduce TP communication for draft token generation (#46448)
Signed-off-by: EanWang211123 <wangyiheng@sangfor.com.cn> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
yiheng committed
Jun 26, 2026 at 10:30 UTC
652d962bc9df7e04959e84ce478c3a8d26fe52a7
1 file changed
+34
-3
vllm/v1/worker/gpu/spec_decode/speculator.py
+34
-3
@@ -8,6 +8,7 @@ import torch.nn as nn
8
9
from vllm.config import VllmConfig, get_layers_from_vllm_config
10
from vllm.config.compilation import CUDAGraphMode
11
+from vllm.logger import init_logger
12
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
13
from vllm.v1.kv_cache_interface import KVCacheConfig
14
from vllm.v1.worker.gpu.attn_utils import (
@@ -23,6 +24,8 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
24
from vllm.v1.worker.gpu.model_states.interface import ModelState
25
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
26
27
+logger = init_logger(__name__)
28
+
29
30
class BaseSpeculator(ABC):
31
@abstractmethod
@@ -95,6 +98,9 @@ class DraftModelSpeculator(BaseSpeculator):
98
self.vocab_size = self.draft_model_config.get_vocab_size()
99
self.dtype = vllm_config.model_config.dtype
100
self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
101
+ self.use_local_argmax_reduction = (
102
+ self.speculative_config.use_local_argmax_reduction
103
+ )
104
105
# DP configuration
106
self.dp_size = vllm_config.parallel_config.data_parallel_size
@@ -149,6 +155,7 @@ class DraftModelSpeculator(BaseSpeculator):
155
)
156
157
self.model = self.load_draft_model(target_model, target_attn_layer_names)
158
+ self._validate_local_argmax_reduction()
159
160
all_attn_layers = set[str](
161
get_layers_from_vllm_config(
@@ -211,6 +218,31 @@ class DraftModelSpeculator(BaseSpeculator):
218
)
219
return attn_metadata
220
221
+ def _validate_local_argmax_reduction(self) -> None:
222
+ if not self.use_local_argmax_reduction:
223
+ return
224
+ if self.speculative_config.draft_sample_method == "probabilistic":
225
+ raise ValueError(
226
+ "use_local_argmax_reduction is not compatible with "
227
+ "draft_sample_method='probabilistic'."
228
+ )
229
+ if not hasattr(self.model, "get_top_tokens"):
230
+ raise ValueError(
231
+ "use_local_argmax_reduction is enabled but draft model "
232
+ f"{self.model.__class__.__name__} does not implement "
233
+ "get_top_tokens()."
234
+ )
235
+ logger.info(
236
+ "Using local argmax reduction for draft token generation "
237
+ "(communication: O(2*tp_size) vs O(vocab_size))."
238
+ )
239
+
240
+ def _greedy_sample_draft(self, hidden_states: torch.Tensor) -> torch.Tensor:
241
+ if self.use_local_argmax_reduction:
242
+ return self.model.get_top_tokens(hidden_states)
243
+ logits = self.model.compute_logits(hidden_states)
244
+ return logits.argmax(dim=-1)
245
+
246
def sample_draft(
247
self,
248
hidden_states: torch.Tensor,
@@ -221,8 +253,8 @@ class DraftModelSpeculator(BaseSpeculator):
253
draft_step: torch.Tensor,
254
draft_logits: torch.Tensor | None,
255
) -> torch.Tensor:
224
- logits = self.model.compute_logits(hidden_states)
256
if draft_logits is not None:
257
+ logits = self.model.compute_logits(hidden_states)
258
# NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
259
# used for draft and target sampling.
260
return gumbel_sample(
@@ -236,8 +268,7 @@ class DraftModelSpeculator(BaseSpeculator):
268
output_processed_logits_col=draft_step,
269
use_fp64=self.use_fp64_gumbel,
270
)
239
- else:
240
- return logits.argmax(dim=-1)
271
+ return self._greedy_sample_draft(hidden_states)
272
273
def _copy_request_inputs(
274
self,