[BUG] matrix::select_k outputs 0.0 for negative bfloat16 numbers

[apivovarov]

Describe the bug
When using raft::matrix::select_k with nv_bfloat16 (bfloat16) input values, any negative numbers in the input are incorrectly converted to 0.0 in the output. The corresponding indices are also incorrect, showing all zeros.

Steps/Code to reproduce bug
GTest Code: PR-2795

./build.sh -v tests

./cpp/build/gtests/MATRIX_SELECT_TEST

Running main() from /root/workspace/raft/cpp/build/_deps/gtest-src/googletest/src/gtest_main.cc
[==========] Running 608 tests from 6 test suites.
[----------] Global test environment set-up.
[----------] 1 test from SelectKBF16Negative
[ RUN      ] SelectKBF16Negative.TopK3
Input: -3.5 -1.5 -7.5 0.5 -2.5
Output: 0.5 0.0 0.0

Expected behavior
Negative input values should be correctly reflected in the output, e.g., [0.5, -1.5, -2.5], with the correct corresponding indices.

Environment details (please complete the following information):

• branch: branch-25.10
• Docker container: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest

Additional context
None

[apivovarov]

@cjnolet @divyegala @beckernick Could you take a look at this issue? It’s blocking the integration of XLA hlo.topk with raft::select_k, which noticeably improves LLM model training performance.

[apivovarov]

Solved the issue by registering nv_bfloat16 type in https://github.com/rapidsai/raft/blob/branch-25.10/cpp/include/raft/util/cudart_utils.hpp#L442

// Negative infinity in bfloat16 (sign=1, exp=all ones, mantissa=0)
lower_bound - 0xff80u

// Positive infinity in bfloat16 (sign=0, exp=all ones, mantissa=0)
upper_bound - 0x7f80u

Will open PR with corresponding unit tests

[apivovarov]

#2799 - Fix select_k for negative bfloat16