diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index 051777aad6..1c99cc2533 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/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, diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index fdfc5b8be4..6650a56de1 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/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, )