DeepInfer/tests/neuron/flash_mla/test_deep_select.py

"""CPU correctness tests for the pure-torch deep_select vendor package.

The checks are ported from vendor/DeepSelect/tests/test.py (NaN guard, index range, uniqueness,
gathered value, topk condition, sorted index / sorted value).
"""

import itertools
import sys

import torch

from deep_select.interface import get_stride_requirement, topk as ds_topk

torch.manual_seed(0)
FAILS = []
NAN_SENTINEL = 0x3F3F3F3F


def chk(name, cond):
    print(f"[{'ok ' if cond else 'FAIL'}] {name}")
    if not cond:
        FAILS.append(name)


def run_case(b, vocab, k, dtype, idx_dtype, sorted_value, sorted_index, return_value,
             use_end, use_offset, with_nan=False, use_out=False):
    tag = (f"b{b} v{vocab} k{k} {str(dtype).split('.')[-1]} {str(idx_dtype).split('.')[-1]} "
           f"sv{int(sorted_value)} si{int(sorted_index)} rv{int(return_value)} "
           f"end{int(use_end)} off{int(use_offset)} nan{int(with_nan)} out{int(use_out)}")
    x = (torch.randn(b, vocab) * 3).to(dtype)
    end = torch.randint(1, vocab + 1, (b,), dtype=torch.int32) if use_end else None
    if use_end:
        end[0] = min(k // 2 + 1, vocab)  # a short row, fewer valid entries than k
    offset = torch.randint(0, 1000, (b,), dtype=torch.int32) if use_offset else None
    nan_rows = torch.zeros(b, dtype=torch.bool)
    if with_nan:
        nan_rows[0] = True
        x[0, 0] = float("nan")
    out_buf = torch.empty((b, k), dtype=idx_dtype) if use_out else None

    val, idx = ds_topk(
        x, k, sorted=sorted_value, end=end, indices_type=idx_dtype, sorted_index=sorted_index,
        output_idx=out_buf, output_idx_offset=offset, idx_oob_fill_value=-1,
        value_oob_fill_value=float("-inf"), return_value=return_value, abort_when_nan_found=False,
    )
    chk(f"{tag}: dtypes", idx.dtype is idx_dtype and (val is None) == (not return_value)
        and (val is None or val.dtype is dtype))
    chk(f"{tag}: shapes", tuple(idx.shape) == (b, k) and (val is None or tuple(val.shape) == (b, k)))
    if use_out:
        chk(f"{tag}: writes output_idx in place", idx.data_ptr() == out_buf.data_ptr())

    valid_len = (torch.full((b,), vocab, dtype=torch.int64) if end is None else end.to(torch.int64))
    counts = torch.clamp(valid_len, max=k)
    sel = torch.arange(k).unsqueeze(0) < counts.unsqueeze(1)

    # NaN guard: whole row filled with the sentinel (Ascend behaviour)
    if with_nan:
        chk(f"{tag}: NaN row sentinel", bool((idx[nan_rows] == NAN_SENTINEL).all()))
    sel = sel & ~nan_rows.unsqueeze(1)

    raw = idx.to(torch.int64)
    if offset is not None:
        raw = raw - offset.to(torch.int64).unsqueeze(1)
    in_range = (raw >= 0) & (raw < valid_len.unsqueeze(1))
    chk(f"{tag}: index range", bool((in_range | ~sel).all()))
    chk(f"{tag}: padding fill", bool((idx[~sel & ~nan_rows.unsqueeze(1)] == -1).all()))

    srt = torch.sort(torch.where(sel, raw, torch.iinfo(torch.int64).max), dim=1).values
    dup = (srt[:, 1:] == srt[:, :-1]) & sel[:, 1:]
    chk(f"{tag}: unique index", not bool(dup.any()))

    safe = torch.where(sel & in_range, raw, torch.zeros_like(raw))
    gathered = x.gather(1, safe)
    rows = ~nan_rows
    if rows.any():
        if val is not None:
            exp = gathered[rows].masked_fill(~sel[rows], float("-inf"))
            chk(f"{tag}: gathered value bitwise", torch.equal(val[rows], exp))
        gv = gathered[rows].float()
        smin = gv.masked_fill(~sel[rows], float("inf")).amin(dim=1)
        unsel = x[rows].float().clone()
        unsel.scatter_(1, safe[rows], float("-inf"))
        col = torch.arange(vocab).unsqueeze(0)
        unsel = torch.where(col >= valid_len[rows].unsqueeze(1), torch.full_like(unsel, float("-inf")), unsel)
        unsel = torch.where(unsel != unsel, torch.full_like(unsel, float("-inf")), unsel)
        chk(f"{tag}: topk condition", bool((smin >= unsel.amax(dim=1)).all()))

    ordered = sel[:, 1:]
    if sorted_index:
        ok = (idx[:, 1:].to(torch.int64) >= idx[:, :-1].to(torch.int64)) | ~ordered
        chk(f"{tag}: sorted index", bool(ok.all()))
    if sorted_value and val is not None:
        ok = (val[:, :-1].float() >= val[:, 1:].float()) | ~ordered
        chk(f"{tag}: sorted value", bool(ok.all()))


if __name__ == "__main__":
    chk("stride requirement == (32, 32)", get_stride_requirement() == (32, 32))

    combos = [(False, False, False), (False, False, True), (True, False, True),
              (False, True, True), (False, True, False)]
    for dtype in (torch.float32, torch.bfloat16):
        for idx_dtype in (torch.int32, torch.int64):
            for sv, si, rv in combos:
                if sv and dtype is torch.bfloat16:
                    continue  # sorted_value is fp32-only upstream
                run_case(5, 97, 8, dtype, idx_dtype, sv, si, rv, use_end=True, use_offset=True)
    for use_end, use_offset in itertools.product((False, True), repeat=2):
        run_case(3, 64, 16, torch.float32, torch.int32, False, True, False, use_end, use_offset)
    run_case(4, 33, 8, torch.float32, torch.int32, False, True, False, True, True, with_nan=True)
    run_case(4, 33, 8, torch.float32, torch.int32, True, False, True, False, False, with_nan=True)
    run_case(3, 16, 8, torch.float32, torch.int32, False, True, False, True, True, use_out=True)
    # k larger than the vocab / every row short
    run_case(2, 6, 8, torch.float32, torch.int32, False, True, False, True, False)
    # DeepInfer indexer call shape: big-ish k, end, offset, sorted_index, no values
    run_case(2, 512, 128, torch.float32, torch.int32, False, True, False, True, True)

    print("\nFAILURES:", FAILS if FAILS else "none")
    sys.exit(1 if FAILS else 0)