diff --git a/src/cuda/cuda_topk_distributed_select_sort.cuh b/src/cuda/cuda_topk_distributed_select_sort.cuh index 8a1a01c164..6d91c5a7a5 100644 --- a/src/cuda/cuda_topk_distributed_select_sort.cuh +++ b/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(); } diff --git a/src/cuda/cuda_topk_select_sort.cuh b/src/cuda/cuda_topk_select_sort.cuh index 17553540fa..345ea4bd5a 100644 --- a/src/cuda/cuda_topk_select_sort.cuh +++ b/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(); } diff --git a/test/cuda/cuda_topk_tests.cpp b/test/cuda/cuda_topk_tests.cpp index 3b62caa90c..60e4fe778d 100644 --- a/test/cuda/cuda_topk_tests.cpp +++ b/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(params.batch_size) * params.vocab_size; + size_t topk_size = static_cast(params.batch_size) * params.k; + + auto scores_in_d = Generators::CudaMallocArray(total_vocab_size); + + // Fill input with all NaN values + std::vector scores_in_h(total_vocab_size, std::numeric_limits::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(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 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 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