init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/_310p/sample/sampler.py

@@ -0,0 +1,128 @@
+#
+# 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.
+#
+
+import torch
+import vllm.envs as envs
+
+from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.sample.sampler import (
+    DEFAULT_LOGPROBS_MODE,
+    AscendSampler,
+    AscendTopKTopPSampler,
+)
+from vllm_ascend.utils import global_stream, npu_stream_switch
+
+_CPU_GENERATOR_CACHE_310P: dict[int, tuple[torch.Generator, int]] = {}
+
+
+def _get_cpu_generator_310p(i: int, generator: torch.Generator) -> torch.Generator:
+    cache_entry = _CPU_GENERATOR_CACHE_310P.get(i)
+    if cache_entry is None or cache_entry[1] != id(generator):
+        cpu_generator = torch.Generator(device="cpu")
+        try:
+            # Keep RNG stream consistent with the original generator.
+            cpu_generator.set_state(generator.get_state())
+        except Exception:
+            cpu_generator.manual_seed(generator.initial_seed())
+        cache_entry = (cpu_generator, id(generator))
+        _CPU_GENERATOR_CACHE_310P[i] = cache_entry
+    return cache_entry[0]
+
+
+def _fill_cpu_exponential_310p(
+    q_cpu: torch.Tensor,
+    generators: dict[int, torch.Generator],
+    has_draft_mask: torch.Tensor | None = None,
+) -> None:
+    """Fill a CPU tensor with exponential values for 310P stability."""
+    if has_draft_mask is not None:
+        has_draft_mask = has_draft_mask.cpu()
+        # Prefill all rows so unmasked requests do not keep uninitialized values.
+        q_cpu.exponential_()
+    elif len(generators) != q_cpu.shape[0]:
+        q_cpu.exponential_()
+    if not generators:
+        return
+    for i, generator in generators.items():
+        cpu_gen = _get_cpu_generator_310p(i, generator)
+        if has_draft_mask is not None:
+            temp_q = torch.empty_like(q_cpu[i])
+            temp_q.exponential_(generator=cpu_gen)
+            q_cpu[i] = torch.where(has_draft_mask[i], temp_q, q_cpu[i])
+        else:
+            q_cpu[i].exponential_(generator=cpu_gen)
+
+
+def fill_exponential_310p(
+    q: torch.Tensor,
+    generators: dict[int, torch.Generator],
+    has_draft_mask: torch.Tensor | None = None,
+) -> None:
+    """Fill ``q`` with exponential values using CPU RNG for 310P stability."""
+    q_cpu = q.cpu()
+    _fill_cpu_exponential_310p(q_cpu, generators, has_draft_mask)
+    q.copy_(q_cpu.to(q.device))
+    # Ensure H2D of q is visible before rejection/recover consumes it.
+    torch.npu.current_stream().synchronize()
+
+
+def _random_sample_310p(
+    probs: torch.Tensor,
+    generators: dict[int, torch.Generator],
+) -> torch.Tensor:
+    """310P-specific random sampling with CPU exponential generation for q."""
+    with npu_stream_switch(global_stream()):
+        q = torch.empty_like(probs).cpu()
+        _fill_cpu_exponential_310p(q, generators)
+        q = q.npu()
+    torch.npu.current_stream().wait_stream(global_stream())
+    return probs.div_(q).argmax(dim=-1).view(-1)
+
+
+class AscendTopKTopPSampler310(AscendTopKTopPSampler):
+    def forward_native(self, logits, generators, k, p):
+        if envs.VLLM_BATCH_INVARIANT:
+            return super().forward_native(logits, generators, k, p)
+        if get_ascend_config().enable_reduce_sample:
+            cand_logits, cand_idx = self.apply_top_k_top_p(logits, k, p, self.top_k)
+            logits_to_return = None
+            if self.logprobs_mode == "processed_logits":
+                logits_to_return = cand_logits
+            elif self.logprobs_mode == "processed_logprobs":
+                logits_to_return = cand_logits.log_softmax(dim=-1, dtype=torch.float32)
+
+            probs = torch.softmax(cand_logits, dim=-1)
+            pos = _random_sample_310p(probs, generators)  # [B]
+
+            next_token = cand_idx.gather(dim=1, index=pos.unsqueeze(1)).squeeze(1)  # [B]
+            return next_token, logits_to_return
+        else:
+            logits = self.apply_top_k_top_p(logits, k, p)
+            logits_to_return = None
+            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)
+            return _random_sample_310p(probs, generators), logits_to_return
+
+
+class AscendSampler310(AscendSampler):
+    def __init__(self, logprobs_mode=DEFAULT_LOGPROBS_MODE):
+        super().__init__(logprobs_mode=logprobs_mode)
+        self.topk_topp_sampler = AscendTopKTopPSampler310(logprobs_mode=logprobs_mode)