Repository: vllm-project/vllm
Pull request #18470: [Platform] Move platform check to right place
Author: login: wangxiyuan
Labels: name: tpu
description: Related to Google TPUs
name: ready
description: ONLY add when PR is ready to merge/full CI is needed
Created: 2025-05-21T08:29:30Z | Merged: 2025-05-22T19:11:29Z

platform provide `check_and_update_config` interface for device to set its own config value. The platform specified config check and set logic should not be in config `__post_init__`.  It'll limit other platform to address some case. For example, in some other oot platform, mla and chunked prefill may works well.

This PR move the specified platform check from config `__post_init__` to `check_and_update_config` to make the common code more generic.


Diff:
diff --git a/vllm/config.py b/vllm/config.py
index 3fa1db0e8390..1c916915a046 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -42,7 +42,10 @@
     try_get_generation_config, uses_mrope)
 from vllm.transformers_utils.s3_utils import S3Model
 from vllm.transformers_utils.utils import is_s3, maybe_model_redirect
-from vllm.utils import (GiB_bytes, LayerBlockType, cuda_device_count_stateless,
+from vllm.utils import (DEFAULT_MAX_NUM_BATCHED_TOKENS,
+                        MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,
+                        POOLING_MODEL_MAX_NUM_BATCHED_TOKENS, GiB_bytes,
+                        LayerBlockType, cuda_device_count_stateless,
                         get_cpu_memory, get_open_port, is_torch_equal_or_newer,
                         random_uuid, resolve_obj_by_qualname)
 
@@ -64,12 +67,6 @@
 
 ConfigT = TypeVar("ConfigT", bound=ConfigType)
 
-# This value is chosen to have a balance between ITL and TTFT. Note it is
-# not optimized for throughput.
-_DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
-_POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
-_MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
-
 TaskOption = Literal["auto", "generate", "embedding", "embed", "classify",
                      "score", "reward", "transcription"]
 
@@ -2074,28 +2071,28 @@ def __post_init__(self) -> None:
                     # so we don't reject sequences on account of a short
                     # max_num_batched_tokens.
                     self.max_num_batched_tokens = max(
-                        self.max_model_len, _DEFAULT_MAX_NUM_BATCHED_TOKENS)
+                        self.max_model_len, DEFAULT_MAX_NUM_BATCHED_TOKENS)
                 else:
                     self.max_num_batched_tokens = (
-                        _DEFAULT_MAX_NUM_BATCHED_TOKENS)
+                        DEFAULT_MAX_NUM_BATCHED_TOKENS)
             else:
                 # If max_model_len is too short, use
-                # _DEFAULT_MAX_NUM_BATCHED_TOKENS as the default value
+                # DEFAULT_MAX_NUM_BATCHED_TOKENS as the default value
                 # for higher throughput.
                 self.max_num_batched_tokens = max(
-                    self.max_model_len, _DEFAULT_MAX_NUM_BATCHED_TOKENS)
+                    self.max_model_len, DEFAULT_MAX_NUM_BATCHED_TOKENS)
 
             if self.runner_type == "pooling":
                 # Choose specific value for higher throughput
                 self.max_num_batched_tokens = max(
                     self.max_num_batched_tokens,
-                    _POOLING_MODEL_MAX_NUM_BATCHED_TOKENS,
+                    POOLING_MODEL_MAX_NUM_BATCHED_TOKENS,
                 )
             if self.is_multimodal_model:
                 # The value needs to be at least the number of multimodal tokens
                 self.max_num_batched_tokens = max(
                     self.max_num_batched_tokens,
-                    _MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,
+                    MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS,
                 )
 
             # When using default settings,
@@ -4316,18 +4313,6 @@ def __post_init__(self):
                 "full_cuda_graph is not supported with "
                 "cascade attention. Disabling cascade attention.")
             self.model_config.disable_cascade_attn = True
-
-        if self.model_config and self.model_config.use_mla and \
-            not (current_platform.is_cuda() or current_platform.is_rocm()):
-            logger.info(
-                "MLA is enabled on a non-GPU platform; forcing chunked "
-                "prefill and prefix caching to be disabled.")
-            self.scheduler_config.enable_chunked_prefill = False
-            self.scheduler_config.chunked_prefill_enabled = False
-            self.scheduler_config.max_num_batched_tokens = max(
-                self.scheduler_config.max_model_len,
-                _DEFAULT_MAX_NUM_BATCHED_TOKENS)
-
             if self.cache_config is not None:
                 self.cache_config.enable_prefix_caching = False
 
diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py
index 2d48af397636..5c0c90972b58 100644
--- a/vllm/platforms/cpu.py
+++ b/vllm/platforms/cpu.py
@@ -9,6 +9,7 @@
 import torch
 
 from vllm.logger import init_logger
