init v0.23.0

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

vllm_ascend/worker/v2/sample/bad_words.py

@@ -0,0 +1,183 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/sample/bad_words.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.
+#
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
+
+MAX_BAD_WORDS_TOTAL_TOKENS = 1024  # Max total tokens for all bad words per request
+MAX_NUM_BAD_WORDS = 128  # Max number of bad words per request
+
+
+@triton.jit(do_not_specialize=["num_tokens", "max_num_bad_words"])
+def _bad_words_kernel(
+    logits_ptr,
+    logits_stride,
+    expanded_idx_mapping_ptr,
+    bad_word_token_ids_ptr,
+    bad_word_token_ids_stride,
+    bad_word_offsets_ptr,
+    bad_word_offsets_stride,
+    num_bad_words_ptr,
+    all_token_ids_ptr,
+    all_token_ids_stride,
+    prompt_len_ptr,
+    total_len_ptr,
+    input_ids_ptr,
+    expanded_local_pos_ptr,
+    num_tokens,
+    max_num_bad_words,
+    MAX_PREFIX_LEN: tl.constexpr,
+):
+    """
+    Optimized bad words filtering kernel for Ascend NPU.
+
+    Key optimizations:
+    - Optimized memory access patterns
+    - Reduced redundant calculations
+    - Enhanced data locality
+    - Minimized conditional branches
+    - Improved load balancing
+    """
+    pid = tl.program_id(0)
+    num_cores = tl.num_programs(0)
+
+    # Calculate tokens per core for better load balancing
+    tokens_per_core = (num_tokens + num_cores - 1) // num_cores
+    start_token = pid * tokens_per_core
+    end_token = min(start_token + tokens_per_core, num_tokens)
+
+    # Process each token assigned to this core
+    for token_idx in range(start_token, end_token):
+        # Load request state index
+        req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
+        num_bad_words = tl.load(num_bad_words_ptr + req_state_idx)
+
+        # Only process if there are bad words for this request
+        if num_bad_words > 0:
+            # Load position information
+            pos = tl.load(expanded_local_pos_ptr + token_idx)
+            cur_req_first_pos = token_idx - pos
+
+            # Load length information
+            prompt_len = tl.load(prompt_len_ptr + req_state_idx)
+            total_len = tl.load(total_len_ptr + req_state_idx)
+            output_len = total_len - prompt_len
+            effective_len = output_len + pos
+
+            # Precompute base addresses
+            bd_offsets_base = bad_word_offsets_ptr + req_state_idx * bad_word_offsets_stride
+            bd_tokens_base = bad_word_token_ids_ptr + req_state_idx * bad_word_token_ids_stride
+            output_base = all_token_ids_ptr + req_state_idx * all_token_ids_stride + prompt_len
+
+            # Process each bad word for this token
+            for bw_idx in range(max_num_bad_words):
+                if bw_idx < num_bad_words:
+                    # Load bad word range
+                    start = tl.load(bd_offsets_base + bw_idx)
+                    end = tl.load(bd_offsets_base + bw_idx + 1)
+                    bad_word_len = end - start
+                    prefix_len = bad_word_len - 1
+
+                    # Check prefix length validity
+                    if prefix_len <= effective_len:
+                        # Load last token
+                        last_token = tl.load(bd_tokens_base + end - 1)
+
+                        # Match checking with early termination
+                        match = 1
+                        j = 0
+                        while j < prefix_len and match:
+                            # Load expected token
+                            expected = tl.load(bd_tokens_base + start + j)
+
+                            # Calculate actual position and load actual token
+                            actual_pos = effective_len - prefix_len + j
+                            if actual_pos >= output_len:
+                                spec_offset = actual_pos - output_len
+                                actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset)
+                            else:
+                                actual = tl.load(output_base + actual_pos)
+
+                            # Check for mismatch
+                            if expected != actual:
+                                match = 0
+                            j += 1
+
+                        # Store result if match found
+                        if match:
+                            tl.store(logits_ptr + token_idx * logits_stride + last_token, -float("inf"))
+
+
+def apply_bad_words(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    bad_word_token_ids: torch.Tensor,
+    bad_word_offsets: torch.Tensor,
+    num_bad_words: torch.Tensor,
+    all_token_ids: torch.Tensor,
+    prompt_len: torch.Tensor,
+    total_len: torch.Tensor,
+    input_ids: torch.Tensor,
+    expanded_local_pos: torch.Tensor,
+    max_num_bad_words: int,
+) -> None:
+    """
+    Apply bad words filtering to logits.
+
+    Args:
+        logits: [num_tokens, vocab_size] - Model output logits
+        expanded_idx_mapping: [num_tokens] - Token to request mapping
+        bad_word_token_ids: [max_num_reqs, MAX_BAD_WORDS_TOTAL_TOKENS] - Bad word token IDs
+        bad_word_offsets: [max_num_reqs, MAX_NUM_BAD_WORDS + 1] - Bad word offsets
+        num_bad_words: [max_num_reqs] - Number of bad words per request
+        all_token_ids: [max_num_reqs, max_seq_len] - All token IDs
+        prompt_len: [max_num_reqs] - Prompt length
+        total_len: [max_num_reqs] - Total length
+        input_ids: [num_tokens] - Input IDs
+        expanded_local_pos: [num_tokens] - Expanded local position
+        max_num_bad_words: Maximum number of bad words to check
+    """
+    num_tokens = logits.shape[0]
+
+    core_num = get_vectorcore_num()
+
+    MAX_PREFIX_LEN = 32
+
+    _bad_words_kernel[(core_num,)](
+        logits,
+        logits.stride(0),
+        expanded_idx_mapping,
+        bad_word_token_ids,
+        bad_word_token_ids.stride(0),
+        bad_word_offsets,
+        bad_word_offsets.stride(0),
+        num_bad_words,
+        all_token_ids,
+        all_token_ids.stride(0),
+        prompt_len,
+        total_len,
+        input_ids,
+        expanded_local_pos,
+        num_tokens,
+        max_num_bad_words,
+        MAX_PREFIX_LEN,
+    )

