@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 17871983a2

Fork 0

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 14, 05:30 UTC

[Bugfix] Fix sequence parallelism bug when enable pipeline parallelism (#24021)

Signed-off-by: cascade812 <cascade812@outlook.com>

cascade committed Sep 15, 2025 at 21:32 UTC 17871983a2efde49d93ee8692f17da490da47e95

6 files changed +135 -42

tests/distributed/test_sequence_parallel.py

+2 -1

@@ -235,7 +235,6 @@ def _compare_sp(

235 'level': 3,

236 'custom_ops': ["+rms_norm"],

237 'compile_sizes': [4, 8],

238 - 'splitting_ops': [],

238 'pass_config': {

239 'enable_sequence_parallelism': True,

240 'enable_fusion': enable_fusion,

@@ -251,6 +250,8 @@ def _compare_sp(

250 *common_args,

251 "--tensor-parallel-size",

252 str(tp_size),

253 + "--pipeline-parallel-size",

254 + str(pp_size),

255 "--distributed-executor-backend",

256 distributed_backend,

257 "--compilation_config",

vllm/distributed/parallel_state.py

+46 -5

@@ -663,14 +663,29 @@ class GroupCoordinator:

663 tensor_dict: dict[str, Union[torch.Tensor, Any]],

664 dst: Optional[int] = None,

665 all_gather_group: Optional["GroupCoordinator"] = None,

666 + all_gather_tensors: Optional[dict[str, bool]] = None,

667 ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

668 """Send the input tensor dictionary.

669 NOTE: `dst` is the local rank of the source rank.

670 +

671 + all_gather_group: The group for the all-gather operation. If provided,

672 + an optimization is enabled where each rank in the group sends a

673 + slice of a tensor and the receiver reconstructs it using an

674 + all-gather, which can improve performance. This is typically the

675 + tensor-parallel group.

676 + all_gather_tensors: A dictionary to specify which tensors should use

677 + the all-gather optimization, which is only effective when

678 + `all_gather_group` is provided. By default, this optimization is

679 + on for any tensor whose size is divisible by the

680 + `all_gather_group`'s world size. However, it should be disabled

681 + for tensors that are not fully replicated across the group (e.g.,

682 + the residual tensor when sequence parallelism is enabled). This

683 + dictionary allows overriding the default behavior on a per-tensor

684 + basis.

685 """

686 # Bypass the function if we are using only 1 GPU.

687 if not torch.distributed.is_initialized() or self.world_size == 1:

688 return tensor_dict

673 -

689 all_gather_size = (1 if all_gather_group is None else

690 all_gather_group.world_size)

691 all_gather_rank = (0 if all_gather_group is None else

@@ -699,14 +714,23 @@ class GroupCoordinator:

714 # `send_object_list` has serialization & deserialization,

715 # all happening on CPU. Therefore, we can use the CPU group.

716 self.send_object(metadata_list, dst=dst)

702 - for tensor in tensor_list:

717 +

718 + tensor_keys = [

719 + k for k, v in tensor_dict.items() if isinstance(v, torch.Tensor)

720 + ]

721 + assert len(tensor_keys) == len(tensor_list)

722 +

723 + for key, tensor in zip(tensor_keys, tensor_list):

724 if tensor.numel() == 0:

725 # Skip sending empty tensors.

726 continue

727

728 # send-allgather: send only a slice, then do allgather.

708 - if (all_gather_group is not None

709 - and tensor.numel() % all_gather_size == 0):

729 + use_all_gather = (all_gather_group is not None

730 + and tensor.numel() % all_gather_size == 0)

731 + use_all_gather = all_gather_tensors.get(key, use_all_gather) \

732 + if all_gather_tensors else use_all_gather

733 + if use_all_gather:

734 tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]

735

736 if tensor.is_cpu:

@@ -725,14 +749,29 @@ class GroupCoordinator:

749 self,

750 src: Optional[int] = None,

751 all_gather_group: Optional["GroupCoordinator"] = None,

752 + all_gather_tensors: Optional[dict[str, bool]] = None,

753 ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

754 """Recv the input tensor dictionary.

755 NOTE: `src` is the local rank of the source rank.

756 +

757 + all_gather_group: The group for the all-gather operation. If provided,

758 + an optimization is enabled where each rank in the group sends a

759 + slice of a tensor and the receiver reconstructs it using an

760 + all-gather, which can improve performance. This is typically the

761 + tensor-parallel group.

762 + all_gather_tensors: A dictionary to specify which tensors should use

763 + the all-gather optimization, which is only effective when

764 + `all_gather_group` is provided. By default, this optimization is

765 + on for any tensor whose size is divisible by the

766 + `all_gather_group`'s world size. However, it should be disabled

767 + for tensors that are not fully replicated across the group (e.g.,

768 + the residual tensor when sequence parallelism is enabled). This

769 + dictionary allows overriding the default behavior on a per-tensor

770 + basis.

771 """

772 # Bypass the function if we are using only 1 GPU.

773 if not torch.distributed.is_initialized() or self.world_size == 1:

774 return None

735 -

775 all_gather_size = (1 if all_gather_group is None else

776 all_gather_group.world_size)

777 all_gather_rank = (0 if all_gather_group is None else

@@ -766,6 +805,8 @@ class GroupCoordinator:

805 # send-allgather: send only a slice, then do allgather.

806 use_all_gather = (all_gather_group is not None

807 and tensor.numel() % all_gather_size == 0)

808 + use_all_gather = all_gather_tensors.get(key, use_all_gather) \

809 + if all_gather_tensors else use_all_gather

810

811 if use_all_gather:

812 orig_shape = tensor.shape

vllm/v1/worker/cpu_worker.py

+15 -3

@@ -19,6 +19,7 @@ from vllm.v1.outputs import ModelRunnerOutput

19 from vllm.v1.worker.cpu_model_runner import CPUModelRunner

20 from vllm.v1.worker.gpu_worker import (Worker,

21 init_worker_distributed_environment)

22 +from vllm.v1.worker.utils import is_residual_scattered_for_sp

24 logger = init_logger(__name__)

@@ -107,18 +108,29 @@ class CPUWorker(Worker):

108 scheduler_output: "SchedulerOutput",

109 ) -> Optional[ModelRunnerOutput]:

110 intermediate_tensors = None

111 + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

112 + num_input_tokens = self.model_runner._get_num_input_tokens(

113 + num_scheduled_tokens)

114 + all_gather_tensors = {

115 + "residual":

116 + not is_residual_scattered_for_sp(self.vllm_config,

117 + num_input_tokens)

118 + }

119 if not get_pp_group().is_first_rank:

120 intermediate_tensors = IntermediateTensors(

121 get_pp_group().recv_tensor_dict(

113 - all_gather_group=get_tp_group()))

122 + all_gather_group=get_tp_group(),

123 + all_gather_tensors=all_gather_tensors))

124

125 output = self.model_runner.execute_model(scheduler_output,

126 intermediate_tensors)

127

128 if not get_pp_group().is_last_rank:

129 assert isinstance(output, IntermediateTensors)

120 - get_pp_group().send_tensor_dict(output.tensors,

121 - all_gather_group=get_tp_group())

130 + get_pp_group().send_tensor_dict(

131 + output.tensors,

132 + all_gather_group=get_tp_group(),

133 + all_gather_tensors=all_gather_tensors)

134 return None

135

136 assert isinstance(output, ModelRunnerOutput)

vllm/v1/worker/gpu_model_runner.py

+33 -30

@@ -88,6 +88,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

88 from vllm.v1.worker.kv_connector_model_runner_mixin import (

89 KVConnectorModelRunnerMixin)

90 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

91 +from vllm.v1.worker.utils import is_residual_scattered_for_sp

93 from .utils import (AttentionGroup, MultiModalBudget,

94 add_kv_sharing_layers_to_kv_cache_groups, bind_kv_cache,

@@ -1633,21 +1634,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1634 assert self.intermediate_tensors is not None

1635

1636 tp = self.vllm_config.parallel_config.tensor_parallel_size

1636 - enabled_sp = self.compilation_config.pass_config. \

1637 - enable_sequence_parallelism

1638 - if enabled_sp:

1639 - # When sequence parallelism is enabled, we always pad num_tokens

1640 - # to be a multiple of tensor_parallel_size (tp) earlier

1641 - assert num_tokens % tp == 0

1642 - is_residual_scattered = tp > 1 and enabled_sp \

1643 - and num_tokens % tp == 0

1637 + is_rs = is_residual_scattered_for_sp(self.vllm_config, num_tokens)

1638

1639 # When sequence parallelism is enabled, the "residual" tensor is sharded

1640 # across tensor parallel ranks, so each rank only needs its own slice.

1641 if sync_self:

1642 assert intermediate_tensors is not None

1643 for k, v in intermediate_tensors.items():

1650 - is_scattered = k == "residual" and is_residual_scattered

1644 + is_scattered = k == "residual" and is_rs

1645 copy_len = num_tokens // tp if is_scattered else \

1646 num_tokens

1647 self.intermediate_tensors[k][:copy_len].copy_(

@@ -1655,8 +1649,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1649

1650 return IntermediateTensors({

1651 k:

1658 - v[:num_tokens // tp]

1659 - if k == "residual" and is_residual_scattered else v[:num_tokens]

1652 + v[:num_tokens //

1653 + tp] if k == "residual" and is_rs else v[:num_tokens]

1654 for k, v in self.intermediate_tensors.items()

1655 })

1656

@@ -1741,6 +1735,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1735 pooler_output=pooler_output,

1736 )

1737

1738 + def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:

1739 + if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

1740 + and not envs.VLLM_DISABLE_PAD_FOR_CUDAGRAPH

1741 + and hasattr(self, "cudagraph_batch_sizes")

1742 + and self.cudagraph_batch_sizes

1743 + and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):

1744 + # Use CUDA graphs.

1745 + # Add padding to the batch size.

1746 + return self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)

1747 +

1748 + # Eager mode.

1749 + # Pad tokens to multiple of tensor_parallel_size when

1750 + # enabled collective fusion for SP

1751 + tp_size = self.vllm_config.parallel_config.tensor_parallel_size

1752 + if (self.compilation_config.pass_config.enable_sequence_parallelism

1753 + and tp_size > 1):

1754 + return round_up(num_scheduled_tokens, tp_size)

1755 + return num_scheduled_tokens

1756 +

1757 def _preprocess(

1758 self,

1759 scheduler_output: "SchedulerOutput",

@@ -1750,24 +1763,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1763 Optional[IntermediateTensors], dict[str, Any]]:

1764

1765 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

1753 - if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

1754 - and not envs.VLLM_DISABLE_PAD_FOR_CUDAGRAPH

1755 - and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):

1756 - # Use CUDA graphs.

1757 - # Add padding to the batch size.

1758 - num_input_tokens = self.vllm_config.pad_for_cudagraph(

1759 - num_scheduled_tokens)

1760 - else:

1761 - # Eager mode.

1762 - # Pad tokens to multiple of tensor_parallel_size when

1763 - # enabled collective fusion for SP

1764 - tp_size = self.vllm_config.parallel_config.tensor_parallel_size

1765 - if self.compilation_config.pass_config. \

1766 - enable_sequence_parallelism and tp_size > 1:

1767 - num_input_tokens = round_up(num_scheduled_tokens, tp_size)

1768 - else:

1769 - num_input_tokens = num_scheduled_tokens

1770 -

1766 + num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)

