diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh
--- a/cpp/include/raft/matrix/detail/select_warpsort.cuh
+++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh
@@ -31,6 +31,8 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cub/util_type.cuh>  // cub::Traits
+
 #include <algorithm>
 #include <functional>
 #include <type_traits>
@@ -765,7 +767,13 @@ template <template <int, bool, typename,
                                                 IdxT* out_idx)
 {
   extern __shared__ __align__(256) uint8_t smem_buf_bytes[];
-  using bq_t         = block_sort<WarpSortClass, Capacity, Ascending, T, IdxT>;
+  // Map the value type T to an appropriate uint type using cub traits.
+  // This gives us two advantages:
+  //   1. More stable behavior for IEEE 754 floating-point values w.r.t. infinity and NaN
+  //      (the behavior matches the select_radix implementation)
+  //   2. Fewer expensive template instantiations
+  using bits_t       = typename cub::Traits<T>::UnsignedBits;
+  using bq_t         = block_sort<WarpSortClass, Capacity, Ascending, bits_t, IdxT>;
   uint8_t* warp_smem = bq_t::queue_t::mem_required(blockDim.x) > 0 ? smem_buf_bytes : nullptr;
   bq_t queue(k, warp_smem);
   const size_t batch_id = blockIdx.y;
@@ -779,13 +787,17 @@ template <template <int, bool, typename,
   const IdxT stride         = gridDim.x * blockDim.x;
   const IdxT per_thread_lim = l_len + laneId();
   for (IdxT i = threadIdx.x + blockIdx.x * blockDim.x; i < per_thread_lim; i += stride) {
-    queue.add(i < l_len ? __ldcs(in + i) : WarpSortClass<Capacity, Ascending, T, IdxT>::kDummy,
+    // Twiddle the input value to ensure proper comparison of floating-point and signed int values
+    queue.add(i < l_len ? cub::Traits<T>::TwiddleIn(__ldcs(reinterpret_cast<const bits_t*>(in) + i))
+                        : WarpSortClass<Capacity, Ascending, bits_t, IdxT>::kDummy,
               (i < l_len && in_idx != nullptr) ? __ldcs(in_idx + i) : i);
   }
 
   queue.done(smem_buf_bytes);
   const int block_id = blockIdx.x + gridDim.x * blockIdx.y;
-  queue.store(out + block_id * k, out_idx + block_id * k);
+  queue.store(reinterpret_cast<bits_t*>(out) + block_id * k,
+              out_idx + block_id * k,
+              cub::Traits<T>::TwiddleOut);  // Restore the FP / signed int representation.
 }
 
 struct launch_params {
diff --git a/cpp/tests/matrix/select_k.cuh b/cpp/tests/matrix/select_k.cuh
--- a/cpp/tests/matrix/select_k.cuh
+++ b/cpp/tests/matrix/select_k.cuh
@@ -234,12 +234,17 @@ struct SelectK  // NOLINT
 
     // If the dists (keys) are the same, different corresponding ids may end up in the selection
     // due to non-deterministic nature of some implementations.
-    auto compare_ids = [this](const IdxT& i, const IdxT& j) {
+    auto compare_ids = [this](const IdxT& i, const IdxT& j, const size_t pos) {
       if (i == j) return true;
       auto& in_ids   = ref.get_in_ids();
       auto& in_dists = ref.get_in_dists();
-      auto ix_i = static_cast<int64_t>(std::find(in_ids.begin(), in_ids.end(), i) - in_ids.begin());
-      auto ix_j = static_cast<int64_t>(std::find(in_ids.begin(), in_ids.end(), j) - in_ids.begin());
+      auto batch     = pos / spec.k;
+      auto start_pos = in_ids.begin() + batch * spec.len;
+      // Get the indices of found neighbor ids in the source (input ids)
+      auto ix_i =
+        static_cast<int64_t>(std::find(start_pos, start_pos + spec.len, i) - in_ids.begin());
+      auto ix_j =
+        static_cast<int64_t>(std::find(start_pos, start_pos + spec.len, j) - in_ids.begin());
       auto forgive_i = forgive_algo(ref.algo, i);
       auto forgive_j = forgive_algo(res.algo, j);
       // Some algorithms return invalid indices in special cases.
diff --git a/cpp/tests/sparse/select_k_csr.cu b/cpp/tests/sparse/select_k_csr.cu
--- a/cpp/tests/sparse/select_k_csr.cu
+++ b/cpp/tests/sparse/select_k_csr.cu
@@ -56,7 +56,7 @@ struct CompareApproxWithInf {
   CompareApproxWithInf(T eps_) : eps(eps_) {}
   bool operator()(const T& a, const T& b) const
   {
-    if (std::isinf(a) && std::isinf(b)) return true;
+    if ((std::isinf(a) || std::isnan(a)) && (std::isinf(b) || std::isnan(b))) return true;
     T diff  = std::abs(a - b);
     T m     = std::max(std::abs(a), std::abs(b));
     T ratio = diff > eps ? diff / m : diff;
diff --git a/cpp/tests/test_utils.cuh b/cpp/tests/test_utils.cuh
--- a/cpp/tests/test_utils.cuh
+++ b/cpp/tests/test_utils.cuh
@@ -188,9 +188,16 @@ testing::AssertionResult hostVecMatch(co
            << "vector sizez mismatch: "
            << "actual=" << n << " != expected=" << expected_h.size() << "; ";
   for (size_t i = 0; i < n; ++i) {
-    auto exp = expected_h[i];
-    auto act = actual_h[i];
-    if (!eq_compare(exp, act)) {
+    auto exp       = expected_h[i];
+    auto act       = actual_h[i];
+    bool are_equal = [&]() {
+      if constexpr (std::is_invocable_v<L, decltype(exp), decltype(act), size_t>) {
+        return eq_compare(exp, act, i);
+      } else {
+        return eq_compare(exp, act);
+      }
+    }();
+    if (!are_equal) {
       return testing::AssertionFailure()
              << "actual=" << act << " != expected=" << exp << " @" << i << "; ";
     }