[0.9.1][bugfix] fix ascend_forward_context (#1554)

### What this PR does / why we need it?
1. fix v0_model_runner, pooling_model_runner, draft_model_runner are not
adapted to `ascend_forward_context`
2. fix moe_distributed_combine's param `global_bs` bug.


### Does this PR introduce _any_ user-facing change?


### How was this patch tested?

Signed-off-by: zzzzwwjj <1183291235@qq.com>

Author: zzzzwwjj

vllm_ascend/ops/fused_moe.py

@@ -214,7 +214,7 @@ def fused_experts_with_mc2(
         "expert_shard_type": 0,
         "shared_expert_rank_num": 0,
         "moe_expert_num": moe_expert_num,
-        "global_bs": 0,
+        "global_bs": global_bs,
     }
     tp_recv_counts = output[5]
     stage3_kwargs = {

vllm_ascend/quantization/w8a8_dynamic.py

@@ -218,7 +218,7 @@ def fused_experts_with_mc2(
         "expert_shard_type": 0,
         "shared_expert_rank_num": 0,
         "moe_expert_num": moe_expert_num,
-        "global_bs": 0,
+        "global_bs": global_bs,
     }
     tp_recv_counts = torch.empty(1,
                                  dtype=torch.int32,

vllm_ascend/worker/draft_model_runner.py

@@ -18,7 +18,6 @@
 from typing import List, Optional
 
 import torch
-from vllm.forward_context import set_forward_context
 from vllm.logger import logger
 from vllm.model_executor.layers.sampler import SamplerOutput
 from vllm.multimodal import MultiModalKwargs
@@ -27,6 +26,7 @@
                                            ModelRunnerInputBase,
                                            ModelRunnerWrapperBase)
 
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 from vllm_ascend.attention.attention import AscendMetadata
 
 # A flag to enable debug prints for the updated input tensors
@@ -261,8 +261,8 @@ def execute_model(
                 spec_step_idx = kwargs.get("spec_step_idx", step)
                 model_execute_kwargs["spec_step_idx"] = spec_step_idx
                 compute_logits_kwargs["spec_step_idx"] = spec_step_idx
-            with set_forward_context(model_input.attn_metadata,
-                                     self.vllm_config):
+            with set_ascend_forward_context(model_input.attn_metadata,
+                                            self.vllm_config):
 
                 if model_input.attn_metadata is not None:
                     model_input.attn_metadata.input_positions = model_input.input_positions

vllm_ascend/worker/model_runner.py

@@ -35,7 +35,6 @@
 from vllm.core.scheduler import SchedulerOutputs
 from vllm.distributed import broadcast_tensor_dict, get_dp_group, get_pp_group
 from vllm.distributed.kv_transfer import get_kv_transfer_group
-from vllm.forward_context import set_forward_context
 from vllm.inputs import INPUT_REGISTRY, InputRegistry
 from vllm.logger import logger
 from vllm.lora.layers import LoRAMapping
@@ -66,6 +65,7 @@
     _init_sampling_metadata_from_tensor_dict)
 
 from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionBackend
@@ -1431,8 +1431,12 @@ def execute_model(
             model_forward_start.record()
 
         if not bypass_model_exec:
-            with set_forward_context(model_input.attn_metadata,
-                                     self.vllm_config, virtual_engine):
+            with set_ascend_forward_context(
+                    model_input.attn_metadata,
+                    self.vllm_config,
+                    virtual_engine,
+                    with_prefill=prefill_meta is not None,
+                    in_profile_run=self.in_profile_run):
                 if model_input.attn_metadata is not None:
                     model_input.attn_metadata.input_positions = model_input.input_positions
                 if self.torchair_graph_enabled:

vllm_ascend/worker/pooling_model_runner.py

@@ -21,13 +21,13 @@
 
 import torch
 from vllm.distributed import get_pp_group
-from vllm.forward_context import set_forward_context
 from vllm.model_executor.pooling_metadata import PoolingMetadata
 from vllm.multimodal import MultiModalKwargs
 from vllm.pooling_params import PoolingParams
 from vllm.sequence import (IntermediateTensors, SequenceData,
                            SequenceGroupMetadata)
 
+from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 from vllm_ascend.worker.model_runner import (ModelInputForNPU,
                                              ModelInputForNPUBuilder,
                                              NPUModelRunnerBase)
@@ -142,8 +142,8 @@ def execute_model(
         if model_input.token_types is not None:
             cross_enc_kwargs["token_type_ids"] = model_input.token_types
 
-        with set_forward_context(model_input.attn_metadata, self.vllm_config,
-                                 virtual_engine):
+        with set_ascend_forward_context(model_input.attn_metadata,
+                                        self.vllm_config, virtual_engine):
             hidden_or_intermediate_states = model_executable(
                 input_ids=model_input.input_tokens,
                 positions=model_input.input_positions,