@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 721fb9b181

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 15, 09:45 UTC

[Platform] Move platform check to right place (#18470)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

wangxiyuan committed May 23, 2025 at 03:11 UTC 721fb9b1818ef23c15fd176c7ea49285de544021

7 files changed +71 -25

vllm/config.py

+10 -25

@@ -42,7 +42,10 @@ from vllm.transformers_utils.config import (

42 try_get_generation_config, uses_mrope)

43 from vllm.transformers_utils.s3_utils import S3Model

44 from vllm.transformers_utils.utils import is_s3, maybe_model_redirect

45 -from vllm.utils import (GiB_bytes, LayerBlockType, cuda_device_count_stateless,

45 +from vllm.utils import (DEFAULT_MAX_NUM_BATCHED_TOKENS,

46 + MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,

47 + POOLING_MODEL_MAX_NUM_BATCHED_TOKENS, GiB_bytes,

48 + LayerBlockType, cuda_device_count_stateless,

49 get_cpu_memory, get_open_port, is_torch_equal_or_newer,

50 random_uuid, resolve_obj_by_qualname)

@@ -64,12 +67,6 @@ logger = init_logger(__name__)

68 ConfigT = TypeVar("ConfigT", bound=ConfigType)

67 -# This value is chosen to have a balance between ITL and TTFT. Note it is

68 -# not optimized for throughput.

69 -_DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048

70 -_POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768

71 -_MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120

72 -

70 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",

71 "score", "reward", "transcription"]

@@ -2074,28 +2071,28 @@ class SchedulerConfig:

2071 # so we don't reject sequences on account of a short

2072 # max_num_batched_tokens.

2073 self.max_num_batched_tokens = max(

2077 - self.max_model_len, _DEFAULT_MAX_NUM_BATCHED_TOKENS)

2074 + self.max_model_len, DEFAULT_MAX_NUM_BATCHED_TOKENS)

2075 else:

2076 self.max_num_batched_tokens = (

2080 - _DEFAULT_MAX_NUM_BATCHED_TOKENS)

2077 + DEFAULT_MAX_NUM_BATCHED_TOKENS)

2078 else:

2079 # If max_model_len is too short, use

2083 - # _DEFAULT_MAX_NUM_BATCHED_TOKENS as the default value

2080 + # DEFAULT_MAX_NUM_BATCHED_TOKENS as the default value

2081 # for higher throughput.

2082 self.max_num_batched_tokens = max(

2086 - self.max_model_len, _DEFAULT_MAX_NUM_BATCHED_TOKENS)

2083 + self.max_model_len, DEFAULT_MAX_NUM_BATCHED_TOKENS)

2084

2085 if self.runner_type == "pooling":

2086 # Choose specific value for higher throughput

2087 self.max_num_batched_tokens = max(

2088 self.max_num_batched_tokens,

2092 - _POOLING_MODEL_MAX_NUM_BATCHED_TOKENS,

2089 + POOLING_MODEL_MAX_NUM_BATCHED_TOKENS,

2090 )

2091 if self.is_multimodal_model:

2092 # The value needs to be at least the number of multimodal tokens

2093 self.max_num_batched_tokens = max(

2094 self.max_num_batched_tokens,

2098 - _MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,

2095 + MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,

2096 )

2097

2098 # When using default settings,

@@ -4316,18 +4313,6 @@ class VllmConfig:

4313 "full_cuda_graph is not supported with "

4314 "cascade attention. Disabling cascade attention.")

4315 self.model_config.disable_cascade_attn = True

4319 -

4320 - if self.model_config and self.model_config.use_mla and \

4321 - not (current_platform.is_cuda() or current_platform.is_rocm()):

4322 - logger.info(

4323 - "MLA is enabled on a non-GPU platform; forcing chunked "

4324 - "prefill and prefix caching to be disabled.")

4325 - self.scheduler_config.enable_chunked_prefill = False

4326 - self.scheduler_config.chunked_prefill_enabled = False

4327 - self.scheduler_config.max_num_batched_tokens = max(

4328 - self.scheduler_config.max_model_len,

4329 - _DEFAULT_MAX_NUM_BATCHED_TOKENS)

4330 -

4316 if self.cache_config is not None:

4317 self.cache_config.enable_prefix_caching = False

4318

vllm/platforms/cpu.py

+11

@@ -9,6 +9,7 @@ import psutil

9 import torch

11 from vllm.logger import init_logger

12 +from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS

14 from .interface import CpuArchEnum, Platform, PlatformEnum, _Backend

@@ -177,6 +178,16 @@ class CpuPlatform(Platform):

178 " set VLLM_WORKER_MULTIPROC_METHOD to fork explicitly.")

179 os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'

180

181 + if vllm_config.model_config and vllm_config.model_config.use_mla:

182 + logger.info(

183 + "MLA is enabled on a non-GPU platform; forcing chunked "

184 + "prefill and prefix caching to be disabled.")

185 + vllm_config.scheduler_config.enable_chunked_prefill = False

186 + vllm_config.scheduler_config.chunked_prefill_enabled = False

187 + vllm_config.scheduler_config.max_num_batched_tokens = max(

188 + vllm_config.scheduler_config.max_model_len,

189 + DEFAULT_MAX_NUM_BATCHED_TOKENS)

