speculative decoding under tensor parallelism (TP>1) , workspace creation select max hidden dim of target and draft model (#52193)
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
Khushali Desai committed
Aug 24, 2026 at 14:33 UTC
23ab0cfdbc432888549d6f30bb65906aac5cdcda
2 files changed
+66
-3
tests/compile/passes/distributed/test_fusion_all_reduce.py
+44
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from importlib.util import find_spec
4
+from types import SimpleNamespace
5
6
import pytest
7
import torch
@@ -13,6 +14,7 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
14
from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
15
AllReduceFusionPass,
16
RocmAiterAllReduceFusionPass,
17
+ _fused_ar_workspace_hidden_dim,
18
_select_flashinfer_allreduce_use_oneshot,
19
)
20
from vllm.compilation.passes.fx_utils import find_op_nodes
@@ -79,6 +81,48 @@ def test_select_flashinfer_allreduce_use_oneshot(
81
)
82
83
84
+def _make_workspace_hidden_config(
85
+ target_hidden: int, draft_hidden: int | None
86
+) -> SimpleNamespace:
87
+ """Minimal stand-in exposing only what _fused_ar_workspace_hidden_dim reads."""
88
+ speculative_config = None
89
+ if draft_hidden is not None:
90
+ speculative_config = SimpleNamespace(
91
+ draft_model_config=SimpleNamespace(
92
+ get_hidden_size=lambda: draft_hidden,
93
+ )
94
+ )
95
+ return SimpleNamespace(
96
+ model_config=SimpleNamespace(get_hidden_size=lambda: target_hidden),
97
+ speculative_config=speculative_config,
98
+ )
99
+
100
+
101
+@pytest.mark.parametrize(
102
+ ("target_hidden", "draft_hidden", "expected"),
103
+ [
104
+ (2048, None, 2048), # no speculative decoding -> target only
105
+ (2048, 2560, 2560), # draft wider than target -> draft (issue #52023)
106
+ (2048, 1024, 2048), # draft narrower than target -> target
107
+ (2048, 2048, 2048), # equal -> target
108
+ ],
109
+)
110
+def test_fused_ar_workspace_hidden_dim(
111
+ target_hidden: int, draft_hidden: int | None, expected: int
112
+):
113
+ config = _make_workspace_hidden_config(target_hidden, draft_hidden)
114
+ assert _fused_ar_workspace_hidden_dim(config) == expected
115
+
116
+
117
+def test_fused_ar_workspace_hidden_dim_spec_without_draft():
118
+ # speculative_config present but no draft model (e.g. ngram) -> target only.
119
+ config = SimpleNamespace(
120
+ model_config=SimpleNamespace(get_hidden_size=lambda: 2048),
121
+ speculative_config=SimpleNamespace(draft_model_config=None),
122
+ )
123
+ assert _fused_ar_workspace_hidden_dim(config) == 2048
124
+
125
+
126
class TestAllReduceRMSNormModel(torch.nn.Module):
127
def __init__(
128
self,
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+22
-3
@@ -989,6 +989,23 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
989
)
990
991
992
+def _fused_ar_workspace_hidden_dim(config: VllmConfig) -> int:
993
+ """Widest hidden size across the target and (optional) draft models.
994
+
995
+ The FlashInfer allreduce+RMSNorm workspace is a process-global singleton
996
+ created eagerly at pass construction and reused by every model. Under
997
+ speculative decoding the draft shares it, so it must fit the larger of the
998
+ two hidden sizes; a draft wider than the target otherwise overflows the
999
+ target-sized buffer (vLLM #52023).
1000
+ """
1001
+ hidden_dim = config.model_config.get_hidden_size()
1002
+ spec = config.speculative_config
1003
+ draft_model_config = getattr(spec, "draft_model_config", None) if spec else None
1004
+ if draft_model_config is not None:
1005
+ hidden_dim = max(hidden_dim, draft_model_config.get_hidden_size())
1006
+ return hidden_dim
1007
+
1008
+
1009
class AllReduceFusionPass(VllmPatternMatcherPass):
1010
def __init__(self, config: VllmConfig) -> None:
1011
super().__init__(config)
@@ -1005,7 +1022,9 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1022
"AllReduce fusion pass is disabled for missing model_config."
1023
)
1024
return
1008
- self.hidden_dim = config.model_config.get_hidden_size()
1025
+ # Size the shared workspace for the widest model that reuses it (the
1026
+ # draft under speculative decoding may be wider than the target).
1027
+ self.workspace_hidden_dim = _fused_ar_workspace_hidden_dim(config)
1028
self.group = get_tp_group().cpu_group
1029
rank = get_tensor_model_parallel_rank()
1030
if flashinfer_comm is None:
@@ -1026,7 +1045,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1045
)
1046
return
1047
element_size = torch.tensor([], dtype=self.model_dtype).element_size()
1029
- self.max_token_num = max_size // (self.hidden_dim * element_size)
1048
+ self.max_token_num = max_size // (self.workspace_hidden_dim * element_size)
1049
# take the min to save workspace size and we'll never use more
1050
# than max_num_batched_tokens anyways
1051
self.max_token_num = min(
@@ -1043,7 +1062,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
1062
world_size=self.tp_size,
1063
rank=rank,
1064
max_token_num=self.max_token_num,
1046
- hidden_dim=self.hidden_dim,
1065
+ hidden_dim=self.workspace_hidden_dim,
1066
dtype=self.model_dtype,
1067
group=self.group,
1068
)