vllm_ascend/worker/v2/sample/gumbel.py

@@ -0,0 +1,205 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/sample/gumbel.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.
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit(do_not_specialize=["logits_stride", "vocab_size"])
+def _temperature_kernel(
+    logits_ptr,
+    logits_stride,
+    expanded_idx_mapping_ptr,
+    temperature_ptr,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
+    temperature = tl.load(temperature_ptr + req_state_idx).to(tl.float32)
+    if temperature == 0.0 or temperature == 1.0:
+        # Early return to avoid loading logits
+        return
+
+    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 + token_idx * logits_stride + block, mask=mask)
+    logits = logits.to(tl.float32)
+    logits = logits / temperature
+    tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
+
+
+def apply_temperature(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    temperature: torch.Tensor,
+) -> None:
+    """
+    Args:
+        logits: Tensor of shape (num_tokens, vocab_size) containing the logits.
+        expanded_idx_mapping: Tensor containing the mapping from token index
+            to request index of tensor temperature.
+        temperature: Tensor containing the temperature value for each request.
+    """
+    num_tokens, vocab_size = logits.shape
+    BLOCK_SIZE = 44032
+    num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+    _temperature_kernel[(num_tokens, num_blocks)](
+        logits,
+        logits.stride(0),
+        expanded_idx_mapping,
+        temperature,
+        vocab_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+        multibuffer=False,
+    )
+
+
+@triton.jit(
+    do_not_specialize=[
+        "local_argmax_stride",
+        "local_max_stride",
+        "processed_logits_stride",
+        "logits_stride",
+        "vocab_size",
+    ]
+)
+def _gumbel_sample_kernel(
+    local_argmax_ptr,
+    local_argmax_stride,
+    local_max_ptr,
+    local_max_stride,
+    processed_logits_ptr,
+    processed_logits_stride,
+    processed_logits_col_ptr,
+    logits_ptr,
+    logits_stride,
+    expanded_idx_mapping_ptr,
+    seeds_ptr,
+    pos_ptr,
+    temp_ptr,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+    APPLY_TEMPERATURE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    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 + token_idx * logits_stride + block,
+        mask=mask,
+        other=float("-inf"),
+    )
+    logits = logits.to(tl.float32)
+
+    req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
+    temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
+
+    if temp != 0.0 and APPLY_TEMPERATURE:
+        # NOTE(woosuk): Match the behavior of _temperature_kernel.
+        logits = logits / temp
+
+    if processed_logits_ptr is not None:
+        # Store the temperature-applied logits.
+        if processed_logits_col_ptr is not None:
+            col = tl.load(processed_logits_col_ptr)
+        else:
+            col = 0
+        tl.store(
+            processed_logits_ptr + req_state_idx * processed_logits_stride + col * vocab_size + block,
+            logits,
+            mask=mask,
+        )
+
+    if temp != 0.0:
+        # Calculate the seed for gumbel noise.
+        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 uint64 of pos arg.
+        pos = tl.load(pos_ptr + token_idx).to(tl.int32)
+        gumbel_seed = tl.randint(seed, pos)
+
+        # NOTE(Ronald1995): r is tl.float64 in vllm, change it to tl.float32,
+        # because triton-ascend's compiler does not support float64.
+        r = tl.rand(gumbel_seed, block).to(tl.float32)
+        gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
+
+        # Apply gumbel noise.
+        logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
+
+    idx = tl.argmax(logits, axis=0)
+    token_id = block_idx * BLOCK_SIZE + idx
+    value = tl.max(logits, axis=0)
+    tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)
+    tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
+
+
+def gumbel_sample(
+    logits: torch.Tensor,  # [num_tokens, vocab_size]
+    expanded_idx_mapping: torch.Tensor,  # [num_tokens]
+    temperature: torch.Tensor,  # [max_num_reqs]
+    seed: torch.Tensor,  # [max_num_reqs]
+    pos: torch.Tensor,  # [num_tokens]
+    apply_temperature: bool,
+    output_processed_logits: torch.Tensor | None = None,
+    output_processed_logits_col: torch.Tensor | None = None,
+    use_fp64: bool = False,
+) -> torch.Tensor:
+    if use_fp64:
+        raise NotImplementedError("FP64 Gumbel sampling is not supported on NPU.")
+    num_tokens, vocab_size = logits.shape
+    BLOCK_SIZE = 1024
+    num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+    local_argmax = torch.empty(
+        num_tokens,
+        num_blocks,
+        dtype=torch.int64,
+        device=logits.device,
+    )
+    local_max = torch.empty(
+        num_tokens,
+        num_blocks,
+        dtype=torch.float32,
+        device=logits.device,
+    )
+    _gumbel_sample_kernel[(num_tokens, num_blocks)](
+        local_argmax,
+        local_argmax.stride(0),
+        local_max,
+        local_max.stride(0),
+        output_processed_logits,
+        output_processed_logits.stride(0) if output_processed_logits is not None else 0,
+        output_processed_logits_col,
+        logits,
+        logits.stride(0),
+        expanded_idx_mapping,
+        seed,
+        pos,
+        temperature,
+        vocab_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+        APPLY_TEMPERATURE=apply_temperature,
+    )
+    # NOTE(woosuk): Use int64 for later indexing.
+    max_block_idx = local_max.argmax(dim=-1, keepdim=True)
+    sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
+    return sampled