+from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
 
 from .interface import CpuArchEnum, Platform, PlatformEnum, _Backend
 
@@ -177,6 +178,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                     " set VLLM_WORKER_MULTIPROC_METHOD to fork explicitly.")
                 os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn'
 
+        if vllm_config.model_config and vllm_config.model_config.use_mla:
+            logger.info(
+                "MLA is enabled on a non-GPU platform; forcing chunked "
+                "prefill and prefix caching to be disabled.")
+            vllm_config.scheduler_config.enable_chunked_prefill = False
+            vllm_config.scheduler_config.chunked_prefill_enabled = False
+            vllm_config.scheduler_config.max_num_batched_tokens = max(
+                vllm_config.scheduler_config.max_model_len,
+                DEFAULT_MAX_NUM_BATCHED_TOKENS)
+
     @classmethod
     def is_pin_memory_available(cls) -> bool:
         logger.warning("Pin memory is not supported on CPU.")
diff --git a/vllm/platforms/hpu.py b/vllm/platforms/hpu.py
index 456b054b2b43..6f7c5a6d3cae 100644
--- a/vllm/platforms/hpu.py
+++ b/vllm/platforms/hpu.py
@@ -7,6 +7,7 @@
 
 from vllm import envs
 from vllm.logger import init_logger
+from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
 
 from .interface import Platform, PlatformEnum, _Backend
 
@@ -80,6 +81,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                     "VLLM_WORKER_MULTIPROC_METHOD=fork explicitly.")
                 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
+        if vllm_config.model_config and vllm_config.model_config.use_mla:
+            logger.info(
+                "MLA is enabled on a non-GPU platform; forcing chunked "
+                "prefill and prefix caching to be disabled.")
+            vllm_config.scheduler_config.enable_chunked_prefill = False
+            vllm_config.scheduler_config.chunked_prefill_enabled = False
+            vllm_config.scheduler_config.max_num_batched_tokens = max(
+                vllm_config.scheduler_config.max_model_len,
+                DEFAULT_MAX_NUM_BATCHED_TOKENS)
+
     @classmethod
     def is_pin_memory_available(cls):
         logger.warning("Pin memory is not supported on HPU.")
diff --git a/vllm/platforms/neuron.py b/vllm/platforms/neuron.py
index e08337b8391d..9cd49fd34804 100644
--- a/vllm/platforms/neuron.py
+++ b/vllm/platforms/neuron.py
@@ -6,6 +6,7 @@
 
 from vllm import envs
 from vllm.logger import init_logger
+from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
 
 from .interface import Platform, PlatformEnum
 
@@ -56,6 +57,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
             vllm_config.cache_config.block_size = \
                 vllm_config.model_config.max_model_len  # type: ignore
 
+        if vllm_config.model_config and vllm_config.model_config.use_mla:
+            logger.info(
+                "MLA is enabled on a non-GPU platform; forcing chunked "
+                "prefill and prefix caching to be disabled.")
+            vllm_config.scheduler_config.enable_chunked_prefill = False
+            vllm_config.scheduler_config.chunked_prefill_enabled = False
+            vllm_config.scheduler_config.max_num_batched_tokens = max(
+                vllm_config.scheduler_config.max_model_len,
+                DEFAULT_MAX_NUM_BATCHED_TOKENS)
+
     @classmethod
     def is_pin_memory_available(cls) -> bool:
         logger.warning("Pin memory is not supported on Neuron.")
diff --git a/vllm/platforms/tpu.py b/vllm/platforms/tpu.py
index 6c573c1b3635..0173b15697cf 100644
--- a/vllm/platforms/tpu.py
+++ b/vllm/platforms/tpu.py
@@ -9,6 +9,7 @@
 from vllm.inputs import ProcessorInputs, PromptType
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams, SamplingType
+from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
 
 from .interface import Platform, PlatformEnum, _Backend
 
@@ -161,6 +162,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
             "Forcing --disable_chunked_mm_input.")
             scheduler_config.disable_chunked_mm_input = True
 
+        if vllm_config.model_config and vllm_config.model_config.use_mla:
+            logger.info(
+                "MLA is enabled on a non-GPU platform; forcing chunked "
+                "prefill and prefix caching to be disabled.")
+            vllm_config.scheduler_config.enable_chunked_prefill = False
+            vllm_config.scheduler_config.chunked_prefill_enabled = False
+            vllm_config.scheduler_config.max_num_batched_tokens = max(
+                vllm_config.scheduler_config.max_model_len,
+                DEFAULT_MAX_NUM_BATCHED_TOKENS)
+
     @classmethod
     def is_pin_memory_available(cls):
         logger.warning("Pin memory is not supported on TPU.")
