mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
27 lines
798 B
Python
27 lines
798 B
Python
from typing import Any
|
|
|
|
import torch
|
|
from deep_select.interface import topk as _topk
|
|
|
|
|
|
def _torch_topk(
|
|
input: torch.Tensor,
|
|
topk: int,
|
|
sorted: bool = False,
|
|
indices_type: torch.dtype = torch.int64,
|
|
**kwargs: Any,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
data, index = torch.topk(input, topk, dim=-1, sorted=sorted)
|
|
if index.dtype is not indices_type:
|
|
index = index.to(indices_type)
|
|
return data, index
|
|
|
|
|
|
# TODO(NPU): To be removed when `topk` for sampler is ready
|
|
def topk(*args: Any, **kwargs: Any) -> tuple[torch.Tensor | None, torch.Tensor]:
|
|
if kwargs.get("return_value", True):
|
|
# Sampler is not supported yet, fallback to torch.
|
|
return _torch_topk(*args, **kwargs)
|
|
# Indexer use real topk kernel.
|
|
return _topk(*args, **kwargs)
|