143 lines
5.0 KiB
Python
143 lines
5.0 KiB
Python
# 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.distributed.parallel_state import get_tp_group
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
from vllm_ascend.ascend_config import get_ascend_config
|
|
from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
|
|
|
|
|
|
@triton.jit(do_not_specialize=["batch_size", "seq_len"])
|
|
def token_bin_counts_and_mask_kernel(
|
|
tokens_ptr,
|
|
tokens_batch_stride,
|
|
tokens_seq_stride,
|
|
batch_size,
|
|
seq_len,
|
|
vocab_size,
|
|
bin_counts_ptr,
|
|
tp_rank,
|
|
counts_batch_stride,
|
|
counts_vocab_stride,
|
|
total_blocks,
|
|
SEQ_BLOCK: tl.constexpr,
|
|
):
|
|
"""Count token occurrences per batch row.
|
|
|
|
1D grid with grid-stride loop: each program processes blocks at
|
|
stride=num_programs to stay within the Triton-Ascend coreDim
|
|
limit (65535) while distributing work evenly across cores.
|
|
"""
|
|
pid = tl.program_id(axis=0)
|
|
num_progs = tl.num_programs(axis=0)
|
|
|
|
vocab_start_idx = tp_rank * vocab_size
|
|
n_seq_blocks = tl.cdiv(seq_len, SEQ_BLOCK)
|
|
|
|
for linear_block in tl.range(pid, total_blocks, num_progs):
|
|
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 + vocab_start_idx,
|
|
)
|
|
|
|
local_token = token - vocab_start_idx
|
|
token_in_range = pos_mask & (token >= vocab_start_idx) & (local_token < vocab_size)
|
|
|
|
safe_local_token = tl.where(token_in_range, local_token, 0)
|
|
count_ptr = batch_counts_start + safe_local_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
|
|
|
|
if n_rows == 0 or 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()
|
|
|
|
# 1D grid: distribute all (batch, seq_block) work items across
|
|
# vector cores via a loop inside the kernel. This avoids the
|
|
# Triton-Ascend grid-size limit of 65535.
|
|
SEQ_BLOCK = 256
|
|
n_seq_blocks = triton.cdiv(n_cols, SEQ_BLOCK)
|
|
total_blocks = n_rows * n_seq_blocks
|
|
grid_size = min(core_num, total_blocks)
|
|
|
|
if get_ascend_config().enable_reduce_sample:
|
|
tp_group = get_tp_group()
|
|
tp_rank = tp_group.rank_in_group
|
|
else:
|
|
tp_rank = 0
|
|
token_bin_counts_and_mask_kernel[(grid_size,)](
|
|
tokens,
|
|
tokens.stride(0),
|
|
tokens.stride(1),
|
|
n_rows,
|
|
n_cols,
|
|
vocab_size,
|
|
bin_counts,
|
|
tp_rank,
|
|
bin_counts.stride(0),
|
|
bin_counts.stride(1),
|
|
total_blocks,
|
|
SEQ_BLOCK=SEQ_BLOCK,
|
|
multibuffer=False,
|
|
)
|
|
return bin_counts, bin_counts > 0
|