init v0.11.0rc0

vllm_ascend/sample/logits_processor/__init__.py

@@ -0,0 +1,50 @@
+import itertools
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Union
+
+import torch
+from vllm.logger import init_logger
+from vllm.v1.sample import logits_processor
+from vllm.v1.sample.logits_processor.builtin import (LogitBiasLogitsProcessor,
+                                                     MinTokensLogitsProcessor)
+from vllm.v1.sample.logits_processor.interface import LogitsProcessor
+from vllm.v1.sample.logits_processor.state import LogitsProcessors
+
+from vllm_ascend.sample.logits_processor.builtin import \
+    AscendMinPLogitsProcessor
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
+
+# Error message when the user tries to initialize vLLM with a pooling model
+# and custom logitsproces
+STR_POOLING_REJECTS_LOGITSPROCS = ("Pooling models do not support custom"
+                                   " logits processors.")
+
+BUILTIN_LOGITS_PROCESSORS: list[type[LogitsProcessor]] = [
+    MinTokensLogitsProcessor,
+    LogitBiasLogitsProcessor,
+    AscendMinPLogitsProcessor,
+]
+
+
+def build_logitsprocs(
+    vllm_config: "VllmConfig",
+    device: torch.device,
+    is_pin_memory: bool,
+    is_pooling_model: bool,
+    custom_logitsprocs: Sequence[Union[str, type[LogitsProcessor]]] = (),
+) -> LogitsProcessors:
+    if is_pooling_model:
+        if custom_logitsprocs:
+            raise ValueError(STR_POOLING_REJECTS_LOGITSPROCS)
+        logger.debug("Skipping logits processor loading because pooling models"
+                     " do not support logits processors.")
+        return LogitsProcessors()
+    custom_logitsprocs_classes = logits_processor._load_custom_logitsprocs(
+        custom_logitsprocs)
+    return LogitsProcessors(
+        ctor(vllm_config, device, is_pin_memory) for ctor in itertools.chain(
+            BUILTIN_LOGITS_PROCESSORS, custom_logitsprocs_classes))

vllm_ascend/sample/logits_processor/builtin.py

@@ -0,0 +1,35 @@
+import torch
+from vllm.config import VllmConfig
+from vllm.v1.sample.logits_processor import MinPLogitsProcessor
+
+
+class AscendMinPLogitsProcessor(MinPLogitsProcessor):
+
+    def __init__(self, vllm_config: "VllmConfig", device: torch.device,
+                 is_pin_memory: bool):
+        super().__init__(vllm_config, device, is_pin_memory)
+
+        decode_max_num_seqs = getattr(vllm_config.scheduler_config,
+                                      'decode_max_num_seqs', 0)
+        if decode_max_num_seqs != 0:
+            max_num_reqs = max(vllm_config.scheduler_config.max_num_seqs,
+                               decode_max_num_seqs)
+
+            self.min_p_count: int = 0
+
+            self.min_p_cpu_tensor = torch.zeros((max_num_reqs, ),
+                                                dtype=torch.float32,
+                                                device="cpu",
+                                                pin_memory=is_pin_memory)
+            self.min_p_cpu = self.min_p_cpu_tensor.numpy()
+
+            self.use_double_tensor = torch.device(device).type != "cpu"
+
+            if self.use_double_tensor:
+                # Pre-allocated device tensor
+                self.min_p_device: torch.Tensor = torch.empty(
+                    (max_num_reqs, ), dtype=torch.float32, device=device)
+            else:
+                self.min_p_device = self.min_p_cpu_tensor
+            # Current slice of the device tensor
+            self.min_p: torch.Tensor = self.min_p_device[:0]

vllm_ascend/sample/sampler.py

@@ -5,11 +5,10 @@ from vllm.v1.sample.sampler import Sampler
 
 from vllm_ascend.utils import is_310p, vllm_version_is
 
-if not (vllm_version_is("0.10.1.1") or vllm_version_is("0.10.1")):
+if vllm_version_is("0.10.2"):
     from vllm.config import LogprobsMode
     DEFAULT_LOGPROBS_MODE = LogprobsMode.RAW_LOGPROBS
 else:
-    LogprobsMode = None
     DEFAULT_LOGPROBS_MODE = "raw_logprobs"
 
 
@@ -30,7 +29,8 @@ class AscendTopKTopPSampler(TopKTopPSampler):
         p: torch.Tensor,
     ) -> torch.Tensor:
         # npu_top_k_top_p uses the operator aclnnApplyTopKTopP, but aclnnApplyTopKTopP currently does not support 310P
-        if not is_310p() and p is not None and k is not None:
+        if not is_310p() and p is not None and k is not None and 1 <= int(
+                k.max()) <= 1024:
             # npu_top_k_top_p's parameter order is (logits, p, k), not (logits, k, p)
             return torch_npu.npu_top_k_top_p(logits, p, k)
 
@@ -68,19 +68,19 @@ class AscendTopKTopPSampler(TopKTopPSampler):
     def forward_native(self, logits, generators, k, p):
         """Override pytorch native implementation to torch_npu"""
         logits = self._apply_top_k_top_p(logits, k, p)
-        if not (vllm_version_is("0.10.1.1") or vllm_version_is("0.10.1")):
-
-            logits_to_return = None
+        logits_to_return = None
+        if vllm_version_is("0.10.2"):
             if self.logprobs_mode == LogprobsMode.PROCESSED_LOGITS:
                 logits_to_return = logits
             elif self.logprobs_mode == LogprobsMode.PROCESSED_LOGPROBS:
                 logits_to_return = logits.log_softmax(dim=-1,
                                                       dtype=torch.float32)
+        else:
+            if self.logprobs_mode == "processed_logits":
+                logits_to_return = logits
+            elif self.logprobs_mode == "processed_logprobs":
+                logits_to_return = logits.log_softmax(dim=-1,
+                                                      dtype=torch.float32)
 
         probs = logits.softmax(dim=-1, dtype=torch.float32)
-        output = None
-        if vllm_version_is("0.10.1.1") or vllm_version_is("0.10.1"):
-            output = random_sample(probs, generators)
-        else:
-            output = (random_sample(probs, generators), logits_to_return)
-        return output
+        return random_sample(probs, generators), logits_to_return