vllm_ascend/worker/v2/sample/logprob.py

@@ -0,0 +1,164 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/sample/logprob.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.
+
+import torch
+from vllm.triton_utils import tl, triton
+from vllm.v1.outputs import LogprobsTensors
+
+from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
+
+
+@triton.jit
+def _topk_log_softmax_kernel(
+    output_ptr,
+    logits_ptr,
+    logits_stride,
+    topk_ids_ptr,
+    topk,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+    PADDED_TOPK: tl.constexpr,
+):
+    req_idx = tl.program_id(0)
+    row_ptr = logits_ptr + req_idx * logits_stride
+
+    max_val = float("-inf")
+    for i in range(0, vocab_size, BLOCK_SIZE):
+        block = i + tl.arange(0, BLOCK_SIZE)
+        logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
+        max_val = tl.max(tl.maximum(logits, max_val, propagate_nan=tl.PropagateNan.ALL))
+    max_val = max_val.to(tl.float32)  # type: ignore
+
+    se = 0.0
+    for i in range(0, vocab_size, BLOCK_SIZE):
+        block = i + tl.arange(0, BLOCK_SIZE)
+        logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
+        logits = logits.to(tl.float32)
+        e = tl.exp(logits - max_val)
+        se += tl.sum(e)
+    lse = tl.log(se)
+
+    k_offset = tl.arange(0, PADDED_TOPK)
+    k_mask = k_offset < topk
+    topk_ids = tl.load(topk_ids_ptr + req_idx * topk + k_offset, mask=k_mask, other=0)
+
+    logits = tl.load(row_ptr + topk_ids, mask=k_mask)
+    logits = logits.to(tl.float32)
+    o = logits - lse - max_val
+    tl.store(output_ptr + req_idx * topk + k_offset, o, mask=k_mask)
+
+
+def compute_token_logprobs(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor:
+    batch_size, vocab_size = logits.shape
+    token_ids = token_ids.to(torch.int64)
+    num_logprobs = token_ids.shape[1]
+    logprobs = logits.new_empty((batch_size, num_logprobs), dtype=torch.float32)
+    _topk_log_softmax_kernel[(batch_size,)](
+        logprobs,
+        logits,
+        logits.stride(0),
+        token_ids,
+        num_logprobs,
+        vocab_size,
+        BLOCK_SIZE=12944,
+        PADDED_TOPK=max(triton.next_power_of_2(num_logprobs), 2),
+        multibuffer=False,
+    )
+    return logprobs
+
+
+@triton.jit(do_not_specialize=["batch_size", "rows_per_core"])
+def _ranks_kernel(
+    output_ptr,
+    logits_ptr,
+    logits_stride,
+    token_ids_ptr,
+    vocab_size,
+    batch_size,
+    rows_per_core,
+    BLOCK_SIZE: tl.constexpr,
+):
+    core_id = tl.program_id(0)
+
+    start_row = core_id * rows_per_core
+    end_row = start_row + rows_per_core
+
+    for req_idx in range(start_row, end_row):
+        if req_idx < batch_size:
+            row_ptr = logits_ptr + req_idx * logits_stride
+
+            token_id = tl.load(token_ids_ptr + req_idx)
+            x = tl.load(row_ptr + token_id)
+
+            n_vec = tl.zeros([BLOCK_SIZE], dtype=tl.int32)
+            for i in range(0, vocab_size, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                logits = tl.load(row_ptr + block, mask=block < vocab_size, other=float("-inf"))
+                n_vec += (logits > x).to(tl.int32)
+            n = tl.sum(n_vec)
+            tl.store(output_ptr + req_idx, n)
+
+
+def compute_topk_logprobs(
+    logits: torch.Tensor,
+    num_logprobs: int,
+    sampled_token_ids: torch.Tensor,
+    cu_num_logits: list[int] | None = None,
+) -> LogprobsTensors:
+    assert num_logprobs >= 0
+    batch_size, vocab_size = logits.shape
+    logprob_token_ids = sampled_token_ids.unsqueeze(-1)
+    if num_logprobs > 0:
+        topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices
+        logprob_token_ids = torch.cat((sampled_token_ids.unsqueeze(-1), topk_indices), dim=1)
+
+    # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full
+    # logprobs tensor. Instead, we only compute and return the logprobs of
+    # the topk + 1 tokens.
+    logprobs = compute_token_logprobs(logits, logprob_token_ids)
+    token_ranks = torch.empty(
+        batch_size,
+        dtype=torch.int64,
+        device=logits.device,
+    )
+
+    vec_core = get_vectorcore_num()
+    NUM_CORES = min(batch_size, vec_core)
+
+    rows_per_core = triton.cdiv(batch_size, NUM_CORES)
+    BLOCK_SIZE = 8192
+    grid = (NUM_CORES,)
+    _ranks_kernel[grid](
+        token_ranks,
+        logits,
+        logits.stride(0),
+        sampled_token_ids,
+        vocab_size,
+        batch_size,
+        rows_per_core,
+        BLOCK_SIZE=BLOCK_SIZE,
+        multibuffer=False,
+    )
+
+    return LogprobsTensors(
+        logprob_token_ids=logprob_token_ids,
+        logprobs=logprobs,
+        selected_token_ranks=token_ranks,
+        cu_num_generated_tokens=cu_num_logits,
+    )

vllm_ascend/worker/v2/sample/min_p.py

@@ -0,0 +1,91 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/sample/min_p.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.
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
+
+
+@triton.jit(do_not_specialize=["num_tokens"])
+def _min_p_kernel(
+    in_logits_ptr,
+    out_logits_ptr,
+    logits_stride,
+    expanded_idx_mapping_ptr,
+    min_p_ptr,
+    vocab_size,
+    num_tokens,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    core_num = tl.num_programs(0)
+
+    tokens_per_block = (num_tokens + core_num - 1) // core_num
+    start_token = pid * tokens_per_block
+    end_token = tl.minimum(start_token + tokens_per_block, num_tokens)
+
+    for token_idx in tl.range(start_token, end_token):
+        req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
+        min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32)
+        if min_p != 0.0:
+            max_val = float("-inf")
+            for i in range(0, vocab_size, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                mask = block < vocab_size
+                logits = tl.load(
+                    in_logits_ptr + token_idx * logits_stride + block,
+                    mask=mask,
+                    other=float("-inf"),
+                )
+                max_val = tl.max(tl.maximum(logits, max_val))
+            max_val = max_val.to(tl.float32)  # type: ignore
+
+            threshold = max_val + tl.log(min_p)
+            for i in range(0, vocab_size, BLOCK_SIZE):
+                block = i + tl.arange(0, BLOCK_SIZE)
+                mask = block < vocab_size
+                logits = tl.load(
+                    in_logits_ptr + token_idx * logits_stride + block,
+                    mask=mask,
+                    other=float("-inf"),
+                )
+                logits = tl.where(logits < threshold, float("-inf"), logits)
+                tl.store(out_logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
+
+
+def apply_min_p(logits: torch.Tensor, expanded_idx_mapping: torch.Tensor, min_p: torch.Tensor) -> None:
+    num_tokens, vocab_size = logits.shape
+
+    vec_core = get_vectorcore_num()
+    core_nums = min(num_tokens, vec_core)
+
+    BLOCK_SIZE = min(triton.next_power_of_2(vocab_size), 8192)
+
+    _min_p_kernel[(core_nums,)](
+        logits,
+        logits,
+        logits.stride(0),
+        expanded_idx_mapping,
+        min_p,
+        vocab_size,
+        num_tokens,
+        BLOCK_SIZE=BLOCK_SIZE,
+        multibuffer=False,
+    )

vllm_ascend/worker/v2/sample/penalties.py

@@ -0,0 +1,215 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/sample/penalties.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.
+#
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _penalties_kernel(
+    logits_ptr,
+    logits_stride,
+    expanded_idx_mapping_ptr,
+    token_ids_ptr,
+    expanded_local_pos_ptr,
+    repetition_penalty_ptr,
+    frequency_penalty_ptr,
+    presence_penalty_ptr,
+    prompt_bin_mask_ptr,
+    prompt_bin_mask_stride,
+    output_bin_counts_ptr,
+    output_bin_counts_stride,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
+    rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)
+    freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)
+    pres_penalty = tl.load(presence_penalty_ptr + req_state_idx)
+
+    use_rep_penalty = rep_penalty != 1.0
+    use_freq_penalty = freq_penalty != 0.0
+    use_pres_penalty = pres_penalty != 0.0
+
+    # NPU doesn't support chained 'or' operations like 'A or B or C'
+    use_penalty = use_rep_penalty or use_freq_penalty
+    use_penalty = use_penalty or use_pres_penalty
+    if not use_penalty:
+        # Early return to avoid loading logits.
+        return
+
+    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 + token_idx * logits_stride + block, mask=mask)
+    logits = logits.to(tl.float32)
+
+    base_output_counts = tl.load(
+        output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
+        mask=mask,
+        other=0,
+    )
+
+    # Accumulate draft token counts from previous positions directly into
+    # output_bin_counts (preserves its native tensor layout, avoiding an
+    # expensive shared-memory layout conversion after the loop).
+    pos = tl.load(expanded_local_pos_ptr + token_idx)
+    start_idx = token_idx - pos
+    output_bin_counts = base_output_counts
+    for prev_pos in tl.range(pos):
+        prev_token = tl.load(token_ids_ptr + start_idx + prev_pos + 1)
+        token_match = block == prev_token
+        output_bin_counts = output_bin_counts + token_match.to(tl.int32)
+    output_bin_mask = output_bin_counts != 0
+
+    # Apply repetition penalties.
+    if use_rep_penalty:
+        packed_block = block_idx * BLOCK_SIZE // 32 + tl.arange(0, BLOCK_SIZE // 32)
+        packed_mask = tl.load(
+            prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + packed_block,
+            mask=packed_block < tl.cdiv(vocab_size, 32),
+            other=0,
+        )
+        bit_masks = 1 << tl.arange(0, 32)
+        bit_masks_expanded = bit_masks[None, :]
+        packed_expanded = packed_mask[:, None]
+        bits_matrix = (packed_expanded & bit_masks_expanded) != 0
+        prompt_bin_mask = bits_matrix.reshape(BLOCK_SIZE)
+
+        # If token appears in prompt or output, apply, otherwise use 1.0 for no-op.
+        scale = tl.where(prompt_bin_mask | output_bin_mask, rep_penalty, 1.0)
+        # If logits are positive, divide by penalty, otherwise multiply by penalty.
+        logits *= tl.where(logits > 0, 1.0 / scale, scale)
+
+    # Apply frequency penalties.
+    logits -= freq_penalty * output_bin_counts
+    # Apply presence penalties.
+    logits -= pres_penalty * output_bin_mask
+    # Store back to logits.
+    tl.store(logits_ptr + token_idx * logits_stride + block, logits, mask=mask)
+
+
+def apply_penalties(
+    logits: torch.Tensor,
+    expanded_idx_mapping: torch.Tensor,
+    token_ids: torch.Tensor,
+    expanded_local_pos: torch.Tensor,
+    repetition_penalty: torch.Tensor,
+    frequency_penalty: torch.Tensor,
+    presence_penalty: torch.Tensor,
+    prompt_bin_mask: torch.Tensor,
+    output_bin_counts: torch.Tensor,
+) -> None:
+    num_tokens, vocab_size = logits.shape
+    BLOCK_SIZE = 4096
+    num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
+    _penalties_kernel[(num_tokens, num_blocks)](
+        logits,
+        logits.stride(0),
+        expanded_idx_mapping,
+        token_ids,
+        expanded_local_pos,
+        repetition_penalty,
+        frequency_penalty,
+        presence_penalty,
+        prompt_bin_mask,
+        prompt_bin_mask.stride(0),
+        output_bin_counts,
+        output_bin_counts.stride(0),
+        vocab_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+
+@triton.jit
+def _bincount_kernel(
+    expanded_idx_mapping_ptr,
+    all_token_ids_ptr,
+    all_token_ids_stride,
+    prompt_len_ptr,
+    prefill_len_ptr,
+    prompt_bin_mask_ptr,
+    prompt_bin_mask_stride,
+    output_bin_counts_ptr,
+    output_bin_counts_stride,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    block_idx = tl.program_id(1)
+    req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
+
+    prefill_len = tl.load(prefill_len_ptr + req_state_idx)
+    if block_idx * BLOCK_SIZE >= prefill_len:
+        return
+
+    prompt_len = tl.load(prompt_len_ptr + req_state_idx)
+    block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    if block_idx * BLOCK_SIZE < prompt_len:
+        mask = block < prompt_len
+        prompt_tokens = tl.load(all_token_ids_ptr + req_state_idx * all_token_ids_stride + block, mask=mask)
+        idx = prompt_tokens // 32
+
+        bit_idx = prompt_tokens % 32
+        bit = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_idx
+
+        tl.atomic_or(
+            prompt_bin_mask_ptr + req_state_idx * prompt_bin_mask_stride + idx,
+            bit,
+            mask=mask,
+        )
+
+    if (block_idx + 1) * BLOCK_SIZE >= prompt_len:
+        mask = block < prefill_len
+        mask &= block >= prompt_len
+        output_tokens = tl.load(all_token_ids_ptr + req_state_idx * all_token_ids_stride + block, mask=mask)
+        tl.atomic_add(
+            output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + output_tokens,
+            1,
+            mask=mask,
+        )
+
+
+def bincount(
+    expanded_idx_mapping: torch.Tensor,
+    all_token_ids: torch.Tensor,
+    prompt_len: torch.Tensor,
+    prefill_len: torch.Tensor,
+    prompt_bin_mask: torch.Tensor,
+    output_bin_counts: torch.Tensor,
+    max_prefill_len: int,
+) -> None:
+    prompt_bin_mask[expanded_idx_mapping] = 0
+    output_bin_counts[expanded_idx_mapping] = 0
+    num_tokens = expanded_idx_mapping.shape[0]
+    BLOCK_SIZE = 1024
+    num_blocks = triton.cdiv(max_prefill_len, BLOCK_SIZE)
+    _bincount_kernel[(num_tokens, num_blocks)](
+        expanded_idx_mapping,
+        all_token_ids,
+        all_token_ids.stride(0),
+        prompt_len,
+        prefill_len,
+        prompt_bin_mask,
+        prompt_bin_mask.stride(0),
+        output_bin_counts,
+        output_bin_counts.stride(0),
+        BLOCK_SIZE=BLOCK_SIZE,
+    )