[Misc][MRV2] Align V2 Model Runner with Upstream (vllm-project#12731)

### What this PR does / why we need it?
This PR introduces several updates and alignments for the Ascend backend
in vLLM:
- **Top-K Top-P Sampler Patching**: Overrides `apply_top_k_top_p` to use
the `ascendc` operator instead of Triton on NPU, bypassing the upstream
Triton routing which is broken on Ascend for batch sizes >= 8.
- **Model Runner Alignments**:
- Uses `sort_batch_req_ids` from upstream.
[https://github.com/vllm-project/vllm/pull/47381](url)
- Cleans up unused attributes (`cudagraph_and_dp_padding`,
`self.input_batch`).
- Updates draft token and logit calculation to support speculative
decoding with multiple new sampled tokens per step (`num_bonus_tokens`).

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

### How was this patch tested?
CI/CD testing with existing tests.


- vLLM version: v0.25.1
- vLLM main:
vllm-project/vllm@fe784ff

Signed-off-by: zouzy <zouzongyu@huawei.com>

Author: zouzy5137

vllm_ascend/worker/v2/model_runner.py

@@ -36,6 +36,11 @@
 )
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
+from vllm_ascend.utils import vllm_version_is
+
+if not vllm_version_is("0.25.1"):
+    from vllm.v1.worker.gpu.model_runner import sort_batch_req_ids
+
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import (
     MoECommType,
@@ -118,22 +123,15 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
             pin_memory=True,
         )
 
+        # NOTE: In GPUModelRunner, decode_query_len is initialized in execute_model(),
+        # +1 is hardcoded here but not in vllm.
+        self.decode_query_len = self.num_speculative_steps + 1
         # Set _mc2_tokens_capacity and _reserved_mc2_mask for MoE communication optimization.
         # TODO: remove set_cos_and_sin (together with update_cos_sin) when mla can properly handle cos/sin internally
-        self.decode_query_len = self.num_speculative_steps + 1
         set_cos_and_sin(vllm_config, self.max_num_reqs, self.decode_query_len, self.dtype, self.device)
         set_mc2_tokens_capacity(vllm_config, self.max_num_reqs, self.decode_query_len)
         set_mc2_mask(vllm_config, self.device)
 
-        # we need to use return value of `get_cudagraph_and_dp_padding`
-        # to set forward_context in `run_fullgraph`.
-        # so we can inherit `execute_model` method.
-        self.cudagraph_and_dp_padding: tuple[int, torch.Tensor | None, int] | None = None
-
-        # we need to use input_batch to set forward_context in run_fullgraph.
-        # so we can inherit `execute_model` method.
-        self.input_batch: AscendInputBatch | None = None
-
     def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         with graph_manager_wrapper(self):
             super().initialize_kv_cache(kv_cache_config)
@@ -171,9 +169,13 @@ def prepare_inputs(
         num_tokens_per_req = scheduler_output.num_scheduled_tokens
         num_reqs = len(num_tokens_per_req)
 
-        # Decode first, then prefill.
         # batch_idx -> req_id
-        req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get)  # type: ignore
+        if vllm_version_is("0.25.1"):
+            # vllm 0.25.1 does not have sort_batch_req_ids;
+            # TODO: remove this patch when main2main is applied.
+            req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get)  # type: ignore
+        else:
+            req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
 
         self._update_seq_lens_cpu(scheduler_output, req_ids)
 
