183
vllm_ascend/worker/v2/sample/bad_words.py
Normal file
183
vllm_ascend/worker/v2/sample/bad_words.py
Normal file
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user