diff --git a/vllm/platforms/xpu.py b/vllm/platforms/xpu.py
index 225e756cd7ce..785fb6ce1b79 100644
--- a/vllm/platforms/xpu.py
+++ b/vllm/platforms/xpu.py
@@ -5,6 +5,7 @@
 import torch
 
 from vllm.logger import init_logger
+from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
 
 from .interface import DeviceCapability, Platform, PlatformEnum, _Backend
 
@@ -113,6 +114,16 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
                 parallel_config.distributed_executor_backend)
             parallel_config.distributed_executor_backend = "ray"
 
+        if vllm_config.model_config and vllm_config.model_config.use_mla:
+            logger.info(
+                "MLA is enabled on a non-GPU platform; forcing chunked "
+                "prefill and prefix caching to be disabled.")
+            vllm_config.scheduler_config.enable_chunked_prefill = False
+            vllm_config.scheduler_config.chunked_prefill_enabled = False
+            vllm_config.scheduler_config.max_num_batched_tokens = max(
+                vllm_config.scheduler_config.max_model_len,
+                DEFAULT_MAX_NUM_BATCHED_TOKENS)
+
     @classmethod
     def is_pin_memory_available(cls):
         logger.warning("Pin memory is not supported on XPU.")
diff --git a/vllm/utils.py b/vllm/utils.py
index 0cd90c130d3e..bfc01972bbd2 100644
--- a/vllm/utils.py
+++ b/vllm/utils.py
@@ -77,6 +77,12 @@
 
 logger = init_logger(__name__)
 
+# This value is chosen to have a balance between ITL and TTFT. Note it is
+# not optimized for throughput.
+DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
+POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
+MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
+
 # Exception strings for non-implemented encoder/decoder scenarios
 
 # Reminder: Please update docs/source/features/compatibility_matrix.md


Reviews:
user: login: DarkLight1337
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-21T08:30:57Z
user: login: DarkLight1337
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-22T14:09:54Z
user: login: wangxiyuan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:11:39Z
user: login: wangxiyuan
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:15:05Z
user: login: DarkLight1337
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-22T14:22:19Z

Review comments:
diff_hunk: @@ -66,7 +67,6 @@
 
 # This value is chosen to have a balance between ITL and TTFT. Note it is
 # not optimized for throughput.
-_DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
 _POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
 _MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
path: vllm/config.py
user: login: DarkLight1337
body: Can you move these constants as well to keep everything in one place?
created_at: 2025-05-22T14:09:54Z
author_association: MEMBER
diff_hunk: @@ -66,7 +67,6 @@
 
 # This value is chosen to have a balance between ITL and TTFT. Note it is
 # not optimized for throughput.
-_DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
 _POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
 _MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
path: vllm/config.py
user: login: wangxiyuan
body: I found that these two value is only used in config.py. So I left them here. But it's fine to move to utils as well
created_at: 2025-05-22T14:11:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -66,7 +67,6 @@
 
 # This value is chosen to have a balance between ITL and TTFT. Note it is
 # not optimized for throughput.
-_DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
 _POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
 _MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
path: vllm/config.py
user: login: wangxiyuan
body: Done
created_at: 2025-05-22T14:15:05Z
author_association: CONTRIBUTOR

PR comments:
user: login: DarkLight1337
created_at: 2025-05-21T08:31:24Z
body: Thanks for cleaning this up!
author_association: MEMBER
user: login: github-actions[bot]
created_at: 2025-05-21T08:33:50Z
body: 👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run `fastcheck` CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your `fastcheck` build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping `simon-mo` or `khluu` to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add `ready` label to the PR or enable auto-merge.

🚀
author_association: NONE
user: login: DarkLight1337
created_at: 2025-05-21T08:48:07Z
body: There is a circular import error in Charts CI
author_association: MEMBER
user: login: wangxiyuan
created_at: 2025-05-22T01:29:53Z
body: @DarkLight1337 oh, yeah, `platform` is used everywhere. it should just import some common modules to avoid circle import error.
author_association: CONTRIBUTOR
user: login: wangxiyuan
created_at: 2025-05-22T03:00:18Z
body: Checked the CI failure, it seems doesn't relate to this PR. 
author_association: CONTRIBUTOR
user: login: DarkLight1337
created_at: 2025-05-22T03:01:56Z
body: We are currently pausing non-essential PRs to get CI green for the release.
author_association: MEMBER
user: login: DarkLight1337
created_at: 2025-05-22T14:01:50Z
body: You can now merge from main to fix CI
author_association: MEMBER
user: login: wangxiyuan
created_at: 2025-05-22T14:05:41Z
body: @DarkLight1337 rebased now. If a new release of vllm is coming, I hope  the CI be happy and this PR can be merged, it'll fix the mla + chunked prefill feature for non-gpu device. Thanks.
author_association: CONTRIBUTOR