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>

tests/compile/passes/distributed/test_fusion_all_reduce.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from importlib.util import find_spec
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -13,6 +14,7 @@ from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
     AllReduceFusionPass,
     RocmAiterAllReduceFusionPass,
+    _fused_ar_workspace_hidden_dim,
     _select_flashinfer_allreduce_use_oneshot,
 )
 from vllm.compilation.passes.fx_utils import find_op_nodes
@@ -79,6 +81,48 @@ def test_select_flashinfer_allreduce_use_oneshot(
     )
 
 
+def _make_workspace_hidden_config(
+    target_hidden: int, draft_hidden: int | None
+) -> SimpleNamespace:
+    """Minimal stand-in exposing only what _fused_ar_workspace_hidden_dim reads."""
+    speculative_config = None
+    if draft_hidden is not None:
+        speculative_config = SimpleNamespace(
+            draft_model_config=SimpleNamespace(
+                get_hidden_size=lambda: draft_hidden,
+            )
+        )
+    return SimpleNamespace(
+        model_config=SimpleNamespace(get_hidden_size=lambda: target_hidden),
+        speculative_config=speculative_config,
+    )
+
+
+@pytest.mark.parametrize(
+    ("target_hidden", "draft_hidden", "expected"),
+    [
+        (2048, None, 2048),  # no speculative decoding -> target only
+        (2048, 2560, 2560),  # draft wider than target -> draft (issue #52023)
+        (2048, 1024, 2048),  # draft narrower than target -> target
+        (2048, 2048, 2048),  # equal -> target
+    ],
+)
+def test_fused_ar_workspace_hidden_dim(
+    target_hidden: int, draft_hidden: int | None, expected: int
+):
+    config = _make_workspace_hidden_config(target_hidden, draft_hidden)
+    assert _fused_ar_workspace_hidden_dim(config) == expected
+
+
+def test_fused_ar_workspace_hidden_dim_spec_without_draft():
+    # speculative_config present but no draft model (e.g. ngram) -> target only.
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(get_hidden_size=lambda: 2048),
+        speculative_config=SimpleNamespace(draft_model_config=None),
+    )
+    assert _fused_ar_workspace_hidden_dim(config) == 2048
+
+
 class TestAllReduceRMSNormModel(torch.nn.Module):
     def __init__(
         self,

vllm/compilation/passes/fusion/allreduce_rms_fusion.py

@@ -989,6 +989,23 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
         )
 
 
+def _fused_ar_workspace_hidden_dim(config: VllmConfig) -> int:
+    """Widest hidden size across the target and (optional) draft models.
+
+    The FlashInfer allreduce+RMSNorm workspace is a process-global singleton
+    created eagerly at pass construction and reused by every model. Under
+    speculative decoding the draft shares it, so it must fit the larger of the
+    two hidden sizes; a draft wider than the target otherwise overflows the
+    target-sized buffer (vLLM #52023).
+    """
+    hidden_dim = config.model_config.get_hidden_size()
+    spec = config.speculative_config
+    draft_model_config = getattr(spec, "draft_model_config", None) if spec else None
+    if draft_model_config is not None:
+        hidden_dim = max(hidden_dim, draft_model_config.get_hidden_size())
+    return hidden_dim
+
+
 class AllReduceFusionPass(VllmPatternMatcherPass):
     def __init__(self, config: VllmConfig) -> None:
         super().__init__(config)
@@ -1005,7 +1022,9 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
                 "AllReduce fusion pass is disabled for missing model_config."
             )
             return
-        self.hidden_dim = config.model_config.get_hidden_size()
+        # Size the shared workspace for the widest model that reuses it (the
+        # draft under speculative decoding may be wider than the target).
+        self.workspace_hidden_dim = _fused_ar_workspace_hidden_dim(config)
         self.group = get_tp_group().cpu_group
         rank = get_tensor_model_parallel_rank()
         if flashinfer_comm is None:
@@ -1026,7 +1045,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
             )
             return
         element_size = torch.tensor([], dtype=self.model_dtype).element_size()
-        self.max_token_num = max_size // (self.hidden_dim * element_size)
+        self.max_token_num = max_size // (self.workspace_hidden_dim * element_size)
         # take the min to save workspace size and we'll never use more
         # than max_num_batched_tokens anyways
         self.max_token_num = min(
@@ -1043,7 +1062,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):
             world_size=self.tp_size,
             rank=rank,
             max_token_num=self.max_token_num,
-            hidden_dim=self.hidden_dim,
+            hidden_dim=self.workspace_hidden_dim,
             dtype=self.model_dtype,
             group=self.group,
         )