[releases/v0.18.0][Triton][Sampler] Add penalty-related Triton kernel for better performance of penalties

tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_apply_penalties_triton.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# Compare vllm_ascend.sample.penalties.apply_all_penalties (Triton-Ascend) with
+# vllm.v1.sample.ops.penalties.apply_all_penalties (PyTorch via model_executor).
+# Requires NPU and Triton-Ascend.
+
+import gc
+import pytest
+import torch
+
+from vllm.v1.sample.ops.penalties import apply_all_penalties as v1_apply_all_penalties
+from vllm_ascend.sample.penalties import apply_all_penalties as ascend_apply_all_penalties
+
+# Same scenario grid as test_apply_penalties_model_executor (equivalence + boundaries).
+APPLY_PENALTY_CASES = [
+    pytest.param(0, 0, "mixed", id="empty-both"),
+    pytest.param(0, 16, "mixed", id="empty-prompt"),
+    pytest.param(32, 0, "mixed", id="empty-output"),
+    pytest.param(1, 1, "mixed", id="single-token-each"),
+    pytest.param(32, 16, "mixed", id="typical-small"),
+    pytest.param(128, 64, "mixed", id="typical-large"),
+    pytest.param(128, 64, "all_padding", id="all-padding"),
+]
+
+
+def _make_tokens(
+    num_seqs: int,
+    seq_len: int,
+    vocab_size: int,
+    mode: str,
+    device: str,
+) -> torch.Tensor:
+    if mode == "all_padding":
+        return torch.full(
+            (num_seqs, seq_len), vocab_size, device=device, dtype=torch.int64
+        )
+    if seq_len == 0:
+        return torch.empty((num_seqs, 0), device=device, dtype=torch.int64)
+    tokens = torch.randint(
+        0, vocab_size, (num_seqs, seq_len), device=device, dtype=torch.int64
+    )
+    pad_mask = torch.rand(num_seqs, seq_len, device=device) > 0.7
+    tokens[pad_mask] = vocab_size
+    return tokens
+
+
+@pytest.mark.parametrize("num_seqs", [1, 8, 32, 128])
+@pytest.mark.parametrize("vocab_size", [5120, 151936])
+@pytest.mark.parametrize(
+    "max_prompt_len,max_output_len,token_mode",
+    APPLY_PENALTY_CASES,
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@torch.inference_mode()
+def test_apply_all_penalties_v1_vs_ascend(
+    num_seqs,
+    vocab_size,
+    max_prompt_len,
+    max_output_len,
+    token_mode,
+    dtype,
+    device="npu",
+    seed=42,
+):
+    from vllm_ascend.ops.triton.triton_utils import init_device_properties_triton
+
+    init_device_properties_triton()
+    torch.manual_seed(seed)
+
+    logits_v1 = torch.randn(num_seqs, vocab_size, device=device, dtype=dtype)
+    logits_ascend = logits_v1.clone()
+
+    prompt_tokens = _make_tokens(
+        num_seqs, max_prompt_len, vocab_size, token_mode, device
+    )
+    output_tokens = _make_tokens(
+        num_seqs, max_output_len, vocab_size, token_mode, device
+    )
+    output_token_ids = [row.tolist() for row in output_tokens.cpu()]
+
+    presence_penalties = torch.rand(num_seqs, device=device, dtype=torch.float32) * 0.2
+    frequency_penalties = torch.rand(num_seqs, device=device, dtype=torch.float32) * 0.2
+    repetition_penalties = torch.rand(num_seqs, device=device, dtype=torch.float32) * 0.4 + 1.0
+
+    v1_apply_all_penalties(
+        logits_v1,
+        prompt_tokens,
+        presence_penalties,
+        frequency_penalties,
+        repetition_penalties,
+        output_token_ids,
+    )
+    ascend_apply_all_penalties(
+        logits_ascend,
+        prompt_tokens,
+        presence_penalties,
+        frequency_penalties,
+        repetition_penalties,
+        output_token_ids,
+    )
+
+    atol = 1e-2 if dtype == torch.bfloat16 else 1e-3
+    rtol = 1e-2 if dtype == torch.bfloat16 else 1e-3
+    assert torch.allclose(
+        logits_ascend.float(), logits_v1.float(), atol=atol, rtol=rtol
+    ), (
+        f"Max diff: {(logits_ascend.float() - logits_v1.float()).abs().max().item()}"
+    )
+    gc.collect()
+    torch.npu.empty_cache()
+    torch.npu.reset_peak_memory_stats()

vllm_ascend/ops/triton/bincount.py

@@ -0,0 +1,139 @@
+# 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.
+#
+# Triton-Ascend implementation of get_token_bin_counts_and_mask.
+# Migrated from model_executor/layers/utils.get_token_bin_counts_and_mask.
+# Reference: https://github.com/vllm-project/vllm-ascend/pull/6979
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
+
+
+@triton.jit
+def token_bin_counts_and_mask_kernel(
+    tokens_ptr,
+    tokens_batch_stride,
+    tokens_seq_stride,
+    batch_size,
+    seq_len,
+    vocab_size,
+    bin_counts_ptr,
+    counts_batch_stride,
+    counts_vocab_stride,
+    SEQ_BLOCK: tl.constexpr,
+):
+    """Count token occurrences per batch row.
+
+    2D tiling:
+      - axis=0: core/program group dimension
+      - axis=1: block id dimension
+
+    We linearize (batch_idx, seq_block_id) into a single global block id and
+    distribute blocks across all programs to improve utilization when
+    batch_size is small but seq_len is large (typical prefill).
+
+    Tokens with value >= vocab_size (e.g. padding) are skipped.
+    """
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+    progs = tl.num_programs(axis=0)
+
+    n_seq_blocks = tl.cdiv(seq_len, SEQ_BLOCK)
+    linear_block = pid1 * progs + pid0
+    total_blocks = batch_size * n_seq_blocks
+    if linear_block >= total_blocks:
+        return
+
+    batch_idx = linear_block // n_seq_blocks
+    seq_block_id = linear_block - batch_idx * n_seq_blocks
+    seq_start = seq_block_id * SEQ_BLOCK
+
+    batch_tokens_start = tokens_ptr + batch_idx * tokens_batch_stride
+    batch_counts_start = bin_counts_ptr + batch_idx * counts_batch_stride
+
+    pos_offsets = seq_start + tl.arange(0, SEQ_BLOCK)
+    pos_mask = pos_offsets < seq_len
+    token = tl.load(
+        batch_tokens_start + pos_offsets * tokens_seq_stride,
+        mask=pos_mask,
+        other=vocab_size,  # force invalid
+    )
+    # Only count valid token ids in [0, vocab_size). Padding must use id >= vocab_size
+    # (see vLLM apply_penalties contract); those positions are masked out here.
+    token_in_range = (token >= 0) & (token < vocab_size) & pos_mask
+    count_ptr = batch_counts_start + token * counts_vocab_stride
+    tl.atomic_add(count_ptr, 1, mask=token_in_range)
+
+
+def get_token_bin_counts_and_mask_triton(
+    tokens: torch.Tensor,
+    vocab_size: int,
+    num_seqs: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Triton-Ascend implementation of token bin counting.
+
+    Args:
+        tokens: [num_seqs, seq_len] tensor of token IDs. Padding value
+            should be vocab_size and will be ignored.
+        vocab_size: Vocabulary size.
+        num_seqs: If provided, asserts tokens.shape[0] == num_seqs.
+
+    Returns:
+        bin_counts: [num_seqs, vocab_size] int32 counts.
+        mask: [num_seqs, vocab_size] bool, True where count > 0.
+    """
+    n_rows, n_cols = tokens.shape
+    if num_seqs is not None and num_seqs > 0:
+        assert n_rows == num_seqs, f"tokens rows must match num_seqs: tokens.shape[0]={n_rows}, num_seqs={num_seqs}"
+    n_rows = num_seqs if num_seqs is not None else n_rows
+
+    # seq_len == 0 is valid for empty decode history; return directly.
+    if n_cols == 0:
+        bin_counts = torch.zeros((n_rows, vocab_size), dtype=torch.int32, device=tokens.device)
+        return bin_counts, bin_counts > 0
+
+    core_num = get_vectorcore_num()
+
+    bin_counts = torch.zeros((n_rows, vocab_size), dtype=torch.int32, device=tokens.device)
+    if not tokens.is_contiguous():
+        tokens = tokens.contiguous()
+
+    # 2D grid: (progs, blocks_per_prog_group)
+    # Keep axis-0 bounded by vector core count, and distribute (batch, seq_block)
+    # blocks across all programs to increase utilization when n_rows is small.
+    SEQ_BLOCK = 256
+    n_seq_blocks = triton.cdiv(n_cols, SEQ_BLOCK)
+    total_blocks = n_rows * n_seq_blocks
+    progs = min(core_num, total_blocks)
+    grid = (progs, triton.cdiv(total_blocks, progs))
+
+    token_bin_counts_and_mask_kernel[grid](
+        tokens,
+        tokens.stride(0),
+        tokens.stride(1),
+        n_rows,
+        n_cols,
+        vocab_size,
+        bin_counts,
+        bin_counts.stride(0),
+        bin_counts.stride(1),
+        SEQ_BLOCK=SEQ_BLOCK,
+    )
+    return bin_counts, bin_counts > 0