184 lines
6.7 KiB
Python
184 lines
6.7 KiB
Python
# 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,
|
|
)
|