init v0.23.0

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

vllm_ascend/_310p/sample/__init__.py

@@ -0,0 +1,3 @@
+from vllm_ascend._310p.sample.sampler import AscendSampler310
+
+__all__ = ["AscendSampler310"]

vllm_ascend/_310p/sample/rejection_sampler.py

@@ -0,0 +1,111 @@
+#
+# Copyright (c) 2026 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 contextlib import contextmanager
+
+import torch
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+
+import vllm_ascend.sample.rejection_sampler as rejection_sampler_module
+from vllm_ascend._310p.sample.sampler import fill_exponential_310p
+from vllm_ascend.sample.rejection_sampler import (
+    AscendRejectionSampler,
+    sample_recovered_tokens_blockwise_pytorch,
+    sample_recovered_tokens_pytorch,
+)
+
+
+@contextmanager
+def _bind_sample_recovered_tokens(fn):
+    original = rejection_sampler_module.sample_recovered_tokens
+    rejection_sampler_module.sample_recovered_tokens = fn
+    try:
+        yield
+    finally:
+        rejection_sampler_module.sample_recovered_tokens = original
+
+
+class AscendRejectionSampler310(AscendRejectionSampler):
+    """310P rejection sampler: PyTorch recovered-token path with CPU RNG (no Triton)."""
+
+    def forward(
+        self,
+        metadata: SpecDecodeMetadata,
+        draft_probs: torch.Tensor | None,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput:
+        with _bind_sample_recovered_tokens(self.sample_recovered_tokens):
+            return super().forward(metadata, draft_probs, logits, sampling_metadata)
+
+    def sample_recovered_tokens(
+        self,
+        max_spec_len: int,
+        num_draft_tokens: list[int],
+        cu_num_draft_tokens: torch.Tensor,
+        draft_token_ids: torch.Tensor,
+        draft_probs: torch.Tensor | None,
+        target_probs: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+        device: torch.device,
+        use_block_verify: bool = False,
+        target_indices: torch.Tensor | None = None,
+        global_vocab_size: int | None = None,
+        enable_reduce_sampling: bool = False,
+    ) -> torch.Tensor:
+        batch_size = len(num_draft_tokens)
+        vocab_size = target_probs.shape[-1]
+
+        q = torch.empty(
+            (batch_size, vocab_size),
+            dtype=torch.float32,
+            device=device,
+        )
+        num_draft_tensor = torch.tensor(num_draft_tokens, pin_memory=True).to(device, non_blocking=True)
+        has_draft_mask = num_draft_tensor > 0
+        fill_exponential_310p(q, sampling_metadata.generators, has_draft_mask)
+
+        recovered_token_ids = torch.empty_like(draft_token_ids)
+        if use_block_verify:
+            sample_recovered_tokens_blockwise_pytorch(
+                recovered_token_ids,
+                cu_num_draft_tokens,
+                draft_token_ids,
+                draft_probs,
+                target_probs,
+                q,
+                vocab_size,
+                IS_NGRAM=draft_probs is None,
+                target_indices=target_indices,
+                enable_reduce_sampling=enable_reduce_sampling,
+            )
+        else:
+            sample_recovered_tokens_pytorch(
+                recovered_token_ids,
+                cu_num_draft_tokens,
+                draft_token_ids,
+                draft_probs,
+                target_probs,
+                q,
+                vocab_size,
+                IS_NGRAM=draft_probs is None,
+                target_indices=target_indices,
+                enable_reduce_sampling=enable_reduce_sampling,
+            )
+        return recovered_token_ids

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)