1767 # Padding for DP

1768 num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)

1769 num_input_tokens += num_pad

@@ -2108,8 +2104,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2104 assert not self.is_pooling_model

2105

2106 if not get_pp_group().is_last_rank:

2107 + all_gather_tensors = {

2108 + "residual":

2109 + not is_residual_scattered_for_sp(

2110 + self.vllm_config, num_input_tokens)

2111 + }

2112 get_pp_group().send_tensor_dict(

2112 - hidden_states.tensors, all_gather_group=get_tp_group())

2113 + hidden_states.tensors,

2114 + all_gather_group=get_tp_group(),

2115 + all_gather_tensors=all_gather_tensors)

2116 logits = None

2117 else:

2118 sample_hidden_states = hidden_states[logits_indices]

vllm/v1/worker/gpu_worker.py

+13 -2

@@ -32,6 +32,7 @@ from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput,

32 DraftTokenIds, ModelRunnerOutput)

33 from vllm.v1.utils import report_usage_stats

34 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

35 +from vllm.v1.worker.utils import is_residual_scattered_for_sp

36 from vllm.v1.worker.worker_base import WorkerBase

38 logger = init_logger(__name__)

@@ -428,10 +429,19 @@ class Worker(WorkerBase):

429 ) -> Optional[Union[ModelRunnerOutput, AsyncModelRunnerOutput]]:

