From 84336ef7552b9b9c4e1b0f13c70a3785a8f83ca1 Mon Sep 17 00:00:00 2001 From: liziyu Date: Thu, 4 Sep 2025 11:01:21 +0800 Subject: [PATCH 1/2] PD disaggregate should not force with_prefill true in P and false in D Signed-off-by: liziyu --- vllm_ascend/models/deepseek_v2.py | 13 +------------ vllm_ascend/worker/model_runner_v1.py | 4 ---- 2 files changed, 1 insertion(+), 16 deletions(-) diff --git a/vllm_ascend/models/deepseek_v2.py b/vllm_ascend/models/deepseek_v2.py index 6d0913c18c9..1811255e2c6 100644 --- a/vllm_ascend/models/deepseek_v2.py +++ b/vllm_ascend/models/deepseek_v2.py @@ -32,8 +32,7 @@ from torch import nn from transformers import PretrainedConfig from vllm.attention import Attention, AttentionMetadata -from vllm.config import (CacheConfig, ModelConfig, VllmConfig, - get_current_vllm_config) +from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, get_tp_group, split_tensor_along_last_dim, @@ -374,10 +373,6 @@ def __init__( self.tp_group = get_tp_group().device_group self.tp_rank = get_tp_group().rank_in_group self.ep_group = get_ep_group() - self.kv_consumer = None - transfer_config = get_current_vllm_config().kv_transfer_config - if transfer_config is not None: - self.kv_consumer = transfer_config.kv_role == "kv_consumer" self.params_dtype = torch.get_default_dtype() self.rm_router_logits = self.experts.rm_router_logits @@ -395,12 +390,6 @@ def forward(self, is_prefill = forward_context.with_prefill - # If this node is kv_consumer, we force the moe always runs in decode path to make sure - # the behaviour aligned between dummy_run and normal model_execute. - if self.kv_consumer: - is_prefill = False - enable_force_load_balance = False - # router_logits: (num_tokens, n_experts) router_logits = None if not self.rm_router_logits and not self.enable_multistream_moe: diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 7a9fe1b8dce..4d09f046019 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -2058,10 +2058,6 @@ def _dummy_run( num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) - # Force dummy run on prefill stage when this node is deemed as kv producer. - if self.is_kv_producer: - with_prefill = True - attn_metadata = self._build_attention_metadata(with_prefill, num_reqs, skip_attn=True) From 1b1da9227b293111e5efb47302f4d433fdf17347 Mon Sep 17 00:00:00 2001 From: liziyu Date: Thu, 4 Sep 2025 16:59:08 +0800 Subject: [PATCH 2/2] PD disaggregate should not force with_prefill false in D Signed-off-by: liziyu --- vllm_ascend/worker/model_runner_v1.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 4d09f046019..7a9fe1b8dce 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -2058,6 +2058,10 @@ def _dummy_run( num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + # Force dummy run on prefill stage when this node is deemed as kv producer. + if self.is_kv_producer: + with_prefill = True + attn_metadata = self._build_attention_metadata(with_prefill, num_reqs, skip_attn=True)