45
vllm_ascend/sample/penalties.py
Normal file
45
vllm_ascend/sample/penalties.py
Normal file
@@ -0,0 +1,45 @@
|
||||
# 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,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,39 +1,240 @@
|
||||
import torch
|
||||
import torch_npu
|
||||
from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler, random_sample
|
||||
import vllm.envs as envs
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
from vllm.logger import logger
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
|
||||
from vllm.v1.sample.sampler import Sampler
|
||||
|
||||
from vllm_ascend.utils import is_310p, vllm_version_is
|
||||
from vllm_ascend.ascend_config import get_ascend_config
|
||||
from vllm_ascend.sample.penalties import apply_all_penalties
|
||||
from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type, global_stream, npu_stream_switch
|
||||
|
||||
if vllm_version_is("0.10.2"):
|
||||
from vllm.config import LogprobsMode
|
||||
DEFAULT_LOGPROBS_MODE = LogprobsMode.RAW_LOGPROBS
|
||||
else:
|
||||
DEFAULT_LOGPROBS_MODE = "raw_logprobs"
|
||||
DEFAULT_LOGPROBS_MODE = "raw_logprobs"
|
||||
|
||||
_SAMPLING_EPS = 1e-5
|
||||
|
||||
|
||||
def random_sample(
|
||||
probs: torch.Tensor,
|
||||
generators: dict[int, torch.Generator],
|
||||
) -> torch.Tensor:
|
||||
"""Randomly sample from the probabilities.
|
||||
|
||||
We use this function instead of torch.multinomial because torch.multinomial
|
||||
causes CPU-NPU synchronization.
|
||||
"""
|
||||
# NOTE(woosuk): To batch-process the requests without their own seeds,
|
||||
# which is the common case, we first assume that every request does
|
||||
# not have its own seed. Then, we overwrite the values for the requests
|
||||
# that have their own seeds.
|
||||
with npu_stream_switch(global_stream()):
|
||||
q = torch.empty_like(probs)
|
||||
if len(generators) != probs.shape[0]:
|
||||
q.exponential_()
|
||||
if generators:
|
||||
# TODO(woosuk): This can be slow because we handle each request
|
||||
# one by one. Optimize this.
|
||||
for i, generator in generators.items():
|
||||
q[i].exponential_(generator=generator)
|
||||
torch.npu.current_stream().wait_stream(global_stream())
|
||||
return probs.div_(q).argmax(dim=-1).view(-1)
|
||||
|
||||
|
||||
class AscendSampler(Sampler):
|
||||
@staticmethod
|
||||
def apply_penalties(
|
||||
logits: torch.Tensor,
|
||||
sampling_metadata: SamplingMetadata,
|
||||
output_token_ids: list[list[int]],
|
||||
) -> torch.Tensor:
|
||||
"""Use Triton-Ascend penalties on NPU when Triton is available; else vLLM default."""
|
||||
if not HAS_TRITON:
|
||||
logger.warning_once(
|
||||
"[sample/sampler] Triton not available, falling back to vLLM default "
|
||||
"penalty implementation. Penalty performance may be degraded on NPU. "
|
||||
)
|
||||
return Sampler.apply_penalties(logits, sampling_metadata, output_token_ids)
|
||||
|
||||
if sampling_metadata.no_penalties:
|
||||
return logits
|
||||
assert sampling_metadata.prompt_token_ids is not None
|
||||
return apply_all_penalties(
|
||||
logits,
|
||||
sampling_metadata.prompt_token_ids,
|
||||
sampling_metadata.presence_penalties,
|
||||
sampling_metadata.frequency_penalties,
|
||||
sampling_metadata.repetition_penalties,
|
||||
output_token_ids,
|
||||
)
|
||||
|
||||
def __init__(self, logprobs_mode=DEFAULT_LOGPROBS_MODE):
|
||||
# TODO: support logprobs_mode in vllm-ascend
|
||||
super().__init__(logprobs_mode=logprobs_mode)
|
||||
self.topk_topp_sampler = AscendTopKTopPSampler()
|
||||
self.topk_topp_sampler = AscendTopKTopPSampler(logprobs_mode=logprobs_mode)
|
||||
self.async_exponential_event = torch.npu.Event()
|
||||
logger.debug(
|
||||
"[sample/sampler] AscendSampler initialized. logprobs_mode=%s, triton_available=%s",
|
||||
logprobs_mode,
|
||||
HAS_TRITON,
|
||||
)
|
||||
|
||||
def set_q_event(self, q, event):
|
||||
self.topk_topp_sampler.set_q_event(q, event)
|
||||
|
||||
def prepare_sampling(self, top_k):
|
||||
self.topk_topp_sampler.prepare_sampling(top_k)
|
||||
|
||||
def do_async_exponential(self, b_s, head_dim, generators):
|
||||
# Calculating exponential randoms in a different stream
|
||||
# and overlapping with model executing.
|
||||
with torch.npu.stream(global_stream()):
|
||||
global_stream().wait_stream(torch.npu.current_stream())
|
||||
q = torch.empty((b_s, head_dim), device="npu", dtype=torch.float32)
|
||||
# Goes to async exponential with AI-CPU exponential or default exponential.
|
||||
if len(generators) != q.shape[0]:
|
||||
q.exponential_()
|
||||
if generators:
|
||||
for i, generator in generators.items():
|
||||
q[i].exponential_(generator=generator)
|
||||
self.async_exponential_event.record()
|
||||
self.set_q_event(q, self.async_exponential_event)
|
||||
|
||||
@staticmethod
|
||||
def greedy_sample(logits: torch.Tensor) -> torch.Tensor:
|
||||
if get_ascend_config().enable_reduce_sample:
|
||||
logger.debug_once(
|
||||
"[sample/sampler] Using reduce-sample greedy sampling. "
|
||||
"TP all-gather will be performed to find global argmax.",
|
||||
)
|
||||
tp_group = get_tp_group()
|
||||
B, V_local = logits.shape
|
||||
rank = tp_group.rank_in_group
|
||||
|
||||
local_max_logits, local_max_indices = logits.max(dim=-1)
|
||||
local_global_idx = local_max_indices + rank * V_local # [B]
|
||||
# [B, world_size]
|
||||
gathered_logits = tp_group.all_gather(local_max_logits.unsqueeze(-1), dim=-1)
|
||||
gathered_global_idx = tp_group.all_gather(local_global_idx.unsqueeze(-1), dim=-1) # [B, world_size]
|
||||
global_max_rank = gathered_logits.argmax(dim=-1) # [B]
|
||||
target_argmax = gathered_global_idx.gather(dim=-1, index=global_max_rank.unsqueeze(-1)).squeeze(-1) # [B]
|
||||
return target_argmax
|
||||
else:
|
||||
return logits.argmax(dim=-1).view(-1)
|
||||
|
||||
|
||||
class AscendTopKTopPSampler(TopKTopPSampler):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.apply_top_k_top_p = apply_top_k_top_p
|
||||
self.top_k = None
|
||||
|
||||
def _apply_top_k_top_p(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
p: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# npu_top_k_top_p uses the operator aclnnApplyTopKTopP, but aclnnApplyTopKTopP currently does not support 310P
|
||||
if not is_310p() and p is not None and k is not None and 1 <= int(
|
||||
k.max()) <= 1024:
|
||||
# npu_top_k_top_p's parameter order is (logits, p, k), not (logits, k, p)
|
||||
return torch_npu.npu_top_k_top_p(logits, p, k)
|
||||
def set_q_event(self, q, event):
|
||||
# Pass in async exponential results.
|
||||
# Also pass in event to prevent synchronize errors.
|
||||
self.q = q
|
||||
self.async_event = event
|
||||
|
||||
def prepare_sampling(self, top_k):
|
||||
if top_k is not None:
|
||||
self.top_k = top_k
|
||||
else:
|
||||
self.top_k = None
|
||||
|
||||
def forward_native(self, logits, generators, k, p):
|
||||
"""Override pytorch native implementation to torch_npu"""
|
||||
# when batch_invariant mode is enabled, we should use vllm's implementation.
|
||||
# or it will make batch_invariant mode not working.
|
||||
if envs.VLLM_BATCH_INVARIANT:
|
||||
logger.debug_once(
|
||||
"[sample/sampler] BATCH_INVARIANT mode enabled, "
|
||||
"falling back to vLLM native top-k/top-p implementation.",
|
||||
)
|
||||
return super().forward_native(logits, generators, k, p)
|
||||
|
||||
if get_ascend_config().enable_reduce_sample:
|
||||
logger.debug_once(
|
||||
"[sample/sampler] Using reduce-sample path in forward_native. "
|
||||
"top-k/top-p with TP all-gather for distributed sampling.",
|
||||
)
|
||||
cand_logits, cand_idx = self.apply_top_k_top_p(logits, k, p, self.top_k)
|
||||
logits_to_return = None
|
||||
if self.logprobs_mode == "processed_logits":
|
||||
logits_to_return = cand_logits
|
||||
elif self.logprobs_mode == "processed_logprobs":
|
||||
logits_to_return = cand_logits.log_softmax(dim=-1, dtype=torch.float32)
|
||||
|
||||
probs = cand_logits.softmax(dim=-1, dtype=torch.float32)
|
||||
pos = random_sample(probs, generators) # [B]
|
||||
|
||||
next_token = cand_idx.gather(dim=1, index=pos.unsqueeze(1)).squeeze(1) # [B]
|
||||
return next_token, logits_to_return
|
||||
else:
|
||||
logits = self.apply_top_k_top_p(logits, k, p)
|
||||
logits_to_return = None
|
||||
if self.logprobs_mode == "processed_logits":
|
||||
logits_to_return = logits
|
||||
elif self.logprobs_mode == "processed_logprobs":
|
||||
logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
|
||||
|
||||
probs = logits.softmax(dim=-1, dtype=torch.float32)
|
||||
if get_ascend_config().enable_async_exponential:
|
||||
# Add synchronize to prevent synchronize error.
|
||||
logger.debug_once(
|
||||
"[sample/sampler] Using async-exponential sampling path. "
|
||||
"Pre-computed exponential randoms from separate stream will be used.",
|
||||
)
|
||||
self.async_event.synchronize()
|
||||
return probs.div_(self.q).argmax(dim=-1).view(-1), logits_to_return
|
||||
return random_sample(probs, generators), logits_to_return
|
||||
|
||||
|
||||
def _apply_top_k_top_p_pytorch(
|
||||
logits: torch.Tensor, # [B, V_local]
|
||||
k: torch.Tensor, # [B] or None
|
||||
p: torch.Tensor, # [B] or None
|
||||
top_k: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
if get_ascend_config().enable_reduce_sample:
|
||||
tp_group = get_tp_group()
|
||||
B, V_local = logits.shape
|
||||
rank = tp_group.rank_in_group
|
||||
|
||||
if top_k is None or (p is None and k is None):
|
||||
k_for_topk = V_local
|
||||
else:
|
||||
k_for_topk = min(top_k, V_local)
|
||||
|
||||
local_vals, local_idx = torch.topk(logits, k=k_for_topk, dim=-1)
|
||||
local_global_idx = local_idx + rank * V_local
|
||||
gathered_vals = tp_group.all_gather(local_vals, dim=-1)
|
||||
gathered_idx = tp_group.all_gather(local_global_idx, dim=-1)
|
||||
|
||||
if p is None and k is None:
|
||||
return gathered_vals, gathered_idx
|
||||
|
||||
probs = gathered_vals.softmax(dim=-1)
|
||||
probs_sort, _ = probs.sort(dim=-1, descending=False)
|
||||
if k is not None:
|
||||
kk = k.to(torch.long).clamp(min=1, max=V_local)
|
||||
top_k_count = (probs_sort.size(1) - kk).unsqueeze(1) # [B,1]
|
||||
top_k_cutoff = probs_sort.gather(-1, top_k_count)
|
||||
no_top_k_mask = (kk == V_local).unsqueeze(1)
|
||||
top_k_cutoff.masked_fill_(no_top_k_mask, -float("inf"))
|
||||
elements_to_discard = probs < top_k_cutoff
|
||||
gathered_vals.masked_fill_(elements_to_discard, -float("inf"))
|
||||
if p is not None:
|
||||
cumprob = torch.cumsum(probs_sort, dim=-1)
|
||||
top_p_mask = cumprob <= (1 - p.unsqueeze(1))
|
||||
top_p_mask[:, -1] = False # at least one
|
||||
top_p_count = top_p_mask.sum(dim=-1, keepdim=True)
|
||||
top_p_cutoff = probs_sort.gather(-1, top_p_count)
|
||||
elements_to_discard = probs < top_p_cutoff
|
||||
gathered_vals.masked_fill_(elements_to_discard, -float("inf"))
|
||||
return gathered_vals, gathered_idx
|
||||
else:
|
||||
if p is None and k is None:
|
||||
return logits
|
||||
|
||||
@@ -41,8 +242,7 @@ class AscendTopKTopPSampler(TopKTopPSampler):
|
||||
probs_sort, _ = probs.sort(dim=-1, descending=False)
|
||||
|
||||
if k is not None:
|
||||
top_k_count = probs_sort.size(1) - k.to(
|
||||
torch.long) # shape: (batch, )
|
||||
top_k_count = probs_sort.size(1) - k.to(torch.long) # shape: (batch, )
|
||||
top_k_count = top_k_count.unsqueeze(dim=1)
|
||||
top_k_cutoff = probs_sort.gather(-1, top_k_count)
|
||||
|
||||
@@ -65,22 +265,39 @@ class AscendTopKTopPSampler(TopKTopPSampler):
|
||||
|
||||
return logits
|
||||
|
||||
def forward_native(self, logits, generators, k, p):
|
||||
"""Override pytorch native implementation to torch_npu"""
|
||||
logits = self._apply_top_k_top_p(logits, k, p)
|
||||
logits_to_return = None
|
||||
if vllm_version_is("0.10.2"):
|
||||
if self.logprobs_mode == LogprobsMode.PROCESSED_LOGITS:
|
||||
logits_to_return = logits
|
||||
elif self.logprobs_mode == LogprobsMode.PROCESSED_LOGPROBS:
|
||||
logits_to_return = logits.log_softmax(dim=-1,
|
||||
dtype=torch.float32)
|
||||
else:
|
||||
if self.logprobs_mode == "processed_logits":
|
||||
logits_to_return = logits
|
||||
elif self.logprobs_mode == "processed_logprobs":
|
||||
logits_to_return = logits.log_softmax(dim=-1,
|
||||
dtype=torch.float32)
|
||||
|
||||
probs = logits.softmax(dim=-1, dtype=torch.float32)
|
||||
return random_sample(probs, generators), logits_to_return
|
||||
def _apply_top_k_top_p_torch_npu(
|
||||
logits: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
p: torch.Tensor,
|
||||
top_k: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
if get_ascend_config().enable_reduce_sample:
|
||||
tp_group = get_tp_group()
|
||||
B, V_local = logits.shape
|
||||
rank = tp_group.rank_in_group
|
||||
|
||||
if top_k is None or (p is None and k is None):
|
||||
k_for_topk = V_local
|
||||
else:
|
||||
k_for_topk = min(top_k, V_local)
|
||||
|
||||
local_vals, local_idx = torch.topk(logits, k=k_for_topk, dim=-1)
|
||||
local_global_idx = local_idx + rank * V_local
|
||||
gathered_vals = tp_group.all_gather(local_vals, dim=-1)
|
||||
gathered_idx = tp_group.all_gather(local_global_idx, dim=-1)
|
||||
|
||||
if not (p is None and k is None):
|
||||
gathered_vals = torch_npu.npu_top_k_top_p(gathered_vals, k=k, p=p)
|
||||
return gathered_vals, gathered_idx
|
||||
|
||||
if p is None and k is None:
|
||||
return logits
|
||||
return torch_npu.npu_top_k_top_p(logits, k=k, p=p)
|
||||
|
||||
|
||||
apply_top_k_top_p = (
|
||||
_apply_top_k_top_p_torch_npu
|
||||
if get_ascend_device_type() in [AscendDeviceType.A2, AscendDeviceType.A3]
|
||||
else _apply_top_k_top_p_pytorch
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user