430 intermediate_tensors = None

431 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

432 + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

433 + num_input_tokens = self.model_runner._get_num_input_tokens(

434 + num_scheduled_tokens)

435 + all_gather_tensors = {

436 + "residual":

437 + not is_residual_scattered_for_sp(self.vllm_config,

438 + num_input_tokens)

439 + }

440 if forward_pass and not get_pp_group().is_first_rank:

441 intermediate_tensors = IntermediateTensors(

442 get_pp_group().recv_tensor_dict(

434 - all_gather_group=get_tp_group()))

443 + all_gather_group=get_tp_group(),

444 + all_gather_tensors=all_gather_tensors))

445

446 output = self.model_runner.execute_model(scheduler_output,

447 intermediate_tensors)

@@ -444,7 +454,8 @@ class Worker(WorkerBase):

454 "external_launcher") and not get_pp_group().is_last_rank

455

456 get_pp_group().send_tensor_dict(output.tensors,

447 - all_gather_group=get_tp_group())

457 + all_gather_group=get_tp_group(),

458 + all_gather_tensors=all_gather_tensors)

459

460 kv_connector_output = output.kv_connector_output

461 if not kv_connector_output:

vllm/v1/worker/utils.py

+26 -1

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Optional

7 import torch

9 from vllm.attention.backends.abstract import AttentionBackend

