Fix CUDA illegal memory access crash in TopK when inputs contain NaN

src/cuda/cuda_topk_distributed_select_sort.cuh

@@ -112,8 +112,13 @@ __global__ void GetTopKKernelDistributedSelectSort(float* scores_in, float* scor
 
     if (tid == 0) {
       distributed_scores_out_curr[ite] = top_k_sequence.u;
-      distributed_indices_out_curr[ite] = top_k_sequence.p;
-      scores_in[top_k_sequence.p] = MIN_FLOAT;
+      // Clamp invalid index to 0 for safety (NaN inputs leave p at INT_MAX)
+      distributed_indices_out_curr[ite] = top_k_sequence.p < vocab_size ? top_k_sequence.p : 0;
+      // Guard against NaN inputs: if all scores are NaN, no valid index is found
+      // and p remains INT_MAX (its init value). Writing to that index would be OOB.
+      if (top_k_sequence.p < vocab_size) {
+        scores_in[top_k_sequence.p] = MIN_FLOAT;
+      }
       __threadfence_block();
     }
     __syncthreads();
@@ -181,8 +186,13 @@ __global__ void GetTopKKernelDistributedSelectSort(float* scores_in, float* scor
         int vocab_index = top_k_sequence_reduced.p_indirection;
 
         scores_out[ite] = top_k_sequence_reduced.u;
-        indices_out[ite] = vocab_index;
-        shared_distributed_scores_out[index] = MIN_FLOAT;
+        // Guard against NaN: if no valid element was found, p_indirection stays at -1 (init value).
+        // Clamp to 0 to avoid downstream OOB in embedding lookups.
+        indices_out[ite] = (vocab_index >= 0 && vocab_index < vocab_size) ? vocab_index : 0;
+        // Guard against NaN inputs: if all scores are NaN, p stays at -1 (init value).
+        if (index >= 0 && index < num_top_k_shards * k) {
+          shared_distributed_scores_out[index] = MIN_FLOAT;
+        }
 
         __threadfence_block();
       }

src/cuda/cuda_topk_select_sort.cuh

@@ -79,8 +79,10 @@ __global__ void GetTop1Kernel(const float* scores_in, float* scores_out, int* in
 
   // Thread 0 writes the final result. No fence or write-back to scores_in is needed.
   if (tid == 0) {
+    // Guard against NaN inputs: clamp invalid index to 0
+    int safe_p = top_k_sequence.p < vocab_size ? top_k_sequence.p : 0;
     scores_out[batch] = top_k_sequence.u;
-    indices_out[batch] = top_k_sequence.p;
+    indices_out[batch] = safe_p;
   }
 }
 
@@ -104,9 +106,14 @@ __global__ void GetTopKKernel(volatile float* scores_in, float* scores_out, int*
     TopK_Pair top_k_sequence = BlockReduce(temp_storage).Reduce(partial, reduce_topk_op);
 
     if (tid == 0) {
+      // Guard against NaN inputs: if all scores are NaN, no valid index is found
+      // and p remains INT_MAX (its init value). Clamp to index 0 to avoid OOB.
+      int safe_p = top_k_sequence.p < vocab_size ? top_k_sequence.p : 0;
       scores_out[ite + batch * k] = top_k_sequence.u;
-      indices_out[ite + batch * k] = top_k_sequence.p;
-      scores_in[batch * vocab_size + top_k_sequence.p] = -FLT_MAX;
+      indices_out[ite + batch * k] = safe_p;
+      if (top_k_sequence.p < vocab_size) {
+        scores_in[batch * vocab_size + top_k_sequence.p] = -FLT_MAX;
+      }
 
       __threadfence_block();
     }

test/cuda/cuda_topk_tests.cpp

@@ -213,4 +213,125 @@ TEST(TopKTests, ParityTests) {
     RunParityTests(params);
   }
 }
+
+// Test that TopK kernels handle all-NaN inputs without crashing or producing OOB indices.
+// This guards against the bug where NaN comparisons always return false, leaving the
+// index at its sentinel value (INT_MAX or -1), which would cause an illegal memory access.
+void RunNaNSafetyTest(const TopKTestParams& params) {
+  std::cout << "\n--- Running NaN Safety Test with batch_size=" << params.batch_size
+            << ", vocab_size=" << params.vocab_size << ", k=" << params.k << " ---\n";
+
+  cudaStream_t stream;
+  CUDA_CHECK(cudaStreamCreate(&stream));
+
+  size_t total_vocab_size = static_cast<size_t>(params.batch_size) * params.vocab_size;
+  size_t topk_size = static_cast<size_t>(params.batch_size) * params.k;
+
+  auto scores_in_d = Generators::CudaMallocArray<float>(total_vocab_size);
+
+  // Fill input with all NaN values
+  std::vector<float> scores_in_h(total_vocab_size, std::numeric_limits<float>::quiet_NaN());
+  CUDA_CHECK(cudaMemcpy(scores_in_d.get(), scores_in_h.data(), scores_in_h.size() * sizeof(float), cudaMemcpyHostToDevice));
+
+  auto topk_data = std::make_unique<Generators::cuda::TopkDataCompact>(params.batch_size, params.vocab_size, stream);
+
+  auto test_algo = [&](const std::string& name, auto func) {
+    // Re-upload NaN data since some algorithms modify the input in-place
+    CUDA_CHECK(cudaMemcpy(scores_in_d.get(), scores_in_h.data(), scores_in_h.size() * sizeof(float), cudaMemcpyHostToDevice));
+
+    func();
+
+    topk_data->CompactOutput(params.batch_size, params.k, stream);
+    CUDA_CHECK(cudaStreamSynchronize(stream));
+
+    // Verify no CUDA errors (the main goal: no illegal memory access)
+    cudaError_t err = cudaGetLastError();
+    ASSERT_EQ(err, cudaSuccess) << "CUDA error after " << name << ": " << cudaGetErrorString(err);
+
+    // Verify all output indices are valid (within [0, vocab_size))
+    std::vector<int> actual_indices_h(topk_size);
+    CUDA_CHECK(cudaMemcpy(actual_indices_h.data(), topk_data->topk_indices_compact.get(),
+                          actual_indices_h.size() * sizeof(int), cudaMemcpyDeviceToHost));
+
+    for (size_t i = 0; i < topk_size; ++i) {
+      ASSERT_GE(actual_indices_h[i], 0)
+          << name << ": output index[" << i << "] = " << actual_indices_h[i] << " is negative";
+      ASSERT_LT(actual_indices_h[i], params.vocab_size)
+          << name << ": output index[" << i << "] = " << actual_indices_h[i]
+          << " is out of bounds (vocab_size=" << params.vocab_size << ")";
+    }
+
+    std::cout << "  [PASS] " << name << " (NaN Safety)" << std::endl;
+  };
+
+  test_algo(Generators::cuda::select_sort::kAlgorithmName, [&]() {
+    Generators::cuda::select_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                           params.vocab_size, params.batch_size, params.k);
+  });
+
+  test_algo(Generators::cuda::per_batch_radix_sort::kAlgorithmName, [&]() {
+    Generators::cuda::per_batch_radix_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                                    params.vocab_size, params.batch_size, params.k);
+  });
+
+  if (Generators::cuda::hybrid_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) {
+    test_algo(Generators::cuda::hybrid_sort::kAlgorithmName, [&]() {
+      Generators::cuda::hybrid_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                             params.vocab_size, params.batch_size, params.k);
+    });
+  }
+
+  if (Generators::cuda::iterative_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) {
+    test_algo(Generators::cuda::iterative_sort::kAlgorithmName, [&]() {
+      Generators::cuda::iterative_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                                params.vocab_size, params.batch_size, params.k);
+    });
+  }
+
+  if (Generators::cuda::cascaded_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) {
+    test_algo(Generators::cuda::cascaded_sort::kAlgorithmName, [&]() {
+      Generators::cuda::cascaded_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                               params.vocab_size, params.batch_size, params.k);
+    });
+  }
+
+  if (Generators::cuda::flash_convergent::IsSupported(params.batch_size, params.vocab_size, params.k)) {
+    test_algo(Generators::cuda::flash_convergent::kAlgorithmName, [&]() {
+      Generators::cuda::flash_convergent::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                                  params.vocab_size, params.batch_size, params.k);
+    });
+  }
+
+  if (Generators::cuda::distributed_select_sort::IsSupported(params.batch_size, params.vocab_size, params.k)) {
+    test_algo(Generators::cuda::distributed_select_sort::kAlgorithmName, [&]() {
+      Generators::cuda::distributed_select_sort::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                                                         params.vocab_size, params.batch_size, params.k);
+    });
+  }
+
+  // Test the default RunTopK dispatcher
+  test_algo("DEFAULT", [&]() {
+    Generators::cuda::RunTopK(topk_data.get(), stream, scores_in_d.get(),
+                              params.vocab_size, params.batch_size, params.k);
+  });
+
+  CUDA_CHECK(cudaStreamDestroy(stream));
+}
+
+TEST(TopKTests, NaNSafetyTests) {
+  // Test representative configurations that exercise different code paths
+  std::vector<TopKTestParams> test_cases = {
+      {1, 200, 1},
+      {1, 200, 4},
+      {4, 2000, 1},
+      {4, 2000, 16},
+      {1, 200000, 1},
+      {1, 200000, 32},
+      {32, 20000, 64},
+  };
+
+  for (const auto& params : test_cases) {
+    RunNaNSafetyTest(params);
+  }
+}
 #endif