[Feature] adapt to uva buffer and main2main (#6657)

### What this PR does / why we need it?
vllm model runner v2 use uva buffer to prepare input data, but npu
doesn't support uva yet, this pr implement a uvawrapper class to mimic
gpu's uva backend. what's more, this pr make some modifications to adapt
to the newer main branch.

### Does this PR introduce _any_ user-facing change?
no
### How was this patch tested?
- vLLM main:
13397841ab

---------

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>

vllm_ascend/patch/__init__.py

@@ -263,3 +263,12 @@
 #       https://gitcode.com/Ascend/torchair/pull/2575
 #    Future Plan:
 #       Remove this patch when the PTA version used by vllm-ascend has been upgraded.
+# ** 14. File: worker/patch_v2_uva.py**
+# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+#   1. `vllm.v1.worker.gpu.states.UvaBuffer`
+#    Why:
+#       ASCEND NPUs do not support UVA yet, so we need to wrap it in vLLM.
+#    How：
+#       make UvaBuffer a dummy class, mimic the interface of vllm UvaBuffer.
+#    Future Plan:
+#       Remove this patch when NPU support UVA.

vllm_ascend/patch/worker/__init__.py

@@ -31,6 +31,7 @@ import vllm_ascend.patch.worker.patch_qwen3_next  # noqa
 import vllm_ascend.patch.worker.patch_qwen3_next_mtp  # noqa
 import vllm_ascend.patch.worker.patch_rejection_sampler  # noqa
 import vllm_ascend.patch.worker.patch_qwen3_next  # noqa
-import vllm_ascend.patch.worker.patch_v2_egale  # noqa
+import vllm_ascend.patch.worker.patch_v2_eagle  # noqa
+import vllm_ascend.patch.worker.patch_v2_uva  # noqa
 import vllm_ascend.patch.worker.patch_huanyuan_vl  # noqa
 import vllm_ascend.patch.worker.patch_npugraph_ex_triton  # noqa

vllm_ascend/patch/worker/patch_v2_eagle.py

@@ -16,10 +16,9 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
-import numpy as np
 import torch
 import vllm
-from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.spec_decode.eagle import prepare_eagle_decode, prepare_eagle_inputs
@@ -31,7 +30,6 @@ from vllm_ascend.worker.v2.attn_utils import build_attn_metadata
 def propose(
     self,
     input_batch: InputBatch,
-    sampling_metadata: SamplingMetadata,
     # [num_tokens, hidden_size]
     last_hidden_states: torch.Tensor,
     # num_layers x [num_tokens, hidden_size]
@@ -40,10 +38,14 @@ def propose(
     num_sampled: torch.Tensor,
     # [num_reqs]
     num_rejected: torch.Tensor,
-    # [num_reqs]
+    # [max_num_reqs]
     last_sampled: torch.Tensor,
-    # [num_reqs]
+    # [max_num_reqs]
     next_prefill_tokens: torch.Tensor,
+    # [max_num_reqs]
+    temperature: torch.Tensor,
+    # [max_num_reqs]
+    seeds: torch.Tensor,
 ) -> torch.Tensor:
     # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
     # number of rejected tokens, we maintain the size of eagle's input_ids and
@@ -74,13 +76,13 @@ def propose(
     last_hidden_states, hidden_states = self.run_model(
         num_tokens,
         input_batch.attn_metadata,
+        input_batch.slot_mappings,
         num_tokens_across_dp=None,  # FIXME
     )
     sample_hidden_states = last_hidden_states[last_token_indices]
     logits = self.model.compute_logits(sample_hidden_states)
 
     num_reqs = input_batch.num_reqs
-    cu_num_logits = input_batch.cu_num_logits[:num_reqs]
     # NOTE(woosuk): For draft sampling, we only consider the temperature
     # and ignore the other sampling parameters such as top_k and top_p,
     # for simplicity and performance.
@@ -89,16 +91,23 @@ def propose(
     # NOTE(Ronald1995): torch.gather will pollute the cache such as self.input_buffers.positions
     # the bug is reported to huawei CANN team, but not fixed yet.
     # So we clone the tensors before calling torch.gather to avoid the issue.
-    temperature = self.temperature[:num_reqs].clone()
-    seeds = self.seeds[:num_reqs].clone()
+    idx_mapping = self.idx_mapping[:num_reqs]
+    idx_mapping.copy_(input_batch.idx_mapping)
+    self.temperature.copy_(temperature)
+    self.seeds.copy_(seeds)
     pos = self.input_buffers.positions[:num_reqs].clone()
     # Gather the values and copy them to the pre-allocated buffers.
-    torch.gather(sampling_metadata.temperature, 0, cu_num_logits, out=temperature)
-    torch.gather(sampling_metadata.seeds, 0, cu_num_logits, out=seeds)
     torch.gather(input_batch.positions, 0, last_token_indices, out=pos)
     # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
     # used for draft and target sampling.
-    draft_tokens = gumbel_sample(logits, temperature, seeds, pos + 1, apply_temperature=True)
+    draft_tokens = gumbel_sample(
+        logits,
+        idx_mapping,
+        self.temperature,
+        self.seeds,
+        pos + 1,
+        apply_temperature=True,
+    )
     if self.num_speculative_steps == 1:
         # Early exit.
         return draft_tokens.view(-1, 1)
@@ -117,9 +126,12 @@ def propose(
         self.max_model_len,
         self.max_num_reqs,
     )
-    query_start_loc = self.input_buffers.query_start_loc
-    query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
-    slot_mappings = self.block_tables.compute_slot_mappings(query_start_loc_gpu, pos)
+    query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
+    slot_mappings = self.block_tables.compute_slot_mappings(
+        idx_mapping,
+        query_start_loc,
+        pos,
+    )
 
     cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
     if cudagraph_size is not None:
@@ -128,10 +140,8 @@ def propose(
         return self.draft_tokens[:num_reqs]
 
     # Run eager mode.
-    query_start_loc.np[: num_reqs + 1] = np.arange(num_reqs + 1)
-    query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
+    query_start_loc_cpu = torch.arange(num_reqs + 1, dtype=torch.int32, device="cpu")
     # HACK(woosuk)
-    seq_lens_np = np.full(num_reqs, self.max_model_len, dtype=np.int32)
     block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
 
     # FIXME(woosuk): This is UNSAFE!!
@@ -139,16 +149,22 @@ def propose(
         attn_metadata_builders=self.attn_metadata_builders,
         num_reqs=num_reqs,
         num_tokens=num_reqs,
-        query_start_loc_gpu=query_start_loc_gpu,
+        query_start_loc_gpu=query_start_loc,
         query_start_loc_cpu=query_start_loc_cpu,
+        max_query_len=1,
         seq_lens=self.input_buffers.seq_lens[:num_reqs],
-        seq_lens_np=seq_lens_np,
-        num_computed_tokens_cpu=None,  # FIXME
+        max_seq_len=self.max_model_len,
         block_tables=block_tables,
         slot_mappings=slot_mappings,
         kv_cache_config=self.kv_cache_config,
     )
-    self.generate_draft(num_reqs, attn_metadata, num_tokens_across_dp=None)  # FIXME
+    slot_mappings_by_layer = build_slot_mappings_by_layer(slot_mappings, self.kv_cache_config)
+    self.generate_draft(
+        num_reqs,
+        attn_metadata,
+        slot_mappings_by_layer,
+        num_tokens_across_dp=None,
+    )  # FIXME
     return self.draft_tokens[:num_reqs]
 
 

vllm_ascend/patch/worker/patch_v2_uva.py

@@ -0,0 +1,125 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/block_table.py
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+from collections.abc import Callable, Sequence
+
+import numpy as np
+import torch
+import vllm.v1.worker.gpu.buffer_utils
+
+
+def get_row_indices_from_key(key: int | slice | tuple, dim_size: int) -> set[int]:
+    """get the set of row indices involved in the given key."""
+    if isinstance(key, int):
+        # parse index such as np[1]
+        key = key if key >= 0 else dim_size + key
+        # handle negative index
+        if key < 0 or key >= dim_size:
+            raise IndexError(f"row index {key} out of [0, {dim_size})")
+        return {key}
+    elif isinstance(key, slice):
+        # parse slice such as np[1:3]
+        start, stop, step = key.indices(dim_size)
+        return set(range(start, stop, step))
+    elif isinstance(key, tuple):
+        # parse row slice such as np[1,:100]
+        if len(key) == 0:
+            return set(range(dim_size))
+        return get_row_indices_from_key(key[0], dim_size)
+    else:
+        # for other types such as list/ndarray, we return all rows.
+        return set(range(dim_size))
+
+
+class MonitoredNumPyArray:
+    """A wrapper around a NumPy array that monitors modifications."""
+
+    def __init__(self, array: np.ndarray, callback: Callable):
+        self._array = array
+        self._callback = callback
+
+    def __setitem__(self, key, value):
+        self._array[key] = value
+        dim_size = self._array.shape[0]
+        row_indices = get_row_indices_from_key(key, dim_size)
+        for row in row_indices:
+            self._callback(row)
+
+    def __getitem__(self, key):
+        return self._array[key]
+
+    def __getattr__(self, name):
+        return getattr(self._array, name)
+
+
+class MonitoredTorchTensor:
+    """A wrapper around a torch tensor that monitors modifications."""
+
+    def __init__(self, tensor: torch.Tensor, callback: Callable):
+        self._tensor = tensor
+        self._callback = callback
+
+    def __setitem__(self, key, value):
+        self._tensor[key] = value
+        dim_size = self._tensor.size(0)
+        row_indices = get_row_indices_from_key(key, dim_size)
+        for row in row_indices:
+            self._callback(row)
+
+    def __getitem__(self, key):
+        return self._tensor[key]
+
+    def __getattr__(self, name):
+        return getattr(self._tensor, name)
+
+
+class UvaBufferWrapper:
+    """Ascend NPU doesn't support UVA tensors directly. This is a wrapper class
+    that provides CPU and NPU views of a UVA tensor."""
+
+    def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
+        self._cpu: torch.Tensor = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True)
+        self._np = self._cpu.numpy()
+        self._uva: torch.Tensor = torch.zeros_like(self._cpu, device="npu")
+        self._modified_indices: set[int] = set()
+
+    def _mark_cpu_modified(self, key: int):
+        self._modified_indices.add(key)
+
+    @property
+    def cpu(self):
+        return MonitoredTorchTensor(self._cpu, self._mark_cpu_modified)
+
+    @property
+    def np(self):
+        return MonitoredNumPyArray(self._np, self._mark_cpu_modified)
+
+    @property
+    def uva(self):
+        """Get the device data of the buffer."""
+        if self._modified_indices:
+            # Sort for better memory access locality
+            dirty_rows = sorted(self._modified_indices)
+            # can't use copy_ method, because copy_ for index tensor
+            #  will malloc new memory.
+            self._uva[dirty_rows] = self._cpu[dirty_rows].to(device="npu", non_blocking=True)
+            self._modified_indices.clear()
+        return self._uva
+
+
+vllm.v1.worker.gpu.buffer_utils.UvaBuffer = UvaBufferWrapper

vllm_ascend/platform.py

@@ -551,7 +551,7 @@ class NPUPlatform(Platform):
         vllm_config: VllmConfig,
         dp_metadata,
         virtual_engine: int = 0,
-        num_tokens: int | None = None,
+        num_tokens: int = 0,
         num_tokens_across_dp: torch.Tensor | None = None,
         cudagraph_runtime_mode=None,
         batch_descriptor=None,
@@ -601,10 +601,6 @@ class NPUPlatform(Platform):
         if not envs_vllm.VLLM_USE_V2_MODEL_RUNNER:
             return {}
 
-        num_actual_tokens = list(attn_metadata.values())[0].num_actual_tokens
-        if num_tokens is None:
-            num_tokens = num_actual_tokens
-
         moe_comm_type = select_moe_comm_method(
             num_tokens,
             vllm_config,
@@ -636,10 +632,13 @@ class NPUPlatform(Platform):
 
         # TODO(Levi-JQ): another PR to normalize the enabling logic for sp/fc2
         flashcomm_v2_enabled = flashcomm2_enable() and tp_world_size > 1 and num_tokens is not None
-        pad_size = 0
+        pad_size = None
+        padded_length = None
         if sp_enabled or flashcomm_v2_enabled:
             pad_size = (tp_world_size - (num_tokens % tp_world_size)) % tp_world_size
 
+        if num_tokens is None and attn_metadata is not None:
+            num_tokens = list(attn_metadata.values())[0].num_actual_tokens
         dp_world_size = get_dp_group().world_size
         if dp_world_size > 1 and dp_metadata is not None:
             max_tokens_across_dp = dp_metadata.max_tokens_across_dp_cpu.item()
@@ -648,8 +647,9 @@ class NPUPlatform(Platform):
                 pad_size = padded_length - num_tokens
         else:
             max_tokens_across_dp = num_tokens
-
+        mc2_mask = None
         if num_tokens is not None:
+            num_actual_tokens = num_tokens
             # NOTE: token num which need to pad to when mc2
             padded_num_tokens = math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size
             reserved_mc2_mask = get_mc2_mask()

vllm_ascend/worker/v2/README.md

@@ -4,3 +4,5 @@ This directory contains the new model runner which is under active development.
 
 please see [Model Runner V2](https://github.com/vllm-project/vllm-ascend/issues/5208)
 to get specific plans.
+
+supported vllm version: main@1339784

vllm_ascend/worker/v2/aclgraph_utils.py

@@ -35,9 +35,14 @@ from vllm_ascend.worker.v2.utils import torch_cuda_wrapper
 class AclGraphManager(CudaGraphManager):
     """ACL Graph Manager for Ascend NPUs."""
 
-    def __init__(self, vllm_config: VllmConfig, device: torch.device):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        use_mrope: bool,
+        device: torch.device,
+    ):
         with torch_cuda_wrapper():
-            super().__init__(vllm_config, device)
+            super().__init__(vllm_config, use_mrope, device)
 
     def capture_graph(
         self,

vllm_ascend/worker/v2/attn_utils.py

@@ -42,17 +42,21 @@ def get_attn_mask_builder(device: torch.device):
 
 
 def build_attn_metadata(
+    *,
     attn_metadata_builders: list[AttentionMetadataBuilder],
     num_reqs: int,
     num_tokens: int,
     query_start_loc_gpu: torch.Tensor,
     query_start_loc_cpu: torch.Tensor,
+    max_query_len: int,
     seq_lens: torch.Tensor,
-    seq_lens_np: np.ndarray,
-    num_computed_tokens_cpu: torch.Tensor | None,
+    max_seq_len: int,
     block_tables: Sequence[torch.Tensor],
     slot_mappings: torch.Tensor,
     kv_cache_config: KVCacheConfig,
+    # extra attributes for ascend npus.
+    seq_lens_np: np.ndarray | None = None,
+    num_computed_tokens_cpu: torch.Tensor | None = None,
     positions: torch.Tensor | None = None,
     attn_state: Any | None = None,
     graph_pad_size: int = -1,
@@ -61,9 +65,13 @@ def build_attn_metadata(
 ) -> dict[str, Any]:
     """Build attention metadata for Ascend NPUs."""
     # TODO(Ronald1995): optimize AscendCommonAttentionMetadata.
-    max_query_len = int(query_start_loc_cpu.max())
-    seq_lens_cpu = torch.from_numpy(seq_lens_np)
-    max_seq_len = int(seq_lens_cpu.max())
+
+    # seq_lens_np is used for ascend npus, it maybe None in spec_decode case,
+    # we fill it with max_seq_len in case `attn_metadata_builder.build` raise
+    # an error.
+    if seq_lens_np is None:
+        seq_lens_np = np.full(num_reqs, max_seq_len, dtype=np.int32)
+    seq_lens_cpu = torch.from_numpy(seq_lens_np)[:num_reqs]
     # torch_npu._reshape_and_cache operator requires slot_mappings to
     # be torch.int32.
     slot_mappings = slot_mappings.to(torch.int32)
@@ -77,7 +85,7 @@ def build_attn_metadata(
         common_attn_metadata = AscendCommonAttentionMetadata(
             query_start_loc=query_start_loc_gpu,
             query_start_loc_cpu=query_start_loc_cpu,
-            seq_lens_cpu=seq_lens_cpu[:num_reqs],
+            seq_lens_cpu=seq_lens_cpu,
             seq_lens=seq_lens[:num_reqs],
             num_reqs=num_reqs,
             num_actual_tokens=num_tokens,

vllm_ascend/worker/v2/input_batch.py

@@ -16,10 +16,11 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
+from dataclasses import asdict, dataclass
 
 import numpy as np
 import torch
-from vllm.v1.worker.gpu.input_batch import InputBuffers
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 
 
 class AscendInputBuffers(InputBuffers):
@@ -29,20 +30,12 @@ class AscendInputBuffers(InputBuffers):
         self,
         max_num_reqs: int,
         max_num_tokens: int,
-        inputs_embeds_size: int,
-        vocab_size: int,
-        dtype: torch.dtype,
         device: torch.device,
-        pin_memory: bool,
     ):
         super().__init__(
             max_num_reqs,
             max_num_tokens,
-            inputs_embeds_size,
-            vocab_size,
-            dtype,
             device,
-            pin_memory,
         )
         # Create seq_lens_cpu and seq_lens_np.
         # npu's attention backend still needs seq_lens on CPU side.
@@ -54,3 +47,36 @@ class AscendInputBuffers(InputBuffers):
         # seq_len_np and seq_lens_cpu share the same memory.
         # define seq_lens_np for easier calculation with numpy.
         self.seq_lens_np: np.ndarray = self.seq_lens_cpu.numpy()
+
+
+@dataclass
+class AscendInputBatch(InputBatch):
+    """Input batch for Ascend NPUs."""
+
+    # Create seq_lens_np.
+    # npu's attention backend still needs seq_lens on CPU side.
+    seq_lens_np: np.ndarray
+
+    @classmethod
+    def make_dummy(
+        cls,
+        num_reqs: int,
+        num_tokens: int,
+        input_buffers: AscendInputBuffers,
+        device: torch.device,
+    ) -> "AscendInputBatch":
+        """Override the make_dummy method to calculate seq_lens_np."""
+        input_batch = InputBatch.make_dummy(
+            num_reqs,
+            num_tokens,
+            input_buffers,
+            device,
+        )
+        # seq_len equals to query_len
+        input_buffers.seq_lens_np[:num_reqs] = num_tokens // num_reqs
+        input_buffers.seq_lens_np[num_reqs - 1] += num_tokens % num_reqs
+        # Pad for full CUDA graph mode.
+        input_buffers.seq_lens_np[num_reqs:] = 0
+        seq_lens_np = input_buffers.seq_lens_np[:num_reqs]
+        input_batch.seq_lens_np = seq_lens_np
+        return cls(**asdict(input_batch), seq_lens_np=seq_lens_np)

vllm_ascend/worker/v2/model_runner.py

@@ -22,9 +22,11 @@ import torch
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
+from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
 from vllm.v1.worker.gpu.input_batch import (
-    InputBatch,
     combine_sampled_and_draft_tokens,
+    expand_idx_mapping,
     prepare_pos_seq_lens,
     prepare_prefill_inputs,
 )
@@ -32,11 +34,11 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
 from vllm_ascend.worker.v2.aclgraph_utils import AclGraphManager
 from vllm_ascend.worker.v2.attn_utils import build_attn_metadata, build_attn_state
-from vllm_ascend.worker.v2.input_batch import AscendInputBuffers
+from vllm_ascend.worker.v2.input_batch import AscendInputBatch, AscendInputBuffers
 from vllm_ascend.worker.v2.sample.sampler import AscendSampler
 from vllm_ascend.worker.v2.spec_decode import init_speculator
 from vllm_ascend.worker.v2.spec_decode.eagle import AscendEagleSpeculator
-from vllm_ascend.worker.v2.states import AscendRequestState, uva_wrapper
+from vllm_ascend.worker.v2.states import AscendRequestState
 from vllm_ascend.worker.v2.utils import torch_cuda_wrapper
 
 logger = init_logger(__name__)
@@ -46,7 +48,7 @@ class NPUModelRunner(GPUModelRunner):
     """Model runner for Ascend NPUs."""
 
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
-        with torch_cuda_wrapper(), uva_wrapper():
+        with torch_cuda_wrapper():
             super().__init__(vllm_config, device)
 
         # because we will override these attribute, delete these attribute to
@@ -59,8 +61,9 @@ class NPUModelRunner(GPUModelRunner):
 
         # NPU specific initializations can be added below.
         self.cudagraph_manager: AclGraphManager = AclGraphManager(
-            vllm_config,
-            device,
+            self.vllm_config,
+            self.uses_mrope,
+            self.device,
         )
 
         # we define AscendEagleSpeculator in vllm_ascend.worker.v2.spec_decode.eagle
@@ -79,23 +82,22 @@ class NPUModelRunner(GPUModelRunner):
             num_speculative_steps=self.num_speculative_steps,
             vocab_size=self.vocab_size,
             device=self.device,
-            pin_memory=self.pin_memory,
         )
         # AscendInputBuffers has extra `seq_lens_cpu` attribute.
         # so reinitialize input_buffers here.
         self.input_buffers: AscendInputBuffers = AscendInputBuffers(
             max_num_reqs=self.max_num_reqs,
             max_num_tokens=self.max_num_tokens,
-            inputs_embeds_size=self.inputs_embeds_size,
-            vocab_size=self.vocab_size,
-            dtype=self.dtype,
             device=self.device,
-            pin_memory=self.pin_memory,
         )
         # we need to adjust triton operators in sampler,
         # so reinitialize sampler here.
         self.sampler: AscendSampler = AscendSampler(
+            max_num_reqs=self.max_num_reqs,
+            vocab_size=self.vocab_size,
+            device=self.device,
             logprobs_mode=self.model_config.logprobs_mode,
+            num_speculative_tokens=self.num_speculative_steps + 1,
         )
 
         # we need to copy num_computed_tokens back to cpu to help
@@ -108,32 +110,31 @@ class NPUModelRunner(GPUModelRunner):
             self.max_num_reqs,
             dtype=torch.int32,
             device="cpu",
-            pin_memory=self.pin_memory,
+            pin_memory=True,
         )
 
     def prepare_inputs(
         self,
         scheduler_output: SchedulerOutput,
         num_tokens_after_padding: int,
-    ) -> InputBatch:
+    ) -> AscendInputBatch:
         """Override GPUModelRunner.prepare_inputs for Ascend NPUs.
         npu attention backends need seq_lens_cpu to work.
         so we need to prepare seq_lens_cpu here.
         """
         num_tokens = scheduler_output.total_num_scheduled_tokens
         assert num_tokens > 0
-        num_reqs = len(scheduler_output.num_scheduled_tokens)
+        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(
-            scheduler_output.num_scheduled_tokens.keys(),
-            key=lambda k: scheduler_output.num_scheduled_tokens[k],
-        )
+        req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get)  # type: ignore
 
         self._update_seq_lens_cpu(scheduler_output, req_ids)
 
-        num_scheduled_tokens = np.array([scheduler_output.num_scheduled_tokens[i] for i in req_ids], dtype=np.int32)
+        numtoks_iter = map(num_tokens_per_req.get, req_ids)
+        num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
         num_valid_tokens = num_scheduled_tokens
         if scheduler_output.scheduled_spec_decode_tokens:
             num_valid_tokens = np.array(
@@ -150,81 +151,94 @@ class NPUModelRunner(GPUModelRunner):
             num_scheduled_tokens,
             num_valid_tokens,
         )
-
-        idx_mapping_list = [self.req_states.req_id_to_index[req_id] for req_id in req_ids]
-        idx_mapping = self.input_buffers.idx_mapping
-        idx_mapping.np[:num_reqs] = idx_mapping_list
-        idx_mapping_np = idx_mapping.np[:num_reqs]
-        idx_mapping_cpu = idx_mapping.cpu[:num_reqs]
-        idx_mapping_npu = idx_mapping.copy_to_gpu(num_reqs)
+        idx_mapping_iter = map(self.req_states.req_id_to_index.get, req_ids)
+        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.int32, count=num_reqs)
+        idx_mapping_cpu = torch.from_numpy(idx_mapping_np)
+        idx_mapping = async_copy_to_gpu(idx_mapping_cpu, device=self.device)
 
         # Get the number of draft tokens for each request.
-        if not scheduler_output.scheduled_spec_decode_tokens:
+        draft_tokens = scheduler_output.scheduled_spec_decode_tokens
+        if not draft_tokens:
             # No draft token scheduled (common case).
             total_num_draft_tokens = 0
             total_num_logits = num_reqs
+            cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
             cu_num_logits = torch.arange(num_reqs + 1, device=self.device, dtype=torch.int32)
+            expanded_idx_mapping = idx_mapping
+            expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
         else:
-            draft_tokens = scheduler_output.scheduled_spec_decode_tokens
             num_draft_tokens = np.array(
-                [len(draft_tokens[req_id]) if req_id in draft_tokens else 0 for req_id in req_ids],
+                [len(draft_tokens.get(req_id, ())) for req_id in req_ids],
                 dtype=np.int32,
             )
             total_num_draft_tokens = int(num_draft_tokens.sum())
             total_num_logits = num_reqs + total_num_draft_tokens
 
-            np.cumsum(
-                num_draft_tokens + 1,
-                out=self.input_buffers.cu_num_logits.np[1 : num_reqs + 1],
+            num_logits = num_draft_tokens + 1
+            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
+            expanded_idx_mapping, expanded_local_pos = expand_idx_mapping(
+                idx_mapping, total_num_logits, cu_num_logits, max_expand_len
             )
-            cu_num_logits = self.input_buffers.cu_num_logits.copy_to_gpu(num_reqs + 1)
 
         # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
-        block_tables = self.block_tables.gather_block_tables(idx_mapping_npu)
+        block_tables = self.block_tables.gather_block_tables(idx_mapping)
 
         # Get query_start_loc.
-        np.cumsum(
-            num_scheduled_tokens,
-            out=self.input_buffers.query_start_loc.np[1 : num_reqs + 1],
-        )
+        query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
+        query_start_loc_np[0] = 0
+        np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1 : num_reqs + 1])
         # Pad for full CUDA graph mode.
         # Some attention backends like FA3 require query_start_loc to be non-decreasing.
-        self.input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
-        self.input_buffers.query_start_loc.copy_to_gpu()
-        query_start_loc_gpu = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
-        query_start_loc_cpu = self.input_buffers.query_start_loc.cpu[: num_reqs + 1]
-        query_start_loc_np = self.input_buffers.query_start_loc.np[: num_reqs + 1]
+        query_start_loc_np[num_reqs + 1 :] = num_tokens
+        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 + 1]
+        query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
+        query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
+        max_query_len = num_scheduled_tokens.max().item()
 
         # Get prefill tokens.
         prepare_prefill_inputs(
             self.input_buffers.input_ids,
             self.req_states.next_prefill_tokens,
-            idx_mapping_npu,
-            query_start_loc_gpu,
-            # use prefill_token_ids.copy_to_gpu() because npu doesn't
-            # support uva buffer.
-            self.req_states.prefill_token_ids.copy_to_gpu(),
+            idx_mapping,
+            query_start_loc,
+            self.req_states.prefill_token_ids.gpu,
             self.req_states.prefill_len.gpu,
-            self.req_states.num_computed_tokens,
+            self.req_states.num_computed_tokens.gpu,
         )
 
         # Prepare positions and seq_lens.
         prepare_pos_seq_lens(
-            idx_mapping_npu,
-            query_start_loc_gpu,
-            self.req_states.num_computed_tokens,
+            idx_mapping,
+            query_start_loc,
+            self.req_states.num_computed_tokens.gpu,
             self.input_buffers.positions,
             self.input_buffers.seq_lens,
         )
         seq_lens = self.input_buffers.seq_lens[:num_reqs]
 
+        # Prepare M-RoPE positions.
+        if self.uses_mrope:
+            self.mrope_states.prepare_mrope_positions(
+                idx_mapping,
+                query_start_loc,
+                self.req_states.prefill_len.gpu,
+                self.req_states.num_computed_tokens.gpu,
+            )
+
         # Some input token ids are directly read from the last sampled tokens
         # and draft tokens. Also, get the logits indices to sample tokens from.
         logits_indices = combine_sampled_and_draft_tokens(
             self.input_buffers.input_ids,
-            idx_mapping_npu,
+            idx_mapping,
             self.req_states.last_sampled_tokens,
-            query_start_loc_gpu,
+            query_start_loc,
             seq_lens,
             self.req_states.prefill_len.gpu,
             self.req_states.draft_tokens,
@@ -234,9 +248,10 @@ class NPUModelRunner(GPUModelRunner):
 
         # Compute slot mappings: [num_kv_cache_groups, num_tokens]
         slot_mappings = self.block_tables.compute_slot_mappings(
-            query_start_loc_gpu, self.input_buffers.positions[:num_tokens]
+            idx_mapping, query_start_loc, self.input_buffers.positions[:num_tokens]
         )
-
+        # Layer name -> slot mapping.
+        slot_mappings_by_layer = build_slot_mappings_by_layer(slot_mappings, self.kv_cache_config)
         # Layer name -> attention metadata.
         # TODO(Ronald1995): try to add a new method `build_attn_metadata` in
         # vllm gpu_model_runner_v2, maybe we don't overwrite `prepare_inputs`
@@ -245,37 +260,51 @@ class NPUModelRunner(GPUModelRunner):
             attn_metadata_builders=self.attn_metadata_builders,
             num_reqs=num_reqs,
             num_tokens=num_tokens,
-            query_start_loc_gpu=query_start_loc_gpu,
+            query_start_loc_gpu=query_start_loc,
             query_start_loc_cpu=query_start_loc_cpu,
+            max_query_len=max_query_len,
             seq_lens=self.input_buffers.seq_lens,
-            seq_lens_np=self.input_buffers.seq_lens_np,
-            num_computed_tokens_cpu=self.req_states.num_computed_tokens_cpu[idx_mapping_cpu],
+            max_seq_len=self.max_model_len,
             block_tables=block_tables,
             slot_mappings=slot_mappings,
             kv_cache_config=self.kv_cache_config,
+            # extra attributes for ascend npus.
+            seq_lens_np=self.input_buffers.seq_lens_np,
+            num_computed_tokens_cpu=self.req_states.num_computed_tokens_cpu[idx_mapping_cpu],
             attn_state=attn_state,
         )
 
         input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
         positions = self.input_buffers.positions[:num_tokens_after_padding]
-        return InputBatch(
+        mrope_positions = None
+        if self.uses_mrope:
+            mrope_positions = self.mrope_states.mrope_positions
+            mrope_positions = mrope_positions[:, :num_tokens_after_padding]
+        return AscendInputBatch(
             req_ids=req_ids,
             num_reqs=num_reqs,
-            idx_mapping=idx_mapping_npu,
+            idx_mapping=idx_mapping,
             idx_mapping_np=idx_mapping_np,
+            expanded_idx_mapping=expanded_idx_mapping,
+            expanded_local_pos=expanded_local_pos,
             num_scheduled_tokens=num_scheduled_tokens,
             num_tokens=num_tokens,
             num_tokens_after_padding=num_tokens_after_padding,
             num_draft_tokens=total_num_draft_tokens,
-            query_start_loc=query_start_loc_gpu,
+            query_start_loc=query_start_loc,
             query_start_loc_np=query_start_loc_np,
             seq_lens=seq_lens,
-            seq_lens_np=self.input_buffers.seq_lens_np,
             input_ids=input_ids,
             positions=positions,
+            mrope_positions=mrope_positions,
+            inputs_embeds=None,
             attn_metadata=attn_metadata,
+            slot_mappings=slot_mappings_by_layer,
             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,
+            seq_lens_np=self.input_buffers.seq_lens_np,
         )
 
     def postprocess(
@@ -303,7 +332,7 @@ class NPUModelRunner(GPUModelRunner):
         with torch.npu.stream(self.num_computed_tokens_stream):
             self.num_computed_tokens_stream.wait_stream(default_stream)
             self.num_computed_tokens_cpu.copy_(
-                self.req_states.num_computed_tokens,
+                self.req_states.num_computed_tokens.gpu,
                 non_blocking=True,
             )
             self.num_computed_tokens_event.record()

vllm_ascend/worker/v2/sample/gumbel.py

@@ -30,6 +30,7 @@ def _gumbel_sample_kernel(
     local_max_stride,
     logits_ptr,
     logits_stride,
+    idx_mapping_ptr,
     seeds_ptr,
     pos_ptr,
     temp_ptr,
@@ -37,24 +38,26 @@ def _gumbel_sample_kernel(
     BLOCK_SIZE: tl.constexpr,
     APPLY_TEMPERATURE: tl.constexpr,
 ):
-    req_idx = tl.program_id(0)
+    batch_idx = tl.program_id(0)
+    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
+
     block_idx = tl.program_id(1)
     block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = block < vocab_size
     logits = tl.load(
-        logits_ptr + req_idx * logits_stride + block,
+        logits_ptr + batch_idx * logits_stride + block,
         mask=mask,
         other=float("-inf"),
     )
     logits = logits.to(tl.float32)
 
-    temp = tl.load(temp_ptr + req_idx).to(tl.float32)
+    temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
     if temp != 0.0:
         # Calculate the seed for gumbel noise.
-        seed = tl.load(seeds_ptr + req_idx)
+        seed = tl.load(seeds_ptr + req_state_idx)
         # NOTE(Ronald1995): change pos's dtype to tl.int32, because triton-ascend's
         # compiler doesn't support unint64 of pos arg.
-        pos = tl.load(pos_ptr + req_idx).to(tl.int32)
+        pos = tl.load(pos_ptr + batch_idx).to(tl.int32)
         gumbel_seed = tl.randint(seed, pos)
 
         # Generate gumbel noise.
@@ -66,7 +69,7 @@ def _gumbel_sample_kernel(
 
         # Apply temperature.
         if APPLY_TEMPERATURE:
-            # NOTE(woosuk): Match the behavior of _penalties_and_temperature_kernel.
+            # NOTE(woosuk): Match the behavior of _temperature_kernel.
             # E.g., if the kernel uses tl.div_rn, we should use tl.div_rn here too.
             logits = logits / temp
 
@@ -76,21 +79,18 @@ def _gumbel_sample_kernel(
     idx = tl.argmax(logits, axis=0)
     token_id = block_idx * BLOCK_SIZE + idx
     value = tl.max(logits, axis=0)
-    tl.store(local_argmax_ptr + req_idx * local_argmax_stride + block_idx, token_id)
-    tl.store(local_max_ptr + req_idx * local_max_stride + block_idx, value)
+    tl.store(local_argmax_ptr + batch_idx * local_argmax_stride + block_idx, token_id)
+    tl.store(local_max_ptr + batch_idx * local_max_stride + block_idx, value)
 
 
 def gumbel_sample(
     logits: torch.Tensor,  # [num_reqs, vocab_size]
+    idx_mapping: torch.Tensor,  # [num_reqs]
     temperature: torch.Tensor,  # [num_reqs]
     seed: torch.Tensor,  # [num_reqs]
     pos: torch.Tensor,  # [num_reqs]
     apply_temperature: bool,
 ) -> torch.Tensor:
-    """Override the function because there are some bugs
-    when _gumbel_sample_kernel runs on npu, we need to make some fixes.
-    you could read NOTE(Ronald1995) comments to understand.
-    """
     num_reqs, vocab_size = logits.shape
     BLOCK_SIZE = 1024
     num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
@@ -114,6 +114,7 @@ def gumbel_sample(
         local_max.stride(0),
         logits,
         logits.stride(0),
+        idx_mapping,
         seed,
         pos,
         temperature,

vllm_ascend/worker/v2/sample/sampler.py

@@ -14,22 +14,25 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
-
+import numpy as np
 import torch
-from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.worker.gpu.sample.gumbel import apply_temperature
 from vllm.v1.worker.gpu.sample.min_p import apply_min_p
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 from vllm_ascend.worker.v2.sample.gumbel import gumbel_sample
-from vllm_ascend.worker.v2.sample.penalties import apply_penalties_and_temperature
 
 
 class AscendSampler(Sampler):
     def sample(
         self,
         logits: torch.Tensor,
-        sampling_metadata: SamplingMetadata,
+        idx_mapping: torch.Tensor,
+        idx_mapping_np: np.ndarray,
+        pos: torch.Tensor,
+        input_ids: torch.Tensor,
+        expanded_local_pos: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Override sample method because we need to override triton operators
         called in the method.
@@ -37,19 +40,42 @@ class AscendSampler(Sampler):
         # Copy logits to a new FP32 tensor.
         logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
 
-        # Apply penalties and temperature in place.
-        apply_penalties_and_temperature(logits, sampling_metadata)
-        # Apply min_p in place.
-        if sampling_metadata.min_p is not None:
-            apply_min_p(logits, sampling_metadata.min_p)
-        # Apply top_k and/or top_p. This might return a new tensor.
-        logits = apply_top_k_top_p(logits, sampling_metadata.top_k, sampling_metadata.top_p)
+        # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
+        self.logit_bias_state.apply_logit_bias(logits, idx_mapping, idx_mapping_np, pos)
 
+        # Apply penalties in place.
+        self.penalties_state.apply_penalties(
+            logits,
+            idx_mapping,
+            idx_mapping_np,
+            input_ids,
+            expanded_local_pos,
+            self.num_speculative_tokens,
+        )
+
+        # Apply temperature in place.
+        apply_temperature(logits, idx_mapping, self.sampling_states.temperature.gpu)
+
+        # Apply min_p in place if any request has a non-zero min_p.
+        do_min_p = self.sampling_states.do_min_p(idx_mapping_np)
+        if do_min_p:
+            apply_min_p(logits, idx_mapping, self.sampling_states.min_p.gpu)
+
+        # Apply top_k and/or top_p. This might return a new tensor.
+        do_top_k = self.sampling_states.do_top_k(idx_mapping_np)
+        top_k = self.sampling_states.top_k.gpu[idx_mapping] if do_top_k else None
+        do_top_p = self.sampling_states.do_top_p(idx_mapping_np)
+        top_p = self.sampling_states.top_p.gpu[idx_mapping] if do_top_p else None
+        if do_top_k or do_top_p:
+            logits = apply_top_k_top_p(logits, top_k, top_p)
+
+        # Sample the next token.
         sampled = gumbel_sample(
             logits,
-            sampling_metadata.temperature,
-            sampling_metadata.seeds,
-            sampling_metadata.pos,
+            idx_mapping,
+            self.sampling_states.temperature.gpu,
+            self.sampling_states.seeds.gpu,
+            pos,
             apply_temperature=False,
         )
         return sampled, logits

vllm_ascend/worker/v2/spec_decode/eagle.py

@@ -42,14 +42,23 @@ class AscendEagleSpeculator(EagleSpeculator):
 
     def propose(
         self,
-        input_batch,
-        sampling_metadata,
-        last_hidden_states,
-        aux_hidden_states,
-        num_sampled,
-        num_rejected,
-        last_sampled,
-        next_prefill_tokens,
+        input_batch: InputBatch,
+        # [num_tokens, hidden_size]
+        last_hidden_states: torch.Tensor,
+        # num_layers x [num_tokens, hidden_size]
+        aux_hidden_states: list[torch.Tensor] | None,
+        # [num_reqs]
+        num_sampled: torch.Tensor,
+        # [num_reqs]
+        num_rejected: torch.Tensor,
+        # [max_num_reqs]
+        last_sampled: torch.Tensor,
+        # [max_num_reqs]
+        next_prefill_tokens: torch.Tensor,
+        # [max_num_reqs]
+        temperature: torch.Tensor,
+        # [max_num_reqs]
+        seeds: torch.Tensor,
     ):
         """Override GPU EagleSpeculator.propose for Ascend NPUs,
         because npu attention metadata needs more information,
@@ -62,19 +71,21 @@ class AscendEagleSpeculator(EagleSpeculator):
         with build_attn_metadata_wrapper():
             return super().propose(
                 input_batch,
-                sampling_metadata,
                 last_hidden_states,
                 aux_hidden_states,
                 num_sampled,
                 num_rejected,
                 last_sampled,
                 next_prefill_tokens,
+                temperature,
+                seeds,
             )
 
     def generate_draft(
         self,
         num_reqs: int,
         attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor],
         num_tokens_across_dp,
     ):
         """Override GPU EagleSpeculator.generate_draft for Ascend NPUs, because
@@ -86,6 +97,7 @@ class AscendEagleSpeculator(EagleSpeculator):
         return super().generate_draft(
             num_reqs,
             attn_metadata,
+            slot_mappings,
             num_tokens_across_dp,
         )
 
@@ -94,6 +106,7 @@ class AscendEagleSpeculator(EagleSpeculator):
         self,
         num_tokens: int,
         attn_metadata: dict[str, Any],
+        slot_mappings: dict[str, torch.Tensor] | None,
         num_tokens_across_dp: torch.Tensor | None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Override GPU EagleSpeculator.run_model for Ascend NPUs, because
@@ -103,6 +116,7 @@ class AscendEagleSpeculator(EagleSpeculator):
         last_hidden_states, hidden_states = super().run_model(
             num_tokens,
             attn_metadata,
+            slot_mappings,
             num_tokens_across_dp,
         )
 

vllm_ascend/worker/v2/states.py

@@ -17,11 +17,7 @@
 # This file is a part of the vllm-ascend project.
 #
 
-from contextlib import contextmanager
-
 import torch
-import vllm
-from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker.gpu.states import RequestState
 
 
@@ -36,7 +32,6 @@ class AscendRequestState(RequestState):
         num_speculative_steps: int,
         vocab_size: int,
         device: torch.device,
-        pin_memory: bool,
     ):
         super().__init__(
             max_num_reqs,
@@ -45,11 +40,7 @@ class AscendRequestState(RequestState):
             num_speculative_steps,
             vocab_size,
             device,
-            pin_memory,
         )
-        # because we will override these attribute, delete these attribute to
-        # make sure it's collected by python gc immediately.
-        del self.prefill_token_ids
         # vllm gpu_model_runner_v2 deprecate the seqs_lens_cpu attribute,
         # because they think most attention backends do not need it.
         # However, Ascend attention backend muse uses seqs_lens_cpu,
@@ -60,11 +51,6 @@ class AscendRequestState(RequestState):
             dtype=torch.int32,
             device="cpu",
         )
-        # NOTE(Ronald1995): Ascend NPUs do not support UVA yet,
-        # so we use CpuGpuBuffer to allocate prefill_token_ids buffer.
-        self.prefill_token_ids: CpuGpuBuffer = self._make_buffer(  # type: ignore
-            (self.max_num_reqs, self.max_model_len), dtype=torch.int32
-        )
 
     def add_request(
         self,
@@ -72,32 +58,12 @@ class AscendRequestState(RequestState):
         prompt_len,
         prefill_token_ids,
         num_computed_tokens,
-        sampling_params,
-        lora_request,
     ):
         super().add_request(
             req_id,
             prompt_len,
             prefill_token_ids,
             num_computed_tokens,
-            sampling_params,
-            lora_request,
         )
         req_idx = self.req_id_to_index[req_id]
         self.num_computed_tokens_cpu[req_idx] = num_computed_tokens
-
-
-@contextmanager
-def uva_wrapper():
-    """Context manager to disable UVA for Ascend NPUs."""
-
-    class UvaBufferWrapper:
-        def __init__(self, *args, **kwargs):
-            pass
-
-    try:
-        # TODO(Ronald1995): rectify this when NPU support uva.
-        vllm.v1.worker.gpu.states.UvaBuffer = UvaBufferWrapper
-        yield
-    finally:
-        pass