190 +

191 @classmethod

192 def is_pin_memory_available(cls) -> bool:

193 logger.warning("Pin memory is not supported on CPU.")

vllm/platforms/hpu.py

+11

@@ -7,6 +7,7 @@ import torch

8 from vllm import envs

9 from vllm.logger import init_logger

10 +from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS

12 from .interface import Platform, PlatformEnum, _Backend

@@ -80,6 +81,16 @@ class HpuPlatform(Platform):

81 "VLLM_WORKER_MULTIPROC_METHOD=fork explicitly.")

82 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

84 + if vllm_config.model_config and vllm_config.model_config.use_mla:

85 + logger.info(

86 + "MLA is enabled on a non-GPU platform; forcing chunked "

87 + "prefill and prefix caching to be disabled.")

88 + vllm_config.scheduler_config.enable_chunked_prefill = False

89 + vllm_config.scheduler_config.chunked_prefill_enabled = False

90 + vllm_config.scheduler_config.max_num_batched_tokens = max(

91 + vllm_config.scheduler_config.max_model_len,

92 + DEFAULT_MAX_NUM_BATCHED_TOKENS)

93 +

94 @classmethod

95 def is_pin_memory_available(cls):

96 logger.warning("Pin memory is not supported on HPU.")

vllm/platforms/neuron.py

+11

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

7 from vllm import envs

8 from vllm.logger import init_logger

9 +from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS

11 from .interface import Platform, PlatformEnum

@@ -56,6 +57,16 @@ class NeuronPlatform(Platform):

57 vllm_config.cache_config.block_size = \

58 vllm_config.model_config.max_model_len # type: ignore

60 + if vllm_config.model_config and vllm_config.model_config.use_mla:

61 + logger.info(

62 + "MLA is enabled on a non-GPU platform; forcing chunked "

63 + "prefill and prefix caching to be disabled.")

64 + vllm_config.scheduler_config.enable_chunked_prefill = False

65 + vllm_config.scheduler_config.chunked_prefill_enabled = False

66 + vllm_config.scheduler_config.max_num_batched_tokens = max(

67 + vllm_config.scheduler_config.max_model_len,

68 + DEFAULT_MAX_NUM_BATCHED_TOKENS)

69 +

70 @classmethod

71 def is_pin_memory_available(cls) -> bool:

72 logger.warning("Pin memory is not supported on Neuron.")

vllm/platforms/tpu.py

+11

@@ -9,6 +9,7 @@ import vllm.envs as envs

9 from vllm.inputs import ProcessorInputs, PromptType

10 from vllm.logger import init_logger

11 from vllm.sampling_params import SamplingParams, SamplingType

12 +from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS

14 from .interface import Platform, PlatformEnum, _Backend

@@ -161,6 +162,16 @@ class TpuPlatform(Platform):

162 "Forcing --disable_chunked_mm_input.")

163 scheduler_config.disable_chunked_mm_input = True

164

165 + if vllm_config.model_config and vllm_config.model_config.use_mla:

166 + logger.info(

167 + "MLA is enabled on a non-GPU platform; forcing chunked "

168 + "prefill and prefix caching to be disabled.")

169 + vllm_config.scheduler_config.enable_chunked_prefill = False

170 + vllm_config.scheduler_config.chunked_prefill_enabled = False

171 + vllm_config.scheduler_config.max_num_batched_tokens = max(

172 + vllm_config.scheduler_config.max_model_len,

173 + DEFAULT_MAX_NUM_BATCHED_TOKENS)

174 +

175 @classmethod

176 def is_pin_memory_available(cls):

177 logger.warning("Pin memory is not supported on TPU.")

vllm/platforms/xpu.py

+11

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

5 import torch

7 from vllm.logger import init_logger

8 +from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS

10 from .interface import DeviceCapability, Platform, PlatformEnum, _Backend

@@ -113,6 +114,16 @@ class XPUPlatform(Platform):

114 parallel_config.distributed_executor_backend)

115 parallel_config.distributed_executor_backend = "ray"

116

117 + if vllm_config.model_config and vllm_config.model_config.use_mla:

118 + logger.info(

119 + "MLA is enabled on a non-GPU platform; forcing chunked "

120 + "prefill and prefix caching to be disabled.")

121 + vllm_config.scheduler_config.enable_chunked_prefill = False

122 + vllm_config.scheduler_config.chunked_prefill_enabled = False

123 + vllm_config.scheduler_config.max_num_batched_tokens = max(

124 + vllm_config.scheduler_config.max_model_len,

125 + DEFAULT_MAX_NUM_BATCHED_TOKENS)

126 +

127 @classmethod

128 def is_pin_memory_available(cls):

129 logger.warning("Pin memory is not supported on XPU.")

vllm/utils.py

@@ -77,6 +77,12 @@ if TYPE_CHECKING:

78 logger = init_logger(__name__)

80 +# This value is chosen to have a balance between ITL and TTFT. Note it is

81 +# not optimized for throughput.

82 +DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048

83 +POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768

84 +MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120

85 +

86 # Exception strings for non-implemented encoder/decoder scenarios

88 # Reminder: Please update docs/source/features/compatibility_matrix.md