Skip to content

Commit 620555f

Browse files
authored
[BugFix]: multiple config reading issues for env var and additional_config (vllm-project#10786)
### What this PR does / why we need it? Fixes three configuration-related bugs: 1. VLLM_ASCEND_BALANCE_SCHEDULING env var not taking effect when --additional-config is not set — the fallback in _balance_scheduling_enabled now properly reads from ascend_config first, then additional_config , then the env var. 2. Removed an unnecessary and (not is_draft_model) condition in dispatch_ffn_combine_enable that caused a performance regression in draft model scenarios. 3. mooncake_hybrid_connector.py was reading enable_transpose_kv_cache_by_block directly from the envs module, bypassing additional_config . Fixed to use get_ascend_config() so that --additional-config settings take effect. ### Does this PR introduce _any_ user-facing change? No breaking changes. Environment variables VLLM_ASCEND_BALANCE_SCHEDULING and VLLM_ASCEND_FUSION_OP_TRANSPOSE_KV_CACHE_BY_BLOCK now work correctly again, while --additional-config settings continue to work as expected. ### How was this patch tested? Pending CI verification. - vLLM version: v0.22.1 - vLLM main: vllm-project/vllm@967c5c3 --------- Signed-off-by: lizy124 <1950471827@qq.com>
1 parent c6343b1 commit 620555f

3 files changed

Lines changed: 15 additions & 5 deletions

File tree

vllm_ascend/ascend_forward_context.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ def select_moe_comm_method(num_tokens: int, vllm_config: VllmConfig, is_draft_mo
253253
Args:
254254
num_tokens (int): The number of tokens in the current batch.
255255
vllm_config (VllmConfig): Runtime configuration for the model.
256-
is_draft_model (bool): Whether the model runs in MTP mode (disables fused MC2).
256+
is_draft_model (bool): Whether the model runs in MTP mode.
257257
258258
Raises:
259259
ValueError: If the soc version is unsupported.
@@ -288,7 +288,7 @@ def select_moe_comm_method(num_tokens: int, vllm_config: VllmConfig, is_draft_mo
288288
# TODO: drop the EP-size guard when dispatch_ffn_combine supports larger EP sizes
289289
# TODO: drop speculative method guard when dispatch_gmm_combine_decode supports w16a16
290290
fused_mc2_enable = get_ascend_config().enable_fused_mc2
291-
dispatch_ffn_combine_enable = get_ep_group().world_size <= 32 and (not is_draft_model)
291+
dispatch_ffn_combine_enable = get_ep_group().world_size <= 32
292292
if num_tokens <= mc2_tokens_capacity:
293293
fused_decode_enable = fused_mc2_enable
294294
if fused_mc2_enable == 1:

vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,6 @@
5151
)
5252
from vllm.v1.request import RequestStatus
5353

54-
from vllm_ascend import envs as ascend_envs
5554
from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config
5655
from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import global_te
5756
from vllm_ascend.distributed.kv_transfer.utils.utils import get_transfer_timeout_value
@@ -724,7 +723,7 @@ def _transfer_kv_cache(self, req_meta: dict[str, Any]):
724723
is_kv_transfer_end = global_offset == tp_num_need_pulls * self._prefill_pp_size - 1
725724
need_cat_cache = tp_num_need_pulls > 1 and is_kv_transfer_end
726725
need_nz_cache = get_ascend_config().enable_kv_nz and is_kv_transfer_end
727-
use_fused_op = ascend_envs.VLLM_ASCEND_FUSION_OP_TRANSPOSE_KV_CACHE_BY_BLOCK
726+
use_fused_op = get_ascend_config().enable_transpose_kv_cache_by_block
728727
if need_nz_cache or need_cat_cache:
729728
# use fused op to reformat kv cache, we keep original implementation to provide ability to disable it.
730729
if use_fused_op and enable_custom_op():

vllm_ascend/patch/platform/patch_balance_schedule.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
# mypy: ignore-errors
2+
import os
23
import signal
34
import time
45

@@ -29,8 +30,18 @@
2930

3031

3132
def _balance_scheduling_enabled(vllm_config) -> bool:
33+
# TODO: Unify this path with AscendConfig once AscendConfig initialization
34+
# is moved earlier in the startup flow.
35+
try:
36+
from vllm_ascend.ascend_config import get_ascend_config
37+
38+
return bool(get_ascend_config().enable_balance_scheduling)
39+
except Exception:
40+
pass
3241
additional_config = getattr(vllm_config, "additional_config", None) or {}
33-
return bool(additional_config.get("enable_balance_scheduling", False))
42+
if "enable_balance_scheduling" in additional_config:
43+
return bool(additional_config["enable_balance_scheduling"])
44+
return bool(int(os.getenv("VLLM_ASCEND_BALANCE_SCHEDULING", "0")))
3445

3546

3647
class BalanceScheduler(Scheduler):

0 commit comments

Comments
 (0)