From eb92e7d50efac8a3a35dfe0aeb2fc04f9e2faff6 Mon Sep 17 00:00:00 2001
From: SILONG ZENG <2609716663@qq.com>
Date: Thu, 19 Mar 2026 20:12:57 +0800
Subject: [PATCH] [Bugfix] Restore balance scheduling patch for v0.17.0 (#7479)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

### What this PR does / why we need it?
Restore previously introduced patches：
- https://github.com/vllm-project/vllm-ascend/pull/5212

- vLLM version: v0.17.0
- vLLM main:
https://github.com/vllm-project/vllm/commit/8b6325758cce5f9c36d38f2462edbd368b97a07c

Signed-off-by: MrZ20 <2609716663@qq.com>
---
 vllm_ascend/patch/platform/__init__.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/vllm_ascend/patch/platform/__init__.py b/vllm_ascend/patch/platform/__init__.py
index 1e4e5b49..d304af1c 100644
--- a/vllm_ascend/patch/platform/__init__.py
+++ b/vllm_ascend/patch/platform/__init__.py
@@ -19,7 +19,8 @@ import os
 import vllm_ascend.patch.platform.patch_distributed  # noqa
 import vllm_ascend.patch.platform.patch_fusion_matcher_compat_ops  # noqa
 import vllm_ascend.patch.platform.patch_kv_cache_interface  # noqa
-from vllm_ascend.utils import is_310p
+from vllm_ascend import envs
+from vllm_ascend.utils import is_310p, vllm_version_is
 
 if not is_310p():
     import vllm_ascend.patch.platform.patch_mamba_config  # noqa
@@ -31,3 +32,5 @@ import vllm_ascend.patch.platform.patch_torch_accelerator  # noqa
 
 if os.getenv("DYNAMIC_EPLB", "false").lower() in ("true", "1") or os.getenv("EXPERT_MAP_RECORD", "false") == "true":
     import vllm_ascend.patch.platform.patch_multiproc_executor  # noqa
+if envs.VLLM_ASCEND_BALANCE_SCHEDULING and vllm_version_is("0.17.0"):
+    import vllm_ascend.patch.platform.patch_balance_schedule  # noqa