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 15, 21:00 UTC
[Misc][Platform] Move use allgather to platform (#14010)
Signed-off-by: Mengqing Cao <cmq0113@163.com>
Mengqing Cao committed
Mar 3, 2025 at 15:40 UTC
b87c21fc89c772d231cae97346e0457ef3bb1bf9
4 files changed
+24
-7
vllm/model_executor/layers/logits_processor.py
+3
-7
@@ -8,7 +8,6 @@ import torch
8
import torch.nn as nn
9
10
import vllm.envs as envs
11
-from vllm.config import get_current_vllm_config
11
from vllm.distributed import (tensor_model_parallel_all_gather,
12
tensor_model_parallel_gather)
13
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -51,11 +50,7 @@ class LogitsProcessor(nn.Module):
50
# Soft cap the logits. Used in Gemma 2.
51
self.soft_cap = soft_cap
52
# Whether to use gather or all-gather to gather the logits.
54
- parallel_config = get_current_vllm_config().parallel_config
55
- self.use_all_gather = current_platform.is_tpu() \
56
- or current_platform.is_neuron() \
57
- or envs.VLLM_USE_V1 \
58
- or parallel_config.distributed_executor_backend == "external_launcher" # noqa
53
+ self.use_all_gather = current_platform.use_all_gather()
54
55
def forward(
56
self,
@@ -83,7 +78,8 @@ class LogitsProcessor(nn.Module):
78
logits *= self.scale
79
80
# Apply logits processors (if any).
86
- if sampling_metadata is not None:
81
+ if sampling_metadata is not None and \
82
+ sampling_metadata.seq_groups is not None:
83
logits = _apply_logits_processors(logits, sampling_metadata)
84
85
return logits
vllm/platforms/interface.py
+13
@@ -330,6 +330,19 @@ class Platform:
330
"""
331
return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa
332
333
+ @classmethod
334
+ def use_all_gather(cls) -> bool:
335
+ """
336
+ Whether to use allgather in LogitsProcessor to gather the logits.
337
+ """
338
+ import vllm.envs as envs
339
+ from vllm.config import get_current_vllm_config
340
+
341
+ parallel_config = get_current_vllm_config().parallel_config
342
+ return (envs.VLLM_USE_V1
343
+ or parallel_config.distributed_executor_backend
344
+ == "external_launcher")
345
+
346
347
class UnspecifiedPlatform(Platform):
348
_enum = PlatformEnum.UNSPECIFIED
vllm/platforms/neuron.py
+4
@@ -55,3 +55,7 @@ class NeuronPlatform(Platform):
55
def is_pin_memory_available(cls) -> bool:
56
logger.warning("Pin memory is not supported on Neuron.")
57
return False
58
+
59
+ @classmethod
60
+ def use_all_gather(cls) -> bool:
61
+ return True
vllm/platforms/tpu.py
+4
@@ -119,3 +119,7 @@ class TpuPlatform(Platform):
119
@classmethod
120
def get_device_communicator_cls(cls) -> str:
121
return "vllm.distributed.device_communicators.tpu_communicator.TpuCommunicator" # noqa
122
+
123
+ @classmethod
124
+ def use_all_gather(cls) -> bool:
125
+ return True