Fix off-by-one num_segments in matrix::sort_cols_per_row (#3010)

Closes [#2049](NVIDIA/cuvs#2049)

`raft::matrix::detail::sortColumnsPerRow` passed `n_rows + 1` to [cub as `num_segments`](https://github.com/NVIDIA/cccl/blob/01b00d00edff5e5ed01fa4921ed9a5ad703c3219/cub/cub/device/device_segmented_radix_sort.cuh#L217). Per the cub contract, an aliased CSR offsets array must have length `num_segments + 1`, so cub read one int past the offsets allocation. This crashes with `cudaErrorIllegalAddress` when `n_rows + 1` is a power of two ≥ 64 and `n_columns ≥ ~16384` (silent otherwise).

Surfaces in `cuvs::stats::trustworthiness_score` whenever `n % batch_size ∈ {63, 127, 255, …}` (e.g. `n=76927, batch_size=512`).

Fix : pass `n_rows` to cub; size the offsets array as `n_rows + 1`.

Authors:
  - Victor Lafargue (https://github.com/viclafargue)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - Corey J. Nolet (https://github.com/cjnolet)
  - Tarang Jain (https://github.com/tarang-jain)

URL: #3010

Author: viclafargue

cpp/include/raft/matrix/detail/columnWiseSort.cuh

@@ -221,7 +221,12 @@ void sortColumnsPerRow(const InType* in,
     // device Segmented radix sort
     // 2^18 column cap to restrict size of workspace ~512 MB
     // will give better perf than below deviceWide Sort for even larger dims
-    int numSegments = n_rows + 1;
+
+    // n_rows CSR-style segments aliased onto an n_rows + 1 offsets array,
+    // with begin = &offsets[0] and end = &offsets[1]. cub's `num_segments`
+    // is the actual segment count (n_rows).
+    int numSegments = n_rows;
+    int numOffsets  = n_rows + 1;
 
     // need auxiliary storage: cub sorting + keys (if user not passing) +
     // staging for values out + segment partition
@@ -248,8 +253,8 @@ void sortColumnsPerRow(const InType* in,
       // value in KV pair need to be passed in, out buffer is separate
       workspaceSize += raft::alignTo(sizeof(OutType) * (size_t)totalElements, memAlignWidth);
 
-      // for segment offsets
-      workspaceSize += raft::alignTo(sizeof(int) * (size_t)numSegments, memAlignWidth);
+      // for segment offsets (numOffsets = numSegments + 1, see above)
+      workspaceSize += raft::alignTo(sizeof(int) * (size_t)numOffsets, memAlignWidth);
     } else {
       size_t workspaceOffset = 0;
 
@@ -264,14 +269,14 @@ void sortColumnsPerRow(const InType* in,
       workspacePtr       = (void*)((size_t)workspacePtr + workspaceOffset);
 
       int* dSegmentOffsets = reinterpret_cast<int*>(workspacePtr);
-      workspaceOffset      = raft::alignTo(sizeof(int) * (size_t)numSegments, memAlignWidth);
+      workspaceOffset      = raft::alignTo(sizeof(int) * (size_t)numOffsets, memAlignWidth);
       workspacePtr         = (void*)((size_t)workspacePtr + workspaceOffset);
 
       // layout idx
       RAFT_CUDA_TRY(layoutIdx(dValuesIn, n_rows, n_columns, stream));
 
       // layout segment lengths - spread out column length
-      RAFT_CUDA_TRY(layoutSortOffset(dSegmentOffsets, n_columns, numSegments, stream));
+      RAFT_CUDA_TRY(layoutSortOffset(dSegmentOffsets, n_columns, numOffsets, stream));
 
       RAFT_CUDA_TRY(cub::DeviceSegmentedRadixSort::SortPairs(workspacePtr,
                                                              workspaceSize,