10 -from vllm.config import ModelConfig, SchedulerConfig

10 +from vllm.config import ModelConfig, SchedulerConfig, VllmConfig

11 from vllm.model_executor.models.interfaces import MultiModalEmbeddings

12 from vllm.model_executor.models.utils import extract_layer_index

13 from vllm.multimodal.cache import processor_only_cache_from_config

@@ -288,3 +288,28 @@ def bind_kv_cache(

288 for layer_name, kv_cache in kv_caches.items():

289 # NOTE: Use list because of v0 PP virtual engine.

290 forward_context[layer_name].kv_cache = [kv_cache]

291 +

292 +

293 +def is_residual_scattered_for_sp(vllm_config: VllmConfig,

294 + num_input_tokens: int) -> bool:

295 + """Check if the residual tensor is scattered for sequence parallelism.

296 +

297 + The residual tensor is scattered across tensor parallel ranks when sequence

298 + parallelism and tensor parallelism is enabled, and the number of

299 + input tokens is one of the compilation sizes.

300 + """

301 + if not vllm_config.compilation_config.pass_config.\

302 + enable_sequence_parallelism:

303 + return False

304 +

305 + tp = vllm_config.parallel_config.tensor_parallel_size

306 +

307 + if tp == 1:

308 + return False

309 +

310 + # When sequence parallelism is enabled, we always pad num_input_tokens

311 + # to be a multiple of tensor_parallel_size (tp) earlier.

312 + assert num_input_tokens % tp == 0

313 +

314 + # Currently, SP is only enabled for static size fx graphs.

315 + return (num_input_tokens in vllm_config.compilation_config.compile_sizes)