DeepInfer/deepinfer/neuron_ops/vendor/deep_select/interface.py

"""Pure-torch drop-in for ``deep_select.interface`` (upstream:
``vendor/DeepSelect/deep_select/interface.py``).

Output contract taken from the upstream docstring and the checks in
``vendor/DeepSelect/tests/test.py`` (lines 70-148):

- ``output_val[i][j] == input[i][output_idx[i][j] - output_idx_offset[i]]`` bitwise, for the
  ``min(end[i], topk)`` valid positions of a row;
- indices are unique and inside ``[0, end[i])``;
- ``min(selected) >= max(unselected)`` within the valid range;
- rows with fewer valid entries than ``topk`` gather their valid results at the beginning and pad
  the rest with ``value_oob_fill_value`` / ``idx_oob_fill_value`` (the fill index is *not*
  shifted by ``output_idx_offset``);
- ``sorted=True`` returns values in descending order, ``sorted_index=True`` returns indices in
  ascending order (upstream never sets both);
- a row containing NaN inside its valid range has its whole ``output_idx`` row filled with
  ``0x3F3F3F3F`` (the Ascend behaviour, ``csrc/ascend_kernels/kernel.asc:190``).

Deviation: ``abort_when_nan_found=True`` aborts the kernel upstream; aborting here would need a
host sync, so the NaN sentinel fill is used for both settings.
"""

from __future__ import annotations

import functools
from typing import Optional, Tuple

import torch

_NAN_SENTINEL = 0x3F3F3F3F


@functools.lru_cache(maxsize=1)
def get_stride_requirement() -> Tuple[int, int]:
    """(input, output) stride requirement in bytes. Ascend values
    (``vendor/DeepSelect/csrc/structs.h:39-40``); CUDA requires 1024 for the input."""
    return (32, 32)


def topk(
    input: torch.Tensor,
    topk: int,
    sorted: bool = False,
    begin: Optional[torch.Tensor] = None,
    end: Optional[torch.Tensor] = None,
    indices_type: torch.dtype = torch.int64,
    sorted_index: bool = False,
    hint: Optional[torch.Tensor] = None,
    output_idx: Optional[torch.Tensor] = None,
    output_idx_offset: Optional[torch.Tensor] = None,
    idx_oob_fill_value: int = 2147483647,
    value_oob_fill_value: float = float("-inf"),
    return_value: bool = True,
    abort_when_nan_found: bool = True,
) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
    """Row-wise top-k selection; see the module docstring for the output contract.

    Args:
        input: (b, vocab_size), bfloat16 or float32.
        topk: number of elements selected per row.
        sorted: return ``output_val`` sorted descending.
        begin: not supported (must be None), as upstream.
        end: (b,) int32, exclusive right bound of each row's valid range (default vocab_size).
        indices_type: torch.int32 or torch.int64.
        sorted_index: return ``output_idx`` sorted ascending.
        hint: not supported (must be None), as upstream.
        output_idx: (b, topk) tensor of ``indices_type`` to write the indices into, in place.
        output_idx_offset: (b,) int32 added to every valid index.
        idx_oob_fill_value / value_oob_fill_value: padding for rows with fewer than ``topk`` valid
            entries.
        return_value: when False the first element of the result tuple is None.
        abort_when_nan_found: see the module docstring (sentinel fill either way here).
    """
    assert begin is None, "`begin` is not supported now"
    assert hint is None, "`hint` is not supported now"
    assert input.dim() == 2, f"input must be 2D, got {tuple(input.shape)}"
    assert indices_type in (torch.int32, torch.int64), f"unsupported indices_type {indices_type}"
    assert topk > 0
    device = input.device
    b, vocab_size = input.shape

    # --- valid range -------------------------------------------------------
    cols = torch.arange(vocab_size, device=device, dtype=torch.int32).view(1, vocab_size)
    if end is None:
        valid_len = torch.full((b, 1), vocab_size, device=device, dtype=torch.int32)
        oob_cols = None
    else:
        assert end.numel() == b, (end.shape, b)
        valid_len = end.reshape(b, 1).to(torch.int32)
        oob_cols = cols >= valid_len  # [b, vocab_size]

    # --- select ------------------------------------------------------------
    # Out-of-range columns must never be selected: push them below every real value.
    scores = input.float()
    if oob_cols is not None:
        scores = torch.where(oob_cols, torch.full_like(scores, float("-inf")), scores)
    # NaN would otherwise win every comparison in torch.topk; keep NaN rows detectable but
    # deterministic by sorting NaN to the bottom (the row is overwritten with the sentinel below).
    is_nan = scores != scores
    scores = torch.where(is_nan, torch.full_like(scores, float("-inf")), scores)

    k_eff = min(topk, vocab_size)
    # torch.topk(sorted=True) returns values descending, which is what `sorted=True` asks for; it
    # is used unconditionally so the result is deterministic for every flag combination.
    from deepinfer.neuron_ops import nki_topk  # NEURON: torch.topk with a large k does not compile on the device

    if nki_topk.supported(scores, k_eff):
        idx = nki_topk.topk(scores, k_eff)
    else:
        idx = torch.topk(scores, k_eff, dim=-1, sorted=True).indices
    idx = idx.to(torch.int64)

    counts = torch.clamp(valid_len.to(torch.int64), max=topk)  # [b, 1]
    kcols = torch.arange(topk, device=device, dtype=torch.int64).view(1, topk)
    selected_mask = kcols < counts  # [b, topk]
    if k_eff < topk:
        pad = torch.zeros((b, topk - k_eff), dtype=torch.int64, device=device)
        idx = torch.cat([idx, pad], dim=1)

    if sorted_index:
        # Keep the valid prefix in ascending index order; the padding stays at the end.
        # Any value above every real index works as the sort sentinel. Keep it in int32 range: the neuron
        # compiler narrows int64 (--implicit-integer-downcast), which turned int64.max into -1 and sorted the
        # padding in front of the selected indices for rows with end < topk.
        big = torch.full_like(idx, vocab_size)
        idx = torch.sort(torch.where(selected_mask, idx, big), dim=1).values
        idx = torch.where(selected_mask, idx, torch.zeros_like(idx))

    # --- values ------------------------------------------------------------
    out_val = None
    if return_value:
        gathered = input.gather(1, idx)  # bitwise equal to the input elements
        fill = torch.full_like(gathered, value_oob_fill_value)
        out_val = torch.where(selected_mask, gathered, fill)

    # --- indices -----------------------------------------------------------
    out_idx = idx
    if output_idx_offset is not None:
        out_idx = out_idx + output_idx_offset.reshape(b, 1).to(torch.int64)
    out_idx = torch.where(selected_mask, out_idx, torch.full_like(out_idx, idx_oob_fill_value))

    # NaN rows: the whole index row becomes the sentinel.
    nan_rows = is_nan.any(dim=1, keepdim=True)
    out_idx = torch.where(nan_rows, torch.full_like(out_idx, _NAN_SENTINEL), out_idx)

    out_idx = out_idx.to(indices_type)
    if output_idx is not None:
        assert output_idx.dtype == indices_type
        assert tuple(output_idx.shape) == (b, topk), (output_idx.shape, (b, topk))
        output_idx.copy_(out_idx)
        out_idx = output_idx
    return out_val, out_idx