73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
################################################################################
|
|
# Copyright(c)2020-2025 Shanghai Biren Technology 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.
|
|
#
|
|
################################################################################
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch_br
|
|
from fastcore.basics import patch_to
|
|
|
|
from vllm.distributed import (get_tensor_model_parallel_rank,
|
|
get_tensor_model_parallel_world_size,
|
|
tensor_model_parallel_all_reduce)
|
|
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
|
from vllm.model_executor.layers.vocab_parallel_embedding import (
|
|
VocabParallelEmbedding)
|
|
from vllm_br import envs
|
|
|
|
|
|
# TODO(shouqing): need to open this patch when fix hang in mtp
|
|
@patch_to(LogitsProcessor)
|
|
def _get_logits(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
lm_head: VocabParallelEmbedding,
|
|
embedding_bias: Optional[torch.Tensor],
|
|
) -> Optional[torch.Tensor]:
|
|
# Get the logits for the next tokens.
|
|
logits = lm_head.quant_method.apply(lm_head,
|
|
hidden_states,
|
|
bias=embedding_bias)
|
|
|
|
spc_num = envs.VLLM_BR_DEVICE_SPC_NUM
|
|
if spc_num > 16:
|
|
bb_input = torch_br._empty_ut_only(size=logits.shape,
|
|
dtype=logits.dtype,
|
|
is_numa=False,
|
|
device=logits.device,
|
|
tensor_type="colmajor")
|
|
|
|
# work around the hang in s1b copy to bb
|
|
bb_input.copy_(logits)
|
|
logits = bb_input
|
|
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
tp_rank = get_tensor_model_parallel_rank()
|
|
|
|
logits_ = torch.zeros((logits.shape[0], logits.shape[-1] * tp_size),
|
|
dtype=logits.dtype,
|
|
device=logits.device)
|
|
|
|
start = logits.shape[-1] * tp_rank
|
|
end = start + logits.shape[-1]
|
|
logits_[:, start:end].copy_(logits)
|
|
logits = tensor_model_parallel_all_reduce(logits_)
|
|
|
|
# Remove paddings in vocab (if any).
|
|
if logits is not None:
|
|
logits = logits[..., :self.org_vocab_size]
|
|
return logits
|