@@ -202,7 +204,7 @@ def prepare_inputs(
 
         # Get the number of draft tokens for each request.
         draft_tokens = scheduler_output.scheduled_spec_decode_tokens
-        num_draft_tokens_per_req: np.ndarray | None = None
+        num_draft_tokens_per_req = None
         if not draft_tokens:
             # No draft token scheduled (common case).
             total_num_draft_tokens = 0
@@ -212,21 +214,21 @@ def prepare_inputs(
             expanded_idx_mapping = idx_mapping
             expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
         else:
-            num_draft_tokens_arr = np.array(
-                [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
+            num_draft_tokens_per_req = np.fromiter(
+                (len(draft_tokens.get(req_id, ())) for req_id in req_ids),
                 dtype=np.int32,
+                count=num_reqs,
             )
-            num_draft_tokens_per_req = num_draft_tokens_arr
-            total_num_draft_tokens = int(num_draft_tokens_arr.sum())
-            total_num_logits = num_reqs + total_num_draft_tokens
-
-            num_logits = num_draft_tokens_arr + 1
+            num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step
+            total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
+            total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens
+            num_logits = num_draft_tokens_per_req + num_bonus_tokens
             cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
             cu_num_logits_np[0] = 0
             np.cumsum(num_logits, out=cu_num_logits_np[1:])
             cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device)
 
-            max_expand_len = self.num_speculative_steps + 1
+            max_expand_len = self.decode_query_len
             expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
                 idx_mapping, total_num_logits, cu_num_logits, max_expand_len
             )
@@ -256,7 +258,7 @@ def prepare_inputs(
         async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
 
         query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
-        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
+        query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
         prefill_len_np = self.req_states.prefill_len.np[idx_mapping_np]
         num_computed_prefill_tokens_np = self.req_states.num_computed_prefill_tokens[idx_mapping_np]
         is_prefilling_np = num_computed_prefill_tokens_np < prefill_len_np
@@ -281,7 +283,7 @@ def prepare_inputs(
             self.input_buffers.positions,
             self.input_buffers.seq_lens,
         )
-        seq_lens = self.input_buffers.seq_lens[:num_reqs]
+        seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
 
         # Pad for full CUDA graph mode.
         self.input_buffers.seq_lens_np[num_reqs_padded:] = 0
@@ -298,11 +300,9 @@ def prepare_inputs(
             self.req_states.draft_tokens,
             cu_num_logits,
             total_num_logits,
+            self.model_state.num_new_sampled_tokens_per_step,
         )
 
-        input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
-        positions = self.input_buffers.positions[:num_tokens_after_padding]
-
         # CPU upper bound on seq_lens (num_computed_tokens + num_scheduled_tokens).
         # Added by vLLM PR #40654 to avoid GPU->CPU sync for seq_lens.
         seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_padded, dtype=np.int32)
@@ -313,12 +313,18 @@ def prepare_inputs(
         )
         seq_lens_cpu_upper_bound = torch.from_numpy(seq_lens_cpu_upper_bound_np)
         num_computed_tokens_np = self.req_states.num_computed_tokens_np[idx_mapping_np]
+
         max_seq_len_np = None
-        if getattr(self, "use_pp", False):
-            # max_seq_len is only consumed by the PP `compute_need_sampled_mask`.
+        if self.use_pp:
+            # max_seq_len is only consumed by the PP `compute_need_sampled_mask`
             max_seq_len_np = self.req_states.max_seq_len[idx_mapping_np]
 
-        self.input_batch = AscendInputBatch(
+        prompt_lens = None
+        if self.model_config.rswa_window is not None:
+            # prompt_lens is only used in R-SWA case.
+            prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]
+
+        input_batch = AscendInputBatch(
             req_ids=req_ids,
             num_reqs=num_reqs,
             num_reqs_after_padding=num_reqs_padded,
@@ -341,24 +347,24 @@ def prepare_inputs(
             prefill_len_np=prefill_len_np,
             num_computed_prefill_tokens_np=num_computed_prefill_tokens_np,
             max_seq_len_np=max_seq_len_np,
-            input_ids=input_ids,
-            positions=positions,
+            input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
+            positions=self.input_buffers.positions[:num_tokens_after_padding],
             is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
             logits_indices=logits_indices,
             cu_num_logits=cu_num_logits,
             cu_num_logits_np=cu_num_logits_np,
             has_structured_output_reqs=scheduler_output.has_structured_output_requests,
             # TODO: only populated for R-SWA (not supported yet).
-            prompt_lens=None,
+            prompt_lens=prompt_lens,
             # extra attributes for ascend npus.
             seq_lens_np=self.input_buffers.seq_lens_np,
             attn_state=attn_state,
         )
 
         # For mla/sfa, update cos/sin. Here is for execute_model.
-        update_cos_sin(self.input_batch.positions)
+        update_cos_sin(input_batch.positions)
 
-        return self.input_batch
+        return input_batch
 
     def postprocess(
         self,

vllm_ascend/worker/worker.py

@@ -761,6 +761,7 @@ def compile_or_warm_up_model(self) -> CompilationTimes:
                 bind_cpus(self.local_rank)
             except Exception as e:
                 logger.warning("Bind cpus failed in rank%s: %s Skip binding cpu.", self.local_rank, e)
+
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)