ed4ef1f4e7
### What this PR does / why we need it? Implement get_token_bin_counts_and_mask and apply_penalties with Triton-Ascend kernels. This significantly reduces latency of the sampling process when repetition/frequency/presence penalties are enabled. Cherry-pick from main PR #7569 ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? CI passed. Signed-off-by: linfeng-yuan <1102311262@qq.com> Co-authored-by: realliujiaxu <realliujiaxu@163.com>
46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
#
|
|
# apply_all_penalties for AscendSampler - uses Triton-Ascend kernels.
|
|
|
|
import torch
|
|
from vllm.utils.platform_utils import is_pin_memory_available
|
|
from vllm.utils.torch_utils import make_tensor_with_pad
|
|
|
|
from vllm_ascend.ops.triton.penalty import apply_penalties_triton
|
|
|
|
|
|
def _convert_to_tensors(output_token_ids: list[list[int]], vocab_size: int, device: torch.device) -> torch.Tensor:
|
|
"""Convert output_token_ids (list of lists) to padded tensor."""
|
|
output_tokens_tensor = make_tensor_with_pad(
|
|
output_token_ids,
|
|
pad=vocab_size,
|
|
device="cpu",
|
|
dtype=torch.int64,
|
|
pin_memory=is_pin_memory_available(),
|
|
)
|
|
return output_tokens_tensor.to(device, non_blocking=True)
|
|
|
|
|
|
def apply_all_penalties(
|
|
logits: torch.Tensor,
|
|
prompt_token_ids: torch.Tensor,
|
|
presence_penalties: torch.Tensor,
|
|
frequency_penalties: torch.Tensor,
|
|
repetition_penalties: torch.Tensor,
|
|
output_token_ids: list[list[int]],
|
|
) -> torch.Tensor:
|
|
"""Apply penalties to logits via Triton-Ascend."""
|
|
_, vocab_size = logits.shape
|
|
output_tokens_t = _convert_to_tensors(output_token_ids, vocab_size, logits.device)
|
|
output_tokens_t.masked_fill_(output_tokens_t == -1, vocab_size)
|
|
|
|
return apply_penalties_triton(
|
|
logits,
|
|
prompt_token_ids,
|
|
output_tokens_t,
|
|
presence_penalties,
|
|
frequency_penalties,
|
|
repetition_penalties,
|
|
)
|